diff --git a/.buildkite/bootstrap-amd-omni.sh b/.buildkite/amd/scripts/bootstrap-amd-omni.sh similarity index 94% rename from .buildkite/bootstrap-amd-omni.sh rename to .buildkite/amd/scripts/bootstrap-amd-omni.sh index 16ed02b266c..ee858068a23 100755 --- a/.buildkite/bootstrap-amd-omni.sh +++ b/.buildkite/amd/scripts/bootstrap-amd-omni.sh @@ -154,14 +154,14 @@ upload_pipeline() { fi # Use local template file for vllm-omni - cp .buildkite/test-template-amd-omni.j2 .buildkite/test-template.j2 + cp .buildkite/amd/test-template-amd-omni.j2 .buildkite/amd/test-template.j2 # (WIP) Use pipeline generator instead of jinja template - if [ -e ".buildkite/pipeline_generator/pipeline_generator.py" ]; then + if [ -e ".buildkite/amd/pipeline_generator/pipeline_generator.py" ]; then python -m pip install click pydantic - python .buildkite/pipeline_generator/pipeline_generator.py --run_all=$RUN_ALL --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW" - buildkite-agent pipeline upload .buildkite/pipeline.yaml + python .buildkite/amd/pipeline_generator/pipeline_generator.py --run_all=$RUN_ALL --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW" + buildkite-agent pipeline upload .buildkite/amd/pipeline.yaml exit 0 fi echo "List file diff: $LIST_FILE_DIFF" @@ -171,7 +171,7 @@ upload_pipeline() { FAIL_FAST=$(fail_fast) - cd .buildkite + cd .buildkite/amd # Select test definition file: merge suite for main, ready suite for PRs. # For debugging, DEBUG_TEST_YAML can override the selection — accepts @@ -182,7 +182,7 @@ upload_pipeline() { TEST_YAML="test-amd-merge.yml" ;; ready) - TEST_YAML="test-amd-ready.yaml" + TEST_YAML="test-amd-ready.yml" ;; *) echo "ERROR: DEBUG_TEST_YAML must be 'merge' or 'ready', got '$DEBUG_TEST_YAML'" >&2 @@ -193,7 +193,7 @@ upload_pipeline() { elif [[ $BUILDKITE_BRANCH == "main" ]]; then TEST_YAML="test-amd-merge.yml" else - TEST_YAML="test-amd-ready.yaml" + TEST_YAML="test-amd-ready.yml" fi ( diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/amd/scripts/run-amd-test.sh similarity index 98% rename from .buildkite/scripts/hardware_ci/run-amd-test.sh rename to .buildkite/amd/scripts/run-amd-test.sh index f880455d479..c49614e70e3 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/amd/scripts/run-amd-test.sh @@ -56,7 +56,7 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea # # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images # sudo apt-get update && sudo apt-get install -y awscli ## Use safe docker login helper to prevent race conditions -# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh" +# source "$(dirname "${BASH_SOURCE[0]}")/../../common/scripts/docker_login_ecr_public.sh" # safe_docker_login_ecr_public ## Pull the container from AMD Docker Hub diff --git a/.buildkite/test-amd-merge.yml b/.buildkite/amd/test-amd-merge.yml similarity index 89% rename from .buildkite/test-amd-merge.yml rename to .buildkite/amd/test-amd-merge.yml index 1a5217b84ae..63433cc12e3 100644 --- a/.buildkite/test-amd-merge.yml +++ b/.buildkite/amd/test-amd-merge.yml @@ -7,17 +7,29 @@ steps: - group: ":card_index_dividers: Simple Test" steps: - - label: "Simple · Diffusion & Model Executor Test" + - label: "Simple · Diffusion & Model Executor Test (1/2)" agent_pool: mi325_1 depends_on: amd-build - timeout_in_minutes: 50 + timeout_in_minutes: 70 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" + # CPU tests on ROCm gpu machines are very slow, so we give them more time. + - "timeout 60m pytest -sv tests/diffusion -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" + + - label: "Simple · Diffusion & Model Executor Test (2/2)" + agent_pool: mi325_1 + depends_on: amd-build + timeout_in_minutes: 70 + mirror_hardwares: [amdproduction] + grade: Blocking + commands: + - export VLLM_ROCM_USE_AITER=0 + # CPU tests on ROCm gpu machines are very slow, so we give them more time. + - "timeout 60m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml" - label: "Simple · Other Test" agent_pool: mi325_1 @@ -158,7 +170,7 @@ steps: grade: Blocking commands: - | - timeout 20m bash -c ' + timeout 50m bash -c ' export VLLM_LOGGING_LEVEL=DEBUG export VLLM_WORKER_MULTIPROC_METHOD=spawn export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" @@ -170,26 +182,14 @@ steps: depends_on: amd-build mirror_hardwares: [amdproduction] grade: Blocking + timeout_in_minutes: 50 commands: - | - timeout 30m bash -c ' + timeout 40m 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" + pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model' ' # TODO: Bagel test on ROCm is very unstable. @tjtanaa diff --git a/.buildkite/test-amd-ready.yaml b/.buildkite/amd/test-amd-ready.yml similarity index 88% rename from .buildkite/test-amd-ready.yaml rename to .buildkite/amd/test-amd-ready.yml index 9da9507c8de..8c645fff002 100644 --- a/.buildkite/test-amd-ready.yaml +++ b/.buildkite/amd/test-amd-ready.yml @@ -161,19 +161,6 @@ steps: - 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 @@ -185,20 +172,6 @@ steps: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model" - - label: "Voxtral-TTS E2E Test" - agent_pool: mi325_1 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - | - timeout 20m bash -c ' - export VLLM_LOGGING_LEVEL=DEBUG - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" - pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model" - ' - - label: "CosyVoice3-TTS E2E Test" agent_pool: mi325_1 depends_on: amd-build @@ -207,7 +180,7 @@ steps: commands: - | timeout 20m bash -c ' - pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model" + pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "slow" --run-level "core_model" ' # TODO: Bagel test on ROCm is very unstable. @tjtanaa diff --git a/.buildkite/test-template-amd-omni.j2 b/.buildkite/amd/test-template-amd-omni.j2 similarity index 96% rename from .buildkite/test-template-amd-omni.j2 rename to .buildkite/amd/test-template-amd-omni.j2 index aa5ca8bcd5d..2f1a9aed31a 100644 --- a/.buildkite/test-template-amd-omni.j2 +++ b/.buildkite/amd/test-template-amd-omni.j2 @@ -57,7 +57,7 @@ steps: {% 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 + command: bash .buildkite/amd/scripts/run-amd-test.sh env: DOCKER_BUILDKIT: "1" {% if step.env %} @@ -104,7 +104,7 @@ steps: {% 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 + command: bash .buildkite/amd/scripts/run-amd-test.sh env: DOCKER_BUILDKIT: "1" {% if step.env %} diff --git a/.buildkite/pipeline.yml b/.buildkite/cuda/pipeline.yml similarity index 75% rename from .buildkite/pipeline.yml rename to .buildkite/cuda/pipeline.yml index b4f5aa9df13..3b460c1c6bd 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/cuda/pipeline.yml @@ -9,7 +9,7 @@ steps: - label: ":github: Resolve skip-ci (docs / skip marks) & upload pipeline" key: upload-ci-pipeline commands: - - "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml" + - "python3 .buildkite/cuda/scripts/upload_pipeline.py --upload .buildkite/cuda/pipeline.yml" agents: queue: "cpu_queue_premerge" @@ -26,23 +26,33 @@ steps: agents: queue: "cpu_queue_premerge" - # L2 Test + # L2 Test — PR ready label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group) - label: "Upload Ready Pipeline" depends_on: image-build key: upload-ready-pipeline if: __UPLOAD_READY_IF__ commands: - - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml + - | + EXTRA="" + if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--e2e" + fi + python3 .buildkite/cuda/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-ready.yml agents: queue: "cpu_queue_premerge" - # L3 Test + # L3 Test — PR merge-test label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group) - label: "Upload Merge Pipeline" depends_on: image-build key: upload-merge-pipeline if: __UPLOAD_MERGE_IF__ commands: - - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml + - | + EXTRA="" + if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then + EXTRA="--e2e" + fi + python3 .buildkite/cuda/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/cuda/test-merge.yml agents: queue: "cpu_queue_premerge" @@ -60,7 +70,7 @@ steps: build.pull_request.labels includes "diffusion-x2v-test" )) commands: - - buildkite-agent pipeline upload .buildkite/test-nightly.yml + - buildkite-agent pipeline upload .buildkite/cuda/test-nightly.yml agents: queue: "cpu_queue_premerge" @@ -72,6 +82,6 @@ steps: (build.branch == "main" && build.env("WEEKLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "weekly-test") commands: - - buildkite-agent pipeline upload .buildkite/test-weekly.yml + - buildkite-agent pipeline upload .buildkite/cuda/test-weekly.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/rebase-pipeline.yaml b/.buildkite/cuda/rebase-pipeline.yml similarity index 76% rename from .buildkite/rebase-pipeline.yaml rename to .buildkite/cuda/rebase-pipeline.yml index f8f4aa217fb..514c9cc1e0c 100644 --- a/.buildkite/rebase-pipeline.yaml +++ b/.buildkite/cuda/rebase-pipeline.yml @@ -14,7 +14,7 @@ steps: depends_on: image-build key: upload-ready-pipeline commands: - - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml + - python3 .buildkite/cuda/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-ready.yml agents: queue: "cpu_queue_premerge" @@ -23,7 +23,7 @@ steps: depends_on: image-build key: upload-merge-pipeline commands: - - python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml + - python3 .buildkite/cuda/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-merge.yml agents: queue: "cpu_queue_premerge" @@ -33,6 +33,6 @@ steps: key: upload-nightly-pipeline if: build.env("NIGHTLY") == "1" commands: - - buildkite-agent pipeline upload .buildkite/test-nightly.yml + - python3 .buildkite/cuda/scripts/upload_pipeline.py --upload --all .buildkite/cuda/test-nightly.yml agents: queue: "cpu_queue_premerge" diff --git a/.buildkite/scripts/upload_pipeline.py b/.buildkite/cuda/scripts/upload_pipeline.py similarity index 83% rename from .buildkite/scripts/upload_pipeline.py rename to .buildkite/cuda/scripts/upload_pipeline.py index 7356144a69c..dc233af7a46 100644 --- a/.buildkite/scripts/upload_pipeline.py +++ b/.buildkite/cuda/scripts/upload_pipeline.py @@ -13,13 +13,19 @@ recognize this uploader-only key). Usage: - python3 upload_pipeline.py [--upload] + python3 upload_pipeline.py [--upload] [--all | --e2e] # Bootstrap (replaces upload_pipeline_with_skip_ci.sh): - python3 upload_pipeline.py --upload .buildkite/pipeline.yml + python3 upload_pipeline.py --upload .buildkite/cuda/pipeline.yml # Test pipeline (replaces upload_test_pipeline_with_diff_skip.py): - python3 upload_pipeline.py --upload .buildkite/test-merge.yml + python3 upload_pipeline.py --upload .buildkite/cuda/test-merge.yml + + # Full suite (rebase pipeline): keep all steps, still strip the uploader-only key: + python3 upload_pipeline.py --upload --all .buildkite/cuda/test-merge.yml + + # Nightly L2 E2E only: keep the E2E Test group, ignore source_file_dependencies: + python3 upload_pipeline.py --upload --e2e .buildkite/cuda/test-merge.yml """ from __future__ import annotations @@ -41,9 +47,10 @@ import yaml LOG = "upload_pipeline" -ROOT = Path(__file__).resolve().parent.parent.parent +ROOT = Path(__file__).resolve().parent.parent.parent.parent DOC_SEP = "\n---\n" BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__" +E2E_GROUP_MARKER = "E2E Test" def _log(message: str) -> None: @@ -319,18 +326,18 @@ def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str: nightly_only = ( '(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")' ) + nightly_main = 'build.branch == "main" && build.env("NIGHTLY") == "1"' + ready_pr = 'build.branch != "main" && build.pull_request.labels includes "ready"' + merge_main = 'build.branch == "main" && build.env("NIGHTLY") != "1" && build.env("WEEKLY") != "1"' + merge_pr = 'build.branch != "main" && build.pull_request.labels includes "merge-test"' if skip_ci: image_if = f"'{nightly_only}'" - ready_if = "'false'" - merge_if = "'false'" + ready_if = f"'({nightly_main})'" + merge_if = f"'({nightly_main})'" else: image_if = "'true'" - ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\'' - merge_if = ( - '\'(build.branch == "main" && build.env("NIGHTLY") != "1" ' - '&& build.env("WEEKLY") != "1") || ' - '(build.branch != "main" && build.pull_request.labels includes "merge-test")\'' - ) + ready_if = f"'(({nightly_main}) || ({ready_pr}))'" + merge_if = f"'(({nightly_main}) || (({merge_main}) || ({merge_pr})))'" return ( continuation.replace("__IMAGE_BUILD_IF__", image_if) @@ -393,6 +400,20 @@ def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]: return filtered +def _is_e2e_group(step: dict[str, Any]) -> bool: + group = step.get("group") + return isinstance(group, str) and E2E_GROUP_MARKER in group + + +def _select_e2e_group_steps(steps: list[Any]) -> list[Any]: + selected = [step for step in steps if isinstance(step, dict) and _is_e2e_group(step)] + if not selected: + _log(f"no group matching {E2E_GROUP_MARKER!r} found") + else: + _log(f"keep {len(selected)} group(s) matching {E2E_GROUP_MARKER!r}") + return selected + + def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]: """Remove uploader-only keys while keeping all steps (no diff filtering).""" stripped: list[Any] = [] @@ -419,11 +440,16 @@ def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]: def render_test_pipeline( doc: dict[str, Any], changed_files: list[str] | None, + *, + e2e_only: bool = False, ) -> dict[str, Any]: steps = doc.get("steps") if not isinstance(steps, list): return doc - if changed_files is not None: + if e2e_only: + steps = _select_e2e_group_steps(steps) + steps = _strip_uploader_metadata_from_steps(steps) + elif changed_files is not None: steps = _filter_steps(steps, changed_files) else: steps = _strip_uploader_metadata_from_steps(steps) @@ -437,11 +463,22 @@ def resolve_pipeline_path(arg: str) -> Path: return ROOT / path -def render_pipeline(path: Path) -> str: +def render_pipeline( + path: Path, + *, + force_all: bool = False, + e2e_only: bool = False, +) -> str: text = path.read_text(encoding="utf-8") diff_range = resolve_diff_range() changed_files = resolve_changed_files() + # ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still + # stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds + # run the full e2e suite (see .buildkite/cuda/rebase-pipeline.yml). + if force_all or e2e_only: + changed_files = None + if BOOTSTRAP_MARKER in text: rendered = render_bootstrap_pipeline( text, @@ -453,7 +490,7 @@ def render_pipeline(path: Path) -> str: if not isinstance(doc, dict): raise ValueError(f"invalid pipeline YAML: {path}") - doc = render_test_pipeline(doc, changed_files) + doc = render_test_pipeline(doc, changed_files, e2e_only=e2e_only) return yaml.safe_dump(doc, sort_keys=False) @@ -472,14 +509,25 @@ def main() -> int: parser.add_argument( "pipeline", nargs="?", - default=".buildkite/pipeline.yml", - help="Pipeline YAML path (default: .buildkite/pipeline.yml)", + default=".buildkite/cuda/pipeline.yml", + help="Pipeline YAML path (default: .buildkite/cuda/pipeline.yml)", ) parser.add_argument( "--upload", action="store_true", help="Pipe rendered YAML to buildkite-agent pipeline upload", ) + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--all", + action="store_true", + help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies", + ) + mode.add_argument( + "--e2e", + action="store_true", + help="Keep only the E2E Test group; disable diff-aware skipping for those steps", + ) args = parser.parse_args() path = resolve_pipeline_path(args.pipeline) @@ -487,7 +535,7 @@ def main() -> int: _log(f"missing pipeline file: {path}") return 1 - rendered = render_pipeline(path) + rendered = render_pipeline(path, force_all=args.all, e2e_only=args.e2e) if args.upload: upload_to_buildkite(rendered) else: diff --git a/.buildkite/test-merge.yml b/.buildkite/cuda/test-merge.yml similarity index 62% rename from .buildkite/test-merge.yml rename to .buildkite/cuda/test-merge.yml index 84c93a2f3a1..7a78eda4af7 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/cuda/test-merge.yml @@ -60,14 +60,32 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Diffusion · Offloader Test" + timeout_in_minutes: 40 + commands: + - timeout 40m pytest -s -v tests/diffusion/offloader -m "advanced_model and cuda" + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Distributed Test" + depends_on: upload-merge-pipeline timeout_in_minutes: 60 commands: - | timeout 55m bash -c " set -e export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/distributed -m 'advanced_model and cuda' --run-level 'advanced_model' + pytest -s -v tests/distributed -m 'advanced_model and cuda and H100' --run-level 'advanced_model' " agents: queue: "mithril-h100-pool" @@ -149,7 +167,7 @@ steps: commands: - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model" agents: - queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU plugins: - docker#v5.2.0: image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT @@ -164,28 +182,6 @@ steps: - group: ":card_index_dividers: E2E Test" depends_on: upload-merge-pipeline steps: - - label: "TTS · Step-Audio2 E2E 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_step_audio2.py -m 'advanced_model' --run-level 'advanced_model' - " - agents: - queue: "gpu_4_queue" - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - label: "TTS · Qwen3-TTS CustomVoice Test" source_file_dependencies: - tests/e2e/online_serving/test_qwen3_tts_customvoice.py @@ -195,7 +191,7 @@ steps: - vllm_omni/deploy/qwen3_tts.yaml commands: - | - timeout 20m bash -c " + timeout 40m bash -c " export VLLM_LOGGING_LEVEL=DEBUG 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 and cuda' --run-level 'advanced_model' @@ -222,7 +218,7 @@ steps: - vllm_omni/deploy/qwen3_tts.yaml commands: - | - timeout 20m bash -c " + timeout 40m bash -c " export VLLM_LOGGING_LEVEL=DEBUG export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model' @@ -240,19 +236,19 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "TTS · Ming-Omni-TTS Test" + - label: "TTS · Higgs-Audio-V3 Test" source_file_dependencies: - - tests/e2e/online_serving/test_ming_tts.py - - tests/e2e/offline_inference/test_ming_tts.py - - vllm_omni/model_executor/models/ming_tts/ - - vllm_omni/model_executor/stage_input_processors/ming_tts.py - - vllm_omni/deploy/ming_tts.yaml - timeout_in_minutes: 20 + - tests/e2e/online_serving/test_higgs_audio_v3.py + - vllm_omni/model_executor/models/higgs_audio_v3/ + - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py + - vllm_omni/deploy/higgs_multimodal_qwen3.yaml 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' + 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 'advanced_model' --run-level 'advanced_model' " agents: queue: "gpu_1_queue" @@ -264,105 +260,10 @@ steps: shm-size: "8gb" environment: - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" 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: - - | - timeout 20m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py tests/e2e/online_serving/test_voxtral_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - 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: - - | - timeout 20m bash -c " - export VLLM_LOGGING_LEVEL=DEBUG - pytest -s -v tests/e2e/online_serving/test_glm_tts.py tests/e2e/offline_inference/test_glm_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model' - " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - label: "Omni · Qwen3-Omni Test" source_file_dependencies: - tests/e2e/offline_inference/test_qwen3_omni.py @@ -370,7 +271,7 @@ steps: - 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 + timeout_in_minutes: 50 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' agents: @@ -407,49 +308,6 @@ steps: path: /mnt/hf-cache 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' - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - label: "TTS · VoxCPM2 Test" source_file_dependencies: - tests/e2e/offline_inference/test_voxcpm2_tts.py @@ -478,22 +336,19 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Omni · Higgs-Audio V2 Test" + - label: "TTS · SoulX-Singer Test" source_file_dependencies: - - tests/e2e/online_serving/test_higgs_audio_v2.py - - tests/e2e/offline_inference/test_higgs_audio_v2.py - - vllm_omni/model_executor/models/higgs_audio_v2/ - - vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py - - vllm_omni/deploy/higgs_audio_v2.yaml + - tests/e2e/offline_inference/test_soulxsinger.py + - vllm_omni/diffusion/models/soulx_singer/ + - vllm_omni/deploy/soulxsinger_svc.yaml + - vllm_omni/deploy/soulxsinger_svs.yaml timeout_in_minutes: 30 commands: - | timeout 30m bash -c " export VLLM_LOGGING_LEVEL=DEBUG export VLLM_WORKER_MULTIPROC_METHOD=spawn - export VLLM_USE_DEEP_GEMM=0 - export VLLM_MOE_USE_DEEP_GEMM=0 - pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model' + pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model' " agents: queue: "gpu_1_queue" @@ -550,88 +405,6 @@ steps: path: /mnt/hf-cache 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' - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: "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' - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - label: "Diffusion · Bagel Test" source_file_dependencies: - tests/e2e/offline_inference/test_bagel.py @@ -811,50 +584,6 @@ steps: path: /mnt/hf-cache 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: - - | - timeout 40m bash -c " - pytest -s -v tests/e2e/online_serving/test_hunyuan_video_15.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' - " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - label: "Diffusion · Z Image Test" source_file_dependencies: - tests/e2e/offline_inference/test_z_image.py @@ -879,11 +608,10 @@ steps: - label: "Diffusion · FLUX.2 Klein Test" source_file_dependencies: - tests/e2e/offline_inference/test_flux2_klein.py - - tests/e2e/online_serving/test_flux2_klein.py - vllm_omni/diffusion/models/flux2_klein/ timeout_in_minutes: 30 commands: - - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' + - pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: @@ -896,48 +624,3 @@ steps: - "HF_TOKEN" volumes: - "/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: - - | - timeout 45m bash -c " - export VLLM_IMAGE_FETCH_TIMEOUT=60 - pytest -s -v tests/e2e/online_serving/test_sd3.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model' - " - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 4 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate diff --git a/.buildkite/test-nightly.yml b/.buildkite/cuda/test-nightly.yml similarity index 84% rename from .buildkite/test-nightly.yml rename to .buildkite/cuda/test-nightly.yml index dd6d5ec31c1..dd62ee68b67 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -12,7 +12,7 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test" steps: - label: ":full_moon: Omni · Function Test with H100" - timeout_in_minutes: 90 + timeout_in_minutes: 120 commands: - pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py agents: @@ -49,25 +49,6 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Omni · Function Test with L4" - timeout_in_minutes: 90 - commands: - - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" - agents: - queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - - "HF_TOKEN" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - label: ":full_moon: Omni · Doc Test with L4" timeout_in_minutes: 90 commands: @@ -172,7 +153,7 @@ steps: - label: ":full_moon: Omni · Perf Test · No Async Chunk" key: nightly-omni-performance-no-async-chunk - timeout_in_minutes: 180 + timeout_in_minutes: 300 commands: - export BENCHMARK_DIR=tests/dfx/perf/results - | @@ -395,9 +376,11 @@ steps: if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" steps: - label: ":full_moon: TTS · Function Test with L4" - timeout_in_minutes: 90 + timeout_in_minutes: 120 commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - export VLLM_USE_DEEP_GEMM="0" + - export VLLM_MOE_USE_DEEP_GEMM="0" - pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py agents: queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU @@ -568,81 +551,19 @@ 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 - depends_on: upload-nightly-pipeline - if: >- - build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2iat-test" - steps: - # Use explicit file shards so single-card cases do not reserve 2+ GPUs, - # and Bagel feature coverage is not duplicated via a broad tests/e2e sweep. - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU" - timeout_in_minutes: 120 + - label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test" + key: nightly-tts-performance-higgs-audio-v3 + timeout_in_minutes: 180 commands: - - >- - pytest -sv - tests/e2e/online_serving/test_qwen_image_expansion.py - tests/e2e/online_serving/test_qwen_image_edit_expansion.py - tests/e2e/online_serving/test_qwen_image_layered_expansion.py - tests/e2e/online_serving/test_longcat_image_expansion.py - 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" + - 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: @@ -677,14 +598,61 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image" + - group: ":card_index_dividers: Tiny Model Tests [Multi-GPU]" + key: nightly-tiny-model-test-group + depends_on: upload-nightly-pipeline + if: >- + build.env("NIGHTLY") == "1" || + build.pull_request.labels includes "nightly-test" + steps: + - label: ":full_moon: Tiny Model Tests · Diffusion (Multi-GPU accelerations)" + timeout_in_minutes: 30 + # NOTE: Single-GPU (core_model) tiny model tests run on the ready label, + # so nightly only adds the multi-GPU parallelism configurations. + # + # This is a bit confusing, but `full_model` is used to mark the multigpu tests, and + # --run-level `core_model` is used to indicate that we should use tiny model weights + # for the CI, since these tests should pass irrespective of whether or not we use + # the full weights or not. + commands: + - timeout 25m pytest -sv tests/model_tests/diffusion -m 'full_model and cuda' --run-level "core_model" + agents: + queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + # 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 + depends_on: upload-nightly-pipeline + if: >- + build.env("NIGHTLY") == "1" || + build.pull_request.labels includes "nightly-test" || + build.pull_request.labels includes "diffusion-x2iat-test" + steps: + # Use explicit file shards so single-card cases do not reserve 2+ GPUs, + # and Bagel feature coverage is not duplicated via a broad tests/e2e sweep. + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Single-GPU" timeout_in_minutes: 120 commands: - >- pytest -sv tests/e2e/online_serving/test_qwen_image_expansion.py tests/e2e/online_serving/test_qwen_image_edit_expansion.py - -m "full_model and diffusion and H100 and distributed_cuda" + tests/e2e/online_serving/test_qwen_image_layered_expansion.py + tests/e2e/online_serving/test_bagel_expansion.py + tests/e2e/online_serving/test_krea2_expansion.py + tests/e2e/offline_inference/test_krea2_expansion.py + -m "full_model and diffusion and H100 and not distributed_cuda" --run-level "full_model" agents: queue: "mithril-h100-pool" @@ -695,7 +663,7 @@ steps: - 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 + nvidia.com/gpu: 1 volumeMounts: - name: devshm mountPath: /dev/shm @@ -720,16 +688,14 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX" - timeout_in_minutes: 90 + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Qwen-Image" + timeout_in_minutes: 120 commands: - >- pytest -sv - tests/e2e/online_serving/test_qwen_image_layered_expansion.py - tests/e2e/online_serving/test_longcat_image_expansion.py - tests/e2e/online_serving/test_longcat_image_edit_expansion.py - tests/e2e/online_serving/test_flux_2_dev_expansion.py - tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py + tests/e2e/online_serving/test_qwen_image_expansion.py + tests/e2e/online_serving/test_qwen_image_edit_expansion.py + tests/e2e/offline_inference/test_krea2_expansion.py -m "full_model and diffusion and H100 and distributed_cuda" --run-level "full_model" agents: @@ -766,12 +732,12 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel" - timeout_in_minutes: 120 + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Layered/LongCat/FLUX" + timeout_in_minutes: 90 commands: - >- pytest -sv - tests/e2e/online_serving/test_bagel_expansion.py + tests/e2e/online_serving/test_qwen_image_layered_expansion.py -m "full_model and diffusion and H100 and distributed_cuda" --run-level "full_model" agents: @@ -808,12 +774,12 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU DreamZero" + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU Bagel" timeout_in_minutes: 120 commands: - >- pytest -sv - tests/e2e/online_serving/test_dreamzero_expansion.py + tests/e2e/online_serving/test_bagel_expansion.py -m "full_model and diffusion and H100 and distributed_cuda" --run-level "full_model" agents: @@ -850,11 +816,15 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100" timeout_in_minutes: 120 commands: - - pytest -sv tests/e2e/online_serving/test_glm_image_expansion.py -m "full_model and diffusion and H100" --run-level "full_model" + - >- + pytest -sv + tests/e2e/online_serving/test_glm_image_expansion.py + tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py + -m "full_model and diffusion and H100" + --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: @@ -891,7 +861,7 @@ steps: - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4" - timeout_in_minutes: 60 + timeout_in_minutes: 120 commands: - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy agents: @@ -1276,10 +1246,48 @@ steps: build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2v-test" steps: - - label: ":full_moon: Diffusion X2V · Function Test" + - label: ":full_moon: Diffusion X2V · Wan2.2 T2V Function Test" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "t2v" -m "full_model and cuda" --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2V · Wan2.2 I2V/TI2V Function Test" timeout_in_minutes: 90 commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k "i2v" -m "full_model and cuda" --run-level "full_model" agents: queue: "mithril-h100-pool" plugins: @@ -1314,14 +1322,129 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Diffusion X2V · Accuracy Test" + - label: ":full_moon: Diffusion X2V · Doc Test" + timeout_in_minutes: 60 + commands: + - pytest -s -v tests/examples/offline_inference/test_image_to_video.py -m "full_model and example and H100" --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2V · Other Function Test" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model" + - pytest -s -v tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2V · Wan2.2 Accuracy Test" timeout_in_minutes: 180 commands: - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2V · Other Accuracy Test" + timeout_in_minutes: 180 + commands: - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model - pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model - - pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model + - pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model agents: queue: "mithril-h100-pool" plugins: @@ -1366,9 +1489,12 @@ steps: set +e pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json EXIT1=$$? + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json + EXIT2=$$? buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - exit $$EXIT1 + if [ $$EXIT1 -ne 0 ]; then exit $$EXIT1; fi + exit $$EXIT2 agents: queue: "mithril-h100-pool" plugins: @@ -1403,11 +1529,28 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - group: ":card_index_dividers: Diffusion · Quantization Test" + - group: ":card_index_dividers: Diffusion Test" depends_on: upload-nightly-pipeline if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test" steps: - - label: ":full_moon: Quantization Test with H100" + - label: ":full_moon: Diffusion · Offloader Test" + timeout_in_minutes: 40 + commands: + - timeout 40m pytest -s -v tests/diffusion/offloader -m "full_model and cuda" + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + - label: ":full_moon: Diffusion · Quantization Test with H100" commands: - timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model" agents: @@ -1444,7 +1587,7 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Quantization Test with L4" + - label: ":full_moon: Diffusion · Quantization Test with L4" timeout_in_minutes: 60 commands: - pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model" @@ -1461,11 +1604,52 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: ":full_moon: Diffusion · Wan VAE Spatial Shard Correctness Test" + timeout_in_minutes: 60 + commands: + - >- + pytest -sv tests/diffusion/distributed/test_wan_spatial_shard.py + -m "full_model and diffusion and H100 and distributed_cuda" + --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: ":bar_chart: Testcase Statistics" key: nightly-testcase-statistics timeout_in_minutes: 120 depends_on: upload-nightly-pipeline - if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" + if: build.env("NIGHTLY") == "1" commands: - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html - buildkite-agent artifact upload "tests/dfx/perf/results/*.html" diff --git a/.buildkite/test-ready.yml b/.buildkite/cuda/test-ready.yml similarity index 87% rename from .buildkite/test-ready.yml rename to .buildkite/cuda/test-ready.yml index 99d87447568..75e834be96e 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -61,7 +61,7 @@ steps: steps: - label: "Diffusion · Offloader Test" commands: - - timeout 10m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model" + - timeout 40m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model" agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: @@ -108,6 +108,12 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" - label: "Diffusion · Distributed Test" + source_file_dependencies: + - tests/diffusion/distributed/ + - vllm_omni/diffusion/distributed/ + - vllm_omni/diffusion/attention/parallel/ + - vllm_omni/diffusion/data.py + - vllm_omni/diffusion/models/ commands: - timeout 20m pytest -sv tests/diffusion/distributed -m 'core_model and cuda' --run-level "core_model" agents: @@ -157,6 +163,27 @@ steps: - "/fsx/hf_cache:/fsx/hf_cache" + - group: ":card_index_dividers: Tiny Model Tests [Single GPU]" + depends_on: upload-ready-pipeline + steps: + - label: "Tiny Model Tests · Diffusion (base)" + commands: + - timeout 15m pytest -sv tests/model_tests/diffusion -m 'core_model and cuda' --run-level "core_model" + agents: + # Ready label only runs the cases that don't need parallelism for tiny + # model tests, which is usually just a smoke test for TTI and I2I + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "Engine&Model Executor Test" depends_on: upload-ready-pipeline commands: @@ -184,96 +211,43 @@ steps: commands: - | timeout 20m bash -c " - pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model' + pytest -sv tests/distributed/ -m 'core_model and cuda and L4' --run-level 'core_model' " agents: - queue: "mithril-h100-pool" + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU 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 - 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 + - 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" - - group: ":card_index_dividers: Entrypoints Test" + - label: "Entrypoints Test" depends_on: upload-ready-pipeline - steps: - - label: "Entrypoints Test with H100" - commands: - - | - timeout 30m bash -c " - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -sv tests/entrypoints/ -m 'core_model and cuda and H100' --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: "Entrypoints Test with L4" - commands: - - | - timeout 30m bash -c " - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -sv tests/entrypoints/ -m 'core_model and cuda and L4' --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" + 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: "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" - group: ":card_index_dividers: E2E Test" depends_on: upload-ready-pipeline @@ -584,6 +558,50 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: "Diffusion · Cosmos3 Test" + source_file_dependencies: + - tests/e2e/online_serving/test_cosmos3.py + - vllm_omni/diffusion/models/cosmos3/ + commands: + - | + timeout 40m bash -c " + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_cosmos3.py -m 'core_model' --run-level 'core_model' + " + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: "Diffusion · Qwen Image Test" source_file_dependencies: @@ -672,3 +690,24 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + - label: "Diffusion · GGUF Plugin Test" + soft_fail: true + source_file_dependencies: + - tests/plugin/test_gguf_diffusion.py + - vllm_omni/diffusion/model_loader + commands: + - pip install "vllm-gguf-plugin>=0.0.3" + - timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" diff --git a/.buildkite/cuda/test-weekly.yml b/.buildkite/cuda/test-weekly.yml new file mode 100644 index 00000000000..d71bdae0a36 --- /dev/null +++ b/.buildkite/cuda/test-weekly.yml @@ -0,0 +1,451 @@ +env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_USE_FLASHINFER_MOE_FP16: 0 + HF_HUB_DOWNLOAD_TIMEOUT: 300 + HF_HUB_ETAG_TIMEOUT: 60 + +steps: + - label: "Reliability Test - qwen3-omni" + timeout_in_minutes: 180 + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + commands: + - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: "Reliability Test - wan22" + timeout_in_minutes: 180 + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + commands: + - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: "Reliability Test - hunyuan-image" + timeout_in_minutes: 180 + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + commands: + - pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 4 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: "Reliability Test - voxcpm2" + timeout_in_minutes: 180 + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + commands: + - pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow" + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + # Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips). + - group: "Reliability Test - Invalid parameters Test" + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + steps: + - label: "Invalid parameters Test · H100" + timeout_in_minutes: 120 + commands: + - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + # Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``. + - label: "Invalid parameters Test · L4" + timeout_in_minutes: 120 + commands: + - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4" + 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" + + - group: ":card_index_dividers: E2E Tests" + depends_on: upload-weekly-pipeline + if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" + steps: + - label: "Omni · L4" + timeout_in_minutes: 90 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model" + - >- + pytest -s -v tests/e2e/ + --ignore=tests/e2e/online_serving/test_qwen2_5_omni_expansion.py + --ignore=tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py + -m "slow and L4 and omni" + --run-level "full_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: "Omni · H100" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and H100 and omni" + --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: "TTS · L4" + timeout_in_minutes: 90 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - export VLLM_USE_DEEP_GEMM="0" + - export VLLM_MOE_USE_DEEP_GEMM="0" + - | + failed=0 + for f in $(pytest --collect-only -q tests/e2e/ -m "slow and L4 and tts" --run-level "full_model" | grep '\.py::' | sed -E 's|::.*||' | sort -u); do + pytest -s -v "$$f" -m "slow and L4 and tts" --run-level "full_model" || failed=1 + done + exit $$failed + 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 · H100" + timeout_in_minutes: 90 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - export VLLM_USE_DEEP_GEMM="0" + - export VLLM_MOE_USE_DEEP_GEMM="0" + - >- + pytest -sv + tests/e2e/ -m "slow and H100 and tts" + --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: "Diffusion · L4" + timeout_in_minutes: 60 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and diffusion and L4" + --run-level "full_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 · H100 · Single-GPU" + timeout_in_minutes: 120 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and diffusion and H100 and not distributed_cuda" + --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: "Diffusion · H100 · Multi-GPU" + timeout_in_minutes: 90 + commands: + - >- + pytest -sv + tests/e2e/ -m "slow and diffusion and H100 and distributed_cuda" + --run-level "full_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate diff --git a/.buildkite/pipeline-intel.yaml b/.buildkite/intel/pipeline-intel.yml similarity index 73% rename from .buildkite/pipeline-intel.yaml rename to .buildkite/intel/pipeline-intel.yml index cb1d76a1cfa..e68f59e5321 100644 --- a/.buildkite/pipeline-intel.yaml +++ b/.buildkite/intel/pipeline-intel.yml @@ -1,16 +1,16 @@ -# .buildkite/pipeline-intel.yaml +# .buildkite/intel/pipeline-intel.yml steps: - group: "INTEL Tests" steps: - label: "Intel: XPU Test" - command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh" + command: "bash .buildkite/intel/scripts/run-xpu-test.sh" agents: queue: "intel-gpu-omni" env: DOCKER_BUILDKIT: "1" # Buildkite will automatically replace this with the actual commit hash VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}" - VLLM_VERSION: "v0.23.0" + VLLM_VERSION: "v0.25.0" priority: 100 timeout_in_minutes: 60 soft_fail: false diff --git a/.buildkite/bootstrap-intel-omni.sh b/.buildkite/intel/scripts/bootstrap-intel-omni.sh similarity index 94% rename from .buildkite/bootstrap-intel-omni.sh rename to .buildkite/intel/scripts/bootstrap-intel-omni.sh index b99568141b2..e51252bd982 100755 --- a/.buildkite/bootstrap-intel-omni.sh +++ b/.buildkite/intel/scripts/bootstrap-intel-omni.sh @@ -11,7 +11,7 @@ fi upload_pipeline() { echo "--- 🛠 Preparing Intel pipeline" - buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml + buildkite-agent pipeline upload .buildkite/intel/pipeline-intel.yml } get_diff() { diff --git a/.buildkite/scripts/hardware_ci/run-xpu-test.sh b/.buildkite/intel/scripts/run-xpu-test.sh similarity index 59% rename from .buildkite/scripts/hardware_ci/run-xpu-test.sh rename to .buildkite/intel/scripts/run-xpu-test.sh index fa4f00e2970..36860522bc0 100755 --- a/.buildkite/scripts/hardware_ci/run-xpu-test.sh +++ b/.buildkite/intel/scripts/run-xpu-test.sh @@ -13,12 +13,33 @@ container_name="xpu_${BUILDKITE_COMMIT}_$( )" cd "${omni_source_dir}" + +# The XPU base image is ~37GB; the default gzip layer exporter is single-threaded +# and takes ~16min to compress it. zstd is multi-threaded (uses all cores) at a +# similar ratio; set EXPORT_COMPRESSION=uncompressed to skip compression entirely +# for local-only images. This requires the containerd image store (buildx docker +# driver), which is the default here. +EXPORT_COMPRESSION="${EXPORT_COMPRESSION:-zstd}" +if [ "${EXPORT_COMPRESSION}" = "uncompressed" ]; then + export_args=(--output "type=image,name={{IMAGE}},compression=uncompressed") +else + export_args=(--output "type=image,name={{IMAGE}},compression=${EXPORT_COMPRESSION},compression-level=3,force-compression=true") +fi + +docker_build() { + # $1 = image name; remaining args passed through to docker build. + local image="$1" + shift + local out=("${export_args[@]/'{{IMAGE}}'/${image}}") + docker build "${out[@]}" "$@" -f docker/Dockerfile.xpu . +} + if [ -z "$(docker images -q "${base_image_name}")" ]; then - docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu . + docker_build "${base_image_name}" --target vllm-base --build-arg "VLLM_VERSION=${VLLM_VERSION}" fi # Try building the docker image -docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu . +docker_build "${image_name}" --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" # Setup cleanup remove_docker_container() { diff --git a/.buildkite/pipeline-npu-a3.yaml b/.buildkite/npu/pipeline-npu-a3.yml similarity index 97% rename from .buildkite/pipeline-npu-a3.yaml rename to .buildkite/npu/pipeline-npu-a3.yml index 8776219e15a..a57f5f8e569 100644 --- a/.buildkite/pipeline-npu-a3.yaml +++ b/.buildkite/npu/pipeline-npu-a3.yml @@ -88,7 +88,7 @@ # # queue: "ascend-a3" # # resource_class: "npu-2" # # command: | -# # buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml +# # buildkite-agent pipeline upload .buildkite/npu/test-npu-ready.yml # # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) # - label: "Upload Nightly Pipeline" @@ -104,7 +104,7 @@ # build.pull_request.labels includes "diffusion-x2v-test" # )) # commands: - # - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml + # - buildkite-agent pipeline upload .buildkite/npu/test-npu-nightly.yml # agents: # queue: "ascend-a3" # resource_class: "npu-2" diff --git a/.buildkite/pipeline-npu.yaml b/.buildkite/npu/pipeline-npu.yml similarity index 53% rename from .buildkite/pipeline-npu.yaml rename to .buildkite/npu/pipeline-npu.yml index 161b0e94476..fd67711f07b 100644 --- a/.buildkite/pipeline-npu.yaml +++ b/.buildkite/npu/pipeline-npu.yml @@ -78,8 +78,89 @@ steps: echo "--- Image pushed successfully" echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" + - label: ":buildkit: Build and Push NPU Test Image - A3" + key: image-build-a3 + agents: + queue: "ascend-a2b3" + resource_class: "npu-1" + plugins: + - kubernetes: + metadata: + annotations: + vault.hashicorp.com/agent-init-first: "true" + vault.hashicorp.com/agent-inject: "true" + vault.hashicorp.com/agent-inject-perms-ca.pem: "0400" + vault.hashicorp.com/agent-inject-perms-cert.pem: "0400" + vault.hashicorp.com/agent-inject-perms-config.json: "0400" + vault.hashicorp.com/agent-inject-perms-key.pem: "0400" + vault.hashicorp.com/agent-inject-secret-ca.pem: internal/data/ascend/buildkitd + vault.hashicorp.com/agent-inject-secret-cert.pem: internal/data/ascend/buildkitd + vault.hashicorp.com/agent-inject-secret-config.json: internal/data/ascend/buildkitd + vault.hashicorp.com/agent-inject-secret-key.pem: internal/data/ascend/buildkitd + vault.hashicorp.com/agent-inject-template-ca.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.RootCA }}\n{{- end }}" + vault.hashicorp.com/agent-inject-template-cert.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaCert }}\n{{- end }}" + vault.hashicorp.com/agent-inject-template-config.json: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.dockerConfig }}\n{{- end }}" + vault.hashicorp.com/agent-inject-template-key.pem: "{{- with secret \"internal/data/ascend/buildkitd\" -}}\n{{ .Data.data.ClientCaKey }}\n{{- end }}" + vault.hashicorp.com/agent-pre-populate-only: "true" + vault.hashicorp.com/agent-run-as-group: "1000" + vault.hashicorp.com/agent-run-as-user: "1000" + vault.hashicorp.com/agent-service-account-token-volume-name: token-vol + vault.hashicorp.com/role: ascend-gha-runners + vault.hashicorp.com/secret-volume-path: /home/user/.docker/ + vault.hashicorp.com/tls-skip-verify: "true" + podSpecPatch: + volumes: + - name: token-vol + projected: + defaultMode: 420 + sources: + - serviceAccountToken: + audience: api + expirationSeconds: 600 + path: token + env: + VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}-a3" + IMAGE_NAME: "vllm-omni-ci-npu" + IMAGE_REGISTRY: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry" + BUILDKITD_ADDR: "tcp://buildkitd-service.buildkitd:1234" + command: | + set -ex + + echo "--- Building and pushing NPU Test Image" + echo "Image: $${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" + echo "buildkitd address: $${BUILDKITD_ADDR}" + + if ! command -v buildctl &> /dev/null; then + echo "Installing buildctl..." + mkdir -p /tmp/buildkit + BUILDKIT_VERSION="v0.29.0" + wget -q "https://github.com/moby/buildkit/releases/download/$${BUILDKIT_VERSION}/buildkit-$${BUILDKIT_VERSION}.linux-arm64.tar.gz" -O /tmp/buildkit.tar.gz + tar -xzf /tmp/buildkit.tar.gz -C /tmp/buildkit + cp /tmp/buildkit/bin/buildctl /usr/local/bin/ + fi + + export DOCKER_CONFIG=/home/user/.docker + buildctl \ + --addr="$${BUILDKITD_ADDR}" \ + --tlscacert=/home/user/.docker/ca.pem \ + --tlscert=/home/user/.docker/cert.pem \ + --tlskey=/home/user/.docker/key.pem \ + build \ + --frontend dockerfile.v0 \ + --local context=. \ + --local dockerfile=./docker \ + --opt filename=Dockerfile.npu.ci.a3 \ + --secret id=dockerconfig,src=/home/user/.docker/config.json \ + --output type=image,name=$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG},push=true \ + --progress=plain + + echo "--- Image pushed successfully" + echo "$${IMAGE_REGISTRY}/$${IMAGE_NAME}:$${VLLM_IMAGE_TAG}" + - label: "Upload Ready Pipeline" - depends_on: image-build + depends_on: + - image-build + - image-build-a3 key: upload-ready-pipeline if: >- build.branch != "main" && ( @@ -90,11 +171,13 @@ steps: queue: "ascend-a2b3" resource_class: "npu-1" command: | - buildkite-agent pipeline upload .buildkite/test-npu-ready.yaml + buildkite-agent pipeline upload .buildkite/npu/test-npu-ready.yml # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) - label: "Upload Nightly Pipeline" - depends_on: image-build + depends_on: + - image-build + - image-build-a3 key: upload-nightly-pipeline if: >- (build.branch == "main" && build.env("NIGHTLY") == "1") || @@ -106,7 +189,7 @@ steps: build.pull_request.labels includes "diffusion-x2v-test" )) commands: - - buildkite-agent pipeline upload .buildkite/test-npu-nightly.yaml + - buildkite-agent pipeline upload .buildkite/npu/test-npu-nightly.yml agents: queue: "ascend-a2b3" resource_class: "npu-1" diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml new file mode 100644 index 00000000000..4dbda17f94b --- /dev/null +++ b/.buildkite/npu/test-npu-nightly.yml @@ -0,0 +1,52 @@ +steps: + - group: ":card_index_dividers: Diffusion X2V Model Test" + key: nightly-diffusion-x2v-group + depends_on: upload-nightly-pipeline + if: >- + build.env("NIGHTLY") == "1" || + build.pull_request.labels includes "nightly-test" || + build.pull_request.labels includes "diffusion-x2v-test" + steps: + - label: ":full_moon: Diffusion X2V · Function Test" + agents: + queue: "ascend-a2b3" + resource_class: "npu-4" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}" + plugins: + - kubernetes: + podSpecPatch: + imagePullSecrets: + - name: swr-secret + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + HF_TOKEN: "${HF_TOKEN}" + commands: + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model" + - group: ":card_index_dividers: TTS Model Test" + key: nightly-tts-test-group + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + steps: + - label: ":full_moon: TTS · Perf Test" + agents: + queue: "ascend-a3" + resource_class: "npu-2" + image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3" + plugins: + - kubernetes: + podSpecPatch: + imagePullSecrets: + - name: swr-secret + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_tts.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT diff --git a/.buildkite/test-npu-ready.yaml b/.buildkite/npu/test-npu-ready.yml similarity index 100% rename from .buildkite/test-npu-ready.yaml rename to .buildkite/npu/test-npu-ready.yml diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release/release-pipeline.yml similarity index 92% rename from .buildkite/release-pipeline.yaml rename to .buildkite/release/release-pipeline.yml index c560fefe13e..351e105cc29 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release/release-pipeline.yml @@ -16,7 +16,7 @@ steps: - "uv venv --python=3.12 && source .venv/bin/activate" - "uv pip install --upgrade build" - "python3 -m build" - - "bash .buildkite/scripts/upload-nightly-wheels.sh" + - "bash .buildkite/release/scripts/upload-nightly-wheels.sh" - label: "Generate and upload wheel indices" depends_on: "build-wheel" @@ -24,7 +24,7 @@ steps: agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/generate-and-upload-nightly-index.sh" + - "bash .buildkite/release/scripts/generate-and-upload-nightly-index.sh" - block: "Upload release wheels to PyPI" depends_on: "build-wheel" @@ -38,7 +38,7 @@ steps: agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/publish-release-wheels-pypi.sh" + - "bash .buildkite/release/scripts/publish-release-wheels-pypi.sh" - block: "Unblock to build release Docker images" depends_on: ~ @@ -100,7 +100,7 @@ steps: agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/publish-release-images.sh" + - "bash .buildkite/release/scripts/publish-release-images.sh" plugins: - docker-login#v3.0.0: username: vllmbot diff --git a/.buildkite/scripts/generate-and-upload-nightly-index.sh b/.buildkite/release/scripts/generate-and-upload-nightly-index.sh similarity index 92% rename from .buildkite/scripts/generate-and-upload-nightly-index.sh rename to .buildkite/release/scripts/generate-and-upload-nightly-index.sh index b09c13f5cf9..63f5da59627 100755 --- a/.buildkite/scripts/generate-and-upload-nightly-index.sh +++ b/.buildkite/release/scripts/generate-and-upload-nightly-index.sh @@ -36,12 +36,12 @@ mkdir -p "$INDICES_OUTPUT_DIR" # HACK: we do not need regex module here, but it is required by pre-commit hook # To avoid any external dependency, we simply replace it back to the stdlib re module -sed -i.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak +sed -i.bak 's/import regex as re/import re/g' .buildkite/release/scripts/generate-nightly-index.py && rm -f .buildkite/release/scripts/generate-nightly-index.py.bak # Generate indices -- the version is just the commit hash (not omni/{commit}) # because relative paths are computed between the index and wheel directories, # both of which live under the omni/ prefix in S3. -$PYTHON .buildkite/scripts/generate-nightly-index.py \ +$PYTHON .buildkite/release/scripts/generate-nightly-index.py \ --version "$BUILDKITE_COMMIT" \ --current-objects "$obj_json" \ --output-dir "$INDICES_OUTPUT_DIR" \ @@ -78,7 +78,7 @@ if [[ "$version" != *"dev"* ]]; then rm -rf "${INDICES_OUTPUT_DIR:?}" mkdir -p "$INDICES_OUTPUT_DIR" # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path - $PYTHON .buildkite/scripts/generate-nightly-index.py \ + $PYTHON .buildkite/release/scripts/generate-nightly-index.py \ --version "$s3_version" \ --wheel-dir "$BUILDKITE_COMMIT" \ --current-objects "$obj_json" \ diff --git a/.buildkite/scripts/generate-nightly-index.py b/.buildkite/release/scripts/generate-nightly-index.py similarity index 100% rename from .buildkite/scripts/generate-nightly-index.py rename to .buildkite/release/scripts/generate-nightly-index.py diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/release/scripts/publish-release-images.sh similarity index 100% rename from .buildkite/scripts/publish-release-images.sh rename to .buildkite/release/scripts/publish-release-images.sh diff --git a/.buildkite/scripts/publish-release-wheels-pypi.sh b/.buildkite/release/scripts/publish-release-wheels-pypi.sh similarity index 100% rename from .buildkite/scripts/publish-release-wheels-pypi.sh rename to .buildkite/release/scripts/publish-release-wheels-pypi.sh diff --git a/.buildkite/scripts/upload-nightly-wheels.sh b/.buildkite/release/scripts/upload-nightly-wheels.sh similarity index 100% rename from .buildkite/scripts/upload-nightly-wheels.sh rename to .buildkite/release/scripts/upload-nightly-wheels.sh diff --git a/.buildkite/scripts/docker_login_ecr_public.sh b/.buildkite/scripts/docker_login_ecr_public.sh deleted file mode 100755 index 51c5e1a5d5d..00000000000 --- a/.buildkite/scripts/docker_login_ecr_public.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/bin/bash -# Helper function to safely login to ECR Public with per-job config isolation -# Uses DOCKER_CONFIG environment variable to prevent race conditions -# -# This script prevents the "device or resource busy" error by giving each -# Buildkite job its own isolated Docker config directory. -# -# Usage: -# source docker_login_ecr_public.sh && safe_docker_login_ecr_public - -set -euo pipefail - -# Configuration -ECR_REGISTRY="public.ecr.aws" - -setup_isolated_docker_config() { - # Use BUILDKITE_JOB_ID for job-specific isolation - # Fallback to PID if running outside Buildkite - local job_id="${BUILDKITE_JOB_ID:-$$}" - - # Set Docker config to job-specific directory - export DOCKER_CONFIG="/tmp/docker-config-${job_id}" - - # Create directory if it doesn't exist - mkdir -p "$DOCKER_CONFIG" - - echo "[docker-config] Using isolated Docker config: $DOCKER_CONFIG" -} - -check_docker_auth() { - # Check if already authenticated to the given registry - # Returns 0 if authenticated, 1 if not - local registry="$1" - - # Check if credentials exist in the isolated config - if [[ -f "$DOCKER_CONFIG/config.json" ]]; then - # Check if registry is present in config - if grep -q "$registry" "$DOCKER_CONFIG/config.json" 2>/dev/null; then - return 0 - fi - fi - - return 1 -} - -safe_docker_login_ecr_public() { - # Setup isolated config first - setup_isolated_docker_config - - local registry="$ECR_REGISTRY" - - # Check if already authenticated (within this job) - if check_docker_auth "$registry"; then - echo "[docker-login] Already authenticated to $registry in this job" - return 0 - fi - - # Perform login to isolated config directory - echo "[docker-login] Logging in to $ECR_REGISTRY (isolated config)..." - if aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$ECR_REGISTRY"; then - echo "[docker-login] Login successful (config: $DOCKER_CONFIG)" - return 0 - else - local exit_code=$? - echo "[docker-login] ERROR: Login failed with exit code $exit_code" >&2 - return $exit_code - fi -} - -# Execute if run as script (not sourced) -if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then - safe_docker_login_ecr_public -fi diff --git a/.buildkite/scripts/hardware_ci/run_npu_test.sh b/.buildkite/scripts/hardware_ci/run_npu_test.sh deleted file mode 100644 index 492699c5ffe..00000000000 --- a/.buildkite/scripts/hardware_ci/run_npu_test.sh +++ /dev/null @@ -1,145 +0,0 @@ -#!/bin/bash - -# This script build the Ascend NPU docker image and run the offline inference inside the container. -# It serves a sanity check for compilation and basic model usage. -set -ex - -# Base ubuntu image with basic ascend development libraries and python installed -VLLM_OMNI_REPO="https://github.com/vllm-project/vllm-omni.git" -BASE_IMAGE_NAME="quay.nju.edu.cn/ascend/vllm-ascend:v0.11.0rc2" -image_name="npu/vllm-omni-ci:${BUILDKITE_COMMIT}_${EPOCHSECONDS}" -# image_name="npu/vllm-ci:${BUILDKITE_COMMIT}_${EPOCHSECONDS}" -container_name="npu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)" - -# BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards -agent_idx=$(echo "${BUILDKITE_AGENT_NAME}" | awk -F'-' '{print $(NF-1)}') -echo "agent_idx: ${agent_idx}" -builder_name="cachebuilder${agent_idx}" -builder_cache_dir="/mnt/docker-cache${agent_idx}" -mkdir -p ${builder_cache_dir} - -# Try building the docker image -cat <=6.0 pytest-cov modelscope - -COPY . . - -# Install vllm-omni -WORKDIR /workspace -ARG VLLM_OMNI_REPO=https://github.com/vllm-project/vllm-omni.git -ARG VLLM_OMNI_TAG=main -ARG BUILDKITE_PULL_REQUEST -ARG BUILDKITE_PULL_REQUEST_REPO -RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/" && \ - if [ "\$BUILDKITE_PULL_REQUEST" != "false" ] && [ -n "\$BUILDKITE_PULL_REQUEST" ]; then \ - echo "Cloning and checking out PR #\$BUILDKITE_PULL_REQUEST..." && \ - git clone \$VLLM_OMNI_REPO /workspace/vllm-omni && \ - cd /workspace/vllm-omni && \ - git fetch origin pull/\$BUILDKITE_PULL_REQUEST/head:pr-\$BUILDKITE_PULL_REQUEST && \ - git checkout pr-\$BUILDKITE_PULL_REQUEST; \ - else \ - echo "Not a PR build, using main branch" && \ - git clone --depth 1 \$VLLM_OMNI_REPO /workspace/vllm-omni; \ - fi - -RUN --mount=type=cache,target=/root/.cache/pip \ - export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \ - source /usr/local/Ascend/ascend-toolkit/set_env.sh && \ - source /usr/local/Ascend/nnal/atb/set_env.sh && \ - export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \ - python3 -m pip install -v -e /workspace/vllm-omni/ - -ENV VLLM_WORKER_MULTIPROC_METHOD=spawn - -WORKDIR /workspace/vllm-omni -CMD ["/bin/bash"] - -EOF - -# Setup cleanup -remove_docker_container() { - docker rm -f "${container_name}" || true; - docker image rm -f "${image_name}" || true; - docker system prune -f || true; -} -trap remove_docker_container EXIT - -# Generate corresponding --device args based on BUILDKITE_AGENT_NAME -# Ascend NPU BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards, and agent_idx starts from 1. -# e.g. atlas-a2-001-1-2cards means this is the 1-th agent on atlas-a2-001 host, and it has 2 NPU cards. -# returns --device /dev/davinci0 --device /dev/davinci1 -parse_and_gen_devices() { - local input="$1" - local index cards_num - if [[ "$input" =~ ([0-9]+)-([0-9]+)cards$ ]]; then - index="${BASH_REMATCH[1]}" - cards_num="${BASH_REMATCH[2]}" - else - echo "parse error" >&2 - return 1 - fi - - local devices="" - local i=0 - while (( i < cards_num )); do - local dev_idx=$(((index - 1)*cards_num + i )) - devices="$devices --device /dev/davinci${dev_idx}" - ((i++)) - done - - # trim leading space - devices="${devices#"${devices%%[![:space:]]*}"}" - # Output devices: assigned to the caller variable - printf '%s' "$devices" -} - -devices=$(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1 - -# Run the image and execute the Out-Of-Tree (OOT) platform interface test case on Ascend NPU hardware. -# This test checks whether the OOT platform interface is functioning properly in conjunction with -# the hardware plugin vllm-ascend. -hf_model_cache_dir=/mnt/hf_cache${agent_idx} -ms_model_cache_dir=/mnt/modelscope${agent_idx} -mkdir -p ${hf_model_cache_dir} -mkdir -p ${ms_model_cache_dir} -docker run \ - --init \ - ${devices} \ - --device /dev/davinci_manager \ - --device /dev/devmm_svm \ - --device /dev/hisi_hdc \ - -v /usr/local/dcmi:/usr/local/dcmi \ - -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ - -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ - -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ - -v /etc/ascend_install.info:/etc/ascend_install.info \ - -v ${hf_model_cache_dir}:/root/.cache/huggingface \ - -v ${ms_model_cache_dir}:/root/.cache/modelscope \ - --network host \ - --entrypoint="" \ - --name "${container_name}" \ - "${image_name}" \ - bash -c ' - set -e - VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py -' diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml deleted file mode 100644 index 520c1bbae5d..00000000000 --- a/.buildkite/test-amd.yaml +++ /dev/null @@ -1,118 +0,0 @@ -steps: - -- label: "Diffusion Model Test" - timeout_in_minutes: 30 - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - export GPU_ARCHS=gfx942 - - pytest -s -v tests/e2e/offline_inference/test_z_image.py - -- label: "Diffusion Images API LoRA E2E" - timeout_in_minutes: 30 - 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 - - pytest -s -v tests/diffusion/lora/test_images_generations_lora.py - -- label: "Diffusion Model CPU offloading Test" - timeout_in_minutes: 20 - 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 - - pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py - - pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py - -- label: "Diffusion Cache Backend Test" - timeout_in_minutes: 15 - 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 - - pytest -s -v -m 'core_model and cache and diffusion and not distributed_rocm and MI325' - -- label: "Diffusion Sequence Parallelism Test" - timeout_in_minutes: 20 - agent_pool: mi325_2 - 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 - - pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model - -- label: "Diffusion Tensor Parallelism Test" - timeout_in_minutes: 20 - agent_pool: mi325_2 - depends_on: amd-build - commands: - - export GPU_ARCHS=gfx942 - - export VLLM_LOGGING_LEVEL=DEBUG - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py - -- label: "Diffusion GPU Worker Test" - timeout_in_minutes: 20 - agent_pool: mi325_2 - depends_on: amd-build - mirror_hardwares: [amdproduction] - grade: Blocking - commands: - - pytest -s -v tests/diffusion/test_diffusion_worker.py - -- label: "Omni Model Test Qwen2-5-Omni" - timeout_in_minutes: 15 - agent_pool: mi325_2 - 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 - - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py - - pytest -s -v tests/engine/test_async_omni_engine_abort.py - -- label: "Omni Model Test Qwen3-Omni" - timeout_in_minutes: 15 - 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 - - export VLLM_ROCM_USE_AITER=0 - - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py - - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py - - -- label: "Omni Sleep Mode Test" - timeout_in_minutes: 40 - agent_pool: mi325_2 - 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 - - pytest -s -v tests/entrypoints/test_omni_sleep_mode.py -m "advanced_model and omni and MI325" --run-level "advanced_model" diff --git a/.buildkite/test-npu-nightly.yaml b/.buildkite/test-npu-nightly.yaml deleted file mode 100644 index 2c2372f4c6e..00000000000 --- a/.buildkite/test-npu-nightly.yaml +++ /dev/null @@ -1,24 +0,0 @@ -steps: - - group: ":card_index_dividers: Diffusion X2V Model Test" - key: nightly-diffusion-x2v-group - depends_on: upload-nightly-pipeline - if: >- - build.env("NIGHTLY") == "1" || - build.pull_request.labels includes "nightly-test" || - build.pull_request.labels includes "diffusion-x2v-test" - steps: - - label: ":full_moon: Diffusion X2V · Function Test" - agents: - queue: "ascend-a2b3" - resource_class: "npu-4" - image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}" - plugins: - - kubernetes: - podSpecPatch: - imagePullSecrets: - - name: swr-secret - env: - VLLM_WORKER_MULTIPROC_METHOD: spawn - HF_TOKEN: "${HF_TOKEN}" - commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model" diff --git a/.buildkite/test-weekly.yml b/.buildkite/test-weekly.yml deleted file mode 100644 index 57eadeafd88..00000000000 --- a/.buildkite/test-weekly.yml +++ /dev/null @@ -1,148 +0,0 @@ -env: - VLLM_WORKER_MULTIPROC_METHOD: spawn - VLLM_USE_FLASHINFER_MOE_FP16: 0 - HF_HUB_DOWNLOAD_TIMEOUT: 300 - HF_HUB_ETAG_TIMEOUT: 60 - -steps: - - label: "Reliability Test - qwen3-omni" - timeout_in_minutes: 180 - depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" - commands: - - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: "Reliability Test - wan22" - timeout_in_minutes: 180 - depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" - commands: - - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - # Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips). - - group: "Reliability Test - Invalid parameters Test" - depends_on: upload-weekly-pipeline - if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test" - steps: - - label: "Invalid parameters Test · H100" - timeout_in_minutes: 120 - commands: - - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - # Stable Audio Open, Qwen3-TTS speech/batch/voices paths marked ``L4``. - - label: "Invalid parameters Test · L4" - timeout_in_minutes: 120 - commands: - - pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4" - 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" diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index c080ca41217..76c89447305 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -24,7 +24,7 @@ Check the model's HF repo for `model_index.json`. This determines your path: | Scenario | How to identify | Migration path | |----------|----------------|----------------| -| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) | +| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip implementation, then validate model-specific examples, tests, and docs as needed | | **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below | | **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below | | **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B | @@ -52,7 +52,7 @@ vllm_omni/diffusion/models/your_model_name/ 2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`). 3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`. 4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names. -5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`. +5. Add class attributes for acceleration features such as `_repeated_blocks` and `_layerwise_offload_blocks_attrs` (see [references/transformer-adaptation.md](references/transformer-adaptation.md) for examples). ### A4. Adapt pipeline @@ -78,13 +78,34 @@ class YourPipeline(nn.Module): Add post/pre-process functions in the same pipeline file. Register them in `registry.py`. +### A4.1 Add progress bar support (recommended) + +For pipelines with a standard denoising loop, prefer the existing progress bar pattern instead of hand-rolled logging. + +```python +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin + +class YourPipeline(nn.Module, ProgressBarMixin): + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + # ... prepare timesteps / latents ... + with self.progress_bar(total=len(timesteps)) as progress_bar: + for i, t in enumerate(timesteps): + # predict noise / scheduler step + latents = ... + progress_bar.update() + + return DiffusionOutput(output=output) +``` + +For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` and existing pipelines using `ProgressBarMixin`. + ### A5. Register, test, docs → continue at Step 4 below. --- ## Path B: Custom/Private Repo Model -For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. +For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns. ### B1. Understand the reference repo @@ -134,7 +155,7 @@ Copy the essential model files into `vllm_omni/diffusion/models//` and ada ### B4. Handle custom weight loading -Custom models have two patterns for weight loading: +Custom models have two common patterns for weight loading: **Pattern 1: Bypass standard loader** (DreamID-Omni style) @@ -149,8 +170,10 @@ class CustomPipeline(nn.Module): self.vae = custom_init_vae(model, device=self.device) self.text_encoder = custom_init_text_encoder(model, device=self.device) self.transformer = CustomFusionModel(CONFIG) - load_custom_checkpoint(self.transformer, - checkpoint_path=os.path.join(model, "model.safetensors")) + load_custom_checkpoint( + self.transformer, + checkpoint_path=os.path.join(model, "model.safetensors"), + ) # NO weights_sources defined — bypasses standard loader def load_weights(self, weights): @@ -162,6 +185,8 @@ class CustomPipeline(nn.Module): When weights are in safetensors format but need name remapping: ```python +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + class CustomPipeline(nn.Module): def __init__(self, *, od_config, prefix=""): super().__init__() @@ -187,7 +212,9 @@ class CustomPipeline(nn.Module): # Remap original weight names to vllm-omni module names name = self._remap_weight_name(name) if name in params: - default_weight_loader(params[name], tensor) + param = params[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor) loaded.add(name) return loaded ``` @@ -281,16 +308,43 @@ For Omni or custom models, create: ### Step 7: Update Documentation Required updates: -1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table +1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table 2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table) -3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported -4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported -5. `examples/offline_inference/xxx/README.md` — offline example docs -6. `examples/online_serve/xxx/README.md` — online serve docs +3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported +4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported +5. Offline example docs under `examples/offline_inference//` (`README.md` or category-specific `.md`) +6. `examples/online_serving//README.md` — online serving docs + +### Step 8: Add E2E Tests + +**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming, Buildkite wiring, and run commands. Also read [l4_functionality_tests.inc.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md) and [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md). + +Classify the model's **CI priority** first: + +| Priority | Required test levels | Files & markers | +|----------|---------------------|-----------------| +| **High** (listed in [#1832](https://github.com/vllm-project/vllm-omni/issues/1832) or on the diffusion hot path) | **L1** · **L2** online · **L3** online + offline · **L4** feature + performance | See table below | +| **Medium** (*normal priority* in L4 docs) | **L3** online + offline · **L4** feature only | Fewer L4 parametrized rows | +| **Low** | **L4** feature only | One or two `*_expansion.py` cases | + +**Per-level deliverables (diffusion / `pytest.mark.diffusion`):** + +| Level | Location | Marker | CI pipeline | Notes | +|-------|----------|--------|-------------|-------| +| **L1** | `tests/diffusion/models/{slug}/`, `tests/diffusion/cache/`, transformer unit tests | `core_model` + `cpu` | `test-ready.yml` | Weight remap, `_sp_plan`, cache enabler registration, shape contracts | +| **L2** | `tests/e2e/online_serving/test_{slug}.py` (and offline if the category is offline-first) | **`core_model` + `advanced_model`** (both on baseline smoke) + `diffusion` + `@hardware_test` / `hardware_marks` | `test-ready.yml` | Default deploy smoke — minimal `num_inference_steps`, single prompt | +| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` when offline matters | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `diffusion`) | `test-merge.yml` **or** merged into nightly diffusion function job | Real weights, streaming/API paths, LoRA/offload smoke | +| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` with per-case **`mark`** (`hardware_marks` + `full_model` + `diffusion`) | + +**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`). + +**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`). + +**L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority. -### Step 8: Add E2E Tests (Recommended) +**Reference implementations:** `tests/e2e/online_serving/test_qwen_image_edit_expansion.py`, `tests/e2e/online_serving/test_longcat_image_expansion.py`, `tests/e2e/online_serving/test_hunyuan_video_15_expansion.py`. -Create `tests/e2e/online_serving/test_your_model_expansion.py`. +**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** deploy constants, prompts, sampling dicts, and inline `request_config` / `form_data` belong in each `test_{slug}.py` and `test_{slug}_expansion.py`. Do not add per-model files under `tests/helpers/`; reuse only repo-wide harness (`mark`, `media`, `runtime`, `stage_config`, `assertions`). **L2+ online/offline e2e:** reuse or add `send_*_request` in `tests/helpers/runtime.py` — tests call the handler, not raw `omni.generate` / HTTP. See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers**. ### Step 9: Add Cache-DiT Acceleration @@ -510,7 +564,7 @@ class YourTransformer(nn.Module): #### 10e. Update parallelism documentation After adding parallelism support, update: -1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table +1. `docs/user_guide/diffusion/parallelism/overview.md` — add your model to the support overview/table 2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch) ### Step 11: Add CPU Offload Support @@ -568,6 +622,7 @@ See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling. ## Reference Files +- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands - [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers - [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models - [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details diff --git a/.claude/skills/add-diffusion-model/references/parallelism-patterns.md b/.claude/skills/add-diffusion-model/references/parallelism-patterns.md index 56a21d36ba8..8a6b7fbd857 100644 --- a/.claude/skills/add-diffusion-model/references/parallelism-patterns.md +++ b/.claude/skills/add-diffusion-model/references/parallelism-patterns.md @@ -569,4 +569,4 @@ Check which parallelism methods a model supports: | **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models//` | | **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` | -The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`. +The canonical per-model support overview/table is in `docs/user_guide/diffusion/parallelism/overview.md`. diff --git a/.claude/skills/add-diffusion-model/references/transformer-adaptation.md b/.claude/skills/add-diffusion-model/references/transformer-adaptation.md index 64abc77127d..75622478bd4 100644 --- a/.claude/skills/add-diffusion-model/references/transformer-adaptation.md +++ b/.claude/skills/add-diffusion-model/references/transformer-adaptation.md @@ -219,13 +219,14 @@ def load_weights(self, weights): ### Class-Level Attributes for Features + ```python class MyTransformer(nn.Module): # torch.compile: list block class names that repeat and can be compiled _repeated_blocks = ["MyTransformerBlock"] # CPU offload: attribute name of the nn.ModuleList containing blocks - _layerwise_offload_blocks_attr = "blocks" + _layerwise_offload_blocks_attrs = ["blocks"] # LoRA: mapping of fused param names to original param names packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]} diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index 811a359d0e4..260a7e924f9 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -305,13 +305,72 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list: ] ``` +### Test Case Writing (CI Levels) + +**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming (`test_{slug}.py` / `test_{slug}_expansion.py`), Buildkite wiring, and copy-paste run commands. Also read [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md) and [tests_style.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/tests_style.md). + +Classify the model's **CI priority** first (high / medium / low). High-priority TTS models are typically those on the integration hot path or listed in tracking issues such as [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); medium and low tiers cover the long tail. When unsure, ask the reviewer which tier applies. + +| Priority | Required test levels | Files & markers | +|----------|---------------------|-----------------| +| **High** | **L1** unit/logic · **L2** online smoke · **L3** online + offline integration · **L4** feature + performance | See table below | +| **Medium** | **L3** online + offline · **L4** feature only | Skip dedicated L1/L2 unless fixing a logic bug | +| **Low** | **L4** feature only | One or two `*_expansion.py` parametrized cases | + +**Per-level deliverables (TTS / `pytest.mark.tts`):** + +| Level | Location | Marker | CI pipeline | Notes | +|-------|----------|--------|-------------|-------| +| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU | +| **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path | +| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk | +| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` | + +**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, single non-streaming `/v1/audio/speech` or equivalent offline path) should carry **both** `@pytest.mark.core_model` **and** `@pytest.mark.advanced_model` on the **same** function so it runs in L2 (`test-ready.yml`, `--run-level core_model`, basic validation) and L3 (`test-merge.yml`, `--run-level advanced_model`, deeper validation) without duplicating the test. In-tree examples: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. + +Heavier scenarios in the same file use **`advanced_model` only** (streaming, extra languages, concurrency, async_chunk, batch). Example: `test_voice_clone_en_streaming_001` → `advanced_model` only. When migrating L3 to nightly, move those heavier cases into `test_{slug}_expansion.py` with `full_model` and drop the dedicated merge job (see `test_ming_tts_expansion.py`, `test_glm_tts_expansion.py`). + +```python +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None: + openai_client.send_audio_speech_request({...}) +``` + +**L4 consolidation:** Prefer parametrized `OmniServerParams` rows (`default`, `async_chunk`, feature flags) in one expansion module rather than many merge-only files ([#1832](https://github.com/vllm-project/vllm-omni/issues/1832)). + +**L4 performance (high-priority models):** Add latency / throughput / stress rows in `tests/dfx/perf/tests/test_tts.json`, or a dedicated `tests/dfx/perf/tests/test_{slug}.json` when the model must not join the shared nightly server matrix before integration lands (see VoxCPM2 / Coqui XTTS pattern). Register the model in `benchmarks/tts/model_configs.yaml` for local `bench_tts.py`. Wire a **separate** `test-nightly.yml` Perf Test step when the JSON is not merged into `test_tts.json` yet. + +**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** + +- Put `MODEL`, deploy path, vendored `REF_AUDIO_URL`, `get_prompt()`, and inline `request_config` dicts **in each** `test_{slug}.py`, `test_{slug}_expansion.py`, offline `test_{slug}.py`, and L1 `test_{slug}_*.py` as needed. +- **Do not** add `tests/helpers/{slug}.py` (or `tests/helpers/{model_name}.py`) to deduplicate constants or request builders across those files. A little duplication is intentional; follow in-tree references such as `tests/e2e/online_serving/test_glm_tts.py` and `tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py`. +- `tests/helpers/` is for **repo-wide** harness code only (`mark.py`, `media.py`, `runtime.py`, `stage_config.py`, `assertions.py`, `fixtures/`). Import those; do not extend the tree with per-model modules. + +**Runtime send helpers (`tests/helpers/runtime.py`) — online and offline e2e:** + +| Path | Fixture | Call | +|------|---------|------| +| Online `/v1/*` | `openai_client` | `openai_client.send_*_request(request_config)` | +| Offline inference | `omni_runner_handler` | `omni_runner_handler.send_*_request(request_config)` | + +1. **Grep `runtime.py` first** — reuse `send_omni_request`, `send_diffusion_request`, `send_audio_speech_request` (online + offline Qwen-style TTS), `send_single_stage_tts_request` (Coqui XTTS / MOSS-TTS-Nano offline), etc. +2. **No matching helper** → add `send__request` (or `send__http_request` for negative/dfx) in **`runtime.py`** with general `assert_*` bundled inside, **then** call it from the test. +3. **Test file** holds `request_config` dicts only — not `omni.generate`, not `_collect_audio()`, not raw HTTP/SDK. +4. Wire `tests/helpers/runtime.py` into Buildkite `source_file_dependencies` when you add helpers. + +See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers** for full tables and exceptions. + ### Deliverables - Updated `serving_speech.py` with all 5 integration points (single commit) - Client scripts and server launcher under `examples/online_serving/text_to_speech//` - Gradio demo with streaming and voice cloning UI in the same dir -- E2E online serving test (`tests/e2e/online_serving/test_.py`) -- Buildkite CI entry in `.buildkite/test-merge.yml` +- E2E tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope) +- Buildkite wired per level: `test-ready.yml` (L1/L2), `test-merge.yml` or nightly function job (L3), `test-nightly.yml` (L4) — see [vllm-omni-test skill](../vllm-omni-test/SKILL.md) - New section in `examples/online_serving/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/online_serving//` dir or a per-model `README.md` inside `text_to_speech//`. ### E2E test pitfalls to avoid @@ -326,9 +385,9 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list: - **Use the harness readiness gate.** The fixture waits for HTTP 200 on `/health`; don't add `time.sleep` in tests. If warmup is incomplete, make `/health` return non-200 until you're actually ready. -- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`** - so the test lands in `test-ready.yml` (triggered by the `ready` label) rather - than only nightly. +- **Mark tests per the CI Levels table** — baseline smoke: `core_model` + `advanced_model`; heavier cases: `advanced_model` only; L4 expansion: `full_model` +- **No per-model helper modules** — do not create `tests/helpers/{slug}.py`; keep constants and `request_config` payloads in the test file +- **Online and offline e2e go through `runtime.py`** — `openai_client.send_audio_speech_request` (online); `omni_runner_handler.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests ## Phase 4: Async Chunk (Streaming) @@ -461,8 +520,8 @@ Use this checklist when integrating a new TTS model: - [ ] Voice cloning works (if supported) - [ ] All response formats work (wav, mp3, flac, pcm) - [ ] Client scripts and server launcher created -- [ ] E2E online serving test written (`tests/e2e/online_serving/test_.py`) -- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml` +- [ ] E2E tests added per model priority tier (see **Test Case Writing (CI Levels)**) +- [ ] Buildkite entries match level: `test-ready.yml` / `test-merge.yml` or nightly TTS job / `test-nightly.yml` - [ ] Gradio demo working - [ ] Documentation added (offline + online docs, nav, supported models) @@ -500,5 +559,6 @@ Project docs and adjacent skills: - [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage - [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3 - [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases +- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands - [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md) - `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above) diff --git a/.claude/skills/precheck-pr/SKILL.md b/.claude/skills/precheck-pr/SKILL.md index ab6034a48f5..cd86ec4aa46 100644 --- a/.claude/skills/precheck-pr/SKILL.md +++ b/.claude/skills/precheck-pr/SKILL.md @@ -63,6 +63,8 @@ If multiple rows apply (e.g., a diffusion model is also a new model), union the Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠. +**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope. + ### Step 5: Print Report ``` @@ -74,6 +76,7 @@ Pre-check report for Dimension Result ───────────────── ────── PR title format ✓ + Code quality ⚠ 1 broad except, 2 Any hints PR desc integrity ✓ Registry/config ✓ Dead code ⚠ 2 warnings diff --git a/.claude/skills/precheck-pr/references/checklists.md b/.claude/skills/precheck-pr/references/checklists.md index 5b9472d9621..a232403a7a1 100644 --- a/.claude/skills/precheck-pr/references/checklists.md +++ b/.claude/skills/precheck-pr/references/checklists.md @@ -35,6 +35,31 @@ PR title must follow the project convention documented in `docs/contributing/REA --- +## All PRs: Code Quality (run on every PR, quick & full) + +Diff-scoped sweep for five fragility patterns that are pervasive in the existing codebase — **only lines the PR adds count**, not the pre-existing backlog. Detection commands, codebase examples, and full severity rules live in [code-quality.md](code-quality.md). + +- [ ] **No new `**kwargs` string-lookup plumbing:** the diff adds no `**kwargs` + `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`. ⚠ any new instance; ✗ if a new string key is duplicated across ≥2 files without a shared constant, or unknown keys are silently dropped on a fail-fast path. +- [ ] **No new broad exception swallow:** the diff adds no `except Exception:` / bare `except:`. ⚠ any new instance; ✗ for `except: pass` / `except Exception: return None` / `continue` on a fail-fast path (init, config validation, weight loading, request handling). +- [ ] **No new `Any` / wrong type hints:** the diff adds no `: Any` / `-> Any` / untyped production signature, and no new `SimpleNamespace` in tests faking an object that already has a real typed stub. ⚠ any new instance; ✗ for a *wrong* (actively misleading) annotation or a `SimpleNamespace` test fake of a typed object. +- [ ] **No new hot-path copy:** the diff adds no `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path. ⚠ any new instance; ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without a comment, or `deepcopy` of a scheduler / sampling-config on the request path. +- [ ] **No new event-loop blocking:** the diff adds no `time.sleep` / blocking HTTP (`requests`, `urllib`) / `await` inside a held lock in `async def` code. ⚠ any new instance; ✗ for blocking sleep or HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock held across an `await`. Blocking inside a dedicated worker thread is fine. + +Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row. + +### Conventions (eyeball — no grep) + +Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check: + +- **Right log level** — `info` is for end-user-visible events; routine internal events are `debug`; hot paths are `trace` or removed. Logging is not free. +- **Structured fields over f-strings** — pass dynamic values as structured fields (`logger.info("msg", extra={...})`), not interpolated strings, so they stay greppable and parseable. +- **No reflexive synchronization** — don't add `threading.Lock` / `asyncio.Lock` / `deepcopy` / re-wrap shared state unless there is real concurrent mutation; the GIL covers most Python code. Owners decide their own synchronization — don't pre-lock in a constructor. +- **Context managers for cleanup** — use `with` / `contextlib` for files, locks, CUDA streams, and profilers; not raw `open()` + manual `.close()` or ad-hoc `finally`. +- **No new dependencies without justification** — prefer the stdlib and existing vllm-omni deps; a new top-level dep needs a reason (check `pyproject.toml` / `requirements*.txt`). +- **Naming / interface discipline** — don't rename or change interfaces just for taste in this PR, but do fix names that actively mislead (e.g. `serve` implying a long-running server, `Instance` where multiple instance kinds coexist). + +--- + ## Bug Fix PRs ### Quick @@ -49,7 +74,7 @@ PR title must follow the project convention documented in `docs/contributing/REA - [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file - [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch - [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting -- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added +- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added. See the broad-except pattern in [code-quality.md](code-quality.md#2-broad-exception-swallow) for the full severity rule (⚠ any new broad catch; ✗ for a swallow on a fail-fast path). - [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.) - [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent diff --git a/.claude/skills/precheck-pr/references/code-quality.md b/.claude/skills/precheck-pr/references/code-quality.md new file mode 100644 index 00000000000..cb4b4ee443d --- /dev/null +++ b/.claude/skills/precheck-pr/references/code-quality.md @@ -0,0 +1,179 @@ +# Code-Quality Patterns + +Five fragility patterns to sweep on **every** PR, in both quick and full mode. These are pervasive in the existing codebase (hundreds of pre-existing sites) — so the checks are **diff-scoped**: only lines the PR *adds* count. Whole-repo greps light up every PR with the backlog and are useless. + +All commands assume: +```bash +BASE="$(git merge-base HEAD origin/main)" +``` + +**Severity:** ⚠ for any *new* instance introduced by the diff; ✗ for the acute sub-cases defined per pattern. Acute ✗ findings belong in the report's blocking count. + +--- + +## 1. `**kwargs` string-lookup passthrough + +### Anti-pattern +```python +def forward(self, **kwargs): + # BLOCKER-prone: raw dict + string keys duplicated across files + if "runtime_additional_information" in kwargs and \ + "model_intermediate_buffer" not in kwargs: + ... +``` +Real instances of this exact guard are copy-pasted across `cosyvoice3.py`, `indextts2_talker.py` (×2), `indextts2_s2mel_decoder.py`, `qwen3_tts_talker.py`, `qwen3_omni.py`. A typo in any string key silently no-ops; there is no type check, no IDE completion, and unknown keys are silently dropped. + +### Why dangerous +Raw `dict` + `**kwargs` + lists of param-name strings is **fragile by construction**: the compiler can't see a missing or misspelled key, so failures are silent. Each duplicated copy is a new place for the strings to drift. + +### Severity +- ⚠ any *new* `**kwargs` signature whose body does `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`. +- ✗ when the diff introduces a **new** string key that is duplicated across ≥2 files without a shared module constant, or when `**kwargs` plumbing **silently drops unknown keys** on a fail-fast path (init, config validation, weight loading). + +### Detect (diff-scoped) +```bash +git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \ + | grep -E '\*\*kwargs|kwargs(\.get|\[)|"[a-z_]+"\s+in\s+kwargs' +``` + +### Fix +Prefer explicit typed params. The repo already has the right primitive: `msgspec.Struct` with `forbid_unknown_fields=True` rejects unknown keys instead of silently dropping them — see `_StructBase` in `vllm_omni/data_entry_keys.py:105` and `MoriPullRequest` in `vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py:81`. A `@dataclass` or `TypedDict` works too. + +If `**kwargs` is unavoidable (vLLM base-class compat), centralize the key strings as module-level constants behind **one** typed accessor, instead of repeating the literal across files. + +--- + +## 2. Broad exception swallow + +### Anti-pattern +```python +# BLOCKER-prone: catches everything, returns None on any failure +try: + evt = build_event(...) + return evt +except Exception: + return None +``` +Real clusters: `vllm_omni/metrics/stats.py:240,274,367,450` and `vllm_omni/metrics/{modality,utils}.py` — all `except Exception: return None`. + +### Why dangerous +A generic `except Exception:` is **too flexible**. If we catch an exception type we don't expect, it is almost always an unhandled edge case that needs to be *fixed* — not swallowed into `None`/`pass`/`continue`. Broad catches turn fixable bugs into silent wrong behavior and make debugging impossible. + +The repo backs this with ruff `BLE001` (broad-exception catch); it is currently suppressed via `# noqa: BLE001` in `vllm_omni/patch.py:413`. Prefer fixing over suppressing. + +### Severity +- ⚠ any *new* `except Exception:` / `except BaseException:` / bare `except:`. +- ✗ for `except:` / `except Exception: pass|return None|continue` on a **fail-fast path** (init, config validation, weight loading, request handling, connector setup). This extends the Bug-Fix "No silent failure risk" ✗ in [checklists.md](checklists.md). + +### Detect (diff-scoped) +```bash +# Any new broad catch (⚠ baseline): +git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \ + | grep -E 'except\s*(Exception|BaseException)?\s*:' +# Note: this matches only the immediate-colon swallow forms (`except :`, +# `except Exception:`) — it deliberately excludes specific catches like +# `except ValueError:` and the `except Exception as e:` variant. Eyeball +# the latter; `as e` usually means it is logged, but still verify. +``` + +### Fix +Catch the **specific** types you actually expect (`ValueError`, `KeyError`, `AttributeError`, `OSError`, `torch.cuda.OutOfMemoryError`, `msgspec.ValidationError`, …). At a genuine top-level / best-effort boundary (metrics, signal handlers), at minimum **log** the exception; never swallow into a bare `pass`/`return None` on a path that should fail loudly. + +--- + +## 3. `Any` / wrong type annotations + +### Anti-pattern +```python +# ⚠ Any param + Any return — no contract, no narrowing +def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]: + ... +``` + +### Why dangerous (root cause) +`Any` disables the type checker for that position, so typos and wrong shapes pass silently. The repo has 636 `: Any` params, 244 `-> Any` returns, and **1,173 `SimpleNamespace` in `tests/`**. The `Any` leak is largely test-driven: vibe-coded unit tests fake objects with `SimpleNamespace`, which forces the production code consuming them to be typed `Any`, and the `Any` then propagates outward into public signatures. + +### Severity +- ⚠ any *new* `: Any` / `-> Any` / fully-untyped signature in production code. +- ✗ for a **wrong** annotation (actively misleading — e.g. `-> bool` that returns `Optional[bool]`), or a *new* `SimpleNamespace` in a test that mimics an object which already has a real typed stub / `@dataclass` / `TypedDict` / `Protocol`. + +### Detect (diff-scoped) +```bash +# Any leaking into new signatures: +git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \ + | grep -E ':\s*Any\b|->\s*Any\b' +# SimpleNamespace newly added in tests (the Any-leak source): +git diff ${BASE}...HEAD -- "tests/" | grep "^+[^+]" | grep "SimpleNamespace" +``` + +### Fix +Replace `Any` with the concrete type, a `Protocol`, or a `Union`. If a type is genuinely dynamic, prefer `object` + `isinstance` narrowing over `Any` — `object` keeps the checker engaged. In tests, replace `SimpleNamespace` with the real class, a small `@dataclass`/`TypedDict`, or a `Protocol` + fake, so the system under test stays fully typed. + +--- + +## 4. Unnecessary copy — `.clone()` / `copy.deepcopy` in hot paths + +### Anti-pattern +```python +# Per-step in a diffusion/AR loop: clones the full latent every iteration +for t in timesteps: + trajectory_latents.append(x_t.clone()) # O(N) GPU alloc + memcpy each step +``` +Real clusters: `vllm_omni/diffusion/models/bagel/bagel_transformer.py` calls `x_t.clone()` ~9× across a trajectory loop (lines 1820–2186); the cache hooks `vllm_omni/diffusion/cache/teacache/hook.py:155` and `magcache/hook.py:182` clone `hidden_states` per step. Separately, `copy.deepcopy(self.scheduler)` / `copy.deepcopy(req.sampling_params)` is copy-pasted across diffusion pipelines — `pipeline_qwen_image.py:786`, `pipeline_hunyuan_image3.py:1904`, `pipeline_ltx2*.py` (×3), `pipeline_dreamzero.py:1085`, `diffusion_model_runner.py:398` — a full object-graph deepcopy on the request path. (Contrast `dreamzero/causal_wan_model.py:203`, which carries a `# No .clone() needed.` comment — the aware pattern.) + +### Why dangerous +A `.clone()` of a GPU tensor or a `copy.deepcopy` of a config object inside a per-step / per-request path is a hidden allocator: it reads as one Python line but triggers a GPU memcpy or a deep object-graph walk every iteration. Across 50 diffusion steps or hundreds of AR steps it dominates latency and VRAM bandwidth. `deepcopy` additionally ignores any custom device / `__copy__` logic and can drag tensors through CPU. + +### Severity +- ⚠ any *new* `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path. +- ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without an explanatory comment, or `copy.deepcopy` of a scheduler / sampling-config on the request hot path. + +### Detect (diff-scoped) +```bash +git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \ + | grep -E '\.clone\(\)|\.copy_\(|copy\.deepcopy|deepcopy\(' +``` + +### Fix +Tensors: prefer a view (`x[..., :k]`, `torch.narrow`, `as_strided`) or write into a pre-allocated buffer with `dst.copy_(src)` instead of `.clone()`; if the caller no longer needs the original, `move`/`del` it rather than copy. Schedulers / configs: build a fresh lightweight instance per request, or split out the mutable per-request bits so the shared immutable part isn't deepcopied. If a copy is genuinely required (e.g. the source is mutated in place later), keep it and add a one-line comment saying why. + +--- + +## 5. Blocking the asyncio event loop / lock held across `await` + +### Anti-pattern +```python +# BLOCKER-prone: blocking sleep / blocking HTTP / lock held across an await, +# all inside an async (event-loop) function +async def handle(self, req): + with self._lock: + resp = await self._client.call(req) # lock held across await -> serializes all requests + time.sleep(0.1) # stalls the single loop for every concurrent request +``` + +### Why dangerous +vLLM-Omni serves over a single asyncio event loop. Any synchronous blocking work on an `async def` path — `time.sleep`, blocking HTTP (`requests`/`urllib`), heavy CPU work, or a lock held across an `await` — stalls the loop for *every* concurrent request, not just the blocking one. A lock held across `await` silently serializes the whole pipeline; blocking I/O tanks throughput and inflates tail latency. + +### Severity +- ⚠ any *new* `time.sleep(`, blocking HTTP call (`requests.get/post`, `urllib.request`, `urlopen`), or `await` inside a `with Lock` / `async with` critical section, in `async def` code. +- ✗ for `time.sleep` or blocking HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock acquired and held across an `await` on the request path. + +### Detect (diff-scoped) +```bash +# New blocking primitives introduced by the diff — then eyeball whether each +# sits in an async function on the loop vs a dedicated worker thread: +git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \ + | grep -E 'time\.sleep|requests\.(get|post|put|delete)|urllib\.request|urlopen' +# Lock acquired with an await potentially inside the critical section (manual eyeball): +git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" | grep -E 'async with|\.acquire\(\)' +``` +Note: blocking inside a **dedicated worker thread is correct** — the repo already does this (`diffusion/model_loader/hub_prefetch.py`, `distributed/omni_connectors/kv_transfer_manager.py`, the `mooncake`/`mori` connectors poll with `time.sleep` on their own threads). The check targets the *event-loop* path only. + +### Fix +Move blocking work off the loop with `asyncio.to_thread(...)` / `run_in_executor`, or use a native async client (`httpx.AsyncClient` / `aiohttp`). Never `await` while holding a lock — drop the lock before the `await` and re-acquire after, or restructure so the critical section contains no `await`. + +--- + +## Running the sweep + +Run all five detections against `${BASE}...HEAD`, then for each hit decide ⚠ vs ✗ using the severity rules above. Roll the results into the report as a single **Code quality** dimension row (count of ⚠ and ✗), alongside the type-specific checklist from [checklists.md](checklists.md). diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md index f42fcb6abc0..b8e9bbb1389 100644 --- a/.claude/skills/readme.md +++ b/.claude/skills/readme.md @@ -30,6 +30,12 @@ include: - `precheck-pr`: self-check a branch before creating a PR — validates PR title format, catches dead code, verifies accuracy/perf claims, and confirms merge readiness +- `vllm-omni-test`: guides generation and execution of CI-aligned tests (L1–L4), + pytest marker selection (`core_model` / `advanced_model` / `full_model`, + `omni` / `tts` / `diffusion`), Buildkite wiring (`test-ready.yml`, + `test-merge.yml`, `test-nightly.yml`, `test-weekly.yml`), and copy-paste + local plus CI-like `pytest` commands; see `references/test-routing.md` for + level-to-command mapping - `review-pr`: provides a structured workflow for reviewing pull requests - `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the latest vllm-ascend NPUModelRunner diff --git a/.claude/skills/vllm-omni-test/SKILL.md b/.claude/skills/vllm-omni-test/SKILL.md new file mode 100644 index 00000000000..db5068a507d --- /dev/null +++ b/.claude/skills/vllm-omni-test/SKILL.md @@ -0,0 +1,894 @@ +--- +name: vllm-omni-test +description: Generate and run tests for vllm-project/vllm-omni with CI-aligned levels and markers; wire new tests into Buildkite (test-ready.yml for L1/L2, test-merge.yml for L3, test-nightly.yml for L4). On completion, always provide copy-paste local and CI-like pytest commands plus prerequisites. Use when creating regression tests, adding L1-L4 coverage, selecting pytest markers, or validating fixes from issues/PRs. +--- + +# vLLM-Omni Test Generator & Runner + +## Purpose + +Use this skill to generate minimal, stable test cases and run them with the correct marker/level strategy for [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni). + +**Link convention:** Paths such as `.buildkite/` and `docs/contributing/` live at the **vllm-omni repo root**. Markdown links use repo-relative paths from this skill file (e.g. `../../../.buildkite/cuda/test-ready.yml`, `../../../docs/contributing/ci/CI_5levels.md`). + +Default priorities: + +1. Reproducible regression coverage for bug fixes +2. Correct test level and marker selection +3. Low flake, low dependency tests first +4. CI-compatible run commands +5. **Actionable run commands for the human**: whenever you add or change tests, always finish with **copy-paste-ready** `pytest` lines (local: single file and/or single test; CI-like: markers + `--run-level`), plus short **prerequisites** (GPU tier, HF cache, optional `model_prefix`). Do not assume the reader will infer commands from `test-routing.md` alone. + +## Inputs + +- Issue/PR link and summary +- Changed files or suspected code path +- Whether the user wants local quick validation or CI-equivalent validation +- Hardware constraints (CPU only / CUDA / ROCm / NPU) + +## Workflow + +### Step 1: Classify Test Goal + +- **Bugfix regression**: start from a minimal failing scenario and add assertions that prevent recurrence. Before writing tests, output **`required` / `recommended` / `not_needed`**: **`required`** — stable logic/contract bug that should have been caught; **`recommended`** — environment-sensitive but a small regression still helps; **`not_needed`** — one-off external/config failure or existing tests already cover the path. Prefer the narrowest stable L1 (CPU) case; escalate to L2/L3 only when the bug needs real weights or serving. +- **Feature coverage**: verify new behavior and one negative/boundary case. +- **Perf/benchmark claim**: require benchmark-oriented tests and explicit metrics. + +### Step 2: Select Test Level + +- **L1**: unit/logic, deterministic, CPU-friendly, fastest feedback. +- **L2**: basic e2e/integration and platform-dependent checks. +- **L3/L4**: advanced model/integration/perf validation. + +Use [references/test-routing.md](references/test-routing.md) for level-to-marker and command mapping. + +### Step 3: Pick Markers + +Always attach markers deliberately: + +- **Level**: `core_model` (L1/L2) and/or `advanced_model` (L3) and/or `full_model` (L4 nightly) +- **Model type** (required on model-centric e2e — pick exactly one): + - `omni` — end-to-end multimodal LLM pipelines (thinker/talker/stages; Qwen-Omni family) + - `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.) + - `diffusion` — generative diffusion models (image / audio / text / video from noise) +- **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark` +- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, … +- Optional: `slow`, distributed markers when multi-card is required + +**Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**. + +For hardware-aware tests, prefer `@hardware_test(...)` or `hardware_marks(...)` in `tests/helpers/mark.py`. + +**Diffusion nightly split** (`test-nightly.yml`): all diffusion tests use `pytest.mark.diffusion`, but CI groups them by **output modality**, not by separate pytest markers: + +| Nightly group | Scope | Typical models / paths | +|---------------|--------|-------------------------| +| **Diffusion X2I(&A&T) Model Test** | x2**i** (image), x2**a** (audio), x2**t** (text) and other **non-video** diffusion | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero, … — `test_*_expansion.py` under X2I shards; L4 sweep uses `-k "not test_wan and not test_bagel_expansion and not hunyuan"` for L4 | +| **Diffusion X2V Model Test** | x2**v** (video) only | Wan2.2, HunyuanVideo 1.5, Wan VACE, … — e.g. `test_wan22_expansion.py`, `test_hunyuan_video_15_expansion.py` | + +Wire new **video** diffusion expansion tests into the **X2V** group; wire **image/audio/text** diffusion expansion tests into **X2I(&A&T)**. Do not place x2v modules in X2I shards (see comment in `test-nightly.yml` above the X2I group). + +### Naming: generated test module files (L2–L4, model-centric e2e) + +When adding **new** pytest modules whose primary scope is a **specific model** (typical under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`), use this filename pattern: + +| Level | Filename pattern | Example (model `Qwen/Qwen2.5-Omni-7B`) | +|-------|------------------|----------------------------------------| +| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` | +| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` | + +**Slug rules for `{lowercase_model_slug}`:** + +1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`. +2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree). +3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) only then. +4. **L1** unit tests are **not** bound to this pattern; use `tests//test_.py` as today. + +Routing tables and commands: [references/test-routing.md](references/test-routing.md) § *Model-centric e2e filename convention*. + +Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-style omni), `tests/e2e/offline_inference/test_qwen3_5_9b.py` (L2-style omni, single-stage VL), `tests/e2e/online_serving/test_qwen3_omni_expansion.py` (L4-style omni), `tests/e2e/online_serving/test_qwen_image_edit_expansion.py` / `test_qwen_image_expansion.py` (L4-style diffusion). + +### Step 4: Generate Test Case Skeleton + +**1. Pick the functional scenario** (then choose directory, fixtures, and markers): + +| Scenario | Typical location | Fixtures / runner pattern | Baseline markers & level | +|----------|------------------|---------------------------|---------------------------| +| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required | +| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` | +| **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/cuda/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). | +| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: JSON `mark` with `full_model` + `omni`/`tts`/`diffusion`; wire `test-nightly.yml` Perf steps | +| **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) | + +**If the user’s test plan includes invalid parameter validation / invalid params / negative HTTP / 400 validation:** do **not** add those `test_*` functions to `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move or author them** under `tests/dfx/reliability/invalid_param_test/` in the **endpoint-matching script** (see **Invalid parameter validation** below). Success-path e2e and invalid-param dfx tests must stay in separate modules. + +**1b. Model type — `omni` vs `tts` vs `diffusion`** + +After choosing offline/online/docs/perf, classify the **product under test** and attach **exactly one** model-type marker. All three can live under the same `tests/e2e/` trees; conventions diverge: + +| Dimension | **Omni** (`pytest.mark.omni`) | **TTS** (`pytest.mark.tts`) | **Diffusion** (`pytest.mark.diffusion`) | +|-----------|------------------------------|----------------------------|----------------------------------------| +| **What it is** | Multimodal LLM pipeline (thinker/talker/stages; text + vision + audio I/O) | Speech synthesis / voice models | Generative diffusion (noise → image, audio, text, or video) | +| **Examples** | Qwen2.5-Omni, Qwen3-Omni | Qwen3-TTS, VoxCPM2, Higgs-Audio, Step-Audio2 | Qwen-Image, BAGEL, Wan2.2, HunyuanVideo | +| **Offline runner** | `omni_runner` + `omni_runner_handler`, `generate_multimodal` | `Omni(...).generate` or stage YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` | +| **Online client** | `openai_client.send_omni_request` (chat completions, modalities) | `openai_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) | +| **Typical assertions** | Stage outputs, text/audio keywords via `OmniRunnerHandler` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` | +| **Stage / deploy YAML** | Per-model omni stage configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT | +| **Nightly group** (`test-nightly.yml`) | **Omni Model Test** — `-m "full_model and omni"` | **TTS Model Test** — `-m "full_model and tts"` | **Diffusion X2I(&A&T)** *or* **Diffusion X2V** (see Step 3 table; same `diffusion` marker, different YAML group / file shard) | +| **L4 pressure** | Expansion per modality/model as needed | Expansion + accuracy/perf in TTS group | X2I: merge feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); X2V: separate nightly group | + +Do **not** mix fixtures across types (e.g. do not use `omni_runner` layout for a pure diffusion or TTS model without mirroring an in-tree test in that family). + +**Diffusion only — X2I(&A&T) vs X2V (nightly routing, not extra markers):** + +- **X2I(&A&T)**: image / audio / text generation — Qwen-Image*, FLUX, SD3, Z-Image, BAGEL (expansion), LongCat, audio diffusion, etc. +- **X2V**: video generation only — Wan2.2, HunyuanVideo 1.5, Wan VACE, LTX video similarity paths. + +When adding a new `test_*_expansion.py`, place it in the matching **nightly group** step (explicit file list in `test-nightly.yml`), not only by marker expression. + +**2. Use the narrowest deterministic skeleton for the scenario** + +*L1 unit / logic (CPU-first):* + +**Mocking rule (L1 only):** use **pytest** integration — **`mocker`** (`pytest-mock`) or **`monkeypatch`** (built-in). **Do not** import or call **`unittest.mock`** (`patch`, `MagicMock`, `@patch`, `with patch(...)`, etc.) in L1 tests; patches must auto-revert with the test lifecycle. + +```python +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_(mocker): + # Prefer mocker.patch / mocker.spy / mocker.Mock — not unittest.mock.patch + fake_fn = mocker.patch("vllm_omni.some.module.expensive_call", return_value=...) + # Act + # Assert + fake_fn.assert_called_once() + + +def test_(monkeypatch): + # For simple env / attribute substitution without a Mock object + monkeypatch.setenv("SOME_FLAG", "1") + monkeypatch.setattr("vllm_omni.some.module.CONST", 42) + ... +``` + +**Avoid in L1:** + +```python +# BAD — do not use in L1 unit tests +from unittest.mock import patch, MagicMock + +@patch("vllm_omni.some.module.fn") +def test_bad(mock_fn): ... + +def test_also_bad(): + with patch("...") as m: ... +``` + +See **L1 unit test constraints (mocking)** below for the full do/don't list. + +*Offline multimodal e2e — **Omni** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_(omni_runner, omni_runner_handler) -> None: + request_config = {"prompts": ..., "modalities": [...]} # optional: images, videos, audios + omni_runner_handler.send_omni_request(request_config) +``` + +*Offline generative e2e — **Diffusion** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.diffusion +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_text_to_image_001(omni_runner_handler) -> None: + omni_runner_handler.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}}) +``` + +*Offline TTS e2e — **Qwen3-TTS** (two-stage; representative):* + +```python +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None: + omni_runner_handler.send_audio_speech_request({ + "input": "...", + "task_type": "Base", + "ref_audio": REF_AUDIO_URL, + "ref_text": REF_TEXT, + }) +``` + +*Offline TTS e2e — **single-stage** (Coqui XTTS, MOSS-TTS-Nano; representative):* + +```python +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(...) +@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True) +def test_voice_clone_001(omni_runner_handler) -> None: + omni_runner_handler.send_single_stage_tts_request({ + "input": "...", + "language": "en", + "prompt_audio_path": REF_AUDIO_PATH, + "response_format": "wav", + "run_level": "advanced_model", + }) +``` + +*Online serving e2e — **Omni** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(...) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_text_001(omni_server, openai_client) -> None: + request_config = {"model": omni_server.model, "messages": ..., "modalities": ["text"]} + openai_client.send_omni_request(request_config) +``` + +*Online serving e2e — **TTS** (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(...) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_text_to_audio_001(omni_server, openai_client) -> None: + request_config = {"model": omni_server.model, "input": "...", "response_format": "wav", ...} + openai_client.send_audio_speech_request(request_config) +``` + +*Online serving e2e — **Diffusion** X2I (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True) +def test_text_to_image_001(omni_server, openai_client) -> None: + openai_client.send_diffusion_request({...}) # chat completions + extra_body +``` + +*Online serving e2e — **Diffusion** X2V (representative):* + +```python +@pytest.mark.core_model +@pytest.mark.diffusion +def test_text_to_video_001(omni_server, openai_client) -> None: + openai_client.send_video_diffusion_request({"model": ..., "form_data": {...}}) # /v1/videos +``` + +*Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts). + +*Performance tests:* add or extend entries under `tests/dfx/perf/tests/` (and JSON configs where the project uses them), with **explicit baselines**, **`mark`** on each case (`hardware_marks` + `full_model` + type marker), and the same nightly Perf step pattern as in-tree configs. + +**3. Cross-cutting rules** + +- Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason. +- Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache). +- Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo). +- **E2E test function layout**: **one case → one `test_` function** with a name that states what is validated (endpoint, size/`n`, server flag, or route). **Do not** merge multiple cases into a single test that branches on `request.node.callspec.id`, `param.id`, or `if case_id == ...`. Use `@pytest.mark.parametrize("omni_server", [...], indirect=True)` per function (usually **one** `OmniServerParams` per test). A loop inside one test is OK only when it serves **that function’s single intent** (e.g. three standard sizes in `test_*_sizes_256_512_1024`). +- **Runtime fixture scope** (`tests/helpers/fixtures/runtime.py`): default **`omni_server` / `omni_runner`** (module) + matching client/handler; use **`omni_server_function` / `omni_runner_function`** only when each `test_*` must start a fresh instance (see below). +- **L1 mocks**: never `unittest.mock`; use `mocker` or `monkeypatch` only (see below). +- **API calls (L2+ e2e, online and offline)**: **reuse** `send_*_request` in `tests/helpers/runtime.py` when it exists; **otherwise add** the helper in `runtime.py` first, then call it from the test. General `assert_*` inside `send_*_request`; special `assert_*` only in the test. See **Runtime send helpers** below — do not call `omni.generate`, raw HTTP, or SDK clients from `test_*.py`. +- **Response assertions**: reusable checks on API bodies / decoded media belong in **`tests/helpers/assertions.py`** — not as private `_assert_*` helpers inside `test_*.py` (see below). +- **Model-specific payloads stay in test modules** — per-model `MODEL`, deploy path, `REF_AUDIO_URL`, `get_prompt()`, `_build_request()`, and inline `request_config` dicts live in `test_{slug}.py` / `test_{slug}_expansion.py` (and offline/L1 siblings). **Do not** create `tests/helpers/{slug}.py` to deduplicate them; a little copy across files is preferred (see `test_glm_tts.py`, `test_cosyvoice3_tts_expansion.py`). `tests/helpers/` is repo-wide harness only (`mark`, `media`, `runtime`, `stage_config`, `assertions`, `fixtures/`). + +#### L1 unit test constraints (mocking) + +L1 tests (`core_model and cpu`, under `tests/diffusion/`, `tests/engine/`, `tests/model_executor/`, etc.) must follow the repo’s **pytest-mock** convention (`pytest-mock>=3.10.0` in `pyproject.toml` `[project.optional-dependencies] dev`). + +| Do | Don't | +|----|-------| +| `def test_foo(mocker):` + `mocker.patch(...)`, `mocker.spy(...)`, `mocker.Mock()`, `mocker.MagicMock()`, `mocker.AsyncMock()` | `from unittest.mock import patch, MagicMock, Mock, AsyncMock` | +| `def test_bar(monkeypatch):` + `monkeypatch.setattr` / `setenv` / `delenv` / `setitem` | `@patch(...)` decorator | +| Let `mocker` auto-stop patches after the test | `with patch(...):` / `patch.object(...)` context managers | +| Mirror neighboring L1 tests in the same directory | `unittest.mock.create_autospec` unless an existing file already documents an exception | + +**Rationale:** `mocker` ties patch lifecycle to pytest fixtures (no leaked patches across tests). `unittest.mock` decorators/context managers are easy to compose incorrectly with parametrized or async tests and are inconsistent with in-tree L1 style. + +**Minimal patterns:** + +```python +def test_returns_cached_config(mocker): + loader = mocker.patch( + "vllm_omni.foo.load_yaml", + return_value={"stages": []}, + ) + result = get_deploy_config("ci/foo.yaml") + assert result["stages"] == [] + loader.assert_called_once_with("ci/foo.yaml") + + +def test_skips_when_env_unset(monkeypatch): + monkeypatch.delenv("VLLM_OMNI_FEATURE", raising=False) + assert should_enable_feature() is False +``` + +E2E levels (L2+) generally **avoid mocks**; if a rare L2 stub is unavoidable, still prefer `mocker` over `unittest.mock` for consistency. + +#### Runtime fixtures — scope (`tests/helpers/fixtures/runtime.py`) + +vLLM-Omni e2e tests start a real **OmniServer** (online) or **OmniRunner** (offline). Pick scope by **how often the process must be recreated**, not by test level alone. + +| Scope | Online fixtures | Offline fixtures | When to use | +|-------|-----------------|------------------|-------------| +| **Module** (default) | `omni_server` → `openai_client` | `omni_runner` → `omni_runner_handler` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. | +| **Function** | `omni_server_function` → `openai_client_function` | `omni_runner_function` → `omni_runner_handler_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. | + +**Rules:** + +1. **Default to module scope** (`omni_server` / `omni_runner`) unless the scenario **explicitly** needs a fresh instance per test function. +2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `openai_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `openai_client_function`. Do not mix module fixture with function client (or vice versa). +3. Different `OmniServerParams` per `test_*` (e.g. default vs `--enable-cpu-offload`) is still OK with **module** `omni_server` — pytest parametrizes per test node; only switch to `_function` when isolation between tests matters, not merely because `server_args` differ. +4. One `test_*` with **many** server configs (expansion matrix) → single function + `@pytest.mark.parametrize("omni_server", [...], indirect=True)` + module `omni_server` (see in-tree `test_qwen_image_expansion.py`). + +```python +# Default — module-scoped server (L4 expansion) +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True) +def test_foo_images_generations_default_1024(omni_server, openai_client) -> None: + openai_client.send_images_generations_request({...}) + +# Function-scoped — reliability / per-test clean state only +@pytest.mark.parametrize( + "omni_server_function", + [pytest.param(OmniServerParams(model=MODEL), marks=H100)], + indirect=True, +) +def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> None: + openai_client_function.send_omni_sleep_http_request({...}) + openai_client_function.send_omni_wakeup_http_request({...}) +``` + +#### Runtime send helpers — online **and** offline (`tests/helpers/runtime.py`) + +**L2+ e2e (online serving and offline inference) must call APIs through `tests/helpers/runtime.py`.** Fixtures live in `tests/helpers/fixtures/runtime.py`; **send/assert implementation** lives in `tests/helpers/runtime.py` + `tests/helpers/assertions.py`. + +| Principle | Action | +|-----------|--------| +| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OpenAIClientHandler` (online) or `OmniRunnerHandler` (offline) that matches the endpoint / pipeline shape. | +| **Extend when close** | If an existing `send_*_request` almost fits (e.g. missing one optional field), extend it in `runtime.py` — do not fork logic in the test file. | +| **Add when missing** | No suitable helper → add **`send__request`** (high-level: call + general `assert_*`) or **`send___http_request`** (low-level HTTP for negative/dfx) in **`runtime.py` first**, then call it from tests. | +| **Test module owns payload only** | In `test_*.py`: `MODEL`, deploy path, vendored media, `get_prompt()`, and inline **`request_config` dicts** only. **No** `omni.generate(...)`, raw `requests.post`, OpenAI SDK calls, or `_collect_audio()` / `_process_output()` in e2e tests. | + +**Online** (`openai_client` from `omni_server`): `OpenAIClientHandler.send_*_request`. + +**Offline** (`omni_runner_handler` from `omni_runner`): `OmniRunnerHandler.send_*_request`. + +| Do | Don't | +|----|-------| +| Online: `openai_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test | +| Offline: `omni_runner_handler.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` | +| Add missing `send_*` to **`runtime.py`** first; bundle general `assert_*` inside the send helper | A one-off `def _post_*` / `def _collect_audio` at the bottom of a test module | +| Mirror naming/style of neighboring `send_*` (docstring, `request_config` dict, optional `run_level`, `err_code` / `err_message` for negative cases) | Different parameter shapes per test file for the same endpoint | + +**Workflow when generating online/offline e2e tests:** + +1. Decide whether the needed check is **general** (every success call of this `send_*`) or **special** (one case / one parameter combo only). See **General vs special assert placement** below. +2. **Search `runtime.py`** for a matching **`send_*_request`**; if missing, add it (with general `assert_*` inside) before writing the test body. +3. In the test module: build `request_config` → call **`send_*_request` only** for the general contract; call **`assert_*` in the test file only** for special, case-specific checks (import from `assertions.py`). +4. Reserve **low-level** `send_*_http_request` for negative/dfx tests (`err_code` / `err_message`) — not for ordinary L2+ success-path e2e. +5. When wiring Buildkite `source_file_dependencies`, include **`tests/helpers/runtime.py`** and/or **`tests/helpers/assertions.py`** when new helpers were added. + +**Exceptions (document in the test docstring why):** models whose offline prompt path cannot go through existing handlers yet (e.g. `test_higgs_audio_v2.py`, `test_voxtral_tts.py` with custom tokenizer compose) may call `omni.generate` directly until a `send_*_request` is added to `runtime.py` — treat as **debt**, not the default for new TTS/diffusion/omni e2e. + +**General vs special assert placement:** + +| Kind | Definition | Where it lives | Test file calls | +|------|------------|----------------|-----------------| +| **General** | Default success contract for this endpoint/client method — e.g. HTTP 200, `data[]` shape, `n` count, `size` dimensions, decodable image/audio, bundled omni/diffusion fields | Implement in `assertions.py`, invoke **inside** `send_*_request` in `runtime.py` (after low-level send / SDK call, when `err_code` is not set) | **`send_*_request` only** — do **not** repeat the same `assert_*` | +| **Special** | Extra check tied to **this test case only** — e.g. seed byte-identical replay, accuracy/CLIP threshold, perf ceiling, model-specific optional field | Add or reuse `assert_*` in `assertions.py` (never inline in test) | `send_*_request` **then** `assert_(...)` **once** for that case | + +Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** special validation — the general `assert_*` should read those from `request_config` inside `send_*`. + +**`send_*` ↔ `assert_*` pairing (online `OpenAIClientHandler`):** + +| High-level `send_*` (prefer in L2+ e2e) | Assert already invoked inside `runtime.py` | Low-level HTTP-only sibling | +|----------------------------------------|--------------------------------------------|-----------------------------| +| `send_omni_request` | `assert_omni_response` | `send_chat_completions_http_request` → `assert_http_error` only | +| `send_diffusion_request` | `assert_diffusion_response` | — | +| `send_audio_speech_request` | `assert_audio_speech_response` | `send_audio_speech_http_request` → `assert_http_error` only | +| `send_video_diffusion_request` | `assert_diffusion_response` | `send_videos_*_http_request` → `assert_http_error` only | +| `send_images_generations_request` *(add when needed)* | **`assert_images_generations_response`** *(general — inside send)* | `send_images_generations_http_request` → `assert_http_error` only | +| `send_images_edits_request` *(add when needed)* | **`assert_images_edits_response`** *(general — inside send)* | `send_images_edits_http_request` → `assert_http_error` only | + +**Rule:** Tests call high-level `send_*_request` for the general contract. **Never** call the same bundled `assert_*` again. Call an extra `assert_*` in the test **only** for special, case-specific validation. + +**Common `OpenAIClientHandler` entry points** (non-exhaustive — grep `runtime.py` before adding): + +| Area | High-level (SDK + assert) | Low-level HTTP (`*_http_request`) | +|------|---------------------------|-----------------------------------| +| Omni chat | `send_omni_request` | `send_chat_completions_http_request` | +| Diffusion T2I (chat route) | `send_diffusion_request` | — | +| Diffusion X2V | `send_video_diffusion_request` | `send_videos_create_http_request`, `send_video_content_http_request`, … | +| DALL-E T2I / edit | `send_images_generations_request`, `send_images_edits_request` | `send_images_generations_http_request`, `send_images_edits_http_request` | +| TTS (online) | `send_audio_speech_request` | `send_audio_speech_http_request`, `send_audio_generate_http_request`, … | +| Ops / meta | — | `send_health_http_request`, `send_models_http_request`, `send_omni_sleep_http_request`, … | + +**Common `OmniRunnerHandler` entry points** (offline — grep `runtime.py` before adding): + +| Area | High-level (`send_*_request` + assert) | Notes | +|------|----------------------------------------|-------| +| Omni multimodal | `send_omni_request` | `generate_multimodal` path | +| Diffusion offline | `send_diffusion_request` | chat-route / `OmniTextPrompt` offline | +| Qwen-style TTS | `send_audio_speech_request` | two-stage, `generate_multimodal` + `mm_processor_kwargs` | +| Single-stage TTS (Coqui XTTS, MOSS-TTS-Nano, …) | `send_single_stage_tts_request`, `send_single_stage_tts_batch_request` | `prompt` + `additional_information` + `omni.generate` — **do not** duplicate in tests | +| New model family | **Add** `send__request` here first | Then call from `tests/e2e/offline_inference/test_*.py` | + +L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OpenAIClientHandler` / `OmniRunnerHandler`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers. + +#### Invalid parameter validation (`tests/dfx/reliability/invalid_param_test/`) + +When the user asks for **invalid parameter validation**, **invalid request bodies**, **HTTP 4xx contract tests**, or any case that sends **malformed / out-of-range / mismatched** API payloads against a **live** server: + +1. **Do not** place these in `tests/e2e/online_serving/` or `*_expansion.py`. If already drafted there, **move** the `test_*` into the correct `invalid_param_test` script and delete the duplicate from e2e. +2. **Pick the script by HTTP route** (extend an existing file; add a new `test_invalid_.py` only when no in-tree script covers that route family): + +| API / area | Script | +|------------|--------| +| Omni chat completions, WebSocket video/realtime paths | `test_invalid_omni_chat.py` | +| `POST /v1/audio/speech`, stream, batch, voices | `test_invalid_audio_speech.py` | +| Audio diffusion endpoints | `test_invalid_audio_diffusion.py` | +| `POST /v1/images/generations` | `test_invalid_image_generation.py` | +| `POST /v1/images/edits` | `test_invalid_image_editing.py` | +| `POST/GET/DELETE /v1/videos*` | `test_invalid_video_generation.py` | +| Sleep / wakeup / server control | `test_invalid_server_control.py` | + +3. **Match in-tree style** in the chosen script: + +| Element | Convention | +|---------|------------| +| **Module markers** | `pytestmark = [pytest.mark.slow, pytest.mark.]` | +| **Server fixture** | `_PARAMS` / `_QWEN3_TTS_SPEECH`-style list of `pytest.param(OmniServerParams(...), id="...", marks=hardware_marks(...))`; `@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)` | +| **Hardware** | `hardware_marks(res={"cuda": "H100"})` for heavy diffusion/omni/video; `hardware_marks(res={"cuda": "L4"})` for smaller TTS models (must match weekly `-m "slow and L4"` step) | +| **HTTP client** | **Low-level** `openai_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) | +| **Case shape** | Prefer **one `test_*` per route family** + `@pytest.mark.parametrize("body_spec, err_message", [...])` with stable `id=` per case; or dedicated `test__malformed_json` when not parametrized | +| **Body helpers** | `_minimal__json(omni_server)` / `_minimal_*_form_data()`; merge overrides with `body.update(body_spec)` | +| **Known gaps** | `pytest.mark.skip(reason="…#3649")` as `_SKIP_ISSUE_3649` when server validation is not yet strict (mirror neighboring cases) | +| **Sections** | Route banner comments (`# ─── POST /v1/images/generations ───`) like existing files | +| **Shared fixtures** | Reuse `tests/dfx/reliability/invalid_param_test/conftest.py` (`tiny_png_bytes`, env defaults) | + +4. **Adding a new model** to invalid-param coverage: append a `pytest.param(OmniServerParams(model="...", stage_config_path=..., server_args=...), ...)` entry to the script’s `_PARAMS` list — do **not** create `test_invalid_.py` unless the route is new. + +5. **CI — `.buildkite/cuda/test-weekly.yml` only** (not `test-ready.yml` / `test-merge.yml` / `test-nightly.yml`): + +| Weekly step | Command | When your cases run | +|-------------|---------|---------------------| +| **Invalid parameters Test · H100** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"` | Diffusion / omni / video invalid-param tests with `H100` hardware mark | +| **Invalid parameters Test · L4** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"` | TTS / lighter models marked `L4` | + +- **Trigger:** `build.env("WEEKLY") == "1"` or PR label `weekly-test`. +- **Default:** extending an existing `invalid_param_test` script needs **no YAML edit** — the weekly steps already sweep the whole directory. +- **Edit YAML only** when adding a **new hardware queue**, a **new top-level script** that must run in isolation, or a **model-specific weekly shard** (mirror neighboring reliability steps). +- Weekly steps **do not** use `source_file_dependencies`. + +**Example — append to `test_invalid_image_generation.py`:** + +```python +@pytest.mark.parametrize( + "body_spec, err_message", + [ + pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"), + ], +) +@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True) +def test_images_generations_invalid_requests( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + body_spec: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body = _minimal_images_gen_json(omni_server) + body.update(body_spec) + openai_client.send_images_generations_http_request( + {"json": body, "timeout": 300, "err_code": 400, "err_message": err_message} + ) +``` + +See [references/test-routing.md](references/test-routing.md) **Invalid parameter / weekly CI**. + +#### Assertion helpers (`tests/helpers/assertions.py`) + +**Do not** add module-local helpers such as `_assert_images_generations_payload` or `_send_and_assert_*` in e2e test files. Response/media validation belongs in **`tests/helpers/assertions.py`**, grouped by **category**: + +| Category | Existing anchors | When to extend vs add | +|----------|------------------|------------------------| +| **Image** (chat diffusion + DALL-E JSON) | `assert_image_diffusion_response`, `assert_image_valid` | DALL-E `/v1/images/generations` JSON → add **`assert_images_generations_response`** beside image helpers (reuses `assert_image_valid`); do not fork decode logic into tests | +| **Video** | `assert_video_diffusion_response`, `assert_video_valid` | Extend these for new video contracts | +| **Audio** | `assert_audio_diffusion_response`, `assert_audio_speech_response`, `assert_audio_valid` | Extend for new TTS/audio endpoints | +| **Omni multimodal** | `assert_omni_response` | Extend for new modality combos | +| **HTTP errors** | `assert_http_error`, `assert_err_message_in_text` | Used inside low-level `send_*_http_request` and negative tests | + +| Do | Don't | +|----|-------| +| Put **general** `assert_*` inside `send_*_request` in `runtime.py`; tests call `send_*_request` only | Call `assert_diffusion_response` after `send_diffusion_request` (general assert already bundled) | +| Put **special** `assert_*` in `assertions.py` and call it **in the test** after `send_*_request` when that case needs extra checks | Put general decode/count/size logic in the test because “this case uses `n=4`” (that belongs in general assert reading `request_config`) | +| Extend the matching category function, or add `assert__response` next to its category | `_send_and_assert_*` or PIL/base64 loops in `test_*.py` | +| Use low-level `send_*_http_request` in dfx/negative tests (`err_code`) | Use low-level HTTP send + manual general assert in every L2+ expansion file | + +**Workflow when generating tests:** + +1. Implement or reuse **general** `assert_*` in `assertions.py` (by category). +2. Wire it into **`send_*_request`** in `runtime.py` if not already bundled. +3. In the test: `request_config` → `send_*_request`. Add a separate `assert_*` import/call **only** for special case validation. +4. Update `source_file_dependencies` when `runtime.py` or `assertions.py` changes. + +**Example — DALL-E `/v1/images/generations`:** + +```python +# tests/helpers/assertions.py — general contract (image category) +def assert_images_generations_response(resp_body: dict, request_config: dict, *, run_level: str | None = None) -> None: + ... # data[], n, size → assert_image_valid + +# tests/helpers/runtime.py — general assert INSIDE send +def send_images_generations_request(self, request_config: dict[str, Any], ...) -> list[HttpResponse]: + responses = self.send_images_generations_http_request(request_config, ...) + if request_config.get("err_code") is None: + body = responses[0].json_body + assert isinstance(body, dict) + assert_images_generations_response(body, request_config, run_level=self.run_level) + return responses + +# tests/e2e/.../test_foo_expansion.py — one case per test_* (no case_id branching) +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) +def test_foo_images_generations_default_1024(omni_server, openai_client) -> None: + openai_client.send_images_generations_request({"json": body, "timeout": 300}) + +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) +def test_foo_images_generations_sizes_256_512_1024(omni_server, openai_client) -> None: + for size in ("256x256", "512x512", "1024x1024"): + openai_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300}) + +# chat diffusion — separate test; send_diffusion_request bundles assert_diffusion_response +@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True) +def test_foo_chat_completions_t2i_fallback(omni_server, openai_client) -> None: + openai_client.send_diffusion_request({...}) +``` + +L1 tests under `tests/entrypoints/` may keep **minimal** asserts next to the handler under test when validating a single branch; still prefer `assertions.py` when the same JSON/media check appears in more than one file or level (e.g. L1 protocol test + L4 e2e). + +#### Omni Test Writing Guidance (L1-L4 Layering) + +When the goal is a general Omni/multimodal test case, prioritize mapping the test to the correct purpose, directory, and resource assumptions aligned with the layers (see [CI_5levels.md](../../../docs/contributing/ci/CI_5levels.md)): + +- **L1**: Unit/logic validation on CPU (`core_model and cpu`). Cover input validation, branches, and exception paths (`tests//test_*.py`). **Mock with `mocker` / `monkeypatch` only — not `unittest.mock`.** +- **L2**: Basic e2e (online/offline basic scenarios). Prefer dummy/lightweight models to validate the end-to-end request-to-output-structure/streaming chain (typically `tests/e2e/online_serving/` and `tests/e2e/offline_inference/`). +- **L3/L4**: Important integration, performance, and accuracy validation. L4 emphasizes “full functional scenarios + performance/stress + runnable doc examples” (typically `*_expansion.py` plus related expansion cases). + +Also keep markers consistent with the run level: use `core_model` for L1/L2 and `advanced_model` for L3/L4, and pair with `--run-level` to select the intended CI strategy. + +#### Diffusion Test Writing Guidance (L4 Coverage Combinations) + +When the task involves diffusion models/features, organize L4 test cases following [`#1832`](https://github.com/vllm-project/vllm-omni/issues/1832): combine multiple diffusion features into as few test cases as possible to fit limited CI GPU resources. + +Implementation strategy: + +1. If “full L4 is too heavy”: first provide a **reduced local validation case in L1/L2** (so the key assertion and the contract fix point are covered deterministically). +2. Then provide **CI/nightly-ready L3/L4 high-marked cases** (e.g. `advanced_model`) to broaden coverage under resource constraints. + +#### L4 nightly sub-pillars — Function vs Accuracy vs Perf + +**L4 is not only `test_*_expansion.py`.** In `test-nightly.yml`, each model-type group typically runs **separate jobs**: + +| Sub-pillar | Nightly step label pattern | What you add | Typical paths | +|------------|----------------------------|--------------|---------------| +| **Function** | `· Function Test with …` | E2e expansion / feature matrix | `tests/e2e/online_serving/test__expansion.py` (or offline) | +| **Accuracy** | `· Accuracy Test` | Quality / similarity vs baseline | `tests/e2e/accuracy/test_*.py` | +| **Perf** | `· Perf Test · ` | Throughput / latency / memory benchmarks | `tests/dfx/perf/tests/test__vllm_omni.json` + runner script | +| **Doc** (optional) | `· Doc Test` | Runnable doc examples | `tests/examples/*/test_text_to_image.py`, … | + +**Default when the user asks for “L4 functional cases”:** deliver **Function** pillar only (`*_expansion.py` + Function Test shard in `test-nightly.yml`). **Do not** silently add Perf or Accuracy unless the user also asks for **performance** / **benchmark** / **accuracy** / **full L4** / **full L4 coverage**. + +**When the user explicitly asks for L4 perf (or full L4 including perf):** + +1. **Do not** put throughput/latency assertions inside `test_*_expansion.py` — perf uses the **dfx benchmark harness**, not `omni_server` e2e fixtures. +2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, `test_qwen3_omni_async_chunk.json`). +3. Each JSON **case**: `test_name`, optional **`mark`** (`hardware_marks` required when present; `marks`: `full_model` + `omni`/`tts`/`diffusion`), `server_params` (`model`, `serve_args` or `stage_config_name`), `benchmark_params[]` with **`name`** and explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …). Diffusion cases also set `server_type` (e.g. `"vllm-omni"`) and usually `benchmark_endpoint`. +4. Run locally via the matching script (model type → runner): + +| Model type | Runner script | Example config | +|------------|---------------|----------------| +| **Diffusion X2I/X2V** | `tests/dfx/perf/scripts/run_diffusion_benchmark.py` | `tests/dfx/perf/tests/test_qwen_image_vllm_omni.json` | +| **TTS** | `tests/dfx/perf/scripts/run_benchmark.py` | `tests/dfx/perf/tests/test_tts.json` (shared) or `test_{slug}.json` (dedicated; use when the model must not join the shared nightly matrix before integration — e.g. `test_voxcpm2.json`, `test_coqui_tts.json`) | +| **Omni** | `tests/dfx/perf/scripts/run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` | + +5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with **`--test-config-file`** (nightly perf steps do not use `-m`), **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf. +6. **One benchmark scenario → one `test_name` block** in JSON (same spirit as one `test_*` per function case). Combine server `serve_args` + workload in one entry; use multiple `benchmark_params` rows for size/step sweeps under the same server config. + +**Perf local commands:** + +```bash +cd tests +export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN +# Single file (CI-like) +pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \ + --test-config-file dfx/perf/tests/test__vllm_omni.json +# Bulk load + filter by JSON mark +pytest -sv dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100" +pytest -sv dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100" +``` + +**Clarify in the test plan** which L4 pillars you deliver: `Function only` | `Function + Perf` | `Function + Accuracy + Perf`. + +See [references/test-routing.md](references/test-routing.md) **L4 nightly pillars** for Buildkite step patterns. + +### Step 5: Wire Buildkite (when CI must run the new test) + +If the test is not already collected by an existing pipeline command (for example, L1 tests marked `core_model and cpu` are already covered by the **Simple Unit Test** step in ready/merge), update the appropriate pipeline under `.buildkite/`: + +| Test level | Edit this file | Typical trigger / intent | +|------------|----------------|---------------------------| +| **L1** and **L2** | [`.buildkite/cuda/test-ready.yml`](../../../.buildkite/cuda/test-ready.yml) | PR **ready** label; L1 CPU + L2 GPU/basic e2e (steps labeled **Omni ·**, **TTS ·**, **Diffusion ·**) | +| **L3** | [`.buildkite/cuda/test-merge.yml`](../../../.buildkite/cuda/test-merge.yml) | Post-merge; `advanced_model` integration per model type | +| **L4** | [`.buildkite/cuda/test-nightly.yml`](../../../.buildkite/cuda/test-nightly.yml) | Nightly; grouped by model type (see below) | +| **Invalid param / reliability (weekly)** | [`.buildkite/cuda/test-weekly.yml`](../../../.buildkite/cuda/test-weekly.yml) | Weekly; `tests/dfx/reliability/invalid_param_test/` — **not** L1–L4 e2e pipelines | + +**Level-specific Buildkite delivery (follow the user’s requested level):** + +| User asks for | Document / edit **only** | Do **not** add by default | +|---------------|---------------------------|---------------------------| +| **L4** (`*_expansion.py`, `full_model`) | `test-nightly.yml` — append file to the matching nightly shard (X2I / X2V / Omni / TTS) **or** note it is already collected by an existing `-m` / `-k` sweep | `test-merge.yml` (L3 / `advanced_model`) or `test-ready.yml` (L2) | +| **L3** | `test-merge.yml` with `advanced_model` + `source_file_dependencies` | `test-nightly.yml` unless user also wants nightly | +| **L2** | `test-ready.yml` with `core_model` + `source_file_dependencies` | merge / nightly pipelines | +| **Invalid param** | `test-weekly.yml` — **Invalid parameters Test** group (`-m "slow and H100"` / `-m "slow and L4"`). Usually **no YAML edit** when appending cases to existing `invalid_param_test` scripts | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml` | + +`test-nightly.yml` shards use **explicit pytest file paths** in `commands` (and PR labels like `diffusion-x2iat-test`); they generally **do not** use `source_file_dependencies`. Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI). + +**`test-nightly.yml` top-level groups** (each uses `full_model` + type marker in pytest `-m`): + +| Group | Type marker | Notes | +|-------|-------------|--------| +| **Omni Model Test** | `omni` | Function / doc / accuracy / **perf** for Qwen-Omni family | +| **TTS Model Test** | `tts` | `-m "full_model and L4 and tts"` (and H100 variants when enabled); **Perf Test** uses `run_benchmark.py` + `test_tts.json` | +| **Diffusion X2I(&A&T) Model Test** | `diffusion` | Image/audio/text diffusion; explicit file shards for **Function**; **Perf Test · <Model>** uses `run_diffusion_benchmark.py` + JSON config | +| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); PR label `diffusion-x2v-test`; separate **Perf** steps when in-tree | + +Guidelines when editing YAML: + +- **Match markers and `--run-level`** to the test: L1/L2 in `test-ready.yml` use `core_model` + `omni` / `tts` / `diffusion` as appropriate; L3 merge uses `advanced_model`; L4 nightly uses `full_model` with the same type marker. +- **Diffusion L4**: add x2i/x2a/x2t expansion files to an **X2I(&A&T)** step; add x2v expansion files to **X2V** — never rely on a broad `tests/e2e/` sweep alone (shards exist to save GPU and avoid wrong queue). +- **`source_file_dependencies` (required for new E2E jobs in `test-ready.yml` and `test-merge.yml`)**: every step under the **`:card_index_dividers: E2E Test`** group **must** declare `source_file_dependencies` so Buildkite only schedules the job when relevant paths change. Omitting it breaks path-based CI skipping for GPU e2e jobs. +- **Reuse or extend an existing step** when the new test shares the same marker expression, queue, and timeout; otherwise add a new `steps:` entry with the correct `agents.queue`, `timeout_in_minutes`, and docker/kubernetes plugin blocks consistent with neighboring jobs. +- **Job timeout — use `timeout_in_minutes` on the step, not `timeout` before pytest**: when adding or documenting a new E2E job in `test-ready.yml` / `test-merge.yml`, set `timeout_in_minutes` on the step and run `pytest` directly in `commands`. **Do not** wrap pytest in `timeout 40m bash -c "..."` / `timeout 20m bash -c "..."`; Buildkite already enforces the step deadline via `timeout_in_minutes`. +- **Never omit `plugins` in skill examples or PR YAML snippets**: H100 E2E steps on `mithril-h100-pool` use the full `kubernetes` block (`resources`, `volumeMounts`, `env` with `HF_TOKEN` secret, `nodeSelector: gpu-h100-sxm`, `devshm` + `hf-cache` volumes). L4/docker steps on `gpu_1_queue` use the full `docker#v5.2.0` block (`shm-size`, `HF_HOME`, volumes). Do not write `# ... plugins unchanged` — paste the complete block from a sibling step. +- **Platform forks** (e.g. AMD-ready / AMD-merge) live alongside these files; apply the same level → file mapping for those pipelines when the test is platform-specific. + +#### `source_file_dependencies` for E2E Test jobs (`test-ready.yml`, `test-merge.yml`) + +When adding or splitting a step inside the **E2E Test** group, **always** include `source_file_dependencies` on that step. List paths the job actually depends on — at minimum: + +1. **Test module(s)** — every `tests/e2e/...` file the `pytest` command runs (online and offline if both are in one step). +2. **Model implementation** — `vllm_omni/model_executor/models//` and/or `vllm_omni/diffusion/models//` as applicable. +3. **Stage / input plumbing** (omni & multi-stage TTS) — `vllm_omni/model_executor/stage_input_processors/.py` when present. +4. **Deploy / stage YAML** — `vllm_omni/deploy/.yaml` or `vllm_omni/deploy/ci/.yaml` referenced by the test’s `OmniServerParams` / `get_deploy_config_path(...)`. +5. **Test helpers** (when changed in the same PR) — `tests/helpers/runtime.py` for new `send_*` client methods; `tests/helpers/assertions.py` for new shared `assert_*` helpers. + +Use **repo-relative paths** (no leading `./`). Prefer **directory** entries (`.../models/bagel/`) when the whole package matters; use **file** entries for single test modules and YAML configs. Mirror a neighboring step for the same model family (e.g. **Diffusion · Bagel Test** in `test-ready.yml`). + +**Example** (L2 ready — diffusion online smoke, H100 kubernetes): + +```yaml +- label: "Diffusion · Bagel Test" + source_file_dependencies: + - tests/e2e/online_serving/test_bagel.py + - vllm_omni/model_executor/models/bagel/ + - vllm_omni/diffusion/models/bagel/ + - vllm_omni/model_executor/stage_input_processors/bagel.py + - vllm_omni/deploy/bagel.yaml + timeout_in_minutes: 40 + commands: + - | + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate +``` + +**Example** (L3 merge — TTS online + offline, L4 docker): + +```yaml +- label: "TTS · Qwen3-TTS Base Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen3_tts_base.py + - tests/e2e/offline_inference/test_qwen3_tts_base.py + - vllm_omni/model_executor/models/qwen3_tts/ + - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py + - vllm_omni/deploy/qwen3_tts.yaml + timeout_in_minutes: 20 + commands: + - | + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 + pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model' + agents: + queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" +``` + +**Example** (L3 merge — diffusion online + offline, H100 kubernetes): + +```yaml +- label: "Diffusion · Qwen Image Test" + source_file_dependencies: + - tests/e2e/online_serving/test_qwen_image.py + - tests/e2e/offline_inference/test_qwen_image.py + - vllm_omni/diffusion/models/qwen_image/ + timeout_in_minutes: 40 + commands: + - | + export VLLM_IMAGE_FETCH_TIMEOUT=60 + pytest -s -v \ + tests/e2e/online_serving/test_qwen_image.py \ + tests/e2e/offline_inference/test_qwen_image.py \ + -m 'advanced_model and cuda' --run-level 'advanced_model' + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate +``` + +When extending an **existing** E2E step to run an additional test file, **append** that file (and any new model/deploy paths) to `source_file_dependencies` on the same step. When authoring a PR, mention the dependency list in the **Buildkite change** section of your test plan. + +#### L3 / L4: validating Buildkite from a feature branch (`cuda/pipeline.yml`) + +Production behavior is driven by [`.buildkite/cuda/pipeline.yml`](../../../.buildkite/cuda/pipeline.yml): it builds the CI image, then uploads **L2** (`test-ready.yml` on non-`main`), **L3** (`test-merge.yml` when `build.branch == "main"`), and **L4** (`test-nightly.yml` when `build.env("NIGHTLY") == "1"`). + +When you add or debug **L3** or **L4** tests and need those child pipelines to run **off `main`** or **without** `NIGHTLY=1`, apply **temporary** edits (revert before merge): + +1. **Point the upload at the right file (optional shortcut)** + In the **Upload Ready Pipeline** step, the command is `buildkite-agent pipeline upload .buildkite/cuda/test-ready.yml`. For a one-off run you may change that path to `.buildkite/cuda/test-merge.yml` (L3) or `.buildkite/cuda/test-nightly.yml` (L4), and adjust the step’s `if:` so it does not conflict with the other upload steps (avoid double-uploading unless intentional). + +2. **L3 — merge pipeline** + The gate `if: build.branch == "main"` lives on the **Upload Merge Pipeline** step in `cuda/pipeline.yml` (not inside `test-merge.yml`). **Comment out that `if` line** so `test-merge.yml` is uploaded on your feature branch. Alternatively rely on step (1) instead of the dedicated merge upload step. + +3. **L4 — nightly pipeline** + **Comment out** `if: build.env("NIGHTLY") == "1"` on the **Upload Nightly Pipeline** step in `cuda/pipeline.yml` so the nightly definition is uploaded without setting the env var. + Child steps in [`test-nightly.yml`](../../../.buildkite/cuda/test-nightly.yml) often repeat `if: build.env("NIGHTLY") == "1"`; **comment out those lines on the steps you need to run**, otherwise they will still be skipped after upload. + +### Step 6: Run Tests + +Pick one command path: + +- **Quick local regression (preferred first)**: single file or `file.py::test_name` +- **CI-like level run**: marker expression + `--run-level` + +Full templates live in [references/test-routing.md](references/test-routing.md). **After authoring tests, always emit concrete commands** (see **Output Format**). + +**Typical copy-paste examples** (run from repo `tests/` directory; requires matching vLLM/vllm-omni install and hardware): + +| Area | Example | +|------|---------| +| **Omni offline L2** | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` | +| **Omni online L2** | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` | +| **TTS online L2** | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` | +| **Diffusion X2I L2 online** | `pytest -s -v e2e/online_serving/test_qwen_image.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2V L2 online** | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2I L4 nightly** | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` | +| **Diffusion X2I L4 perf (nightly)** | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` | +| **Diffusion X2V L4 nightly** | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` | +| **Invalid param (weekly H100)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` | +| **Invalid param (weekly L4)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` | +| **L1 CPU** | `pytest -s -v -m "core_model and cpu"` | + +**Prerequisites to mention when relevant**: GPU model (e.g. L4 vs H100), `HF_HOME` / token for hub weights, module-level `skipif` (NPU/XPU-only gaps), and whether CI already collects the path (e.g. `test_*_expansion.py` glob in `test-nightly.yml`). + +### Step 7: Validate Result Quality + +Before finishing: + +- Is the new assertion directly tied to the bug/feature contract? +- Are API calls made via **`runtime.py` `send_*_request`** (not inline HTTP/SDK in the test file)? +- Is the **general** `assert_*` bundled inside `send_*_request` (not duplicated in the test)? +- Is each scenario a **separate `test_*` function** (no `case_id` / `if param.id` branching across cases)? +- **Fixture scope**: default **`omni_server` / `omni_runner`** (module); **`_function`** variants only when each test must restart the instance? +- Is the test deterministic (no fragile timing/network coupling)? +- Is runtime appropriate for the selected level? +- Are markers and `--run-level` consistent? +- **Invalid-param cases** in `tests/dfx/reliability/invalid_param_test/` (not e2e), using `send_*_http_request` + `err_code`, with `pytest.mark.slow` and correct `H100` / `L4` hardware mark? + +## Output Format + +When completing a request, return: + +1. **Test plan** (level, markers, file target, **L4 pillar(s)** if applicable: Function / Accuracy / Perf / Doc; **or** **Invalid param** pillar with target `invalid_param_test/test_invalid_.py`; module basename for e2e: `test_{slug}.py` / `test_{slug}_expansion.py`) +2. **Generated/updated test file(s)** — model-specific constants and request payloads stay **inside** those test modules (**never** `tests/helpers/{slug}.py`). Extend **`tests/helpers/runtime.py`** / **`tests/helpers/assertions.py`** only when new **repo-wide** `send_*` or shared assert helpers are required (do not leave ad-hoc HTTP or `_assert_*` in test modules). **If L4 Perf was requested**, also list **`tests/dfx/perf/tests/*.json`**. **If invalid-param cases were requested**, list the **`invalid_param_test/` script** and parametrized `test_*` / new `id=` rows — **not** e2e paths. +3. **Buildkite change** — **match the requested level and pillar**: + - **L4 Function**: `test-nightly.yml` **Function Test** shard (file list / note existing sweep). **No** `test-merge.yml` unless user also asked for L3. + - **L4 Perf**: `test-nightly.yml` **Perf Test · <Model>** step (new step or extend commands); include env exports + artifact upload pattern from a sibling perf job. + - **L4 Accuracy**: `test-nightly.yml` **Accuracy Test** step under the same model-type group. + - **L3**: `test-merge.yml` + full `source_file_dependencies` + `agents` + `plugins`. + - **L2**: `test-ready.yml` + same E2E block requirements. + - **Invalid param**: `test-weekly.yml` — note **Invalid parameters Test · H100/L4** group; usually **no YAML change** when only extending existing scripts. + If the user asked only for **L4 functional cases**, state explicitly that **Perf / Accuracy / Invalid param** were not included (offer to add if needed). If the plan mixes **success e2e** and **invalid param**, split deliverables across e2e vs `invalid_param_test/` and call out both CI files. +4. **Run commands (required)** — always include, in fenced `bash` blocks: + - **Local — whole file**: `cd tests` then `pytest -s -v …` + - **Local — single test** (optional but preferred when the change is one function): `pytest -s -v path::test_func …` + - **CI-like** (when not L1 CPU): the same **marker + `--run-level`** pairing the level uses (see [references/test-routing.md](references/test-routing.md) and **Step 6** table above) + - **Prerequisites** (one line): e.g. “Requires CUDA L4 + weights cached locally”; “Requires H100 + `HF_TOKEN`”; “Syntax-only validation when vllm is not installed locally” +5. **Result summary** (pass/fail if you executed tests; if not executed, state that explicitly and what the user should run) + +## Additional Resources + +- Marker and command routing: [references/test-routing.md](references/test-routing.md) +- CI pipelines (vllm-omni): [test-ready.yml](../../../.buildkite/cuda/test-ready.yml) (L1/L2), [test-merge.yml](../../../.buildkite/cuda/test-merge.yml) (L3), [test-nightly.yml](../../../.buildkite/cuda/test-nightly.yml) (L4), [test-weekly.yml](../../../.buildkite/cuda/test-weekly.yml) (invalid param / reliability) +- L4 documentation example tests (naming, extraction vs copied scripts, output dirs, skips): [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) — see also [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910) diff --git a/.claude/skills/vllm-omni-test/references/test-routing.md b/.claude/skills/vllm-omni-test/references/test-routing.md new file mode 100644 index 00000000000..83b9b200718 --- /dev/null +++ b/.claude/skills/vllm-omni-test/references/test-routing.md @@ -0,0 +1,395 @@ +# Test Routing Reference + +Use this reference to map testing goals to levels, markers, and runnable commands. + +**Repo paths** (`.buildkite/`, `docs/contributing/…`): link with repo-relative paths from this file (e.g. `../../../../.buildkite/cuda/test-ready.yml`, `../../../../docs/contributing/ci/CI_5levels.md`). + +## Model-centric e2e filename convention (L2–L4) + +When generating a **new** test module tied to a **specific model** under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`: + +| Level | Pattern | Example (`Qwen/Qwen2.5-Omni-7B`) | +|-------|---------|----------------------------------| +| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` | +| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` | + +**Slug rules for `{lowercase_model_slug}`** (canonical: [SKILL.md](../SKILL.md) § *Naming: generated test module files*): + +1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`. +2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree). +3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) **only then**. +4. **L1** unit tests are **not** bound to this pattern; use `tests//test_.py` as today. + +## Model type markers (`omni` / `tts` / `diffusion`) + +Every model-centric e2e test must declare **exactly one** type marker: + +| Marker | Model family | Typical APIs | +|--------|--------------|--------------| +| `pytest.mark.omni` | Multimodal LLM (Qwen-Omni, …) | `send_omni_request`, `omni_runner`, stage YAML | +| `pytest.mark.tts` | TTS / speech synthesis | `send_audio_speech_request`, `/v1/audio/speech` | +| `pytest.mark.diffusion` | Diffusion generative models | `send_diffusion_request`, `send_video_diffusion_request`, `send_images_generations_http_request`, `OmniDiffusionSamplingParams` | + +**Diffusion** uses one marker for all output modalities. Nightly CI splits diffusion by **group**, not by extra markers: + +| Nightly group (`test-nightly.yml`) | Diffusion scope | Examples | +|-----------------------------------|-----------------|----------| +| **Diffusion X2I(&A&T) Model Test** | x2**i** / x2**a** / x2**t** — image, audio, text (non-video) | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero | +| **Diffusion X2V Model Test** | x2**v** — video only | Wan2.2, HunyuanVideo 1.5, Wan VACE | + +PR labels for selective nightly runs: `diffusion-x2iat-test`, `diffusion-x2v-test` (plus `nightly-test` / `NIGHTLY=1`). + +## Level and Marker Mapping + +| Goal | Suggested Level | Marker baseline | Typical location | +|------|------------------|-----------------|------------------| +| Unit logic, regression on pure code path | L1 | `core_model and cpu` | `tests//test_*.py` | + +### L1 unit tests — mocking (`mocker`, not `unittest.mock`) + +L1 modules run in **Simple Test** (`test-ready.yml` / `test-merge.yml`: `-m 'core_model and cpu'`). When isolation requires doubles: + +- **Use** the **`mocker`** fixture (`pytest-mock`): `mocker.patch`, `mocker.spy`, `mocker.Mock`, `mocker.MagicMock`, `mocker.AsyncMock`. +- **Use** **`monkeypatch`** for env vars and simple `setattr` / `delattr` without mock objects. +- **Do not use** `unittest.mock` — no `from unittest.mock import ...`, no `@patch`, no `with patch(...)`. + +```python +# Good (L1) +def test_handler(mocker): + mocker.patch("vllm_omni.pkg.fn", return_value=0) + +# Bad (L1) +from unittest.mock import patch + +@patch("vllm_omni.pkg.fn") +def test_handler(mock_fn): ... +``` + +E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer real (or lightweight) execution paths. +| Basic integration/e2e | L2 | `core_model` + **one of** `omni` / `tts` / `diffusion` | `tests/e2e/...` | +| Advanced integration | L3 | `advanced_model` + type marker | `tests/e2e/...` | +| Full function / nightly | L4 | `full_model` (nightly) + type marker | **Function:** `tests/e2e/*_expansion.py`; **Perf:** `tests/dfx/perf/tests/*.json`; **Accuracy:** `tests/e2e/accuracy/` | +| Invalid HTTP / param validation | Weekly (dfx) | `pytest.mark.slow` + type marker + `H100` or `L4` | `tests/dfx/reliability/invalid_param_test/test_invalid_*.py` | + +## Marker Selection Rules + +1. **Level** (pick one): + - `core_model` — L1/L2 (`test-ready.yml`) + - `advanced_model` — L3 (`test-merge.yml`) + - `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration +2. **Model type** (pick one): `omni`, `tts`, or `diffusion` +3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark` +4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, … +5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py` + +## Command Templates + +### Quick local checks + +```bash +cd tests +pytest -s -v test_xxxx.py +``` + +### L1 + +```bash +cd tests +pytest -s -v -m "core_model and cpu" +``` + +### L2 + +```bash +cd tests +pytest -s -v -m "core_model and not cpu" --run-level=core_model +``` + +### L3 (merge) + +```bash +cd tests +pytest -s -v -m "advanced_model" --run-level=advanced_model +``` + +### L4 (nightly) + +**Function** (e2e expansion — default when user asks for “L4 functional cases”): + +```bash +cd tests +pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model +``` + +**Perf** (dfx benchmark harness — separate pillar; not `omni_server` fixtures): + +```bash +cd tests +export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN +# CI-like: single JSON file (nightly Perf Test steps) +pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \ + --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json +# Local bulk load: all *.json under tests/dfx/perf/tests/, filter by JSON mark +pytest -sv dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100" +pytest -sv dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100" +``` + +Broad marker sweep (when no explicit file shard): + +```bash +cd tests +pytest -s -v -m "full_model" --run-level=full_model +``` + +## L4 nightly pillars (Function / Accuracy / Perf) + +Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not a single pytest module: + +| Pillar | User says | Deliver | Nightly step | +|--------|-----------|---------|--------------| +| **Function** | “L4 functional cases” / “functional” / `*_expansion.py` | `tests/e2e/.../test__expansion.py` | `· Function Test with H100/L4` | +| **Perf** | “performance” / “perf” / “benchmark” / “full L4” | `tests/dfx/perf/tests/test__vllm_omni.json` + runner | `· Perf Test · ` | +| **Accuracy** | “accuracy” / “similarity” | `tests/e2e/accuracy/test_*.py` | `· Accuracy Test` | + +**Rule:** “L4 functional cases” → **Function only** by default. State that Perf/Accuracy are separate pillars; add them only when requested. + +**Diffusion perf JSON** (`tests/dfx/perf/tests/test__vllm_omni.json`): array of **case objects**, each with: + +| Field | Required | Notes | +|-------|----------|-------| +| `test_name` | Yes | Unique server/workload id | +| `mark` | Recommended | When present: **`hardware_marks` required** + optional `marks` (`full_model`, `diffusion`, `local_model`, …). Parsed by `resolve_pytest_marks`; enables local `-m` filtering. | +| `server_type` | Diffusion | e.g. `"vllm-omni"` — routes case to `run_diffusion_benchmark.py` | +| `benchmark_endpoint` | Diffusion | e.g. `/v1/videos`, `/v1/images/generations` | +| `server_params` | Yes | `model`, `serve_args`, … | +| `benchmark_params[]` | Yes | Each row: **`name`** (pytest id suffix), workload fields, **`baseline`** | + +Copy structure from `test_qwen_image_vllm_omni.json` or `test_cosmos3_vllm_omni.json` (2-GPU mark example). + +**Omni perf JSON**: `test_qwen3_omni_*.json` — same `mark` shape with `"marks": ["full_model", "omni"]`; run via `run_benchmark.py`. + +**Runner split**: `is_diffusion_perf_config()` → diffusion when `server_type` is set or `"diffusion"` ∈ `mark.marks`; otherwise omni/tts → `run_benchmark.py`. + +**Param ids**: `{test_name}-{benchmark_params.name}` (e.g. `test_tts-p0`, not `test_tts-0`). + +**Result filenames**: runtime hardware from `get_runtime_resource_label()`; `H100` omitted on default CI pool. + +**Diffusion perf nightly step** (under **Diffusion X2I(&A&T)**; copy full `kubernetes` plugins from `· Perf Test · Qwen-Image`): + +```yaml + - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · " + key: nightly-diffusion-x2iat-performance- + timeout_in_minutes: 180 + commands: + - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results + - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN + - export CACHE_DIT_VERSION=1.3.0 + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test__vllm_omni.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + exit $$EXIT + agents: + queue: "mithril-h100-pool" + plugins: + # … full kubernetes block (often multi-GPU) — copy from Perf Test · Qwen-Image +``` + +| Model type | Perf runner | Config example | +|------------|-------------|----------------| +| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json` | +| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` | +| Omni | `dfx/perf/scripts/run_benchmark.py` | `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_no_async_chunk.json`, … | + +Do **not** put throughput/latency baselines inside `test_*_expansion.py` — that belongs in the dfx perf JSON + nightly Perf job. + +## Invalid parameter validation (weekly / dfx) + +When the user’s test plan includes **invalid parameter validation**, **invalid request bodies**, or **HTTP 4xx** validation against a live server: + +1. **Do not** author these in `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move** any drafted `test_*` into `tests/dfx/reliability/invalid_param_test/`. +2. **Pick script by route** (extend in-tree file): + +| Route family | Script | +|--------------|--------| +| Omni chat / WS video·realtime | `test_invalid_omni_chat.py` | +| `/v1/audio/speech` (+ stream / batch / voices) | `test_invalid_audio_speech.py` | +| Audio diffusion | `test_invalid_audio_diffusion.py` | +| `/v1/images/generations` | `test_invalid_image_generation.py` | +| `/v1/images/edits` | `test_invalid_image_editing.py` | +| `/v1/videos*` | `test_invalid_video_generation.py` | +| Sleep / wakeup / server control | `test_invalid_server_control.py` | + +3. **Style:** `pytestmark = [pytest.mark.slow, pytest.mark.]`; `_PARAMS` + `hardware_marks`; `send_*_http_request` with `err_code` + `err_message`; parametrized `body_spec` rows with `id=`; `_minimal_*_json()` helpers; `_SKIP_ISSUE_3649` when tracked in [#3649](https://github.com/vllm-project/vllm-omni/issues/3649). +4. **CI:** [`.buildkite/cuda/test-weekly.yml`](../../../../.buildkite/cuda/test-weekly.yml) group **Reliability Test - Invalid parameters Test** — **not** ready/merge/nightly. + +```bash +# Weekly H100 (diffusion / omni / video invalid-param) +cd tests +pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100" + +# Weekly L4 (e.g. Qwen3-TTS speech invalid-param) +pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4" + +# Single script / case +pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py::test_images_generations_invalid_requests -m "slow and H100" +``` + +**Trigger:** `WEEKLY=1` or PR label `weekly-test`. Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). No `source_file_dependencies` on weekly steps. + +### Platform-targeted examples + +```bash +cd tests +pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model +``` + +### Concrete e2e paths (common in-tree) + +Paths are relative to `tests/` after `cd tests`. + +| Scenario | Example command | +|----------|------------------| +| **Omni** offline L2 | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` | +| **Omni** offline L2 — one test | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py::test_text_to_text -m "core_model and omni and not cpu" --run-level=core_model` | +| **Omni** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` | +| **Omni** offline L2 (Qwen3.5-9B VL) | `pytest -s -v e2e/offline_inference/test_qwen3_5_9b.py -m "core_model and omni and not cpu" --run-level=core_model` | +| **TTS** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` | +| **TTS** online L2 — one test | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py::test_text_to_audio_001 -m "core_model and tts" --run-level=core_model` | +| **Diffusion X2I** offline L2 | `pytest -s -v e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion and not cpu" --run-level=core_model` | +| **Diffusion X2I** L2 online smoke | `pytest -s -v e2e/online_serving/test_qwen_image.py e2e/online_serving/test_bagel.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2V** L2 online smoke | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` | +| **Diffusion X2I** L4 online expansion (H100) | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` | +| **Diffusion X2I** L4 perf (nightly) | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` | +| **Diffusion X2I** L4 online expansion (Edit) | `pytest -s -v e2e/online_serving/test_qwen_image_edit_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` | +| **Diffusion X2V** L4 nightly | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` | +| **Invalid param** weekly H100 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` | +| **Invalid param** weekly L4 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` | + +### Agent / author completion checklist + +When adding or modifying tests, do not stop at “where the file lives” — also deliver: + +1. **Local**: `cd tests` + `pytest -s -v ` (and `path::test_func` when a single case is enough). +2. **CI-like**: marker string + `--run-level` matching the pipeline (`core_model` / `advanced_model` / `full_model`). +3. **L1 mocks**: `mocker` / `monkeypatch` only; no `unittest.mock`. +4. **API client + assert placement**: **General** validation → implement in `assertions.py`, call **inside** `send_*_request` in `runtime.py`; tests call **`send_*_request` only**. **Special** case validation → `assert_*` in `assertions.py`, called **in the test** after `send_*_request`. Low-level `send_*_http_request` is for negative/dfx tests (`err_code`), not ordinary L2+ success e2e. +5. **Shared assertions**: logic in **`tests/helpers/assertions.py`**; general checks inside `send_*_request`; special checks only in the test. No `_assert_*` in `test_*.py`. +6. **One case → one `test_*`**: function name reflects what is validated; no `if case_id == ...` mega-test merging multiple scenarios. +7. **Fixture scope**: default **`omni_server` + `openai_client`** / **`omni_runner` + `omni_runner_handler`** (module). Use **`omni_server_function` + `openai_client_function`** / **`omni_runner_function` + `omni_runner_handler_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`). +8. **Type marker**: `omni`, `tts`, or `diffusion` on every model e2e module. +9. **Diffusion L4 Function**: wire `*_expansion.py` into **X2I(&A&T)** or **X2V** **Function Test** in **`test-nightly.yml` only** — do not add `test-merge.yml` unless the user also requested L3. +10. **Diffusion L4 Perf** (only when requested): add `tests/dfx/perf/tests/test__vllm_omni.json` + **Perf Test · <Model>** step (artifact upload); not part of “L4 functional cases” by default. +11. **E2E Buildkite (L2/L3 only)**: `test-ready.yml` / `test-merge.yml` steps need **`source_file_dependencies`** + full **`agents` + `plugins`**. **L4 nightly** uses explicit file lists or perf scripts in `test-nightly.yml` (no merge job). +12. **Invalid param**: cases in **`tests/dfx/reliability/invalid_param_test/`** (route-matching script), `send_*_http_request` + `err_code`, `pytest.mark.slow` + `H100`/`L4`; CI = **`test-weekly.yml`** only — do not put in e2e or nightly. +13. **Prerequisites**: GPU tier, HF cache/token, and any module `skipif` / platform-only YAML. + +## Buildkite pipeline mapping + +| Level | Repo file | Model-type grouping | +|-------|-----------|---------------------| +| L1, L2 | [`.buildkite/cuda/test-ready.yml`](../../../../.buildkite/cuda/test-ready.yml) | Steps prefixed **Omni ·**, **TTS ·**, **Diffusion ·** under **E2E Test** — **`source_file_dependencies` required** | +| L3 | [`.buildkite/cuda/test-merge.yml`](../../../../.buildkite/cuda/test-merge.yml) | Per-model E2E steps; `-m "advanced_model and …"` — **`source_file_dependencies` required** | +| L4 | [`.buildkite/cuda/test-nightly.yml`](../../../../.buildkite/cuda/test-nightly.yml) | **Omni / TTS / Diffusion X2I(&A&T) / X2V** — each group may have **Function**, **Accuracy**, **Perf**, **Doc** steps | +| Invalid param (weekly) | [`.buildkite/cuda/test-weekly.yml`](../../../../.buildkite/cuda/test-weekly.yml) | **Invalid parameters Test · H100** / **· L4** — sweeps `tests/dfx/reliability/invalid_param_test/` | + +### `source_file_dependencies` (E2E Test only — ready & merge) + +Required on each step inside the **E2E Test** group in `test-ready.yml` and `test-merge.yml`. Typical entries: + +| Category | Path pattern | +|----------|----------------| +| Test module | `tests/e2e/online_serving/test_.py`, `tests/e2e/offline_inference/test_.py` | +| API client helpers | `tests/helpers/runtime.py` (when the e2e job uses newly added `send_*` methods) | +| Shared assert helpers | `tests/helpers/assertions.py` (when the e2e job uses newly added `assert_*` helpers) | +| AR / omni model | `vllm_omni/model_executor/models//` | +| Diffusion model | `vllm_omni/diffusion/models//` | +| Stage processor | `vllm_omni/model_executor/stage_input_processors/.py` | +| Deploy config | `vllm_omni/deploy/.yaml` or `vllm_omni/deploy/ci/.yaml` | + +Copy the dependency block from the closest in-tree E2E step for the same model family. Update the list whenever the pytest command or server YAML changes. + +### E2E `agents` + `plugins` blocks (required — do not truncate) + +When generating or documenting E2E steps, **always include the complete `plugins` section**. Set **`timeout_in_minutes`** on the step for the job deadline; run **`pytest` directly** in `commands` (no `timeout 40m bash -c` wrapper around pytest). + +Common patterns in `test-ready.yml` / `test-merge.yml`: + +| Queue | Plugin | Required fields | +|-------|--------|-----------------| +| `mithril-h100-pool` | `kubernetes` | `resources.limits.nvidia.com/gpu`, `volumeMounts` (`devshm`, `hf-cache`), `env` (`HF_HOME`, `HF_TOKEN` via `secretKeyRef`), `nodeSelector: gpu-h100-sxm`, `volumes` (`devshm` emptyDir, `hf-cache` hostPath `/mnt/hf-cache`) | +| `gpu_1_queue` / `gpu_4_queue` | `docker#v5.2.0` | `image`, `always-pull`, `propagate-environment`, `shm-size: "8gb"`, `environment: HF_HOME`, `volumes: /fsx/hf_cache` (add `HF_TOKEN` when the step needs hub access) | + +**H100 kubernetes template** (copy from **Diffusion · Bagel Test** / **Diffusion · Qwen Image Test** in-tree): + +```yaml + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate +``` + +**L4 docker template** (copy from **TTS · Qwen3-TTS Base Test** in `test-merge.yml`): + +```yaml + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" +``` + +When extending an existing step (e.g. add `tests/e2e/offline_inference/test_qwen_image.py` to **Diffusion · Qwen Image Test**), update `source_file_dependencies` and `commands` only; **keep the existing `agents` + `plugins` block unchanged** unless the hardware tier changes. + +[`.buildkite/cuda/pipeline.yml`](../../../../.buildkite/cuda/pipeline.yml) decides **which** child file is uploaded. To run **L3** on a feature branch, comment out `if: build.branch == "main"` on the merge upload step. To run **L4** without `NIGHTLY=1`, comment out the nightly upload step’s `if: build.env("NIGHTLY") == "1"` and the same `if` on relevant steps inside `test-nightly.yml`. Revert such edits before merging. + +Platform-specific pipelines (e.g. AMD) follow the same level → file pairing under `.buildkite/`. + +## Diffusion RFC (#1832) Alignment Tips + +For **X2I(&A&T)** coverage planning: + +- Prioritize high-value feature combinations with minimal case count. +- Split into lightweight L1/L2 validation plus L4 `*_expansion.py` under the **X2I** nightly shards. +- **X2V** models (Wan, HunyuanVideo) stay in the **X2V** group — do not merge into X2I sweeps. + +If hardware is insufficient, provide an executable reduced case plus a deferred full CI/nightly plan. diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 0541f1b8612..6b3ddb4743b 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -22,30 +22,30 @@ /.buildkite/ @yenuo26 @congw729 # Engine and model_executor -/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan -/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan -/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan +/vllm_omni/core/ @tzhouam @yinpeiqi @fake0fan @Sy0307 +/vllm_omni/engine/ @tzhouam @yinpeiqi @fake0fan @Sy0307 +/vllm_omni/entrypoints/ @tzhouam @yinpeiqi @fake0fan @alex-jw-brooks /vllm_omni/model_executor/ @tzhouam @gcanlin /vllm_omni/worker/ @tzhouam @gcanlin # Diffusion -/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 +/vllm_omni/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @david6666666 @xuechendi /vllm_omni/lora/ @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao # Diffusion CPU offloading -/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao +/vllm_omni/diffusion/offloader/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @yuanheng-zhao @xuechendi # Omni and TTS support /vllm_omni/transformers_utils/ @tzhouam @yuanheng-zhao -/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter +/vllm_omni/model_executor/models/ @linyueqian @ZeldaHuang @princepride @yuanheng-zhao @amy-why-3459 @Bounty-hunter @Sy0307 @gcanlin # Benchmarks /benchmarks/accuracy/ @david6666666 /benchmarks/build_dataset/ @linyueqian @yuanheng-zhao -/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa +/benchmarks/diffusion/ @Isotr0py @princepride @SamitHuang @wtomin @ZJY0516 @RuixiangMa @xuechendi /benchmarks/distributed/ @princepride /benchmarks/glm_image/ @princepride @JaredforReal -/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride +/benchmarks/tts/ @linyueqian @ZeldaHuang @princepride @Sy0307 @gcanlin # Usability and Observability /vllm_omni/deploy/ @tzhouam @lishunyang12 @@ -53,13 +53,13 @@ /vllm_omni/profiler/ @gcanlin # ComfyUI -/apps/ @fhfuih +/apps/ @fhfuih @alex-jw-brooks # Large Scale Serving /vllm_omni/distributed/ @princepride @yuanheng-zhao @xuechendi @natureofnature @fake0fan # Quantization and configuration -/vllm_omni/config/ @lishunyang12 +/vllm_omni/config/ @lishunyang12 @alex-jw-brooks /vllm_omni/quantization/ @david6666666 @Isotr0py @lishunyang12 # Hardware plugins and accelerators diff --git a/.github/ISSUE_TEMPLATE/400-bug-report.yml b/.github/ISSUE_TEMPLATE/400-bug-report.yml index c3799b19400..7cbf0ac0e78 100644 --- a/.github/ISSUE_TEMPLATE/400-bug-report.yml +++ b/.github/ISSUE_TEMPLATE/400-bug-report.yml @@ -19,6 +19,7 @@ body: Consider redacting or replacing sensitive values with placeholders like `` when sharing configuration or code examples. - type: textarea + id: current-environment attributes: label: Your current environment description: | @@ -42,6 +43,7 @@ body: required: true - type: textarea + id: code-version attributes: label: Your code version description: | @@ -66,6 +68,7 @@ body: required: true - type: textarea + id: bug-description attributes: label: 🐛 Describe the bug description: | diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 964fb5abcc9..2bd0cc13327 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -24,3 +24,9 @@ jobs: python-version: "3.12" cache: 'pip' - uses: pre-commit/action@v3.0.1 + # Don't run the pytest mark check against everything in the CI + # for now since a lot of files are still missing them. This check + # is used to help contributors remember to add the marks when they + # commit. + env: + SKIP: check-test-ci-coverage diff --git a/.gitignore b/.gitignore index 40a6916b552..224dac8676c 100644 --- a/.gitignore +++ b/.gitignore @@ -208,6 +208,7 @@ cache/ .cache/ diffusion_cache/ kv_cache/ +!vllm_omni/experimental/ar_diffusion/kv_cache/ # Logs logs/ @@ -241,11 +242,13 @@ datasets/ *.json !tests/dfx/perf/tests/*.json !apps/ComfyUI-vLLM-Omni/example_workflows/*.json +!tests/assets/soulxsinger/*.json *.jsonl *.parquet # Output files outputs/ +!vllm_omni/outputs/ results/ generated/ output_*/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6db093e0ee2..be649785d8d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -54,6 +54,13 @@ repos: verbose: true stages: [commit-msg] + - id: check-test-ci-coverage + name: Ensure test files have pytestmarks + entry: python tools/pre_commit/check_test_marks.py + language: python + types: [python] + files: ^tests/ + # Keep `suggestion` last - id: suggestion name: Suggestion diff --git a/README.md b/README.md index cf1658b4574..e7fc6912ebe 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone --- *Latest News* 🔥 +- [2026/07] We released [0.24.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.24.0) - aligned with the vLLM 0.24 release line, expanding production-ready coverage across TTS, speech, diffusion, image/video generation, and robot-policy serving, with major Omni stage runtime refactoring, diffusion request-level batching, async output materialization, quantization/cache/memory improvements, and broad CUDA/ROCm/XPU/NPU support. - [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. @@ -31,9 +32,9 @@ Easy, fast, and cheap omni-modality model serving for everyone [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving: -- **Omni-modality**: Text, image, video, and audio data processing +- **Omni-modality**: Text, image, audio, video, and action data processing - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models -- **Heterogeneous outputs**: from traditional text generation to multimodal outputs +- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs

@@ -57,9 +58,10 @@ vLLM-Omni is flexible and easy to use with: vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including: -- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL) +- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos3, 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) +- **Robot-policy and action models** (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy) ## Getting Started @@ -68,6 +70,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn - [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/) - [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/) - [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/) +- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving ## Contributing diff --git a/apps/lvsa-omni-showcase/README.md b/apps/lvsa-omni-showcase/README.md new file mode 100644 index 00000000000..69596fa3c71 --- /dev/null +++ b/apps/lvsa-omni-showcase/README.md @@ -0,0 +1,87 @@ +# LVSA Long-Video Showcase + +**Training-free block-sparse attention for long-video diffusion**, running on +vLLM-Omni via the public plugin entry points (no core changes). LVSA speeds up +extended-horizon generation and prevents the freeze/loop failure mode dense +attention exhibits beyond a model's training length — with no fine-tuning and no +weight changes. + +Upstream: **[LongVideoSparseAttention](https://github.com/JiusiServe/LongVideoSparseAttention)** +(algorithm, all model adapters, full docs). + +## Results + +Wan 2.1 1.3B, single A100 80 GB, mean over 5 prompts (Dense vs LVSA-FlashInfer): + +| Horizon | 1× | 2× | 3× | 4× | 5× | 6× | +|---|---|---|---|---|---|---| +| **Speedup** | 0.9× | 1.4× | 1.9× | 2.5× | 3.0× | **3.5×** | +| **VQeval composite** (Dense→LVSA) | 63→63 | 59→63 | 61→63 | 61→64 | 59→64 | 59→63 | + +The speedup **grows with length** (LVSA is the dense regime at 1× by design, so +the win starts at 2×), and quality *improves* at extension — LVSA's rotating +keyframes prevent the looping/static output dense produces beyond its horizon +(VQeval loop-quality +27 to +34 at ≥3×). At long horizons dense **OOMs on 80 GB** +(HunyuanVideo ≥2×, Wan2.1-14B and Cosmos at 6×) where LVSA still generates. + +→ Full sweep (5 models × 6 horizons, VQeval + VBench-Long) and the +quality-metric methodology: +**[benchmarks/](https://github.com/JiusiServe/LongVideoSparseAttention/tree/main/benchmarks)**. + +## Quickstart + +```bash +# 1. vLLM-Omni 0.22.0 (stable) + matching vLLM +pip install "vllm==0.22.0" +pip install --no-build-isolation \ + "vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@v0.22.0" + +# 2. LVSA core + plugin +git clone https://github.com/JiusiServe/LongVideoSparseAttention.git +( cd LongVideoSparseAttention && pip install -e . && pip install -e lvsa-vllm-omni/ ) + +# 3. this showcase's client dep +pip install requests +``` + +FlashInfer (used by the default `flashinfer` backend) **ships with vLLM** — step 1 +already pulls the pinned `flashinfer-python` + `flashinfer-cubin`, so the fused +backend works with no extra install. Two optional extras: + +```bash +# (optional) faster startup — prebuilt JIT cache, matched to your CUDA + flashinfer +CU=cu128 # match your CUDA: cu121 / cu124 / cu128 / … +FIVER=$(python -c "import importlib.metadata as m; print(m.version('flashinfer-python'))") +pip install --extra-index-url "https://flashinfer.ai/whl/${CU}" \ + "flashinfer-jit-cache==${FIVER}+${CU}" + +# (fallback) no FlashInfer in your environment? run the SDPA backend instead: +export LVSA_BACKEND=sdpa +``` + +```bash +# Serve an LVSA-enabled endpoint (Wan 2.1 1.3B at 2× horizon): +MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 \ + bash serve_lvsa.sh + +# In another shell — generate: +python generate.py --frames 161 --height 480 --width 832 \ + --prompt "A dog running through a sunlit forest." --out wan_lvsa.mp4 +``` + +`serve_lvsa.sh` sets the LVSA env vars for the chosen `MODEL_FAMILY` +(`wan` / `hunyuan`). Wan and HunyuanVideo run through the LVSA **attention +backend** by default — it stays sparse even under sequence-parallel, unlike the +monkey-patch hooks. + +## Configuration + +The one knob you must get right is the model's training horizon +(`LVSA_REFERENCE_LATENT_FRAMES`: Wan = 21, HunyuanVideo = 33); `serve_lvsa.sh` +sets it per family. For the full environment-variable reference, the multi-GPU +support matrix, and tuning (`sparsity_scale`, rotating keyframes, non-standard +resolutions), see the upstream docs: + +- **[lvsa-vllm-omni plugin reference](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/lvsa-vllm-omni/README.md)** — every `LVSA_*` env var +- **[docs/parallelism.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/parallelism.md)** — TP / Ulysses / HSDP support per model +- **[docs/tuning.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/tuning.md)** — picking knobs for your model and horizon diff --git a/apps/lvsa-omni-showcase/generate.py b/apps/lvsa-omni-showcase/generate.py new file mode 100644 index 00000000000..32de4310a9b --- /dev/null +++ b/apps/lvsa-omni-showcase/generate.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python +"""Generate a long video through an LVSA-enabled vLLM-Omni server. + +A thin client over the Omni video API — start the server first with +``serve_lvsa.sh``. The sparse attention runs server-side; this just submits a +job and saves the result. + +Uses the **asynchronous** endpoint (``POST /v1/videos`` → poll +``GET /v1/videos/{id}`` → download ``GET /v1/videos/{id}/content``) rather than +``POST /v1/videos/sync``. The sync endpoint is capped by a fixed server-side +timeout (``VIDEO_SYNC_TIMEOUT_S``, 600 s) and is documented as +"for benchmark and testing scenarios" — long videos (LVSA's use case) routinely +exceed it. The async job has no such cap. + + python generate.py --frames 481 --height 480 --width 832 \ + --prompt "A golden retriever running through a sunlit forest." \ + --out dog_forest_lvsa.mp4 +""" + +import argparse +import os +import time + +import requests + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--server", default=os.environ.get("OMNI_SERVER", "http://localhost:8000")) + ap.add_argument("--prompt", default="A dog running through a sunlit forest.") + ap.add_argument("--frames", type=int, default=161) # Wan 2x horizon + ap.add_argument("--height", type=int, default=480) + ap.add_argument("--width", type=int, default=832) + ap.add_argument("--fps", type=int, default=16) + ap.add_argument("--steps", type=int, default=50) + ap.add_argument("--guidance", type=float, default=5.0) + ap.add_argument("--out", default="lvsa_long_video.mp4") + ap.add_argument("--poll-interval", type=float, default=5.0, help="Seconds between status polls.") + ap.add_argument( + "--max-wait", type=float, default=7200.0, help="Give up if the job hasn't finished after this many seconds." + ) + args = ap.parse_args() + + base = args.server.rstrip("/") + form = { + "prompt": args.prompt, + "size": f"{args.width}x{args.height}", + "num_frames": str(args.frames), + "fps": str(args.fps), + "num_inference_steps": str(args.steps), + "guidance_scale": str(args.guidance), + } + + # ── 1. submit the async job ────────────────────────────────────────────── + print(f"[lvsa-app] POST {base}/v1/videos {args.width}x{args.height} {args.frames}f steps={args.steps}") + resp = requests.post(f"{base}/v1/videos", data=form, timeout=120) + resp.raise_for_status() + job = resp.json() + vid = job["id"] + print(f"[lvsa-app] job {vid} ({job.get('status', '?')})") + + # ── 2. poll until completed / failed ───────────────────────────────────── + deadline = time.time() + args.max_wait + last = None + while True: + if time.time() > deadline: + raise SystemExit(f"[lvsa-app] timed out after {args.max_wait:.0f}s (job {vid} still running)") + time.sleep(args.poll_interval) + r = requests.get(f"{base}/v1/videos/{vid}", timeout=60) + if r.status_code != 200: + # retrieve_video returns a non-200 JSON when the job has FAILED + raise SystemExit(f"[lvsa-app] job failed: {r.text}") + st = r.json() + status = st.get("status") + progress = st.get("progress", 0) + if (status, progress) != last: + print(f"[lvsa-app] {status} {progress}%") + last = (status, progress) + if status == "completed": + break + if status == "failed": + raise SystemExit(f"[lvsa-app] job failed: {st.get('error')}") + + # ── 3. download the rendered mp4 ───────────────────────────────────────── + c = requests.get(f"{base}/v1/videos/{vid}/content", timeout=600) + c.raise_for_status() + with open(args.out, "wb") as f: + f.write(c.content) + print(f"[lvsa-app] saved {args.out}") + + +if __name__ == "__main__": + main() diff --git a/apps/lvsa-omni-showcase/serve_lvsa.sh b/apps/lvsa-omni-showcase/serve_lvsa.sh new file mode 100644 index 00000000000..58b71d14c8f --- /dev/null +++ b/apps/lvsa-omni-showcase/serve_lvsa.sh @@ -0,0 +1,45 @@ +#!/usr/bin/env bash +# Start an LVSA-enabled vLLM-Omni server for Wan or HunyuanVideo. +# +# LVSA is selected as the attention backend via --diffusion-attention-config (see +# below); the `LVSA_*` env vars only tune it. The `lvsa-vllm-omni` plugin registers +# itself through vLLM-Omni's `general_plugins` entry points — nothing in vLLM-Omni +# core is modified. Watch the log for "[LVSA] Geometry detected" to confirm sparse +# attention is active for your request geometry. +# +# MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 bash serve_lvsa.sh +# MODEL=/path/to/HunyuanVideo-1.5-... MODEL_FAMILY=hunyuan FRAMES=193 bash serve_lvsa.sh +set -euo pipefail + +MODEL=${MODEL:?set MODEL=/path/to/Wan-or-HunyuanVideo-Diffusers} +MODEL_FAMILY=${MODEL_FAMILY:-wan} # wan | hunyuan +PORT=${PORT:-8000} +FRAMES=${FRAMES:-161} # target horizon in frames +VAE_T=${VAE_T:-4} # VAE temporal compression (Wan/Hunyuan = 4) +BACKEND=${LVSA_BACKEND:-flashinfer} # flashinfer (fused, fastest) | sdpa + +T_LAT=$(( (FRAMES - 1) / VAE_T + 1 )) + +# Per-family training reference horizon (Wan train ref = 21, Hunyuan = 33). +case "$MODEL_FAMILY" in + wan) REFERENCE=${REFERENCE:-21} ;; + hunyuan) REFERENCE=${REFERENCE:-33} ;; + *) echo "MODEL_FAMILY must be 'wan' or 'hunyuan'"; exit 1 ;; +esac +export LVSA_BACKEND="$BACKEND" +export LVSA_TOTAL_LATENT_FRAMES="$T_LAT" +export LVSA_REFERENCE_LATENT_FRAMES="$REFERENCE" +export LVSA_ROTATE_KEYFRAMES=${LVSA_ROTATE_KEYFRAMES:-1} + +echo "[lvsa-app] serving $MODEL ($MODEL_FAMILY) on :$PORT" +echo "[lvsa-app] LVSA backend=$BACKEND T_lat=$T_LAT (frames=$FRAMES) reference=$REFERENCE" +echo "[lvsa-app] confirm activation by grep '\\[LVSA\\] Geometry detected' in this log" + +# vLLM-Omni 0.22 selects the LVSA attention backend per-role via +# --diffusion-attention-config. The backend is what drives sparse attention +# here; it engages under tensor-parallel too, unlike the LVSA_*_HOOK monkey-patch +# path (which falls back to dense under sequence-parallel), so we deliberately do +# NOT set a hook. +exec vllm serve "$MODEL" --omni --host 0.0.0.0 --port "$PORT" \ + --diffusion-attention-config '{"per_role": {"self": {"backend": "LVSA"}}}' \ + --init-timeout 1800 ${EXTRA_ARGS:-} diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index 368dc5dcb28..f8ad1a71e96 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -28,6 +28,7 @@ class RequestFuncInput: slo_ms: float | None = None extra_body: dict[str, Any] = field(default_factory=dict) image_paths: list[str] | None = None + video_paths: list[str] | None = None request_id: str = field(default_factory=lambda: str(uuid.uuid4())) default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK @@ -338,6 +339,11 @@ async def async_request_v1_videos( form.add_field(k, str(v)) image_file = None + if input.image_paths and input.video_paths: + output.error = "Only one of image_paths or video_paths can be provided" + output.success = False + return output + if input.image_paths and len(input.image_paths) > 0: image_path = input.image_paths[0] image_file = open(image_path, "rb") @@ -345,7 +351,16 @@ async def async_request_v1_videos( "input_reference", image_file, filename=os.path.basename(image_path), - content_type="application/octet-stream", + content_type=_guess_mime_type(image_path), + ) + elif input.video_paths and len(input.video_paths) > 0: + video_path = input.video_paths[0] + image_file = open(video_path, "rb") + form.add_field( + "input_reference", + image_file, + filename=os.path.basename(video_path), + content_type=_guess_mime_type(video_path), ) job_id = None diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 2f385b53857..8f7a02858ce 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -73,12 +73,17 @@ {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1} ]' + v2v: + python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --endpoint /v1/videos --dataset random --task v2v --num-prompts 1 \ + --height 720 --width 1280 --fps 24 --num-frames 189 --num-inference-steps 35 """ import argparse import ast import asyncio +import base64 import glob import json import logging @@ -106,6 +111,9 @@ logger = logging.getLogger(__name__) +_STAGE_METRICS_ENDPOINTS = {"/v1/chat/completions"} +_RETURN_STAGE_METRICS_FIELD = "return_stage_metrics" + class BaseDataset(ABC): def __init__(self, args, api_url: str, model: str): @@ -538,11 +546,17 @@ def __getitem__(self, idx: int) -> RequestFuncInput: if not image_paths: single = row.get("image_path") image_paths = [single] if single else None + video_paths = row.get("video_paths") + if not video_paths: + single_video = row.get("video_path") + video_paths = [single_video] if single_video else None if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i", "it2i"]: raise ValueError( f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row." ) + if not video_paths and self.args.task == "v2v": + raise ValueError("Task v2v requires video input, but no video_path or video_paths found in trace row.") override_w = self.args.width override_h = self.args.height @@ -566,6 +580,7 @@ def __getitem__(self, idx: int) -> RequestFuncInput: timestamp=timestamp, slo_ms=slo_ms, image_paths=image_paths, + video_paths=video_paths, request_id=str(row.get("request_id")) if row.get("request_id") is not None else str(uuid.uuid4()), ) @@ -585,6 +600,8 @@ class CustomDataset(BaseDataset): - seed (optional): Random seed - image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks - image_urls (optional): List of input image URLs (alternative to image_paths) + - video_paths (optional): List of input video paths for v2v tasks + - video_urls (optional): List of input video URLs (alternative to video_paths) Example JSONL for ti2i: {"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]} @@ -656,6 +673,42 @@ def _resolve_image_paths(self, item: dict) -> list[str] | None: return None + def _resolve_video_paths(self, item: dict) -> list[str] | None: + if "video_paths" in item and item["video_paths"]: + paths = item["video_paths"] + if isinstance(paths, str): + paths = [paths] + + valid_paths = [] + for path in paths: + if os.path.exists(path): + valid_paths.append(path) + else: + raise ValueError(f"Video file not found: {path}") + + return valid_paths if valid_paths else None + + if "video_urls" in item and item["video_urls"]: + urls = item["video_urls"] + if isinstance(urls, str): + urls = [urls] + + downloaded_paths = [] + for url in urls: + try: + response = requests.get(url, timeout=30) + response.raise_for_status() + suffix = os.path.splitext(url)[1] or ".mp4" + with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp: + tmp.write(response.content) + downloaded_paths.append(tmp.name) + except Exception as e: + raise ValueError(f"Failed to download video from {url}: {e}") + + return downloaded_paths if downloaded_paths else None + + return None + def __len__(self) -> int: return len(self.data) @@ -668,11 +721,22 @@ def __getitem__(self, idx: int) -> RequestFuncInput: height = item.get("height", self.args.height) num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps) seed = item.get("seed", self.args.seed) - reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"} + reserved_keys = { + "prompt", + "width", + "height", + "num_inference_steps", + "seed", + "image_paths", + "image_urls", + "video_paths", + "video_urls", + } extra_body = {k: v for k, v in item.items() if k not in reserved_keys} # Handle image paths/URLs image_paths = self._resolve_image_paths(item) + video_paths = self._resolve_video_paths(item) return RequestFuncInput( prompt=prompt, @@ -680,6 +744,7 @@ def __getitem__(self, idx: int) -> RequestFuncInput: model=self.model, seed=seed, image_paths=image_paths, + video_paths=video_paths, extra_body=extra_body, width=width, height=height, @@ -733,6 +798,10 @@ def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool self._random_image_path = self._generate_random_image_paths() else: self._random_image_path = None + if self.args.task == "v2v": + self._random_video_path = self._generate_random_video_paths() + else: + self._random_video_path = None def __len__(self) -> int: return self.num_prompts @@ -752,13 +821,15 @@ def __getitem__(self, idx: int) -> RequestFuncInput: if self._sampled_requests: profile = self._sampled_requests[idx] params.update(profile) + prompt = str(params.pop("prompt", f"Random prompt {idx} for benchmarking diffusion models")) return RequestFuncInput( - prompt=f"Random prompt {idx} for benchmarking diffusion models", + prompt=prompt, api_url=self.api_url, model=self.model, seed=self.args.seed, extra_body=extra_body, image_paths=self._random_image_path, + video_paths=self._random_video_path, **params, ) @@ -777,6 +848,35 @@ def _generate_random_image_paths(self) -> list[str]: image_paths.append(image_path) return image_paths + def _generate_random_video_paths(self) -> list[str]: + video_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_video.mp4") + try: + import cv2 + + profile = self.random_request_config[0] if self.random_request_config else {} + width = int(self.args.width or profile.get("width") or 1280) + height = int(self.args.height or profile.get("height") or 720) + num_frames = int(self.args.num_frames or profile.get("num_frames") or 16) + fps = float(self.args.fps or profile.get("fps") or 8) + writer = cv2.VideoWriter(video_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)) + if not writer.isOpened(): + raise RuntimeError("cv2.VideoWriter failed to open") + for frame_idx in range(num_frames): + frame = np.zeros((height, width, 3), dtype=np.uint8) + x = int((frame_idx / max(num_frames - 1, 1)) * max(width - 96, 1)) + y = height // 2 + frame[:, :, 0] = 24 + frame[:, :, 1] = 32 + frame[:, :, 2] = 48 + frame[max(y - 36, 0) : min(y + 36, height), x : min(x + 96, width), :] = (64, 128, 220) + writer.write(frame) + writer.release() + return [video_path] + except Exception as e: + raise RuntimeError( + "Failed to generate synthetic v2v input video. Install opencv-python or provide a custom dataset." + ) from e + def _compute_expected_latency_ms_from_base(req: RequestFuncInput, args, base_time_ms: float | None) -> float | None: """Compute expected execution time (ms) based on a base per-step-per-frame unit time. @@ -1013,6 +1113,81 @@ def calculate_metrics( return metrics +def _save_generated_outputs( + outputs: list[RequestFuncOutput], + requests_list: list[RequestFuncInput], + save_dir: str, +) -> None: + """Decode and save base64 images/videos from successful responses.""" + os.makedirs(save_dir, exist_ok=True) + saved = 0 + failed = 0 + + for idx, (req, out) in enumerate(zip(requests_list, outputs)): + if not out.success or not out.response_body: + continue + + if isinstance(out.response_body, bytes): + fname = f"req_{idx:04d}.mp4" + fpath = os.path.join(save_dir, fname) + with open(fpath, "wb") as f: + f.write(out.response_body) + saved += 1 + continue + + media_urls: list[str] = [] + + # Chat-completions style: choices[*].message.content[*].image_url.url + choices = out.response_body.get("choices", []) + if isinstance(choices, list): + for choice in choices: + content = (choice or {}).get("message", {}).get("content") + if not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict) or item.get("type") != "image_url": + continue + url = (item.get("image_url") or {}).get("url", "") + if isinstance(url, str) and url.startswith("data:"): + media_urls.append(url) + + # Images endpoint style: data[*].b64_json + data_items = out.response_body.get("data", []) + if isinstance(data_items, list): + for data_item in data_items: + if not isinstance(data_item, dict): + continue + b64_json = data_item.get("b64_json", "") + if isinstance(b64_json, str) and b64_json: + media_urls.append(f"data:image/png;base64,{b64_json}") + + for img_idx, url in enumerate(media_urls): + if "," not in url: + continue + + try: + header, b64_data = url.split(",", 1) + ext = "png" + if "image/jpeg" in header: + ext = "jpg" + elif "image/webp" in header: + ext = "webp" + elif "video/mp4" in header: + ext = "mp4" + + img_bytes = base64.b64decode(b64_data) + fname = f"req_{idx:04d}_{img_idx}.{ext}" + fpath = os.path.join(save_dir, fname) + with open(fpath, "wb") as f: + f.write(img_bytes) + saved += 1 + except Exception as e: + failed += 1 + logger.warning(f"Failed to save image for request {idx}: {e}", exc_info=True) + + logger.info(f"Saved {saved} generated image(s) to {save_dir}. Failed to save {failed} image(s).") + + def wait_for_service(base_url: str, timeout: int = 120) -> None: print(f"Waiting for service at {base_url}...") start_time = time.time() @@ -1033,7 +1208,7 @@ def wait_for_service(base_url: str, timeout: int = 120) -> None: def _default_endpoint_for_task(task: str) -> str: - if task in {"t2v", "i2v", "ti2v"}: + if task in {"t2v", "i2v", "ti2v", "v2v"}: return "/v1/videos" if task in {"i2i", "ti2i", "it2i"}: return "/v1/images/edits" @@ -1047,7 +1222,7 @@ async def benchmark(args): if args.base_url is None: args.base_url = f"http://{args.host}:{args.port}" - VIDEO_TASKS = {"t2v", "i2v", "ti2v"} + VIDEO_TASKS = {"t2v", "i2v", "ti2v", "v2v"} IMAGE_TASKS = {"t2i", "i2i", "ti2i", "it2i"} if args.task in VIDEO_TASKS: @@ -1095,6 +1270,14 @@ async def benchmark(args): requests_list = dataset.get_requests() print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.") + if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS: + for req in requests_list: + req.extra_body.setdefault(_RETURN_STAGE_METRICS_FIELD, True) + if args.extra_body: + extra_body = json.loads(args.extra_body) + for req in requests_list: + req.extra_body.update(extra_body) + if args.endpoint == "/v1/images/edits": for req in requests_list: req.default_bot_task = args.bot_task @@ -1202,6 +1385,9 @@ async def limited_request_func(req, session, pbar): print("\n" + "=" * 60) + if args.save_dir: + _save_generated_outputs(outputs, requests_list, args.save_dir) + if args.output_file: with open(args.output_file, "w") as f: json.dump(metrics, f, indent=2) @@ -1242,7 +1428,7 @@ async def limited_request_func(req, session, pbar): "--task", type=str, default="t2v", - choices=["t2v", "i2v", "ti2v", "ti2i", "i2i", "it2i", "t2i"], + choices=["t2v", "i2v", "ti2v", "v2v", "ti2i", "i2i", "it2i", "t2i"], help="Task type.", ) parser.add_argument( @@ -1278,11 +1464,18 @@ async def limited_request_func(req, session, pbar): default=1, help="Number of warmup requests to run before measurement.", ) + # NOTE Changed default from 1 to 2 because some models (e.g., Bagel) run + # `num_timesteps - 1` denoising iterations. A default of 1 results in 0 steps, + # which causes errors. + # TODO If this slightly longer warmup causes regression issues for other + # diffusion pipelines in the future, consider implementing model-specific + # overrides instead of a global default. parser.add_argument( "--warmup-num-inference-steps", type=int, - default=1, - help="num_inference_steps used for warmup requests.", + default=2, + help="Number of inference steps used for warmup requests. " + "Default is 2 to ensure at least one denoising step is executed.", ) parser.add_argument( "--warmup-concurrency", @@ -1323,6 +1516,13 @@ async def limited_request_func(req, session, pbar): default=3.0, help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).", ) + parser.add_argument( + "--save-dir", + type=str, + default=None, + help="Directory to save generated images/outputs for visual inspection. " + "If not set, generated outputs are discarded after metric collection.", + ) parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.") parser.add_argument( "--enable-negative-prompt", @@ -1343,6 +1543,12 @@ async def limited_request_func(req, session, pbar): '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]' ), ) + parser.add_argument( + "--extra-body", + type=str, + default=None, + help="JSON object merged into every generated request body/form.", + ) parser.add_argument( "--num-input-images", type=int, @@ -1358,6 +1564,11 @@ async def limited_request_func(req, session, pbar): default="think", help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."), ) + parser.add_argument( + "--return-stage-metrics", + action="store_true", + help="Request stage duration metrics from endpoints that support return_stage_metrics.", + ) args = parser.parse_args() asyncio.run(benchmark(args)) diff --git a/benchmarks/fish-speech/bench_speaker_cache.py b/benchmarks/fish-speech/bench_speaker_cache.py index a2ac460afbc..764e4e9df36 100644 --- a/benchmarks/fish-speech/bench_speaker_cache.py +++ b/benchmarks/fish-speech/bench_speaker_cache.py @@ -201,6 +201,7 @@ def make_inline_payload(prompt: str) -> dict: "input": prompt, "voice": "default", "stream": True, + "stream_format": "audio", "response_format": "pcm", "ref_audio": ref_audio_b64, "ref_text": args.ref_text, @@ -240,6 +241,7 @@ def make_uploaded_payload(prompt: str) -> dict: "input": prompt, "voice": args.voice_name, "stream": True, + "stream_format": "audio", "response_format": "pcm", "ref_text": args.ref_text, "max_new_tokens": 2048, diff --git a/benchmarks/kernels/mot_linear_benchmarks.py b/benchmarks/kernels/mot_linear_benchmarks.py new file mode 100644 index 00000000000..20674272e85 --- /dev/null +++ b/benchmarks/kernels/mot_linear_benchmarks.py @@ -0,0 +1,1015 @@ +# ruff: noqa: N803 + +"""MoT (Mixture-of-Tokens) GEMM kernel benchmark and auto-tuning. + +Generates optimal Triton kernel configurations for MoT GEMM operations +across different batch sizes, model shapes, TP configurations, and hardware. + +Usage: + # Auto-tune and save configs: + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --tp-size 1 --dtype w16a16 --tune \ + --save-dir vllm_omni/diffusion/layers/mot/configs/ + + # Auto-tune with local model path (offline clusters): + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model /data/models/BAGEL-7B-MoT \ + --tp-size 2 --tune + + # Benchmark only (measure with default configs, no search): + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --tp-size 1 --dtype w16a16 +""" + +import argparse +import gc +import json +import logging +import math +import os +import time +from datetime import datetime +from itertools import product +from typing import Any + +import ray +import torch +from ray.experimental.tqdm_ray import tqdm +from vllm.platforms import current_platform +from vllm.transformers_utils.config import get_config +from vllm.triton_utils import triton +from vllm.utils.torch_utils import set_random_seed + +# NOTE: you should use the same naming syetem for the kernel to load properly +from vllm_omni.diffusion.layers.mot.ops.mot_gemm import build_config_filename, get_device_name + +# clear the triton cache from time to time, usaully no need to change +_CACHE_CLEAR_INTERVAL_ENV = "VLLM_MOT_TUNE_CACHE_CLEAR_INTERVAL" +TRITON_CACHE_CLEAR_INTERVAL = int(os.environ.get(_CACHE_CLEAR_INTERVAL_ENV, "50")) + +# represent the token number of each generated image +_VAE_CHUNK_SIZE_ENV = "VAE_CHUNK_SIZE" +VAE_CHUNK_SIZE = int(os.environ.get(_VAE_CHUNK_SIZE_ENV, "1024")) + +logger = logging.getLogger(__name__) + +# ===================================================================== +# Utility Functions +# ===================================================================== + + +def clear_triton_cache(): + """Clear Triton JIT compilation cache and Python/CUDA memory.""" + gc.collect() + if torch.cuda.is_available(): + torch.accelerator.empty_cache() + try: + if hasattr(triton, "runtime") and hasattr(triton.runtime, "cache") and hasattr(triton.runtime.cache, "clear"): + triton.runtime.cache.clear() + except Exception: + pass + gc.collect() + + +# TODO: check rocm/npus support +# based on https://docs.nvidia.com/cuda/cuda-runtime-api/ +# structcudaDeviceProp.html#structcudaDeviceProp_16cede1829516e86917f0842a5f6498c8 +def get_max_shared_memory() -> int: + """Return the maximum shared memory per block in bytes.""" + props = torch.cuda.get_device_properties(0) + if hasattr(props, "shared_memory_per_block_option"): + return props.shared_memory_per_block_option + return getattr(props, "shared_memory_per_block", 49152) + + +def get_max_regs() -> int: + """Return the maximum registers per block in bytes.""" + props = torch.cuda.get_device_properties(0) + if hasattr(props, "regs_per_block"): + return props.regs_per_block + return 65536 + + +def get_sm_count() -> int: + """Get the number of physical SMs on the target GPU (A100 = 108, H100 = 132).""" + return torch.cuda.get_device_properties(0).multi_processor_count + + +def get_ab_element_bytes(dtype_str: str) -> tuple[int, int]: + """Return ``(activation_bytes, weight_bytes)`` for a dtype config.""" + if dtype_str == "w16a16": + return 2, 2 + elif dtype_str == "fp8_w8a8": + return 1, 1 + elif dtype_str == "int8_w8a16": + return 2, 1 + return 2, 2 + + +def build_regular_indices( + image_num: int, + vae_chunk_size: int, + device: str = "cuda", +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Build deterministic MoT indices with per-image [Text][VAE...][Text].""" + if image_num <= 0: + raise ValueError(f"image_num must be > 0, got {image_num}") + if vae_chunk_size <= 0: + raise ValueError(f"{_VAE_CHUNK_SIZE_ENV} must be > 0, got {vae_chunk_size}") + + text_idx_list: list[int] = [] + vae_idx_list: list[int] = [] + + current_idx = 0 + for _ in range(image_num): + text_idx_list.append(current_idx) + current_idx += 1 + + vae_idx_list.extend(range(current_idx, current_idx + vae_chunk_size)) + current_idx += vae_chunk_size + + text_idx_list.append(current_idx) + current_idx += 1 + + text_indices = torch.tensor(text_idx_list, dtype=torch.long, device=device) + vae_indices = torch.tensor(vae_idx_list, dtype=torch.long, device=device) + exact_M = current_idx + return text_indices, vae_indices, exact_M + + +def get_exact_m(image_num: int, vae_chunk_size: int) -> int: + if image_num <= 0: + raise ValueError(f"image_num must be > 0, got {image_num}") + if vae_chunk_size <= 0: + raise ValueError(f"{_VAE_CHUNK_SIZE_ENV} must be > 0, got {vae_chunk_size}") + return image_num * (vae_chunk_size + 2) + + +# ===================================================================== +# Model Shape Extraction +# ===================================================================== + + +class MoTShape: + """One unique (K, N) GEMM shape in a MoT model layer.""" + + def __init__(self, K: int, N: int, comment: str): + self.K = K + self.N = N + self.comment = comment + + def config_key(self) -> str: + return f"{self.K}_{self.N}" + + def __repr__(self) -> str: + return f"MoTShape(K={self.K}, N={self.N}, comment='{self.comment}')" + + +def get_mot_shapes( + model: str, + tp_size: int, + trust_remote_code: bool = False, +) -> tuple[list[MoTShape], str]: + """Extract MoT GEMM shapes from a HuggingFace model config. + + Supports both remote HuggingFace model IDs and local checkpoint paths. + + Returns + ------- + shapes : list[MoTShape] + De-duplicated GEMM shapes (K, N) with TP applied. + model_name : str + Cleaned model name for the config filename. + """ + config = get_config(model=model, trust_remote_code=trust_remote_code) + model_name = model.rstrip("/").split("/")[-1] + + text_config = getattr(config, "text_config", config) + + hidden_size: int = text_config.hidden_size + num_attention_heads: int = text_config.num_attention_heads + num_kv_heads: int = getattr(text_config, "num_key_value_heads", num_attention_heads) + head_dim: int = getattr(text_config, "head_dim", hidden_size // num_attention_heads) + intermediate_size: int = text_config.intermediate_size + + # ---- Compute per-TP shapes ---- + + # QKV_PROJ (QKVParallelLinear, output partitioned by TP) + q_out = num_attention_heads * head_dim + kv_out = 2 * num_kv_heads * head_dim + qkv_total = q_out + kv_out + assert qkv_total % tp_size == 0, f"QKV output {qkv_total} not divisible by tp {tp_size}" + qkv_N = qkv_total // tp_size + + # O_PROJ (RowParallelLinear, input partitioned by TP) + assert q_out % tp_size == 0, f"Q output {q_out} not divisible by tp {tp_size}" + o_K = q_out // tp_size + o_N = hidden_size + + # FFN gate+up (MergedColumnParallelLinear, output partitioned by TP) + gate_up_total = 2 * intermediate_size + assert gate_up_total % tp_size == 0, f"Gate-up output {gate_up_total} not divisible by tp {tp_size}" + gate_up_N = gate_up_total // tp_size + + # FFN down (RowParallelLinear, input partitioned by TP) + assert intermediate_size % tp_size == 0, f"Intermediate size {intermediate_size} not divisible by tp {tp_size}" + down_K = intermediate_size // tp_size + down_N = hidden_size + + shapes = [ + MoTShape(K=o_K, N=o_N, comment="O_PROJ"), + MoTShape(K=hidden_size, N=qkv_N, comment="QKV_PROJ"), + MoTShape(K=hidden_size, N=gate_up_N, comment="FFN_GATE_UP_PROJ"), + MoTShape(K=down_K, N=down_N, comment="FFN_DOWN_PROJ"), + ] + + seen: dict[str, MoTShape] = {} + unique: list[MoTShape] = [] + for s in shapes: + key = s.config_key() + if key not in seen: + seen[key] = s + unique.append(s) + else: + seen[key].comment += f" / {s.comment}" + + return unique, model_name + + +# ===================================================================== +# Search Space Generation & Pruning +# ===================================================================== + + +def estimate_sram_bytes(config: dict[str, int], dtype_str: str) -> int: + """Estimate SRAM (shared memory) usage for a Triton tile config. + + Formula: + (BLOCK_M * BLOCK_K * a_bytes + BLOCK_N * BLOCK_K * b_bytes) + * num_stages + """ + bm = config["BLOCK_SIZE_M"] + bn = config["BLOCK_SIZE_N"] + bk = config["BLOCK_SIZE_K"] + stages = config["num_stages"] + a_bytes, b_bytes = get_ab_element_bytes(dtype_str) + return (bm * bk * a_bytes + bn * bk * b_bytes) * stages + + +# TODO: +# granule_size= 256 for nvdia gpus, +# warp_size=32 for nvdia gpus, +# not sure if it is true for rocm/other npus +def estimate_register_pressure( + config: dict[str, int], + dtype_str: str, + max_regs_per_block: int = 65536, + max_regs_per_thread: int = 255, + granule_size: int = 256, + warp_size: int = 32, +) -> bool: + """ + Evaluate register pressure for MoT GEMM based on kernel structure and datatypes. + + Args: + config: Triton tile configuration. + dtype_str: for now only support:"w16a16", "fp8_w8a8","int8_w8a16" + max_regs_per_block: Hardware limit (usually 65536). + max_regs_per_thread: PTX limit (usually 255). + granule_size: register allocation size for one warp. + warp_size: number of threads per warp. + Returns: + True if the config is safe to compile and run efficiently, False if it should be pruned. + """ + bm = config["BLOCK_SIZE_M"] + bn = config["BLOCK_SIZE_N"] + bk = config["BLOCK_SIZE_K"] + warps = config["num_warps"] + num_threads = warps * warp_size + + a_bytes, b_bytes = get_ab_element_bytes(dtype_str) + + # Physical register standard: 1 register = 32-bit (4 bytes) + + # [Accumulator C] + # Triton uses fp32/int32 by default as the accumulator for fp16/int8 to prevent overflow + regs_c = (bm * bn) / num_threads * 1.0 + + # [MMA slices A and B] + # Data is loaded into registers from SRAM to participate in Tensor Core operations + regs_a = ((bm * bk) / num_threads) * (a_bytes / 4.0) + regs_b = ((bk * bn) / num_threads) * (b_bytes / 4.0) + + # [MoT specific routing overhead] + # real_row_idxs is tl.int64 (8 bytes), each element needs 2 32-bit registers + regs_routing = (bm / num_threads) * 2.0 + + # [Quantization specific Epilogue overhead] + # W8A8 needs to load scale_a and scale_b after the loop for de-quantization + regs_epilogue = 0.0 + if dtype_str == "fp8_w8a8": + # fp8*token-wise quant scenario: scale_a length is bm, scale_b length is bn + regs_epilogue = ((bm + bn) / num_threads) * 1.0 + elif dtype_str == "int8_w8a16": + # Weight-Only*token-wise quant scenario: usually only scale_b is needed + regs_epilogue = (bn / num_threads) * 1.0 + + # [Control flow and base pointer constant overhead] + # Includes: loop counter(k), pointer addressing, + # Mask predicate calculation, TMA state machine, etc. + constant_overhead = 35 + # --------------------------------------------------------- + # --------------------------------------------------------- + # Summary and red line intercept + # --------------------------------------------------------- + estimated_regs_per_thread = math.ceil(regs_c + regs_a + regs_b + regs_routing + regs_epilogue + constant_overhead) + # Hardware red line 1: single thread physical limit + # (PTX ISA specifies a maximum of 255 registers per thread, + # leaving 10% as a compiler buffer) + if estimated_regs_per_thread > max_regs_per_thread * 0.9: + return False + + # Hardware red line 2: single block total physical limit + # (each warp allocated registers are rounded up to the nearest multiple of 256) + regs_per_warp_raw = estimated_regs_per_thread * warp_size + regs_per_warp_actual = math.ceil(regs_per_warp_raw / granule_size) * granule_size + + # Calculate the actual physical register consumption for the current block + estimated_regs_per_block = regs_per_warp_actual * warps + if estimated_regs_per_block > max_regs_per_block: + return False + + return True + + +def get_mot_search_space( + M: int, + K: int, + N: int, + dtype_str: str, + max_sram: int, + max_regs: int, + num_sms: int, +) -> list[dict[str, int]]: + """Generate a pruned search space of Triton tile configs for MoT GEMM.""" + + param_ranges = { + "BLOCK_SIZE_M": [32, 64, 128, 256], + "BLOCK_SIZE_N": [32, 64, 128, 256], + "BLOCK_SIZE_K": [32, 64, 128], + "GROUP_SIZE_M": [4, 8, 16], + "num_warps": [4, 8], + "num_stages": [2, 3, 4, 5], + } + + def next_power_of_2(n): + return 1 if n == 0 else 2 ** (n - 1).bit_length() + + padded_M = next_power_of_2(M) + padded_N = next_power_of_2(N) + padded_K = next_power_of_2(K) + + keys, values = zip(*param_ranges.items()) + configs: list[dict[str, int]] = [] + + for vals in product(*values): + cfg = dict(zip(keys, vals)) + bm = cfg["BLOCK_SIZE_M"] + bn = cfg["BLOCK_SIZE_N"] + bk = cfg["BLOCK_SIZE_K"] + + # --- Dimension-based pruning --- + if bm > max(32, padded_M): + continue + if bn > max(32, padded_N): + continue + if bk > max(32, padded_K): + continue + if bm * bn < 64: + continue + + # --- Occupancy-based pruning --- + grid_m = (M + bm - 1) // bm + grid_n = (N + bn - 1) // bn + total_blocks = grid_m * grid_n + + if total_blocks < num_sms // 4: + continue + + # --- SRAM capacity check --- + if estimate_sram_bytes(cfg, dtype_str) > max_sram * 0.9: + continue + + # --- register spilling check --- + if not estimate_register_pressure(cfg, dtype_str, max_regs): + continue + + configs.append(cfg) + + return configs + + +# ===================================================================== +# Single-Config Benchmark +# ===================================================================== + + +def benchmark_config( + config: dict[str, int], + image_num: int, + K: int, + N: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a16: bool, + num_iters: int = 10, + cache_flusher: torch.Tensor | None = None, +) -> float: + """Run a MoT GEMM with the given tile config and return avg latency (us).""" + from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm + + text_indices, vae_indices, M = build_regular_indices( + image_num=image_num, + vae_chunk_size=VAE_CHUNK_SIZE, + device="cuda", + ) + + # ---- Allocate tensors on the current CUDA device ---- + A_scale: torch.Tensor | None = None + B_text_scale: torch.Tensor | None = None + B_vae_scale: torch.Tensor | None = None + + if use_fp8_w8a8: + fp8_dtype = current_platform.fp8_dtype() + A = torch.randn(M, K, dtype=torch.float16, device="cuda").to(fp8_dtype) + B_text = torch.randn(K, N, dtype=torch.float16, device="cuda").to(fp8_dtype) + B_vae = torch.randn(K, N, dtype=torch.float16, device="cuda").to(fp8_dtype) + A_scale = torch.ones(M, dtype=torch.float32, device="cuda") + B_text_scale = torch.ones(1, dtype=torch.float32, device="cuda") + B_vae_scale = torch.ones(1, dtype=torch.float32, device="cuda") + C = torch.empty(M, N, dtype=dtype, device="cuda") + elif use_int8_w8a16: + A = torch.randn(M, K, dtype=dtype, device="cuda") + B_text = torch.randint(-127, 127, (K, N), dtype=torch.int8, device="cuda") + B_vae = torch.randint(-127, 127, (K, N), dtype=torch.int8, device="cuda") + B_text_scale = torch.ones(N, dtype=torch.float32, device="cuda") + B_vae_scale = torch.ones(N, dtype=torch.float32, device="cuda") + C = torch.empty(M, N, dtype=dtype, device="cuda") + else: + A = torch.randn(M, K, dtype=dtype, device="cuda") + B_text = torch.randn(K, N, dtype=dtype, device="cuda") + B_vae = torch.randn(K, N, dtype=dtype, device="cuda") + C = torch.empty(M, N, dtype=dtype, device="cuda") + + def run(): + invoke_mot_gemm( + A=A, + B_text=B_text, + B_vae=B_vae, + C=C, + bias_text=None, + bias_vae=None, + text_indices=text_indices, + vae_indices=vae_indices, + A_scale=A_scale, + B_text_scale=B_text_scale, + B_vae_scale=B_vae_scale, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=False, + use_int8_w8a16=use_int8_w8a16, + use_int4_w4a16=False, + A_per_channel_quant=use_fp8_w8a8, + B_per_channel_quant=use_int8_w8a16, + config=config, + ) + + # JIT warmup + run() + torch.accelerator.synchronize() + + # Capture 1 invocations with CUDA Graph + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + torch.accelerator.synchronize() + + # Warmup replays + for _ in range(5): + graph.replay() + torch.accelerator.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + latencies: list[float] = [] + for _ in range(num_iters): + if cache_flusher is not None: + cache_flusher.zero_() + torch.accelerator.synchronize() + + start_event.record() + graph.replay() + end_event.record() + end_event.synchronize() + + latencies.append(start_event.elapsed_time(end_event)) + + latencies.sort() + valid_latencies = latencies[1:-1] if len(latencies) > 2 else latencies + + avg_us = sum(valid_latencies) / len(valid_latencies) * 1000 # ms → us + graph.reset() + + return avg_us + + +# ===================================================================== +# Ray Worker +# ===================================================================== + + +@ray.remote(num_gpus=1) +class BenchmarkWorker: + def __init__(self, seed: int) -> None: + # Ray will automatically set CUDA_VISIBLE_DEVICES, + # so the GPU seen by the worker is always the logical 0 + self.logical_device_id = 0 + torch.set_default_device(f"cuda:{self.logical_device_id}") + + set_random_seed(seed) + self.seed = seed + + # ---- Benchmark (use default config, report latency) ---- + + def benchmark( + self, + image_num: int, + K: int, + N: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a16: bool, + ) -> tuple[dict[str, int], float]: + set_random_seed(self.seed) + from vllm_omni.diffusion.layers.mot.ops.mot_gemm import ( + get_best_mot_config, + ) + + M = get_exact_m(image_num, VAE_CHUNK_SIZE) + loaded_m_key, config = get_best_mot_config(M, N, K) + if loaded_m_key == -1: + print( + " [config] WARNING: No tuned config found — " + "using conservative default. " + "Performance numbers are NOT representative. " + "Run mot_linear_benchmarks.py --tune to generate configs." + ) + else: + print(f" [config] Tuned config loaded (actual M={M}, loaded M={loaded_m_key}) config = {config})") + kernel_time = benchmark_config( + config, + image_num, + K, + N, + dtype, + use_fp8_w8a8, + use_int8_w8a16, + num_iters=10, + ) + return config, kernel_time + + # ---- Tune (search over all configs, return best) ---- + + def tune( + self, + image_num: int, + K: int, + N: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a16: bool, + search_space: list[dict[str, int]], + ) -> dict[str, int] | None: + set_random_seed(self.seed) + M = get_exact_m(image_num, VAE_CHUNK_SIZE) + + best_config: dict[str, int] | None = None + best_time = float("inf") + + # Diagnosis counters + total_configs = len(search_space) + err_oom = 0 + err_triton_resources = 0 + err_other = 0 + + with torch.cuda.device(self.logical_device_id): + cache_flusher = torch.empty(int(256 * 1024 * 1024 / 4), dtype=torch.int32, device="cuda") + + for idx, config in enumerate(tqdm(search_space)): + try: + kernel_time = benchmark_config( + config, + image_num, + K, + N, + dtype, + use_fp8_w8a8, + use_int8_w8a16, + num_iters=10, + cache_flusher=cache_flusher, + ) + except triton.runtime.autotuner.OutOfResources: + err_triton_resources += 1 + continue + except torch.cuda.OutOfMemoryError: + err_oom += 1 + clear_triton_cache() + continue + except Exception: + err_other += 1 + logger.exception("Config %s failed unexpectedly", config) + clear_triton_cache() + continue + + if kernel_time < best_time: + best_time = kernel_time + best_config = config + + if TRITON_CACHE_CLEAR_INTERVAL > 0 and idx > 0 and idx % TRITON_CACHE_CLEAR_INTERVAL == 0: + clear_triton_cache() + + del cache_flusher + clear_triton_cache() + + if best_config is None: + diag_msg = ( + f"\n🚨 [CRITICAL] TUNING FAILED for M={M}, K={K}, N={N}\n" + f" Total configs tested: {total_configs}\n" + f" - Triton OutOfResources (SRAM/Regs): {err_triton_resources}\n" + f" - CUDA OOM: {err_oom}\n" + f" - Other Errors: {err_other}\n" + f" 💡 DIAGNOSIS:\n" + f" 1. If total configs is 0, your 'get_mot_search_space' pruning is too aggressive.\n" + f" 2. If Triton/OOM errors == total configs, hardware limits (SRAM/Regs) in pruning are too loose.\n" + f" 3. If Other Errors is high, check benchmark_config logic or Triton kernel runtime bugs." + ) + print(diag_msg) + return None + + now = datetime.now() + print(f"[{now.ctime()}] Tuning done: M={M}, K={K}, N={N}, best_time={best_time:.2f} us") + return best_config + + # ---- Device info helpers (called from driver) ---- + + def get_device_name(self) -> str: + return get_device_name() + + def get_max_shared_memory(self) -> int: + return get_max_shared_memory() + + def get_sm_count(self) -> int: + return get_sm_count() + + def get_max_regs(self) -> int: + return get_max_regs() + + +# ===================================================================== +# Config I/O +# ===================================================================== + + +def sort_config(config: dict[str, int]) -> dict[str, int]: + return { + "BLOCK_SIZE_M": config["BLOCK_SIZE_M"], + "BLOCK_SIZE_N": config["BLOCK_SIZE_N"], + "BLOCK_SIZE_K": config["BLOCK_SIZE_K"], + "GROUP_SIZE_M": config["GROUP_SIZE_M"], + "num_warps": config["num_warps"], + "num_stages": config["num_stages"], + } + + +def save_configs( + results: dict[str, dict[int, dict[str, int]]], + shapes: list[MoTShape], + model_name: str, + tp_size: int, + device_name: str, + dtype_str: str, + save_dir: str, +) -> str: + """Merge tuned configs into ``device_name=...,dtype=....json``. + + Behavior: + - Create a new file if it does not exist. + - If it exists, merge by shape key (``K_N``) and M key. + - Existing entries are preserved unless overwritten by current results. + """ + shape_map = {s.config_key(): s for s in shapes} + + current_output: dict[str, Any] = {} + for config_key, m_configs in results.items(): + shape = shape_map[config_key] + # Make comments self-descriptive across mixed model/tp runs. + comment = f"model={model_name}|tp={tp_size}|op={shape.comment}" + entry: dict[str, Any] = {"_comment": comment} + for m_val in sorted(m_configs.keys()): + entry[str(m_val)] = sort_config(m_configs[m_val]) + current_output[config_key] = entry + + filename = f"device_name={device_name},dtype={dtype_str}.json" + os.makedirs(save_dir, exist_ok=True) + filepath = os.path.join(save_dir, filename) + + merged_output: dict[str, Any] = {} + if os.path.isfile(filepath): + try: + with open(filepath) as f: + existing = json.load(f) + if isinstance(existing, dict): + merged_output = existing + else: + print(f"WARNING: Existing config is not a JSON object: {filepath}. Overwrite with newly tuned configs.") + except Exception as exc: + print(f"WARNING: Failed to read existing config {filepath}: {exc}. Overwrite with newly tuned configs.") + + # Merge on two levels: shape key -> M key + for config_key, new_entry in current_output.items(): + old_entry = merged_output.get(config_key, {}) + if not isinstance(old_entry, dict): + old_entry = {} + merged_entry = dict(old_entry) + old_comment = merged_entry.get("_comment") + new_comment = new_entry.get("_comment") + merged_entry.update(new_entry) + if old_comment and new_comment and old_comment != new_comment: + merged_entry["_comment"] = f"{old_comment} / {new_comment}" + merged_output[config_key] = merged_entry + + print(f"Saving merged config to {filepath}") + with open(filepath, "w") as f: + json.dump(merged_output, f, indent=2) + f.write("\n") + + return filepath + + +# ===================================================================== +# Main +# ===================================================================== + + +def main(args: argparse.Namespace): + print(args) + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + + if not current_platform.is_cuda(): + raise NotImplementedError( + "Non-CUDA environment detected!" + "This benchmark script has not been fully tested on" + "AMD GPUs and may produce errors or suboptimal results." + ) + + # ---- 1. Extract model shapes ---- + shapes, model_name = get_mot_shapes(args.model, args.tp_size, args.trust_remote_code) + print(f"\nModel: {model_name} | TP: {args.tp_size}") + print(f"Detected {len(shapes)} unique GEMM shape(s):") + for s in shapes: + print(f" {s}") + + # ---- 2. Determine dtype ---- + dtype_str: str = args.dtype + use_fp8_w8a8 = dtype_str == "fp8_w8a8" + use_int8_w8a16 = dtype_str == "int8_w8a16" + dtype = torch.bfloat16 + + # ---- 3. Image counts ---- + image_nums: list[int] = args.batch_size if args.batch_size is not None else [1, 2, 4, 8, 16] + + # ---- 4. Initialize Ray workers ---- + ray.init() + num_gpus = int(ray.available_resources()["GPU"]) + workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)] + print(f"\nRay initialized with {num_gpus} GPU worker(s)") + + device_name = ray.get(workers[0].get_device_name.remote()) + max_sram = ray.get(workers[0].get_max_shared_memory.remote()) + max_regs = ray.get(workers[0].get_max_regs.remote()) + sm_count = ray.get(workers[0].get_sm_count.remote()) + + print( + f"Device: {device_name} | Max SRAM/Block: {max_sram} bytes\n" + f" Max Regs/Block: {max_regs} 32-bit regs\n" + f" SM on GPU: {sm_count} \n" + ) + + # ---- Helper: round-robin distribute tasks to workers ---- + def distribute(method: str, inputs: list[tuple[Any, ...]]) -> list[Any]: + futures = [] + for i, input_args in enumerate(inputs): + worker = workers[i % num_gpus] + futures.append(getattr(worker, method).remote(*input_args)) + return ray.get(futures) + + # ---- 5. TUNE mode ---- + if args.tune: + start = time.time() + + # 1) Checkpoint loading and resuming + filename = build_config_filename(device_name, dtype_str) + filepath = os.path.join(args.save_dir, filename) + + existing_history: dict[str, Any] = {} + if os.path.isfile(filepath): + try: + with open(filepath) as f: + existing_history = json.load(f) + print(f"Loaded existing checkpoint from {filepath}, resuming...") + except Exception as e: + print(f"WARNING: Failed to load existing checkpoint: {e}") + + # 2) Build task queue and execute checkpoint filtering + pending_futures = {} + task_counter = 0 + + for shape in shapes: + shape_key = shape.config_key() + for image_num in image_nums: + exact_M = get_exact_m(image_num, VAE_CHUNK_SIZE) + + if shape_key in existing_history and str(exact_M) in existing_history[shape_key]: + print(f"Skipping image_num={image_num} (M={exact_M}), Shape={shape_key} (Already tuned)") + continue + + # Only tune parameters that have not been tuned yet + search_space = get_mot_search_space( + M=exact_M, + K=shape.K, + N=shape.N, + dtype_str=dtype_str, + max_sram=max_sram, + max_regs=max_regs, + num_sms=sm_count, + ) + if len(search_space) == 0: + print( + f"WARNING: empty search space for " + f"{shape.config_key()} image_num={image_num} (M={exact_M}), " + "skipping" + ) + continue + + # Round-robin assign to Worker + worker = workers[task_counter % num_gpus] + future = worker.tune.remote( + image_num, shape.K, shape.N, dtype, use_fp8_w8a8, use_int8_w8a16, search_space + ) + + # Bind future with its corresponding metadata + pending_futures[future] = (shape, image_num, exact_M) + task_counter += 1 + + print(f"Starting tuning: {len(pending_futures)} new tasks pending...") + + # 3）Async streaming collect results and incremental checkpoint (Streaming Checkpoint) + results: dict[str, dict[int, dict[str, int]]] = {} + + # ray.wait will return when any task is completed + # file I/O is executed serially + while pending_futures: + done_refs, not_done_refs = ray.wait(list(pending_futures.keys()), num_returns=1) + + for ready_future in done_refs: + shape, image_num, exact_M = pending_futures.pop(ready_future) + config_key = shape.config_key() + try: + best_config = ray.get(ready_future) + + if best_config is None: + print( + f"⚠️ SKIPPING CHECKPOINT for image_num={image_num}, " + f"M={exact_M}, Shape={config_key} due to tuning failure. " + "Please review the worker diagnostics above." + ) + continue + + # Put the temporary results of this run into the result set + results.setdefault(config_key, {})[exact_M] = best_config + save_configs( + results={config_key: {exact_M: best_config}}, + shapes=shapes, + model_name=model_name, + tp_size=args.tp_size, + device_name=device_name, + dtype_str=dtype_str, + save_dir=args.save_dir, + ) + print(f"Checkpoint saved for image_num={image_num}, M={exact_M}, Shape={config_key}") + + except Exception as e: + print( + f"🚨 CRITICAL ERROR: Task failed for image_num={image_num}, " + f"M={exact_M}, Shape={config_key}. Error: {e}" + ) + + elapsed = time.time() - start + print(f"\nTuning completed in {elapsed:.1f}s") + print(f"Complete Configs saved to: {filepath}") + + # ---- 6. BENCHMARK mode ---- + else: + all_tasks = [] + task_keys = [] + + for shape in shapes: + for image_num in image_nums: + exact_M = get_exact_m(image_num, VAE_CHUNK_SIZE) + all_tasks.append( + ( + image_num, + shape.K, + shape.N, + dtype, + use_fp8_w8a8, + use_int8_w8a16, + ) + ) + task_keys.append((shape.config_key(), image_num, exact_M)) + + all_results = distribute("benchmark", all_tasks) + + current_key = None + for (config_key, image_num, exact_M), (config, kernel_time) in zip(task_keys, all_results): + if config_key != current_key: + current_key = config_key + shape = next(s for s in shapes if s.config_key() == config_key) + print(f"\n{'=' * 60}") + print(f"Shape: {shape}") + print(f"{'=' * 60}") + print(f" image_num={image_num:>4d} M={exact_M:>6d} {kernel_time:>8.2f} us config={config}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="MoT GEMM kernel benchmark and auto-tuning", + ) + parser.add_argument( + "--model", + type=str, + required=True, + help="HuggingFace model name or local checkpoint path", + ) + parser.add_argument( + "--tp-size", + "-tp", + type=int, + default=1, + help="Tensor parallel size (default: 1)", + ) + parser.add_argument( + "--dtype", + type=str, + default="w16a16", + choices=["w16a16", "fp8_w8a8", "int8_w8a16"], + help="Weight/activation dtype (default: w16a16)", + ) + parser.add_argument( + "--batch-size", + type=int, + nargs="+", + default=None, + help="Image counts to tune/benchmark, note M=batch_size*(VAE_CHUNK_SIZE+2) (default: 1 2 4 8 16)", + ) + parser.add_argument( + "--tune", + action="store_true", + help="Enable auto-tuning mode (search for best configs)", + ) + parser.add_argument( + "--save-dir", + type=str, + default="./", + help="Directory to save tuned config JSON (default: ./)", + ) + parser.add_argument( + "--seed", + type=int, + default=0, + help="Random seed (default: 0)", + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code when loading HuggingFace config", + ) + + args = parser.parse_args() + main(args) diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 9acf4e47e3e..2ec7addca35 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -80,6 +80,17 @@ models: task_extra_body: voice_clone: {} + # MOSS-TTS-Local-Transformer-v1.5: Qwen3 backbone + 1-layer GPT2-style local + # depth transformer (n_vq=12), paired with the 48 kHz stereo + # MOSS-Audio-Tokenizer-v2 codec. Same voice-clone request shape as the + # other MOSS-TTS variants. + OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5: + supported_tasks: [voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + task_extra_body: + voice_clone: {} + # MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset # has single-speaker rows; bench piggybacks on it (both speakers share the # same ref_audio). For real dialogue evaluation use the ttsd dataset. diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index dd55cd6036c..f4444f8accb 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,8 +1,8 @@ -# Target vllm version: v0.21.0 (releases/v0.21.0 branch). +# Target vllm version: v0.24.0 (releases/v0.24.0 branch). # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must # be updated on every rebase to match the new upstream vllm commit/tag. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.23.0 +ARG VLLM_BASE_TAG=v0.25.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index fdea270f046..4e9075f2a31 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.23.0 +ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app diff --git a/docker/Dockerfile.npu b/docker/Dockerfile.npu index 2e961b89e65..642b89882a0 100644 --- a/docker/Dockerfile.npu +++ b/docker/Dockerfile.npu @@ -1,5 +1,5 @@ -ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.18.0rc1 +ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend +ARG VLLM_ASCEND_TAG=v0.25.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.a3 b/docker/Dockerfile.npu.a3 index e3781fc18f8..4fe1d56fe27 100644 --- a/docker/Dockerfile.npu.a3 +++ b/docker/Dockerfile.npu.a3 @@ -1,5 +1,5 @@ -ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend -ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3 +ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend +ARG VLLM_ASCEND_TAG=v0.25.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} # WORKDIR /vllm-workspace/vllm diff --git a/docker/Dockerfile.npu.ci b/docker/Dockerfile.npu.ci index e6cc3d47211..a3cede3a735 100644 --- a/docker/Dockerfile.npu.ci +++ b/docker/Dockerfile.npu.ci @@ -1,11 +1,13 @@ -ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend -ARG VLLM_ASCEND_TAG=nightly-main +ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend +ARG VLLM_ASCEND_TAG=v0.25.0 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni WORKDIR ${APP_DIR} COPY . . +RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list + RUN apt-get update && \ apt-get install -y --no-install-recommends \ ffmpeg \ diff --git a/docker/Dockerfile.npu.ci.a3 b/docker/Dockerfile.npu.ci.a3 index 9ddb356f9a9..1369d8c8580 100644 --- a/docker/Dockerfile.npu.ci.a3 +++ b/docker/Dockerfile.npu.ci.a3 @@ -1,11 +1,13 @@ -ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend -ARG VLLM_ASCEND_TAG=nightly-main-a3 +ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend +ARG VLLM_ASCEND_TAG=v0.25.0-a3 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG} ARG APP_DIR=/vllm-workspace/vllm-omni WORKDIR ${APP_DIR} COPY . . +RUN sed -i "s@^\(deb.*\)https\?://[a-z0-9.-]*\.ubuntu\.com@\1http://cache-service.nginx-pypi-cache.svc.cluster.local:8081@g" /etc/apt/sources.list + RUN apt-get update && \ apt-get install -y --no-install-recommends \ ffmpeg \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index a8409b9a098..49c6f02aad4 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.23.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0 FROM ${BASE_IMAGE} AS base # Declare a variable to know if we want to use the nightly build or the stable build. diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 0dc703a9865..b70dd274e9d 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE="copy" -ARG VLLM_VERSION=v0.23.0 +ARG VLLM_VERSION=v0.25.0 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm WORKDIR /workspace/vllm @@ -101,13 +101,15 @@ RUN --mount=type=cache,target=/root/.cache/uv \ source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \ export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" -ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/" +ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/" ENV VLLM_TARGET_DEVICE=xpu ENV VLLM_WORKER_MULTIPROC_METHOD=spawn RUN --mount=type=cache,target=/root/.cache/uv \ - uv pip install --no-build-isolation . + uv pip uninstall triton triton-xpu && \ + uv pip install triton-xpu==3.7.1 && \ + uv pip install --no-build-isolation --no-deps . CMD ["/bin/bash"] @@ -143,12 +145,11 @@ COPY . . ENV VLLM_OMNI_TARGET_DEVICE=xpu RUN --mount=type=cache,target=/root/.cache/uv \ - echo "triton @ https://example.com/blocked" > /tmp/vllm_omni_constraints.txt && \ - PIP_CONSTRAINT=/tmp/vllm_omni_constraints.txt \ uv pip install --no-build-isolation \ --extra-index-url https://mirrors.ustc.edu.cn/pypi/simple/ \ ".[dev]" && \ - uv pip uninstall oneccl oneccl-devel + uv pip uninstall triton oneccl oneccl-devel && \ + uv pip install triton-xpu==3.7.1 --reinstall FROM vllm-omni AS vllm-omni-openai diff --git a/docs/.nav.yml b/docs/.nav.yml index 258224e2a7c..6bf161b1bc6 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -18,13 +18,12 @@ nav: - Offline Inference: - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md - - Helios Video Generation: user_guide/examples/offline_inference/helios.md - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md - - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md + - MiMo-Audio: user_guide/examples/offline_inference/mimo_audio.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md @@ -39,7 +38,7 @@ nav: - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md - Image-To-Image: user_guide/examples/online_serving/image_to_image.md - Image-To-Video: user_guide/examples/online_serving/image_to_video.md - - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md + - MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md - Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md @@ -69,10 +68,11 @@ nav: - Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md - Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md - Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md - - VAE Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md + - VAE Parallelism: user_guide/diffusion/parallelism/vae_parallelism.md - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md - LoRA: user_guide/diffusion/lora.md - Custom Pipeline: features/custom_pipeline.md + - Request-Level Batching: user_guide/diffusion/request_batching.md - Step Execution: user_guide/diffusion/step_execution.md - Quantization: - Overview: user_guide/quantization/overview.md @@ -80,6 +80,7 @@ nav: - Quantized KV Cache: user_guide/quantization/quantized_kvcache.md - FP8 W8A8: user_guide/quantization/fp8.md - Int8 W8A8: user_guide/quantization/int8.md + - BitsAndBytes W4: user_guide/quantization/bitsandbytes.md - MXFP8 W8A8: user_guide/quantization/mxfp8.md - MXFP4 W4A4: user_guide/quantization/mxfp4.md - ModelOpt: user_guide/quantization/modelopt.md @@ -119,6 +120,7 @@ nav: - design/feature/async_chunk.md - design/feature/vae_parallel.md - design/feature/diffusion_step_execution.md + - design/feature/diffusion_request_level_batching.md - design/feature/diffusion_continuous_batching.md - Module Design: - design/module/ar_module.md diff --git a/docs/README.md b/docs/README.md index 66fc8ef4668..f6263483b3e 100644 --- a/docs/README.md +++ b/docs/README.md @@ -28,9 +28,9 @@ Easy, fast, and cheap omni-modality model serving for everyone [vLLM](https://github.com/vllm-project/vllm) was originally designed to support large language models for text-based autoregressive generation tasks. vLLM-Omni is a framework that extends its support for omni-modality model inference and serving: -- **Omni-modality**: Text, image, video, and audio data processing +- **Omni-modality**: Text, image, audio, video, and action data processing - **Non-autoregressive Architectures**: extend the AR support of vLLM to Diffusion Transformers (DiT) and other parallel generation models -- **Heterogeneous outputs**: from traditional text generation to multimodal outputs +- **Heterogeneous outputs**: from traditional text generation to multimodal and action outputs

@@ -55,8 +55,10 @@ 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. Qwen2.5-Omni, Qwen3-Omni) -- Multi-modality generation models (e.g. Qwen-Image) +- Omni-modality models (e.g. Qwen3-Omni, Cosmos3, 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) +- Robot-policy and action models (e.g. GR00T-N1.7, DreamZero-DROID, InternVLA-A1, Cosmos3 action policy) For more information, checkout the following: diff --git a/docs/api/README.md b/docs/api/README.md index d4c58062918..cfea6d17563 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -46,17 +46,17 @@ Engine classes for offline and online inference. - [vllm_omni.engine.PromptEmbedsPayload][] - [vllm_omni.engine.arg_utils.OmniEngineArgs][] - [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][] -- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][] -- [vllm_omni.engine.mm_outputs.MultimodalPayload][] - [vllm_omni.engine.orchestrator.Orchestrator][] - [vllm_omni.engine.orchestrator.OrchestratorRequestState][] -- [vllm_omni.engine.output_modality.OutputModality][] -- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][] -- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][] -- [vllm_omni.engine.output_processor.OmniRequestState][] - [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][] - [vllm_omni.engine.stage_init_utils.StageMetadata][] - [vllm_omni.engine.stage_init_utils.StartedLlmStage][] +- [vllm_omni.outputs.mm_outputs.MultimodalCompletionOutput][] +- [vllm_omni.outputs.mm_outputs.MultimodalPayload][] +- [vllm_omni.outputs.output_modality.OutputModality][] +- [vllm_omni.outputs.output_modality.TensorAccumulationStrategy][] +- [vllm_omni.outputs.output_processor.MultimodalOutputProcessor][] +- [vllm_omni.outputs.output_processor.OmniRequestState][] ## Core diff --git a/docs/assets/WeChat.jpg b/docs/assets/WeChat.jpg index 5afd09d39c1..fc5894e7a6f 100644 Binary files a/docs/assets/WeChat.jpg and b/docs/assets/WeChat.jpg differ diff --git a/docs/cli/README.md b/docs/cli/README.md index 1fcfdb14eac..9e3240de444 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -18,11 +18,17 @@ Specify the port: vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 ``` -If you have custom stage configs file, launch the server with command below +For a migrated model, load a custom deploy configuration with `--deploy-config`: + ```bash -vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file +vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml ``` +The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml +``` ## bench diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md index d6b3ef39379..2a1acd1f055 100644 --- a/docs/cli/bench/serve.md +++ b/docs/cli/bench/serve.md @@ -2,7 +2,7 @@ The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models. ## Notes -We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend. +vLLM-Omni registers the `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` serving benchmark backends. ## Basic Parameter Description You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below: @@ -10,7 +10,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio. - `--backend` - Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported. + Specify the backend adapter. vLLM-Omni adds `openai-chat-omni`, `openai-audio-speech`, `openai-image-edits-omni`, and `daily-omni` to the upstream vLLM backend choices. - `--model` The model identifier to load, filled according to the models supported by vLLM-Omni. @@ -55,7 +55,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para 'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", ' '"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". ' -- `-print-stage` +- `--print-stage` Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default. - `--save-result` @@ -312,7 +312,7 @@ vllm bench serve \ --ignore-eos \ --percentile-metrics ttft,tpot,itl \ --random-output-len 2 \ - --extra_body '{"modalities": ["text"]}' + --extra-body '{"modalities": ["text"]}' ``` If successful, you will see the following output: @@ -390,7 +390,7 @@ vllm bench serve --omni \ --ignore-eos \ --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \ --random-output-len 900 \ - --extra_body '{"modalities": ["text","audio"]}' \ + --extra-body '{"modalities": ["text","audio"]}' \ --print-stage ``` diff --git a/docs/cli/serve.md b/docs/cli/serve.md index 9d1747e0be2..ff9e561e046 100644 --- a/docs/cli/serve.md +++ b/docs/cli/serve.md @@ -8,8 +8,7 @@ The stage-based CLI is designed for deployments that require launching each pipe - For **migrated models** that utilize the bundled deployment YAML configurations located in `vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...` automatically loads the bundled deployment configuration. -- For **legacy models** utilizing configuration files located in - `vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory. +- For custom legacy `stage_args` YAMLs, pass the file with `--stage-configs-path`. Example: Initializing Stage 0 (Orchestrator and API Server): The commands below show a common device mapping where Stage 0 uses GPU 0 and diff --git a/docs/community/governance.md b/docs/community/governance.md index 83374a155ad..3b0c62c0911 100644 --- a/docs/community/governance.md +++ b/docs/community/governance.md @@ -33,16 +33,20 @@ Lead maintainers are responsible for the overall direction and strategy of the p Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions: +- [@alex-jw-brooks](https://github.com/alex-jw-brooks): Frontend, Configuration, and Entrypoint Support +- [@amy-why-3459](https://github.com/amy-why-3459): Omni Modality Model and World Model Support - [@david6666666](https://github.com/david6666666): Quantization and Community Relationship -- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration +- [@gcanlin](https://github.com/gcanlin): Hardware plugin, NPU integration, and TTS Model Support - [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization - [@linyueqian](https://github.com/linyueqian): TTS and Omni Support - [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration - [@princepride](https://github.com/princepride): Diffusion and Omni Support - [@RuixiangMa](https://github.com/RuixiangMa): Diffusion models, parallel, cache, and docs - [@SamitHuang](https://github.com/SamitHuang): RL, Diffusion, and cache +- [@Sy0307](https://github.com/Sy0307): TTS Model and AR Engine Support - [@tzhouam](https://github.com/tzhouam): Engine and New Model Support - [@wtomin](https://github.com/wtomin): Diffusion models, parallel, and docs +- [@xuechendi](https://github.com/xuechendi): Diffusion and Intel GPU Support - [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support - [@ZJY0516](https://github.com/ZJY0516): Diffusion attention backend, kernel fusion, and CustomOp - [@yuanheng-zhao](https://github.com/yuanheng-zhao): Diffusion cache, offload, and Omni Support @@ -53,7 +57,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou ### Runtime-related -- Diffusion models: @RuixiangMa, @wtomin +- Diffusion models: @RuixiangMa, @wtomin, @xuechendi ### Optimization @@ -61,6 +65,7 @@ This section breaks down diffusion-related responsibilities. If you have PRs tou - Attention backend: @ZJY0516 - Cache: @yuanheng-zhao, @RuixiangMa, @SamitHuang - Offload: @yuanheng-zhao +- Intel GPU support: @xuechendi - Quantization: @lishunyang12, @david6666666 - Kernel fusion / communication-computation: @ZJY0516 diff --git a/docs/configuration/composable_parallel.md b/docs/configuration/composable_parallel.md new file mode 100644 index 00000000000..a456de843e8 --- /dev/null +++ b/docs/configuration/composable_parallel.md @@ -0,0 +1,239 @@ +# Composable parallel strategies + +In vLLM-Omni, the `composable_parallel` system layers a declarative *parallel strategy* on top of the existing stage configs. A strategy YAML maps each pipeline stage (by its `model_stage` name) to a stack of axis declarations — one per parallelism dimension — that the engine translates into per-stage parallel sizing (`tensor_parallel_size`, `data_parallel_size`, `pipeline_parallel_size`, expert-parallel toggle, replica count) and, when relevant, an orchestrator-wide load-balancer policy. + +Strategies are intentionally narrow: each one only writes the axes it declares, leaving every other engine arg untouched. This makes a strategy file a small, reusable overlay you can swap independently of the deploy YAML (see `configuration/stage_configs.md`). + +!!! note + Composable parallel is **opt-in**. Passing `--strategy-config ` overlays the derived sizing onto the registry-merged stages *before* any CLI overrides are applied. When `--strategy-config` is omitted, stages remain exactly as their deploy YAML produces them. + +## Top-level schema reference + +A strategy file has a single top-level mapping, `strategies`, keyed by the `model_stage` name of each stage you want to overlay. Each value is a list of axis declarations (see [StrategySpec fields](#strategyspec-fields) below). + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `strategies` | `dict[str, list]` | required | — | Maps `model_stage` name (e.g. `thinker`, `talker`, `code2wav`, `dit`) to a list of axis declarations. The key MUST be the stage name, not the integer `stage_id`. See `vllm_omni/config/composable_parallel/strategy_loader.py`. | + +### StrategySpec fields + +Each entry in a stage's list becomes one `StrategySpec` (see `vllm_omni/config/composable_parallel/spec.py`). The loader fills in `routing` and `aggregation` from per-kind defaults, so the file stays compact. + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `axis` | str | required | — | Mesh-axis kind. One of `tp`, `dp`, `pp`, `ep`, `stage_replica` (wired today) or `sp_ulysses`, `sp_ring`, `cfg`, `vae_pp`, `hsdp`, `stage_pp`, `cp` (reserved; declaring one raises `AxisTranslationError` at translate time). See `MeshAxisKind` in `vllm_omni/config/composable_parallel/spec.py`. | +| `size` | int | required | — | Axis degree. Must be `> 0`. | +| `routing` | str | optional | per-kind default | Stateless routing policy. Accepted only for `dp` / `stage_replica`; other kinds reject it. Values: `random`, `round_robin`, `least_queue`. (`hash` is recognised but currently raises `NotImplementedError`.) See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`. | +| `l1_owner` | str | optional | per-kind default | Overrides which layer owns request routing for this axis. `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`. Recorded into `shard_extension`. | +| `name` | str | optional | axis kind | Human-readable label for this spec, used in error messages. | + +Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs` — let advanced authors bind layer-walker hook IDs (`LayerHookSpec`) or kernel IDs (`KernelSpec`) onto an axis. These are not exposed by the YAML loader today and are intended for programmatic construction. See `vllm_omni/config/composable_parallel/spec.py` for the dataclasses. + +### CLI flags + +| Flag | Description | +|------|-------------| +| `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. | +| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). | +| `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. | + +!!! warning + `--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy. + +## Sub-schemas + +### `axis` + +The set of recognised axis kinds is the `MeshAxisKind` enum in `vllm_omni/config/composable_parallel/spec.py`: + +```yaml +# Translatable today (see `_SUPPORTED_KINDS` in translator.py). +axis: tp # tensor parallel — sets tensor_parallel_size +axis: dp # engine data parallel — sets data_parallel_size (vLLM's intra-engine DP LB) +axis: pp # pipeline parallel — sets pipeline_parallel_size +axis: ep # dense expert parallel — sets enable_expert_parallel=True; size must equal tp*dp +axis: stage_replica # omni-coordinator-level replica fan-out — sets num_replicas + omni_lb_policy + +# Reserved (raise AxisTranslationError today; declared in the type system for forward compat). +axis: sp_ulysses +axis: sp_ring +axis: cfg +axis: vae_pp +axis: hsdp +axis: stage_pp +axis: cp +``` + +### `routing` + +The full closed hierarchy lives in `vllm_omni/config/composable_parallel/routing.py` (`Broadcast`, `PartitionByHash`, `ShardSequence`, `ShardSpatial`, `RouteByExpert`, `PipelineMicrobatch`, `DuplicateWithCondTag`, `RouteByStage`). For YAML authors only `RouteByStage` is exposed as a `routing:` string, and only for `dp` / `stage_replica` axes (other axes have fixed routing — `tp`/`ep` broadcast, `pp` microbatch). Recognised policies: + +```yaml +routing: random # uniform-random across replicas/DP ranks +routing: round_robin # cycle through replicas/DP ranks (the default for dp/stage_replica) +routing: least_queue # send to the replica with the smallest backlog +# routing: hash # recognised but raises NotImplementedError; no key-stable balancer yet +``` + +### `aggregation` + +Aggregation is filled in automatically by the loader (`_default_aggregation` in `vllm_omni/config/composable_parallel/strategy_loader.py`) from the axis kind, so users do not author it: + +| Axis kind | Aggregation pattern | +|-----------|---------------------| +| `tp`, `ep` | `TakeRank` (all ranks agree; pick one) | +| `dp`, ep (in some flows) | `Union` (disjoint replica results) | +| `pp` | `StitchPipeline` | +| `stage_replica` | `FanInByStage` | + +Full pattern hierarchy: `TakeRank`, `Union`, `GatherDim`, `AllGather`, `StitchSpatial`, `StitchPipeline`, `Combine`, `FanInByStage` (see `vllm_omni/config/composable_parallel/aggregation.py`). + +## Precedence + +From highest to lowest, when the same axis is touched by multiple layers: + +1. **CLI overrides** (`--tensor-parallel-size`, `--stage-overrides`, `--omni-lb-policy`, etc.). A CLI value always wins; a warning is emitted whenever it overrides a strategy-declared axis (see `_reconcile_strategy_with_cli` in `vllm_omni/config/config_factory.py`). +2. **Strategy YAML** (`--strategy-config`). Derived sizing is written onto the merged stages with *conflict-on-explicit* semantics — if the deploy YAML already set a value to something different, application raises `StrategyApplyError`. +3. **Deploy YAML** (`--deploy-config` or the bundled `vllm_omni/deploy/.yaml`). +4. **Parser defaults.** + +After CLI overrides land, the device-layout guard re-runs (`check_device_layout`, called from `_reconcile_strategy_with_cli`) against the *effective* `tp * dp * pp * num_replicas` and the resolved `devices` string, so a `devices` value passed through `--stage-overrides` or a `--tensor-parallel-size` value that conflicts with the strategy's world size cannot slip past the pre-spawn check. + +## Worked example + +### Example 1 — `strategy_tp2.yaml` + +Shard the `thinker` stage with TP=2 on top of the bundled Qwen2.5-Omni deploy YAML. Routing is fixed (broadcast) for TP, so no `routing:` field is needed. + +```yaml +# examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml +strategies: + thinker: + - axis: tp + size: 2 +``` + +Launch: + +```bash +python3 end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_tp2.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +``` + +Effective `OmniParallelConfig` per stage after the overlay (see `OmniParallelConfig` in `vllm_omni/config/composable_parallel/translator.py`): + +| Stage (`model_stage`) | `tensor_parallel_size` | `data_parallel_size` | `pipeline_parallel_size` | `stage_replica_size` | `omni_lb_policy` | +|-----------------------|------------------------|----------------------|--------------------------|----------------------|------------------| +| `thinker` | **2** (from strategy) | 1 | 1 | 1 | — | +| `talker` | 1 | 1 | 1 | 1 | — | +| `code2wav` | 1 | 1 | 1 | 1 | — | + +### Example 2 — `strategy_stage_replica.yaml` + +Replicate the `thinker` stage across two independent engines and let omni's StagePool round-robin requests across them. + +```yaml +# examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml +strategies: + thinker: + - axis: stage_replica + size: 2 + routing: round_robin +``` + +Launch: + +```bash +python3 end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_stage_replica.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +``` + +Effective per-stage config and the derived orchestrator knob: + +| Stage (`model_stage`) | `tensor_parallel_size` | `stage_replica_size` | `omni_lb_policy` (derived) | +|-----------------------|------------------------|----------------------|----------------------------| +| `thinker` | 1 | **2** | **`round-robin`** | +| `talker` | 1 | 1 | — | +| `code2wav` | 1 | 1 | — | + +The derived `omni_lb_policy = round-robin` is **not** written into any stage's engine args. Instead it is surfaced on `StrategyApplyResult.omni_lb_policy` and applied once at orchestrator construction by `AsyncOmniEngine._apply_strategy_lb_policy` (see `vllm_omni/engine/async_omni_engine.py`, lines 1107–1134). If you also pass `--omni-lb-policy round-robin` the values must match; passing a different value raises `Conflicting load-balancer policy`. + +## Per-field deep-dive + +### `name` + +Default: the axis kind (e.g. `"tp"`). + +Human-readable label that appears in `AxisTranslationError` and `StrategyApplyError` messages. Set it to disambiguate two stages declaring the same axis kind when scanning logs. + +### `mesh_axis` (`axis` + `size`) + +Default: no default — both fields are required. + +`axis` chooses the parallelism dimension; `size` is its degree. Five kinds are wired end-to-end today (`tp`, `dp`, `pp`, `ep`, `stage_replica`); the remaining kinds in `MeshAxisKind` are reserved and raise `AxisTranslationError` at translate time so a typo cannot silently become a no-op (see `_SUPPORTED_KINDS` in `vllm_omni/config/composable_parallel/translator.py`). For `ep`, the size MUST equal `tensor_parallel_size * data_parallel_size`. + +### `routing` + +Default: per-kind — `RouteByStage(round_robin)` for `dp` and `stage_replica`; `PipelineMicrobatch()` for `pp`; `Broadcast()` for `tp`/`ep` and any other kind. + +User-facing string only on `dp` and `stage_replica` axes. The four recognised values are `random`, `round_robin`, `least_queue`, `hash`; `hash` is rejected today with `NotImplementedError`. See `RouteByStage` in `vllm_omni/config/composable_parallel/routing.py`. + +### `l1_owner` + +Default: per-kind — `engine` for `tp`/`dp`/`pp`/`ep`, `delegated` for `stage_replica`. + +Overrides which layer owns request routing along this axis. `engine` means vLLM realizes routing intra-engine; `delegated` means the omni coordinator owns routing across independent replicas. Set this only when overriding the per-kind default in `_DEFAULT_L1_OWNER` (`vllm_omni/config/composable_parallel/translator.py`); invalid values fail fast in `_resolve_l1_owner`. + +### `stage_replica` and `omni_lb_policy` + +`stage_replica` is *not* a vLLM world dimension; it spins up `size` independent engine replicas of one stage, coordinated by the omni `StagePool` load balancer. The axis writes both: + +- The stage's `num_replicas` (a per-stage deploy knob). +- The pipeline-wide `omni_lb_policy` string (an orchestrator-level knob). + +Because `omni_lb_policy` is read once at `AsyncOmniEngine` construction — not from per-stage engine args — only one value is allowed across the whole pipeline. If two stages declare conflicting policies, application raises (see `apply_strategy_specs` in `vllm_omni/config/composable_parallel/apply.py`, lines 253–260). To set this knob manually instead, use `--omni-lb-policy` on the CLI; see [the model stage docs](stage_configs.md) for how stage replicas plug into the orchestrator at runtime. + +### `model_stage` keys + +Default: no default — required for every entry under `strategies:`. + +Strategy keys MUST be `model_stage` *names* (strings), the same labels stage configs use for `engine_args.model_stage` (see [`model_stage` in `stage_configs.md`](stage_configs.md#engine_argsmodel_stage)). Integer `stage_id` values are not accepted; the resolver matches on `model_stage` exactly (`_resolve_stage` in `vllm_omni/config/composable_parallel/apply.py`). + +## Common pitfalls + +!!! warning "Using `stage_id` integers as keys" + Strategy keys are `model_stage` names (str), not integer stage IDs. A file that uses an integer index will raise at apply time: + + ```text + StrategyApplyError: strategy role model_stage=0 did not match any stage; + available stages: 'thinker'(id=0), 'talker'(id=1), 'code2wav'(id=2) + ``` + + Use the `model_stage` name (e.g. `thinker`) instead. + +!!! warning "Mixing CLI overrides with a strategy on the same axis" + A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time. + +!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`" + Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits: + + ```text + --strategy-config () was provided but model '' resolves via the + legacy stage_configs YAML path, which does not support composable-parallel + strategies; the strategy is ignored. Use a registry-based model to apply it. + ``` + + Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy. + +!!! warning "Conflicting `omni_lb_policy` across stages" + `omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises: + + ```text + StrategyApplyError: role 'talker' derives omni_lb_policy='least-queue-length' + but role 'thinker' already derived 'round-robin'. The load-balancer policy is + pipeline-wide; only one value is allowed. + ``` + + Make the policies match, drop one, or supply the value explicitly via `--omni-lb-policy`. diff --git a/docs/configuration/pd_disaggregation.md b/docs/configuration/pd_disaggregation.md index 9196bdb0240..7deb0443ade 100644 --- a/docs/configuration/pd_disaggregation.md +++ b/docs/configuration/pd_disaggregation.md @@ -123,7 +123,7 @@ After inserting the extra thinker stage, renumber the remaining stages: runtime: devices: "2" engine_input_source: [1] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker + sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_token_only - stage_id: 3 runtime: @@ -131,7 +131,7 @@ After inserting the extra thinker stage, renumber the remaining stages: engine_args: max_num_seqs: 1 engine_input_source: [2] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav + sync_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_token_only ``` Compared with the default Qwen3-Omni config: diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 3a17f1c4013..1aa4ca1b329 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -1,11 +1,11 @@ -# Stage configs for vLLM-Omni +# Deploy and legacy stage configurations -In vLLM-Omni, the target model is separated into multiple stages, which are processed by different LLMEngines, DiffusionEngines or other types of engines. Depending on different types of stages, such as Autoregressive (AR) stage or Diffusion transformer (DiT) stage, each can choose corresponding schedulers, model workers to load with the Engines in a plug-in fashion. +In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema. !!! note - Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. The bundled defaults have been verified on 1xH100 for Qwen2.5-Omni and 2xH100 for Qwen3-Omni. Models that have not yet migrated to the new schema continue to use the legacy `vllm_omni/model_executor/stage_configs/.yaml` files via `--stage-configs-path`. + Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. Custom legacy `stage_args` YAMLs are still accepted via `--stage-configs-path`. -## New deploy schema reference +## Deploy configuration schema The new deploy schema lives under `vllm_omni/deploy/` and is paired with a frozen `PipelineConfig` registered by the model's `pipeline.py`. Each deploy YAML has these top-level fields: @@ -82,14 +82,14 @@ stages: input_connectors: {from_stage_0: shm} ``` -### CLI flags introduced in this refactor +### CLI flags | Flag | Description | |------|-------------| -| `--deploy-config PATH` | Load a new-schema deploy YAML. Takes precedence over `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/.yaml` is auto-loaded by the model registry. | +| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/.yaml` is auto-loaded by the model registry. | | `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. | | `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. | -| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. | +| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. | ### Stage-Based CLI Paradigm @@ -146,7 +146,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \ From highest to lowest: -1. Per-stage flags (`--stage-overrides` JSON, `--stage--` if registered) +1. Per-stage overrides (`--stage-overrides` JSON) 2. Explicit global CLI flags (`--gpu-memory-utilization 0.85`, etc.) 3. Platform section (`platforms.npu.stages`, etc.) on top of the base `stages:` 4. Overlay YAML (via `base_config:`) on top of the base YAML @@ -210,7 +210,7 @@ Effective config per stage after the merge: | 1 | `max_model_len` | `16384` | global CLI | | 2 | (all defaults) | — | base YAML (no overrides apply) | -Therefore, as a core part of vLLM-Omni, the stage configs for a model have several main functions: +Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment configurations have several main functions: - Claim partition of stages and their corresponding class implementation in `model_executor/models`. - The disaggregated configuration for each stage and the communication topology among them. @@ -219,16 +219,16 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever - Default input parameters. To override specific parameters, explicitly inject the customized configuration schema -in both online and offline instantiation flows. Prioritize the `--deploy-config` flag -when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter +in both online and offline instantiation flows. Use the `--deploy-config` flag +when loading a deploy configuration, reserving the `--stage-configs-path` parameter exclusively to maintain compatibility with legacy `stage_args` YAML constructs. Examples: -For offline (Assume necessary dependencies have been imported): +For offline inference (assuming the necessary dependencies have been imported): ```python model_name = "Qwen/Qwen2.5-Omni-7B" -omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml") +omni = Omni(model=model_name, deploy_config="/path/to/deploy_config.yaml") ``` For online serving: @@ -242,11 +242,14 @@ Legacy online serving: vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file ``` !!! important - We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development! + We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers. + +## Legacy stage configuration reference + +The remainder of this page documents the deprecated `stage_args` schema for models that have not migrated to deploy configurations. The following historical Qwen2.5-Omni configuration illustrates that schema; current Qwen2.5-Omni deployments should use `vllm_omni/deploy/qwen2_5_omni.yaml` instead. -Below is a specific example of stage_configs.yaml in Qwen2.5-omni. ```python -# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime. +# Historical stage configuration for Qwen2.5-Omni. stage_args: - stage_id: 0 # mark the unique id for each stage runtime: # The disaggregated configuration @@ -290,7 +293,6 @@ stage_args: enable_prefix_caching: false engine_output_type: latent engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker default_sampling_params: temperature: 0.9 top_p: 0.8 @@ -339,7 +341,7 @@ runtime: ``` -## Stage Configuration Arguments +## Legacy stage configuration arguments Each stage in the `stage_args` list contains the following configuration options: @@ -379,6 +381,20 @@ The maximum number of sequences for concurrent processing in this stage. For LLM Default: `1` +#### `engine_args.request_batch_max_wait_ms` + +The maximum time, in milliseconds, that a diffusion request-mode stage may wait +before the first `schedule()` of a new scheduler wave so compatible requests can +accumulate for request-level batching. This only applies to diffusion pipelines +that support request-level batching with `step_execution` disabled. + +Use this together with `max_num_seqs > 1` for bursty serving traffic. `0` +disables admission waiting and preserves the lowest first-request latency. +For diffusion request-level batching tuning, see +[Request-Level Batching](../user_guide/diffusion/request_batching.md). + +Default: `0.0` + ### `engine_args` Engine arguments for configuring the LLM engine, diffusion engine, or other engine types used by this stage. diff --git a/docs/contributing/DOCS_GUIDE.md b/docs/contributing/DOCS_GUIDE.md index 100bac67423..916f5e60fea 100644 --- a/docs/contributing/DOCS_GUIDE.md +++ b/docs/contributing/DOCS_GUIDE.md @@ -67,6 +67,21 @@ docs/ └── stylesheets/ # Custom CSS ``` +## Naming Model Examples + +Offline inference and online serving examples for the same model use the same +directory name and the shared display name in +`examples/model_display_names.yml`. Use these title forms: + +- `# : Offline inference` +- `# : Online serving` + +The documentation generator uses the full title for the page H1 and the shared +display name alone for navigation, and fails the build if a mapped README uses +a different H1. Keep checkpoint identifiers in commands and prose rather than +in the display name, and do not rename an existing example directory solely to +adjust its title because the directory defines its public documentation URL. + ## Publishing Documentation ### GitHub Pages (Recommended) diff --git a/docs/contributing/README.md b/docs/contributing/README.md index e10d8d3b33c..63a62eabf92 100644 --- a/docs/contributing/README.md +++ b/docs/contributing/README.md @@ -131,7 +131,7 @@ Please run the L1 and L2 test cases locally first and attach the results before ### Automatic skip-ci (docs and pytest skip marks) -On pull requests and `main` pushes, the bootstrap step in [`.buildkite/pipeline.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/pipeline.yml) runs [`.buildkite/scripts/upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/scripts/upload_pipeline.py) against the git diff. When every changed file qualifies, **L2 (`ready`) and L3 (`merge-test`) pipelines are not uploaded**, so the default GPU CI jobs are skipped. +On pull requests and `main` pushes, the bootstrap step in [`.buildkite/cuda/pipeline.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/pipeline.yml) runs [`.buildkite/cuda/scripts/upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/scripts/upload_pipeline.py) against the git diff. When every changed file qualifies, **L2 (`ready`) and L3 (`merge-test`) pipelines are not uploaded**, so the default GPU CI jobs are skipped. | Change per file | Examples | | --- | --- | diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 01310606eb5..d3c164d9b90 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -11,6 +11,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Level Scope & Focus + Model Coverage Strategy + Feature Coverage Strategy + Interface Coverage Strategy + Tags Time Cost Test Dir Doc @@ -24,6 +28,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Contribution Guideline & PR checklist / / + / + / + / + / .github/PULL_REQUEST_TEMPLATE.md Test Style (PR Checklist) / / @@ -32,6 +40,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari CI Failure Description / / + / + / + / + / CI Failures / / @@ -39,6 +51,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L1
(Unit & Logic) Unit tests for components like entrypoints, models + / + / + / + core_model and cpu <15min /tests/{component_name}/test_xxx @@ -50,10 +66,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L2
(E2E across models & GPU-required UT) - Online & Offline (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT + Online (basic deployment scenarios):
dummy, normal inference function (output format, stream), some instance startup UT + High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy) + High-priority features (using random lightweight models) + High-priority interfaces (using random lightweight models) + core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion + Model tests:
/tests/e2e/online_serving/test_{model_name}.py
- /tests/e2e/offline_inference/test_{model_name}.py + Feature tests:
+ /tests/{component_name}/test_xxx
+ Interface tests:
+ /tests/entrypoints/test_xxx Chapter 1
@@ -65,10 +89,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L3
(Important Perf & Integration & Accuracy) Online & Offline (multiple deployment scenarios):
real model, normal inference function, normal accuracy + High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy + Medium-priority features (using random lightweight models) + Medium-priority interfaces (using random lightweight models) + advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <30min + Model tests:
/tests/e2e/online_serving/test_{model_name}.py
- /tests/e2e/offline_inference/test_{model_name}.py + /tests/e2e/offline_inference/test_{model_name}.py
+ Feature tests:
+ /tests/{component_name}/test_xxx
+ Interface tests:
+ /tests/entrypoints/test_xxx Chapter 2
@@ -79,18 +112,28 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L4
(Perf & Integration & Accuracy) - Online & Offline: full functional scenarios + performance test + doc test + Online: full functional scenarios + performance test + doc test + accuracy test + High-priority models: function, performance, accuracy, and doc testing
Medium-priority models: function and doc testing + Low-priority features (using real weights) + Low-priority interfaces (using real weights) + full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion <3 hour - Full Function:
+ Model tests:
/tests/e2e/online_serving/test_{model_name}_expansion.py
- /tests/e2e/offline_inference/test_{model_name}_expansion.py
+ Feature tests:
+ /tests/{component_name}/test_xxx
+ Interface tests:
+ /tests/entrypoints/test_xxx
Performance:
- /tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),
- and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
+ /tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),
+ test_voxcpm2.json, test_higgs_audio_v3.json, and
+ /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
Doc Test:
- tests/example/online_serving/test_{model_name}.py
- tests/example/offline_inference/test_{model_name}.py + tests/examples/online_serving/test_{model_name}.py
+ tests/examples/offline_inference/test_{model_name}.py
+ Accuracy Test:
+ /tests/e2e/accuracy/test_{model_name}.py Chapter 3
@@ -101,7 +144,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari L5
(Stability & Reliability) - Online & Offline: long-term stability test + reliability test + Online: long-term stability test + reliability test + Long-term stability and reliability testing for high-priority models
Low-priority models: function and doc testing + / + Invalid-parameter validation for high-priority interfaces + slow and hardware_test(H100, L4, etc.) and omni/tts/diffusion Depends on reality Stability:
@@ -109,7 +156,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari /tests/dfx/stability/tests/test_wan22.json
Reliability:
tests/dfx/reliability/test_reliability_{model_key}.py
- (e.g. test_reliability_qwen3_omni.py, test_reliability_wan22.py) + (e.g. test_reliability_qwen3_omni.py, test_reliability_wan22.py, test_reliability_hunyuan_image.py, test_reliability_voxcpm2.py) Chapter 4
@@ -272,7 +319,7 @@ Before entering specific testing levels, the project establishes two common spec ### 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`. +L2 (`.buildkite/cuda/test-ready.yml`) and L3 (`.buildkite/cuda/test-merge.yml`) pipelines can **skip unrelated GPU jobs at upload time** based on the PR diff. This is implemented by `.buildkite/cuda/scripts/upload_pipeline.py`, which filters steps before calling `buildkite-agent pipeline upload`. #### What `source_file_dependencies` is @@ -287,14 +334,14 @@ L2 (`.buildkite/test-ready.yml`) and L3 (`.buildkite/test-merge.yml`) pipelines | `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. +Docs-only PRs are handled earlier in bootstrap (`.buildkite/cuda/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`) | +| L2 | `.buildkite/cuda/test-ready.yml` | `upload_pipeline.py --upload .buildkite/cuda/test-ready.yml` (from `cuda/pipeline.yml`) | +| L3 | `.buildkite/cuda/test-merge.yml` | `upload_pipeline.py --upload .buildkite/cuda/test-merge.yml` (from `cuda/pipeline.yml`) | Steps **without** `source_file_dependencies` are always uploaded (subject to the usual label conditions: `ready` for L2, `merge-test` for L3). @@ -346,10 +393,10 @@ A **group** may also define `source_file_dependencies`; nested steps inherit fil ```bash # Render filtered YAML to stdout (no upload) -python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-ready.yml +python3 .buildkite/cuda/scripts/upload_pipeline.py .buildkite/cuda/test-ready.yml # Confirm uploader-only keys are stripped -python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-merge.yml | grep source_file_dependencies +python3 .buildkite/cuda/scripts/upload_pipeline.py .buildkite/cuda/test-merge.yml | grep source_file_dependencies # (no output expected) ``` @@ -357,7 +404,7 @@ On a PR build, Buildkite logs from `upload_pipeline.py` include lines such as `s #### Related -- Implementation: `.buildkite/scripts/upload_pipeline.py` +- Implementation: `.buildkite/cuda/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 @@ -400,7 +447,11 @@ L2 level testing builds upon L1 by introducing GPU resources and verifying that - ***Time Cost***: Execution time is controlled within ***15 minutes*** to ensure fast feedback. - ***L2 (Basic End-to-End Testing)***: - - ***Scope***: Covers two basic deployment scenarios: `online` (serving) and `offline` (inference). - - ***Focus***: Uses `dummy` models or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, etc. Also includes some unit tests that require launching independent service instances. + - ***Focus***: Uses `dummy` weights (via deploy YAML patching at `core_model`) or lightweight real models to verify that the entire chain from request input to result output works normally, including output data structure, streaming (stream) support, and **cheap payload checks** at `--run-level core_model` (see below). Also includes some unit tests that require launching independent service instances. + - ***L2 response validation (`core_model`)***: Implemented in `tests/helpers/assertions.py` and invoked by `OpenAIClientHandler` based on `--run-level`. At L2 we require **request success** plus minimal output sanity—not full accuracy: + - **Speech / TTS** (`assert_audio_speech_response`): decoded audio must be present and non-empty (or exceed `min_audio_bytes` when set in `request_config`); `response_format` must match the returned content-type (e.g. `wav`, `pcm`). Whisper transcript similarity, PCM HNR, and preset-voice gender checks run only at L3+. + - **Diffusion** (`assert_diffusion_response`): at least one non-empty image, video, or audio artifact. Resolution/frame-count parity and other deep checks run only at L3+. + - **Omni multimodal** (`assert_omni_response`): L2 asserts successful completion; keyword, transcript, and cross-modal similarity checks run only at L3+. - ***Characteristic***: Requires ***GPU*** resources to perform model computations. ### 1.3 Test Directory and Execution Files @@ -509,6 +560,11 @@ L3 level testing executes after code is merged into the main branch. Its core pu - 1. ***Inference Functionality***: Ensures real models can perform forward computation normally and return results. 2. ***Accuracy Compliance***: Verifies that the model's evaluation metrics (e.g., accuracy) meet the expected baseline, preventing code changes from introducing accuracy issues. 3. ***Important Performance***: Verifies whether performance (e.g., P99 latency, throughput) in core scenarios meets preset thresholds. +- ***L3 response validation (`advanced_model`)***: At `--run-level advanced_model`, `tests/helpers/assertions.py` adds semantic checks on top of L2 payload gates (see Chapter 1): + - **Speech / TTS** (`assert_audio_speech_response`): Whisper transcript of returned audio vs `request_config["input"]` (cosine similarity > 0.9 when input text is set); optional `min_audio_bytes` floor; PCM harmonic-to-noise ratio when `response_format` is `pcm`; **preset voice gender** via `_assert_preset_voice_gender_from_audio` when `voice` matches a known preset in `_PRESET_VOICE_GENDER_MAP` (pitch/F0-based classifier; skipped for unknown voices or `pcm` output). + - **Omni multimodal** (`assert_omni_response`): non-empty text/audio outputs; `key_words` in transcript or text; text–audio similarity / containment; preset **`speaker`** gender check (same helper as TTS). + - **Diffusion** (`assert_*_diffusion_response`): image/video dimension and frame-count parity with request parameters where configured. + - **Accuracy suites**: pixel/video similarity and metric baselines under `/tests/e2e/accuracy/` (separate from inline helper assertions). ### 2.3 Test Directory and Execution Files @@ -536,9 +592,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Explanation**: - @pytest.mark.advanced_model: Marks the test as L3 merge level, indicating deep validation with real models. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples). + @pytest.mark.advanced_model: Marks the test as L3 merge level (`--run-level advanced_model`): real weights, Whisper/similarity/keyword checks, diffusion deep checks, and **preset voice/speaker gender** validation where applicable. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples). - @pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly. + @pytest.mark.core_model: Marks the test as L1 or L2 level. At `--run-level core_model`, validation is limited to request success plus cheap payload checks (e.g. non-empty audio bytes, response format/content-type, non-empty diffusion outputs)—not Whisper, keyword, or accuracy gates. Deploy YAML may use `load_format: dummy` for fast PR feedback. @pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing. @@ -634,7 +690,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu **Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n. - **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. + **Implicit Validation**: `send_omni_request`, `send_audio_speech_request`, and `send_diffusion_request` call `assert_*_response` helpers with the session `--run-level`. At `core_model`, checks are limited to success plus cheap payload sanity (non-empty audio/media, format/content-type, optional `min_audio_bytes`). At `advanced_model` and `full_model`, deep validation adds Whisper transcripts, keyword/similarity, diffusion dimension checks, PCM HNR, etc. **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)). @@ -653,16 +709,16 @@ L4 level testing is a comprehensive quality audit before a version release. It e ### 3.2 Testing Content and Scope - ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling. -- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs in the form `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` via `--test-config-file`) to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization. +- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen3_omni_*.json` (Omni), `test_tts.json` / `test_voxcpm2.json` / `test_higgs_audio_v3.json` (TTS), and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` or `run_diffusion_benchmark.py` via `--test-config-file` in nightly **Perf Test** steps) to drive throughput, latency, and memory benchmarks. Each JSON **case** may declare an optional top-level **`mark`** array: exactly one ``hardware_marks`` object plus pytest marker name strings (`full_model`, `omni` / `tts` / `diffusion`, …). Runners attach those marks to each parametrized `(server, benchmark index)` pair so **local** bulk runs can filter with `-m` (for example `-m "full_model and H100 and diffusion"`). Nightly CI perf jobs select workloads by **`--test-config-file`**, not `-m`. Details are in the Performance Tests example below. - ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description. ### 3.3 Test Directory and Execution Files - ***Functional Testing***: Same directories as L3. -- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`) +- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen3_omni_*.json`, `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`, `test_cosmos3_vllm_omni.json`). Optional per-case `mark` for local `-m` filtering is documented in Section 3.4 Performance Tests. - ***Documentation Example Tests***: -- - `tests/example/online_serving/test_{model_name}.py` - - `tests/example/offline_inference/test_{model_name}.py` + - `tests/examples/online_serving/test_{model_name}.py` + - `tests/examples/offline_inference/test_{model_name}.py` ### 3.4 Execution Method and Example @@ -695,14 +751,16 @@ L5 level testing focuses on the performance of model services under ***long-runn ### 4.2 Testing Content and Scope - ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable. -- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `

` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill). +- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `
` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill). ### 4.3 Test Directory and Execution Files - ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family) - ***Reliability Test Suite*** (`tests/dfx/reliability/`): - `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default). - - `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery). + - `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery). + - `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4). + - `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1). - `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`). - `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready. - `README.md` — Short local run commands for this directory. @@ -800,22 +858,33 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in | `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). | | `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). | | `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). | -| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). | +| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). | +| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. | +| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. | | `README.md` | Minimal run / collect examples. | #### Parametrization and markers - Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`. -- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card). +- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4). - **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**. #### CI trigger -Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example: +Weekly Buildkite (`.buildkite/cuda/test-weekly.yml`) runs one step per model suite (trigger: `WEEKLY=1` or PR label `weekly-test`), for example: + +| Buildkite step | Test file | CI hardware | +| -------------- | --------- | ----------- | +| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) | +| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) | +| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) | +| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) | ```bash pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow" pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" +pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow" +pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow" ``` #### Local commands @@ -824,6 +893,8 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow" pytest --collect-only tests/dfx/reliability pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow +pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow +pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow ``` #### Adding a new model suite @@ -831,12 +902,12 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow 1. Add `test_reliability_.py` under `tests/dfx/reliability/`. 2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files). 3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML. -4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/test-weekly.yml`** when the suite should run in weekly L5. +4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/cuda/test-weekly.yml`** when the suite should run in weekly L5.
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_.py` alongside a matching JSON config) - - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_.py` for new models) + - ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_.py` and a matching step in `.buildkite/cuda/test-weekly.yml` for new models) ## Summary diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md index f1f3073dc52..5c59ba49959 100644 --- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md @@ -1,14 +1,41 @@ -When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, or diffusion configs such as `tests/dfx/perf/tests/test_*_vllm_omni.json` (selected via `pytest ... run_benchmark.py --test-config-file `): +When you want to add L4-level ***performance test*** cases, add entries to JSON files under `tests/dfx/perf/tests/` and run them via `tests/dfx/perf/scripts/run_benchmark.py` (omni / TTS) or `run_diffusion_benchmark.py` (diffusion). + +**Config file layout (in-tree examples)** + +| Model type | Runner | Example JSON files | +| ---------- | ------ | ------------------ | +| Omni | `run_benchmark.py` | `test_qwen3_omni_no_async_chunk.json`, `test_qwen3_omni_async_chunk.json`, `test_qwen3_omni_vllm_text.json`, `test_qwen3_omni_multi_replicas.json` | +| TTS | `run_benchmark.py` | `test_tts.json`, `test_voxcpm2.json`, `test_higgs_audio_v3.json` | +| Diffusion | `run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`, `test_wan22_i2v_vllm_omni.json`, `test_cosmos3_vllm_omni.json`, … | + +**How runners pick cases** + +Without **`--test-config-file`**, each runner scans all `*.json` under `tests/dfx/perf/tests/` but only keeps its own model type: + +- **`run_benchmark.py`**: omni and TTS cases only (skips diffusion JSON). +- **`run_diffusion_benchmark.py`**: diffusion cases only (skips omni / TTS JSON). + +Diffusion cases are detected when the JSON has `server_type` (typically `"vllm-omni"`) or `"diffusion"` in the `mark` array. Omni / TTS JSON has neither. + +**Running perf cases** + +Pass **`--test-config-file`** to run one JSON file as-is, or omit it for the bulk scan above and filter with pytest **`-m`** on each case's JSON `mark` (for example `pytest … run_benchmark.py -m "full_model and tts and H100"`). ```JSON { - "test_name": "test_qwen3_omni", + "test_name": "test_qwen3_omni_async_chunk", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, + "full_model", + "omni" + ], "server_params": { "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "stage_config_name": "qwen3_omni.yaml" }, "benchmark_params": [ { + "name": "random_10p", "dataset_name": "random", "num_prompts": [10, 20], "max_concurrency": [1, 4], @@ -33,8 +60,95 @@ When you want to add L4-level ***performance test*** cases, you can refer to the | Field | Required | Description | | ---------------- | -------- | --------------------------------------------------------------- | | test_name | Yes | Unique identifier for the test case | +| mark | No | Pytest marks for this case (see **`mark` field** below). Omit only for configs not meant to be filtered by `-m`. | | server_params | Yes | Server-side configuration parameters | | benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) | +| server_type | Diffusion only | When set (e.g. `"vllm-omni"`), routes the case to `run_diffusion_benchmark.py`. | +| benchmark_endpoint | Diffusion only | API path for the benchmark client (e.g. `/v1/videos`, `/v1/images/generations`). | + +**`mark` field** + +Optional top-level field on each perf JSON **case object** (one per `test_name`). `run_benchmark.py` and `run_diffusion_benchmark.py` read it via `tests.dfx.conftest.resolve_pytest_marks` and attach the marks to the corresponding `pytest.param`, so you can filter locally with `-m`. + +When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same shape as `@hardware_test` / `hardware_marks()` in `tests/helpers/mark.py`), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, `diffusion`, or `local_model`. + +Supported form: + +| Form | Example | Effect | +| ---- | ------- | ------ | +| Array (single platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "diffusion"]` | Hardware marks from `hardware_marks(...)`; `num_cards` > 1 adds `distributed_*` (+ CUDA `skipif_cuda` when applicable). String entries become extra pytest marks. | +| Array (multi-platform) | `"mark": [{"hardware_marks": {"res": {"cuda": "H100", "rocm": "MI325", "npu": "A2"}, "num_cards": {"cuda": 2, "rocm": 2, "npu": 1}}}, "full_model", "omni"]` | Same as above, but declares **multiple platforms** in one case. Filter examples: `-m "full_model and H100 and cuda"`, `-m "full_model and MI325 and rocm"`. | + +Recommended for L4 perf cases: + +```JSON +{ + "test_name": "test_bagel_single_device_single_stage_t2i", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, + "full_model", + "diffusion" + ], + "server_type": "vllm-omni", + "server_params": { "...": "..." }, + "benchmark_params": [ { "name": "1024x1024_steps20", "...": "..." } ] +} +``` + +Multi-GPU diffusion (example: Cosmos3 with `cfg-parallel-size=2`): + +```JSON +{ + "test_name": "test_cosmos3_t2i_official_demo_2gpu", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, + "full_model", + "diffusion" + ], + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/images/generations", + "server_params": { "...": "..." }, + "benchmark_params": [ { "name": "1024x1024_steps4", "...": "..." } ] +} +``` + +- Omni perf: include `"omni"` in the `mark` array +- TTS perf: include `"tts"` in the `mark` array +- Diffusion perf: include `"diffusion"` in the `mark` array +- HunyuanImage local-weight cases: add `"local_model"` to the `mark` array + +**Parametrization IDs** + +Each `(server, benchmark index)` pair becomes one `pytest.param` with id `{test_name}-{suffix}`. The suffix comes from `benchmark_params[].name` when set (for example `test_tts-p0`, `test_omni-p1`); otherwise it is derived from fields like `task` / `eval_phase`. + +**Benchmark result filenames** + +Result files use the **runtime** hardware label from `get_runtime_resource_label()` (detected GPU/NPU name on the machine that ran the job), **not** `mark.hardware_marks.res`. On the default H100 CI pool, `H100` is **omitted** from filenames (`resource_label_for_filename`). + +Examples: + +- Omni/TTS: `result_{test_name}_{optional_hw}_{dataset}_....json` under `BENCHMARK_DIR` +- Diffusion: one aggregate `diffusion_result_{config_stem}_{optional_hw}_{timestamp}.json` per source JSON file (array of all runs from that file) + +**Local commands** + +```bash +# Bulk load + filter by JSON mark +pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and H100 and diffusion" +pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "full_model and omni and H100" + +# Single file (same as nightly CI Perf steps) +pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \ + --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json +pytest -s -v tests/dfx/perf/scripts/run_benchmark.py \ + --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json + +# Optional baseline assertion (default off) +pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --assert-baseline \ + --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +``` + +See also [Markers for Tests](./tests_markers.md) for registered hardware markers (`H100`, `L4`, `cuda`, `distributed_cuda`, …). **`server_params` Configuration** @@ -79,10 +193,12 @@ You can add any benchmark running parameters you need here. For all optional par 1. Change the --xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file. 2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file. 3. Optionally add a `baseline` object (see **Baseline thresholds** below). If you omit `baseline` or leave it empty, the performance test still runs but does not assert metric thresholds from this field. -4. The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below: +4. Set `"name"` on each `benchmark_params` entry for stable pytest ids and readable result keys. +5. The qps and concurrency modes are recommended to be mutually exclusive. For detailed explanations, see the table below: | Parameter | Type | Required | Example/Values | Description | | --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| name | string | No | `"1024x1024_steps4"` | Stable suffix for pytest param id and logs. | | num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. | | request_rate | float / array | No | 0.5, [0.5, 1, inf] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. | | max_concurrency | int / array | No | 1, [1, 2, 3] | Maximum concurrent in-flight requests. Same array / expansion rules as `request_rate` (mutually exclusive with QPS mode). | diff --git a/docs/contributing/ci/test_guide.md b/docs/contributing/ci/test_guide.md index 36dae445103..d1179e8320d 100644 --- a/docs/contributing/ci/test_guide.md +++ b/docs/contributing/ci/test_guide.md @@ -12,76 +12,155 @@ When running test cases, you may need to install the following dependencies: ```bash uv pip install ".[dev]" -apt-get install -y ffmpeg +apt-get install -y espeak-ng jq ``` ## Running Tests Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases. +### CI job runners (L2–L4): logs, timing, and timeouts + +[`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh), [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh), and [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh) share the same run-time behavior (via [`tools/run_jobs_common.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_jobs_common.sh)): + +**Log layout** (default under `logs/ready_jobs`, `logs/merge_jobs`, or `logs/nightly_jobs`; override with `--log-dir`): + +| Path | Description | +|------|-------------| +| `jobs/*.sh` | Generated per-step bash wrappers | +| `jobs/.job_timeouts` | Job key → `timeout_in_minutes` from Buildkite YAML (when set) | +| `.log` | Tee output for each job | +| `timing_summary.log` | Per-job duration and total wall time after the run completes | + +**Timing:** while each job runs, stderr prints `finished in …` or `failed after …`. When all jobs finish, a summary is printed and written to `timing_summary.log`, for example: + +``` +=== Job timing summary === + TTS_VoxCPM2_Test 25m 40s OK + Omni_Qwen3-Omni_Test 8m 05s TIMED OUT +Total wall time: 33m 45s (2 jobs) +Failed jobs: 1/2 +``` + +**Timeouts:** if a Buildkite step defines `timeout_in_minutes`, the runner wraps that job with `timeout ${N}m` (aligned with CI). Steps without `timeout_in_minutes` rely on any inline `timeout …` already present in the extracted pytest command. On timeout the process exits with code `124` and the summary marks the job as `TIMED OUT`. + === "L1 level" ```bash cd tests pytest -s -v -m "core_model and cpu" ``` - The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml). + The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-ready.yml). === "L2 level" + **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_ready_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_ready_jobs.sh). The script reads [`.buildkite/cuda/test-ready.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-ready.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/ready_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts) for timing and timeout behavior). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`). + ```bash - cd tests - pytest -s -v -m "core_model and not cpu" --run-level=core_model + # All L2 and L1 jobs (default: every pytest step in test-ready.yml) + bash tools/run_ready_jobs.sh + + # Skip L1-style "Simple Test" steps in test-ready.yml (group or "Simple ·" labels) + bash tools/run_ready_jobs.sh --skip-simple + bash tools/run_ready_jobs.sh --skip-simple --model-type diffusion + + # Preview extracted commands without running + bash tools/run_ready_jobs.sh --dry-run + + # Filter by model area (OR semantics): omni | tts | diffusion | all + bash tools/run_ready_jobs.sh --model-type omni + bash tools/run_ready_jobs.sh --model-type tts,diffusion + + # Match a Buildkite label substring (e.g. one model suite) + bash tools/run_ready_jobs.sh --model-type tts --label-substr "VoxCPM2" + ``` - If you only want to run a specific test case, you can use: + + **Ad hoc:** run a single test file or marker expression manually: + ```bash + cd tests pytest -s -v test_xxxx.py --run-level=core_model + pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model ``` - If you only want to run specific test cases on a particular platform, you can use: - ```bash - pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model - ``` - The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml). === "L3 level" + **Recommended:** run CI-aligned jobs from the repo root via [`tools/run_merge_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/run_merge_jobs.sh). The script reads [`.buildkite/cuda/test-merge.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-merge.yml), generates per-step bash wrappers, runs pytest, and tees logs under `logs/merge_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts)). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`). + ```bash - pytest -s -v -m "advanced_model" --run-level=advanced_model + # All L3 and L1 jobs (default: every pytest step in test-merge.yml) + bash tools/run_merge_jobs.sh + + # Skip L1-style "Simple Test" steps in test-merge.yml (group or "Simple ·" labels) + bash tools/run_merge_jobs.sh --skip-simple + bash tools/run_merge_jobs.sh --skip-simple --model-type omni + + # Preview extracted commands without running + bash tools/run_merge_jobs.sh --dry-run + + # Filter by model area (OR semantics): omni | tts | diffusion | all + bash tools/run_merge_jobs.sh --model-type diffusion + bash tools/run_merge_jobs.sh --model-type omni,tts + + # Match a Buildkite label substring + bash tools/run_merge_jobs.sh --model-type diffusion --label-substr "Wan22" + ``` - If you only want to run a specific test case, you can use: + + **Ad hoc:** run a single test file or marker expression manually: + ```bash pytest -s -v test_xxxx.py --run-level=advanced_model + pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model ``` - If you only want to run specific test cases on a particular platform, you can use: - ```bash - pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model - ``` - The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml). === "L4 level" + **Recommended:** run CI-aligned nightly jobs from the repo root via [`tools/nightly/run_nightly_jobs.sh`](https://github.com/vllm-project/vllm-omni/blob/main/tools/nightly/run_nightly_jobs.sh). The script reads [`.buildkite/cuda/test-nightly.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-nightly.yml), generates per-step bash wrappers, runs pytest (perf jobs first, then aggregates perf JSON into Excel), and tees logs under `logs/nightly_jobs/` (see [CI job runners](#ci-job-runners-l2l4-logs-timing-and-timeouts); `generate_nightly_perf_excel` is included in the timing summary). Requires `bash`, `python3`, and PyYAML (`pip install pyyaml`). + ```bash - cd tests - pytest -s -v -m "full_model" --run-level=full_model + # All nightly jobs (default: test-type all, model-type all) + bash tools/nightly/run_nightly_jobs.sh + + # Preview extracted commands without running + bash tools/nightly/run_nightly_jobs.sh --dry-run + + # Test kind (--test-type, OR semantics): perf | acc | function | stability | local | all + bash tools/nightly/run_nightly_jobs.sh --test-type function + bash tools/nightly/run_nightly_jobs.sh --test-type perf,acc + + # Model area (--model-type, OR semantics): omni | tts | diffusion | all + bash tools/nightly/run_nightly_jobs.sh --test-type function --model-type omni + bash tools/nightly/run_nightly_jobs.sh --test-type perf --model-type diffusion + + # Stability scripts (tests/dfx/stability/scripts/) and local marker tests + bash tools/nightly/run_nightly_jobs.sh --test-type stability --model-type omni + bash tools/nightly/run_nightly_jobs.sh --test-type local --model-type tts ``` - If you only want to run a specific test case, you can use: + + Perf steps (label contains ``Perf Test``) run before other jobs; Excel output is written under ``logs/``. Diffusion nightly steps may also appear in [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml) on CI. + + **Ad hoc:** run a single test file or marker expression manually: + ```bash + cd tests pytest -s -v test_xxxx.py --run-level=full_model + pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model + pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model ``` If you only want to run specific test cases on a particular platform, you can use: ```bash pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model ``` - Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS): - ```bash - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py - ``` - The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example: - + Note: ``run_benchmark.py`` and ``run_diffusion_benchmark.py`` accept an optional ``--test-config-file``. If omitted, each loads every ``*.json`` under ``tests/dfx/perf/tests/`` (omni/tts vs diffusion split by ``is_diffusion_perf_config``) and pytest ``-m`` filters by each case's JSON ``mark``: ```bash - cd tests - pytest -s -v -m "full_model and (omni or tts) and H100" --run-level=full_model + pytest -sv tests/dfx/perf/scripts/run_benchmark.py -m "full_model and tts and H100" + pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100" + pytest -sv tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json + pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_cosmos3_vllm_omni.json ``` + Nightly **Perf Test** jobs in ``test-nightly.yml`` use ``--test-config-file`` only (no ``-m``). E2e L4 function tests still use ``full_model`` + ``--run-level full_model`` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml)). Example: === "L5 level" @@ -96,15 +175,21 @@ Our test scripts use the pytest framework. First, please use `git clone https:// # Stability: Wan2.2 (v1/videos diffusion benchmark loop) pytest -s -v dfx/stability/scripts/test_stability_wan22.py -m slow - # Reliability: Qwen3-Omni + # Reliability: Qwen3-Omni (H100 × 2) pytest -s -v dfx/reliability/test_reliability_qwen3_omni.py -m slow - # Reliability: Wan2.2 + # Reliability: Wan2.2 (H100) pytest -s -v dfx/reliability/test_reliability_wan22.py -m slow + # Reliability: HunyuanImage DiT (H100 × 4) + pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow + + # Reliability: VoxCPM2 (L4) + pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow + ``` - The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml). + The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/cuda/test-weekly.yml). You can find more information about markers in the documentation: [marker doc](./tests_markers.md) diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index cc89cee4788..0a0a974e641 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -33,7 +33,8 @@ Defined in `pyproject.toml`: | `skipif_rocm` | Skip if the num of ROCm cards is less than the required * | | `skipif_npu` | Skip if the num of NPU cards is less than the required * | | `slow` | Slow tests (may skip in quick CI) | -| `benchmark` | Benchmark tests | +| `benchmark` | Benchmark tests (decorator on runner test functions; perf JSON uses `full_model` + type marker instead) | +| `local_model` | Tests requiring local / non-HF-hub model weights | \* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility). @@ -94,7 +95,7 @@ This decorator is intended to make hardware-aware, cross-platform test authoring num_cards=2, ) ``` -- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), NPU (A2/A3) +- `res` must be a dict; supported resources: CUDA (L4/H100), ROCm (MI325), XPU (B60), MUSA (S5000), NPU (A2/A3) - `num_cards` can be int (all platforms) or dict (per platform); defaults to 1 when missing - Distributed markers (`distributed_cuda`, `distributed_rocm`, `distributed_npu`) are auto-added for multi-card cases - Filtering examples: @@ -120,6 +121,32 @@ def test_feature(omni_server): ... ``` +### JSON `mark` field (L4 perf configs) + +Perf JSON under `tests/dfx/perf/tests/` can attach pytest marks per **case** (`test_name` block). Parsed by `tests.dfx.conftest.resolve_pytest_marks` and applied to each parametrized run in `run_benchmark.py` / `run_diffusion_benchmark.py`. + +When `mark` is present, it must be an **array** with exactly one ``hardware_marks`` object (same semantics as `hardware_marks()` above), followed by registered pytest marker name strings such as `full_model`, `omni`, `tts`, or `diffusion`. + +```json +{ + "test_name": "test_cosmos3_t2i_official_demo_2gpu", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, + "full_model", + "diffusion" + ], + "server_type": "vllm-omni", + "server_params": { "...": "..." }, + "benchmark_params": [{ "name": "1024x1024_steps4", "...": "..." }] +} +``` + +- Local bulk load: `pytest -sv tests/dfx/perf/scripts/run_diffusion_benchmark.py -m "full_model and diffusion and H100"` +- Nightly CI perf steps: `--test-config-file tests/dfx/perf/tests/test__vllm_omni.json` (file selects cases; no `-m`) +- Result filenames use **runtime** GPU detection (`get_runtime_resource_label`); `H100` is omitted on the default CI pool + +See [L4 performance test examples](./test_examples/l4_performance_tests.inc.md) for the full schema. + ## Add Support for a New Platform If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps: diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 35ff6dae202..b3880511210 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -349,12 +349,14 @@ class YourModelPipeline(nn.Module): - def __call__( + def forward( self, -+ req: OmniDiffusionRequest, # ← Add request parameter here ++ req: DiffusionRequestBatch, # ← Add request-batch parameter here - ): -+ ) -> DiffusionOutput: # ← Add return type ++ ) -> list[DiffusionOutput]: # ← Add return type ``` -[`OmniDiffusionRequest`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/request/#vllm_omni.diffusion.request.OmniDiffusionRequest) is a dataclass that contains the **prompts** and **sampling parameters** [`OmniDiffusionSamplingParams`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/inputs/data/#vllm_omni.inputs.data.OmniDiffusionSamplingParams) for the diffusion pipeline execution. It also contains a request_id for other components to trace this request and its outputs. +[`OmniDiffusionRequest`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/request/#vllm_omni.diffusion.request.OmniDiffusionRequest) is a dataclass that contains one **prompt** and the **sampling parameters** [`OmniDiffusionSamplingParams`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/inputs/data/#vllm_omni.inputs.data.OmniDiffusionSamplingParams) for one logical diffusion request. It also contains a request_id for other components to trace this request and its outputs. Before pipeline execution, the runner wraps one or more independent requests into `DiffusionRequestBatch`. + +[`DiffusionRequestBatch`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/worker/request_batch/#vllm_omni.diffusion.worker.request_batch.DiffusionRequestBatch) exposes compatibility properties such as `prompts`, `sampling_params`, and `request_id`. Pipelines that can execute the whole request batch in one forward pass should set `supports_request_batch = True`; other pipelines still receive a single-request batch and return a one-element output list. See some parameters in `OmniDiffusionSamplingParams` as follows: @@ -367,35 +369,36 @@ See some parameters in `OmniDiffusionSamplingParams` as follows: **Extract parameters from request:** ```python -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch def forward( self, - req: OmniDiffusionRequest, -) -> DiffusionOutput: - # Extract prompts from request - if req.prompts is not None: - prompt = [ - p if isinstance(p, str) else (p.get("prompt") or "") - for p in req.prompts - ] - - # Extract sampling parameters + req: DiffusionRequestBatch, +) -> list[DiffusionOutput]: + # Extract prompts from the request batch + prompts = [ + p if isinstance(p, str) else (p.get("prompt") or "") + for p in req.prompts + ] + + # Extract common sampling parameters sampling_params = req.sampling_params num_inference_steps = sampling_params.num_inference_steps or 50 guidance_scale = sampling_params.guidance_scale or 7.5 height = sampling_params.height or (self.default_sample_size * self.vae_scale_factor) width = sampling_params.width or (self.default_sample_size * self.vae_scale_factor) - # For image editing pipelines, extract images from multi_modal_data - if hasattr(req, 'multi_modal_data') and req.multi_modal_data: - input_images = req.multi_modal_data.get('image', []) + # For image editing pipelines, extract media from each prompt dict + input_images = [] + for p in req.prompts: + multi_modal_data = p.get("multi_modal_data", {}) if isinstance(p, dict) else {} + input_images.append(multi_modal_data.get("image")) # ... rest of generation logic ``` -For an image editing model, an example `OmniDiffusionRequest` is like: +For an image editing model, the request `prompt` can be a dict like: ```python { "prompt": "turn this cat to a dog", @@ -472,12 +475,12 @@ def get_your_model_pre_process_func( def pre_process_func( request: OmniDiffusionRequest, ): - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - # image pre-processing - # after pre-processing, update the request attributes - ... + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + # image pre-processing + # after pre-processing, update the request attributes + ... return request return pre_process_func @@ -635,7 +638,7 @@ mkdir -p examples/online_serving/your_model_name - Script: `examples/offline_inference/your_model_name/end2end.py` - Parse args like BAGEL (`--model`, `--modality`, optional `--image-path`, `--steps`, etc.) - - Use `from vllm_omni.entrypoints.omni import Omni` (or `OmniDiffusion` if your model is diffusion-only) + - Use `from vllm_omni.entrypoints.omni import Omni` for both multi-stage and diffusion-only models - Save outputs (images/audio/video/text) with deterministic filenames (e.g., `output_0_0.png`) - Doc: `examples/offline_inference/your_model_name/README.md` - Include at least one runnable command, e.g.: @@ -923,11 +926,11 @@ When implementing a new pipeline, avoid putting all logic inside a single functi For example: ``` -def forward(self, req: OmniDiffusionRequest): +def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: prompt_embeds = self.encode_prompt(req) latents = self.diffuse(prompt_embeds, req) images = self.vae.decode(latents) - return DiffusionOutput(output=images) + return [DiffusionOutput(output=images)] ``` This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily. diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 1eaff10596c..5032c980521 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/docs/contributing/model/adding_omni_model.md @@ -34,6 +34,7 @@ vllm_omni/model_executor/models/ └── your_model_name/ # Model directory (e.g., qwen3_omni) ├── __init__.py # Exports main model class ├── your_model.py # Main unified model class + ├── pipeline.py # Pipeline topology registry entry ├── your_model_stage1_implementation.py # Stage 1 implementation (e.g., thinker) ├── your_model_stage2_implementation.py # Stage 2 implementation (e.g., talker) └── your_model_stage3_implementation.py # Stage 3 implementation (e.g., code2wav) @@ -42,10 +43,14 @@ vllm_omni/model_executor/models/ vllm_omni/model_executor/stage_input_processors/ └── your_model_name.py # Stage transition processors -vllm_omni/model_executor/stage_configs/ -└── your_model_name.yaml # Stage configuration file +vllm_omni/deploy/ +└── your_model_name.yaml # Deployment configuration file ``` +New in-tree models should register their `PipelineConfig` in +`vllm_omni/config/pipeline_registry.py`; use `--stage-configs-path` only for +custom legacy `stage_args` YAMLs. + ## Step-by-Step Implementation ### Step 1: Create the Model Directory @@ -475,92 +480,64 @@ thinker_hidden_states = output.multimodal_output["24"] ### Implementation Example -Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`: +Create stage transition processors in `vllm_omni/model_executor/stage_input_processors/your_model_name.py`. Each inter-stage edge should provide a **coherent processor set** rather than a single monolithic function: + +| Suffix | Role | Runs when | +|--------|------|-----------| +| `*_full_payload` | Worker-side payload producer | `async_chunk=false`; accumulates tensors and ships via connector | +| `*_async_chunk` | Scheduler-side streaming producer | `async_chunk=true`; emits per-chunk payloads | +| `*_token_only` | Orchestrator placeholder builder | `async_chunk=false`; allocates downstream prompt slots only | ```python -# qwen3_omni.py +# qwen3_omni.py (Thinker → Talker, non-async path) -def thinker2talker( - stage_list: list[Any], - engine_input_source: list[int], +def thinker2talker_token_only( + source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | None = None, requires_multimodal_data: bool = False, + streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: - """ - Process thinker outputs to create talker inputs. - - Args: - stage_list: List of stage objects - engine_input_source: Source stage IDs (typically [0] for thinker) - prompt: Original prompt data - - Returns: - List of OmniTokensPrompt for talker stage - """ - source_stage_id = engine_input_source[0] - thinker_outputs = stage_list[source_stage_id].engine_outputs - talker_inputs = [] - - for thinker_output in thinker_outputs: - output = thinker_output.outputs[0] - # Extract thinker embeddings and hidden states - thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda() - thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda() - - info = { - "thinker_prefill_embeddings": thinker_prefill_embeddings, - "thinker_hidden_states": thinker_hidden_states, - "thinker_sequences": thinker_output.prompt_token_ids + output.token_ids, - "thinker_input_ids": thinker_output.prompt_token_ids, - } - - talker_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0] * computed_length, - additional_information=info, - multi_modal_data=None, - ) - ) - - return talker_inputs - + """Allocate talker prefill slots; bulk tensors arrive via the connector.""" + ... + + +def thinker2talker_full_payload( + transfer_manager: Any, + pooling_output: dict[str, Any], + request: OmniEngineCoreRequest, +) -> dict[str, Any] | None: + """Pack accumulated thinker hidden states into OmniPayload for stage-1.""" + ... + + +def thinker2talker_async_chunk( + transfer_manager: Any, + multimodal_output: OmniPayload | dict[str, Any], + request: OmniEngineCoreRequest, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Stream thinker rows to talker while async_chunk is enabled.""" + ... +``` -def talker2code2wav( - stage_list: list[Any], - engine_input_source: list[int], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """ - Process talker outputs to create code2wav inputs. - """ - source_stage_id = engine_input_source[0] - talker_outputs = stage_list[source_stage_id].engine_outputs - code2wav_inputs = [] - - for talker_output in talker_outputs: - output = talker_output.outputs[0] - # Extract codec codes - codec_codes = ( - output.multimodal_output["code_predictor_codes"] - .to(torch.long) - .transpose(0, 1) - .cpu() - .to(torch.long) - .reshape(-1) - .tolist() - ) - - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=codec_codes, - multi_modal_data=None, - ) - ) +Wire these in `pipeline.py`: - return code2wav_inputs +```python +StagePipelineConfig( + stage_id=0, + custom_process_next_stage_input_func=f"{_PROC}.thinker2talker_full_payload", + async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2talker_async_chunk", + ... +), +StagePipelineConfig( + stage_id=1, + sync_process_input_func=f"{_PROC}.thinker2talker_token_only", + ... +), ``` +Do **not** add a no-suffix `thinker2talker` when a `sync_process_input_func` is already declared — `_select_processor_funcs()` always prefers the `*_token_only` hook in non-async mode, so the bare function would never run. See `docs/design/rfc_stage_input_processors_refactor.md` for the full contract. + ## Testing For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md). diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 3e5ae30df6d..b9bb9484f4d 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -175,15 +175,16 @@ vllm_omni/model_executor/models/ your_model_name/ __init__.py your_model.py # Unified class (stage dispatch) + pipeline.py # Pipeline topology registry entry your_model_ar_stage.py # Stage 0: AR stage your_model_decoder.py # Stage 1: audio decoder vllm_omni/model_executor/stage_input_processors/ your_model_name.py # Stage 0 -> Stage 1 transition -vllm_omni/model_executor/stage_configs/ - your_model_name.yaml # Batch mode config - your_model_name_async_chunk.yaml # Streaming mode config +vllm_omni/deploy/ + your_model_name.yaml # Batch/deploy config + your_model_name_async_chunk.yaml # Streaming/deploy config ``` ### Example placement @@ -445,6 +446,13 @@ Key configuration fields: | `final_output` | Whether this stage produces the final user-facing output | | `final_output_type` | Type of final output (`audio`, `text`, etc.) | +!!! note + New in-tree models should define frozen topology in `models//pipeline.py`, + register it in `vllm_omni/config/pipeline_registry.py`, and put deployment + knobs in `vllm_omni/deploy/.yaml`. Legacy `stage_args` YAMLs are still + accepted for custom configs via `--stage-configs-path`, but should not be used + for new bundled defaults. + ### Batch mode ```yaml @@ -771,8 +779,9 @@ directly from the model's own generator. vllm_omni/model_executor/models/your_model_name/ __init__.py modeling_your_model_name.py # unified class: load_weights + forward + streaming + pipeline.py # single-stage pipeline topology -vllm_omni/model_executor/stage_configs/your_model_name.yaml +vllm_omni/deploy/your_model_name.yaml ``` No stage input processor is needed. diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index 400d2a3c02c..b57de441ed1 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -286,3 +286,54 @@ Recommended viewers: - Nsight Systems GUI for CUDA kernel timelines For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/). + +## 7. Orchestrator Monitor + +> **Warning:** Diagnostic only. Adds lightweight counters on the orchestrator poll loop and per-replica queue depth. Disable in production unless you are actively debugging orchestrator saturation. + +Multi-stage omni pipelines route client-facing outputs and inter-stage connector traffic through a single-process orchestrator. When that loop is saturated (for example by large multimodal payloads on the poll path), TTFT and stability can degrade even when GPU stages are healthy. The orchestrator monitor records time-series signals to confirm whether the orchestrator or downstream stage queues are the bottleneck. + +Enable on the stage-0 (orchestrator) process: + +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ + --stage-id 0 \ + --enable-orch-monitor +``` + +Optional output path override: + +```bash +export VLLM_OMNI_ORCH_MONITOR_PATH=/tmp/vllm_omni_orch_monitor.json +``` + +If unset, the monitor writes `./vllm_omni_orch_monitor_.json` at shutdown. + +### Recorded metrics + +Each 1-second window records: + +| Series | Description | +|---|---| +| `windows.duration_s` | Wall time covered by the window | +| `windows.loop_idle` / `windows.loop_active` | Orchestrator poll-loop iterations with no work vs. active forwarding | +| `replicas..outputs_queue_size` | MP client `outputs_queue` backlog for that replica | +| `replicas..inflight` | Requests currently bound/routed to the replica | + +On shutdown the server also logs a short summary (`loop_active_pct`, per-replica queue averages/maxima). + +### Relationship to other diagnostics + +This monitor is intentionally separate from the existing profiling tools: + +| Tool | Scope | Output | +|---|---|---| +| `--enable-diffusion-pipeline-profiler` | Diffusion stage functions (`vae.decode`, `diffuse`, …) | Per-request `stage_durations` / logs | +| `--enable-ar-profiler` | AR stage generation time | Per-request `stage_durations` | +| `profiler_config` (`torch` / `cuda`) | GPU/CPU kernels inside a stage worker | `trace.json`, Nsight, operator tables | +| Prometheus `/metrics` (`vllm:omni_*`) | Pipeline SLOs and cross-stage transfer | Continuous scrape endpoint | +| `--enable-orch-monitor` | Orchestrator poll loop + replica queue backlog | Single JSON file at shutdown | + +The orchestrator monitor does not use `torch.profiler` because the bottleneck signal lives in the orchestrator process (poll-loop duty cycle and queue depth), not inside stage workers. It complements Prometheus metrics when you need a post-run time series for a benchmark or regression run. + +Process liveness monitors (`monitor_engine_liveness`, worker/engine-core monitors) are unrelated: they detect crashed workers, not orchestrator load. diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 1c38ba67183..20c99fd9c45 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -13,7 +13,7 @@ This document outlines the architectural design for vLLM-Omni. The primary goal of the vLLM-Omni project is to build the fastest and easiest-to-use open-source Omni-Modality model inference & serving engine. vLLM-Omni extends the original vLLM, which was created to support large language models for text-based autoregressive (AR) generation tasks. vLLM-Omni is designed to support: -* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, and video, alongside text. +* **Non-textual Output:** Enables the integration, efficient processing and output of various data types, including but not limited to, images, audio, video, and action trajectories, alongside text. * **Non-Autoregressive Structure:** Support model structure beyond autoregressive, especially Diffusion Transformer (DiT), which is widely used in visual and audio generation. * **Integration with vLLM Core:** Maintain compatibility and leverage existing vLLM key modules and optimizations where applicable. * **Extensibility:** Design a modular and flexible architecture that can easily accommodate new modalities, model architectures, and output formats. diff --git a/docs/design/feature/cfg_parallel.md b/docs/design/feature/cfg_parallel.md index c73a87749f5..0533604941f 100644 --- a/docs/design/feature/cfg_parallel.md +++ b/docs/design/feature/cfg_parallel.md @@ -126,16 +126,16 @@ Call `self.diffuse` in your pipeline's forward function: ```python import torch.nn as nn +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + class YourModelPipeline(nn.Module, CFGParallelMixin): def forward( self, - prompt: str, - negative_prompt: str | None = None, - guidance_scale: float = 3.5, - num_inference_steps: int = 50, - **kwargs, - ): - # Encode prompts, Initialize latents, Get timesteps + req: DiffusionRequestBatch, + ) -> list[DiffusionOutput]: + # Read prompts and sampling params from the request + ... + # Encode prompts, initialize latents, get timesteps ... # Run diffusion loop (calls the mixin's diffuse method) latents = self.diffuse( diff --git a/docs/design/feature/diffusion_continuous_batching.md b/docs/design/feature/diffusion_continuous_batching.md index c443dea78fa..f90969369ed 100644 --- a/docs/design/feature/diffusion_continuous_batching.md +++ b/docs/design/feature/diffusion_continuous_batching.md @@ -38,9 +38,13 @@ With continuous batching enabled: The current implementation is conservative: - only compatible requests are batched together -- request-mode diffusion still runs with `max_num_seqs=1` - per-request progress and completion remain independent +Here, "continuous batching" means the step-wise path enabled by +`step_execution=True`. Request-mode `DiffusionRequestBatch` is static +request-level batching for one full pipeline `forward()` call; it does not +admit or remove requests between denoise steps. + ## Enablement Use `--step-execution` as the feature gate, then increase `--max-num-seqs` @@ -72,7 +76,7 @@ which is built from shape-sensitive and CFG-sensitive sampling fields. This is the core correctness rule for batching: requests are only co-batched when they share the same denoise tensor contract. -There are two important details: +There are three important details: - `num_inference_steps` is not part of the key, so requests with different total step counts can still share a batch @@ -88,8 +92,10 @@ key also covers LoRA identity (`lora_int_id`, `lora_scale`), so requests targeting different adapters or scales run in separate batches and the worker can activate exactly one adapter per step. -The current batching unit is one `OmniDiffusionRequest`. Requests with -multiple prompts do not participate in batching today. +The scheduler batching unit is one logical `OmniDiffusionRequest`. In the +step-wise path, runtime tensor batching is represented as `StepInputBatch`. For +request-mode prompt semantics, see +[Request-Level Batching](../../user_guide/diffusion/request_batching.md). ## Runner @@ -126,17 +132,15 @@ request-local scheduler state and outputs. the background loop and async add-request path needed for multiple requests to accumulate in the scheduler. -This is supporting infrastructure, not the main design point. The batching -behavior is defined by scheduler-side compatibility gating and runner-side -batch packing. +When `step_execution=True`, the engine routes work through the step-wise +executor path. The continuous batching behavior is defined by scheduler-side +compatibility gating and runner-side `StepInputBatch` packing. ## Current Limitations - Experimental feature; use `max_num_seqs=1` for the older conservative path. - Only native pipelines that already support `step_execution=True`. -- Request-mode diffusion still clamps `max_num_seqs` back to `1`. - Only homogeneous batches keyed by `SamplingParamsKey` are supported. -- Multi-prompt requests are not batched. - `cache_backend`, KV transfer, and other request-mode extras are not wired into the batched step-wise path yet. - Future work can relax the current same-shape restriction with richer diff --git a/docs/design/feature/diffusion_request_level_batching.md b/docs/design/feature/diffusion_request_level_batching.md new file mode 100644 index 00000000000..57287ada49d --- /dev/null +++ b/docs/design/feature/diffusion_request_level_batching.md @@ -0,0 +1,175 @@ +# Request-Level Batching for Diffusion + +This document describes the request-mode batching path for diffusion pipelines. +For end-user enablement and tuning, see +[Request-Level Batching](../../user_guide/diffusion/request_batching.md). + +This is separate from +[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md). +Request-level batching runs one full pipeline `forward()` over a static batch of +compatible requests. Step-wise continuous batching admits work between denoise +steps when `step_execution=True`. + +## Why It Helps + +The request-level design avoids coupling several logical requests to one request +object. This keeps request identity, abort/error handling, and per-request +metadata unambiguous while still allowing one fused pipeline forward pass for +bursty or concurrent traffic. + +## Overview + +With request-level batching enabled: + +- each `OmniDiffusionRequest` contains one `prompt` and one `request_id` +- the scheduler groups compatible waiting requests into one scheduler wave +- `DiffusionRequestBatch` wraps the scheduled requests for pipeline `forward()` +- batch-capable pipelines return `list[DiffusionOutput]`, one output per + request +- `BatchRunnerOutput` maps each result back to its original `request_id` + +Pipelines opt in with `supports_request_batch = True` and a `forward()` method +that accepts `DiffusionRequestBatch` and returns `list[DiffusionOutput]`. +Pipelines that do not opt in keep the existing per-request execution path. + +## Enablement + +Request-level batching is the request-mode path, so `step_execution` must remain +disabled. Increase `max_num_seqs` above `1` to let the scheduler keep multiple +compatible requests active: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --max-num-seqs 4 +``` + +For bursty online ingress, `request_batch_max_wait_ms` can add a bounded +admission wait before the first `schedule()` of a scheduler wave: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --max-num-seqs 4 \ + --request-batch-max-wait-ms 20 +``` + +`request_batch_max_wait_ms=0` disables this wait and is the default. + +## Request Contract + +`OmniDiffusionRequest` represents one logical request. It owns one prompt, +sampling parameters, request id, and request-local metadata. Runtime batches are +formed by the scheduler and represented separately from the request payload. + +Runtime batching is represented by: + +- [`DiffusionSchedulerOutput`](gh-file:vllm_omni/diffusion/sched/interface.py) + for scheduled request ids and request payloads +- [`DiffusionRequestBatch`](gh-file:vllm_omni/diffusion/worker/request_batch.py) + for the pipeline-facing request batch +- [`BatchRunnerOutput`](gh-file:vllm_omni/diffusion/worker/utils.py) for + per-request results + +`DiffusionRequestBatch` intentionally exposes compatibility properties such as +`prompts`, `sampling_params`, `request_id`, and `kv_sender_info` so migrated +pipelines can stay close to upstream code while using a batch-aware contract. + +## Scheduler + +The scheduler derives its capacity from `max_num_seqs` through +`max_num_running_reqs`. It exposes waiting/running queue counters so the engine +can decide whether admission wait is useful before scheduling a new wave. + +Batch compatibility is controlled by +[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py). The key +contains shape-sensitive and guidance-sensitive fields, including output count +and LoRA identity. Requests with incompatible shapes, CFG settings, output +counts, LoRA adapters, or LoRA scales are kept in separate batches. + +Admission is conservative: + +- the scheduler only batches compatible requests +- FIFO ordering is preserved +- an incompatible request at the head of the waiting queue blocks later + compatible requests + +## Engine + +[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) resolves +request-batch capability during initialization from the configured pipeline +class, including custom pipeline classes. + +The capability check uses the pipeline class attribute +`supports_request_batch = True`. Pipelines that set this attribute must implement +a request-batch-compatible `forward()` contract and return one +`DiffusionOutput` per request; the runner validates that return shape at runtime. + +When the selected pipeline is batch-capable and `step_execution=False`, request +mode routes scheduler waves through the batch executor path. Otherwise it keeps +the per-request executor path. + +The optional admission wait runs only when: + +- request batching is supported +- `step_execution=False` +- `request_batch_max_wait_ms > 0` +- no requests are currently running + +The wait exits early when the waiting queue reaches capacity, when the queue is +stable for a short window, when the deadline expires, or when the engine stops. + +## Executor And Runner + +The executor exposes two request-mode entries: + +- `execute_request`: one worker call per scheduled request +- `execute_batch`: one worker call for the whole `DiffusionSchedulerOutput` + +On the batch path, the worker builds a `DiffusionRequestBatch` and runs the +pipeline once. Request-local setup remains per request: + +- KV transfer metadata +- random generator and seed handling +- request output/error/abort mapping + +Shared batch setup happens once per batch when possible: + +- cache refresh +- LoRA activation for the homogeneous adapter key +- pipeline `forward(req_batch)` + +Large tensor IPC still uses the shared-memory packing path. The packer traverses +both normal `RunnerOutput.result` wrappers and nested batch results so batched +outputs do not fall back to pickle IPC for tensor payloads. + +## Current Limitations + +- Only pipelines that declare the request-batch contract use fused batch + execution. +- Batches are homogeneous under `SamplingParamsKey`; heterogeneous resolution or + incompatible guidance settings do not co-batch yet. +- FIFO scheduling can reduce batching opportunities when an incompatible + request is at the front of the queue. +- `request_batch_max_wait_ms` improves burst coalescing but can add latency to + the first request in a scheduler wave. Keep it small for latency-sensitive + serving. +- Step-wise continuous batching is documented separately and only applies when + `step_execution=True`. + +## Related Files + +- Request object and request batch: + [`vllm_omni/diffusion/request.py`](gh-file:vllm_omni/diffusion/request.py) +- Scheduler interface: + [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py) +- Scheduler base: + [`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py) +- Engine: + [`vllm_omni/diffusion/diffusion_engine.py`](gh-file:vllm_omni/diffusion/diffusion_engine.py) +- Worker runner: + [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py) +- Executor interface: + [`vllm_omni/diffusion/executor/abstract.py`](gh-file:vllm_omni/diffusion/executor/abstract.py) +- Tests: + [`tests/diffusion/test_diffusion_engine.py`](gh-file:tests/diffusion/test_diffusion_engine.py) diff --git a/docs/design/feature/vae_parallel.md b/docs/design/feature/vae_parallel.md index e330b41a68f..00f41378a6c 100644 --- a/docs/design/feature/vae_parallel.md +++ b/docs/design/feature/vae_parallel.md @@ -448,6 +448,42 @@ Complete examples in the codebase: --- +## Spatially-Sharded Decode (Wan) + +The tile-parallel executor above assigns independent spatial **tiles** to ranks. The Wan VAE additionally supports a **spatially-sharded decode** backend, selected through `DiffusionParallelConfig.vae_parallel_mode` (`"spatial_shard_height"` or `"spatial_shard_width"`; default is `"tile"`). + +### How it differs from tile parallel + +| Aspect | Tile parallel (`"tile"`) | Spatially-sharded (`"spatial_shard_height"`/`"spatial_shard_width"`) | +|--------|--------------------------|-----------------------------------------------| +| Unit of work | Independent overlapping tiles | A single global feature map sharded along H or W | +| Cross-rank communication | Gather tiles to rank 0, stitch + blend | Per-conv **halo exchange** of boundary rows/cols (P2P) | +| Output assembly | Blend overlapping tiles | All-gather shards on rank 0, trim padding (matches `broadcast_result=False`) | +| Scope | Decode + encode | Decode only | + +Spatial-shard decode swaps the decoder's spatial convolutions/padding for halo-exchanging variants (`WanDistConv2d`, `WanDistCausalConv3d`, `WanDistZeroPad2d`) so each rank only holds a shard of the activations but still sees the correct receptive field at shard boundaries. Implementation lives in `vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py`. + +### Wiring + +`vae_parallel_mode` flows through the same path as `vae_patch_parallel_size`: + +```text +serve.py (--vae-parallel-mode) / OmniEngineArgs + -> DiffusionParallelConfig.vae_parallel_mode + -> registry.py: model.vae.set_parallel_size(vae_pp_size, mode=...) + -> DistributedVaeExecutor.parallel_mode + -> DistributedAutoencoderKLWan.tiled_decode dispatch +``` + +`DistributedAutoencoderKLWan._spatial_shard_decode_enabled()` gates the path: it requires distributed decode to be enabled, a 5D latent, and `vae_patch_parallel_size == DiT group size`. Otherwise it logs a warning and falls back to tile-parallel decode. + +### Notes + +- The decoder is patched **in place** the first time spatial-shard decode runs and is bound to a single split dimension for the lifetime of the VAE instance. +- Numerical correctness vs. single-GPU decode is covered by `tests/diffusion/distributed/test_wan_spatial_shard.py::test_spatial_shard_decode_matches_reference` (multi-GPU, nightly `full_model` + `distributed_cuda`). + +--- + ## Summary Adding VAE Patch Parallel support to diffusion model: diff --git a/docs/design/index.md b/docs/design/index.md index 8789f2e1b22..cfd177d65da 100644 --- a/docs/design/index.md +++ b/docs/design/index.md @@ -11,6 +11,7 @@ This section contains design documents and architecture specifications for vLLM- - [Disaggregated Inference](feature/disaggregated_inference.md) - [Ray-based Execution](feature/ray_based_execution.md) - [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md) +- [Request-Level Batching for Diffusion](feature/diffusion_request_level_batching.md) - [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md) ## Infrastructure Design Documents diff --git a/docs/design/module/dit_module.md b/docs/design/module/dit_module.md index 50890c59742..02d50181bfa 100644 --- a/docs/design/module/dit_module.md +++ b/docs/design/module/dit_module.md @@ -199,7 +199,7 @@ class _BaseScheduler(SchedulerInterface): - **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of duplicated in each policy. -- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching. +- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`. Request-mode diffusion can use that capacity for compatible independent requests when the configured pipeline declares `supports_request_batch = True`; step-wise diffusion uses the same scheduler capacity for compatible step batches. #### 2.4 Current `RequestScheduler` Policy @@ -215,7 +215,7 @@ class RequestScheduler(_BaseScheduler): - **FIFO request scheduling**: Waiting requests are promoted in queue order. -- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch. +- **Compatible request admission**: `RequestScheduler` admits waiting requests while capacity remains and the request's `SamplingParamsKey` is compatible with the active batch. Request-mode execution keeps each logical request independent, while batch-capable pipelines receive the scheduled requests as a runner-side `DiffusionRequestBatch`. - **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids. @@ -235,7 +235,12 @@ while True: - **No scheduler-owned IPC**: Scheduler no longer talks to workers directly. -- **Split concurrency model**: Request-mode diffusion remains single-active-request, while the step-wise path can keep multiple compatible requests running and advance them independently between denoise steps. +- **Split concurrency model**: Request-mode diffusion can schedule a static + batch of compatible independent requests for one full pipeline forward when + the pipeline supports request-level batching; the step-wise path can also + admit or remove compatible requests between denoise steps. See + [Request-Level Batching](../feature/diffusion_request_level_batching.md) and + [Continuous Batching for Step-Wise Diffusion](../feature/diffusion_continuous_batching.md). --- @@ -465,7 +470,7 @@ class Attention(nn.Module): role="self", role_category=None, ...): # Resolve backend for this role from the active OmniDiffusionConfig config = get_current_diffusion_config_or_none() - attention_config = config.attention_config if config is not None else None + attention_config = config.diffusion_attention_config if config is not None else None attn_backend_cls, spec = get_attn_backend_for_role( role=role, head_size=head_size, @@ -488,7 +493,7 @@ These backends provide the **kernel implementations** for attention computation. #### Backend Selection Mechanism -Selection is driven by `AttentionConfig` on `OmniDiffusionConfig`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order: +Selection is driven by `OmniDiffusionConfig.diffusion_attention_config`, which carries a global `default` plus a `per_role` map of `AttentionSpec` entries. A role string is resolved in this order: ```python def get_attn_backend_for_role(role, head_size, attention_config=None, role_category=None): @@ -904,7 +909,7 @@ def initialize_model_parallel( ``` 1. User Request - └─> OmniDiffusion.generate(prompt) + └─> Omni.generate(prompt) └─> Prepare OmniDiffusionRequest └─> DiffusionEngine.step(requests) diff --git a/docs/features/custom_pipeline.md b/docs/features/custom_pipeline.md index dc8b645e45a..78d680f26d2 100644 --- a/docs/features/custom_pipeline.md +++ b/docs/features/custom_pipeline.md @@ -10,7 +10,7 @@ This guide demonstrates how to use the newly added features for extending vLLM-O Three main features enable custom pipeline extension: 1. **`WorkerWrapperBase`**: A wrapper class that enables dynamic worker extension with custom functionality -2. **`load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines +2. **`diffusion_load_format`**: A parameter that controls how diffusion models are loaded, including support for custom pipelines 3. **`CustomPipelineWorkerExtension`**: An extension class that enables pipeline re-initialization with custom implementations ## Features @@ -27,15 +27,15 @@ Three main features enable custom pipeline extension: **Location:** `vllm_omni/diffusion/worker/diffusion_worker.py` -### load_format Parameter +### diffusion_load_format parameter -The `load_format` parameter controls how diffusion models are loaded. It supports the following values: +The `diffusion_load_format` parameter controls the initial diffusion model load. The relevant values are: - **`"default"`**: Standard model loading using the model registry (default behavior) -- **`"custom_pipeline"`**: Load a custom pipeline class specified by `custom_pipeline_name` -- **`"dummy"`**: Skip model loading (useful for testing or when pipeline will be initialized separately) +- **`"dummy"`**: Skip the initial model load; use this with `custom_pipeline_args["pipeline_class"]` for the custom pipeline workflow below +- **`"diffusers"`**: Load through the Hugging Face Diffusers adapter -**Location:** `vllm_omni/diffusion/model_loader/diffusers_loader.py` +**Location:** `vllm_omni/diffusion/data.py` (`OmniDiffusionConfig`) ### CustomPipelineWorkerExtension @@ -54,20 +54,24 @@ Create a custom pipeline class that extends an existing pipeline. In this exampl ```python # custom_pipeline.py -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch import torch class CustomPipeline(QwenImageEditPipeline): def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) - def forward(self, req, prompt=None, negative_prompt=None, **kwargs): + def forward(self, req: DiffusionRequestBatch): + # Optionally customize sampling parameters + actual_num_steps = req.sampling_params.num_inference_steps or 50 + req.sampling_params.num_inference_steps = actual_num_steps + # Call parent's forward to get normal output - output = super().forward(req=req, prompt=prompt, negative_prompt=negative_prompt, **kwargs) + output = super().forward(req=req) # Add custom trajectory data - actual_num_steps = req.sampling_params.num_inference_steps or kwargs.get('num_inference_steps', 50) output.trajectory_timesteps = torch.linspace(1000, 0, actual_num_steps, dtype=torch.float32) output.trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32) diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 59d03a3e4c6..e21a11ceefd 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels -Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, and 0.23.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). +Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, and 0.25.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). === "NVIDIA CUDA" @@ -94,6 +94,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc ### Build your own docker image +=== "NVIDIA CUDA" + + --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:build-docker" + === "AMD ROCm" --8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 9e394b758e0..d4808139ef6 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.23.0 --torch-backend=auto +uv pip install vllm==0.25.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]' # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library ```bash -uv pip install vllm==0.23.0 --torch-backend=auto +uv pip install vllm==0.25.0 --torch-backend=auto ``` -The 0.23.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.25.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.23.0 +git checkout v0.25.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.23.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.23.0 \ + vllm/vllm-omni:v0.25.0 \ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -110,3 +110,53 @@ docker run --runtime nvidia --gpus 2 \ The CUDA image does not define a default entrypoint, so include `vllm serve ... --omni` after the image name. # --8<-- [end:pre-built-images] + +# --8<-- [start:build-docker] + +#### Build docker image + +```bash +DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda . +``` + +If you want to specify the base vLLM version: + +```bash +DOCKER_BUILDKIT=1 docker build \ + -f docker/Dockerfile.cuda \ + --build-arg BASE_IMAGE=vllm/vllm-openai:v0.22.1 \ + -t vllm-omni-cuda . +``` + +#### Launch the docker image + +##### Launch with OpenAI API Server + +!!! note + The model `Qwen/Qwen3-Omni-30B-A3B-Instruct` requires significant GPU memory. The example below has been verified on 2 x H100's. + +```bash +docker run --runtime nvidia --gpus 2 \ + -v ${HF_HOME:-$HOME/.cache/huggingface}:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8091:8091 \ + --ipc=host \ + vllm-omni-cuda \ + vllm serve --omni --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091 +``` + +By default, this mounts `$HOME/.cache/huggingface` as the model cache directory. To use a custom location, set the `HF_HOME` environment variable before running the command (e.g., `export HF_HOME=/data/models`). + +##### Launch with interactive session for development + +```bash +docker run --runtime nvidia --gpus all -it --rm \ + -v ${HF_HOME:-$HOME/.cache/huggingface}:/root/.cache/huggingface \ + --env "HF_TOKEN=$HF_TOKEN" \ + -p 8091:8091 \ + --ipc=host \ + --entrypoint bash \ + vllm-omni-cuda +``` + +# --8<-- [end:build-docker] diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index a8344b21bab..4a64e872326 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722 +uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 ``` #### Installation of vLLM-Omni @@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library ```bash -uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722 +uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 ``` -The pre-built 0.23.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.25.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.23.0 +git checkout v0.25.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` @@ -130,7 +130,7 @@ docker run --rm \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ - vllm/vllm-omni-rocm:v0.23.0 \ + vllm/vllm-omni-rocm:v0.25.0 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -149,7 +149,7 @@ docker run --rm -it \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ --entrypoint bash \ - vllm/vllm-omni-rocm:v0.23.0 + vllm/vllm-omni-rocm:v0.25.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index c5b13dd73f7..7381f5b1067 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -1,19 +1,16 @@ # --8<-- [start:requirements] -For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html). +For detailed hardware and software requirements, please refer to the [vLLM-Ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html). # --8<-- [end:requirements] # --8<-- [start:installation-release] -The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images: +The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built Docker images: ```bash -# Update the vllm-ascend image -# Atlas A2: -# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1 -# Atlas A3: -# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3 -export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1 +# vLLM-Ascend image +# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0 +# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3 docker run --rm \ --name vllm-omni-npu \ --shm-size=1g \ @@ -31,11 +28,11 @@ docker run --rm \ -v /etc/ascend_install.info:/etc/ascend_install.info \ -v /root/.cache:/root/.cache \ -p 8000:8000 \ - -it $IMAGE bash + -it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash # Inside the container, install vLLM-Omni from source cd /vllm-workspace -git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git +git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git cd vllm-omni pip install -v -e . --no-build-isolation # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e . @@ -87,14 +84,11 @@ Supported images as following. Here's an example deployment command that has been verified on 4 x NPUs: ```bash -# Atlas A2: -# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0 -# Atlas A3: -# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0-a3 -export IMAGE=quay.io/ascend/vllm-omni:v0.18.0 +# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0 +# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3 docker run --rm \ - --name vllm-omni-npu \ - --shm-size=1g \ + --name vllm-omni-a3 \ + --shm-size=4g \ --device /dev/davinci0 \ --device /dev/davinci1 \ --device /dev/davinci2 \ @@ -106,10 +100,17 @@ docker run --rm \ -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ + --device /dev/davinci_manager \ + --device /dev/devmm_svm \ + --device /dev/hisi_hdc \ + -v /usr/local/dcmi:/usr/local/dcmi \ + -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \ + -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \ + -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \ -v /etc/ascend_install.info:/etc/ascend_install.info \ - -v /root/.cache:/root/.cache \ - -p 8000:8000 \ - -it $IMAGE bash + -v ~/.cache:/root/.cache \ + -p 8091:8091 \ + -it quay.io/ascend/vllm-omni:v0.25.0-a3 bash ``` !!! tip diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 652c512a245..914f1f7b43e 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.23.0 --torch-backend=auto +uv pip install vllm==0.25.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722 +uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -31,11 +31,10 @@ uv pip install -e . For additional installation methods — please see the [installation guide](installation/README.md). - !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. - If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. + If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.25.0` and vLLM Omni `0.25.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference @@ -52,14 +51,14 @@ if __name__ == "__main__": images[0].save("coffee.png") ``` -You can pass a list of prompts and wait for them to process altogether, shown below. +You can pass a list of prompts and wait for the independent requests to finish, +as shown below. !!! info - However, it is not currently recommended to do so - because not all models support batch inference, - and batch requesting mostly does not provide significant performance improvement (despite the impression that it does). - This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements. + For diffusion pipelines, each prompt becomes a separate logical request. + The runtime may automatically batch compatible in-flight requests through + the scheduler and runner. ```python from vllm_omni.entrypoints.omni import Omni @@ -88,9 +87,9 @@ if __name__ == "__main__": !!! info - For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input - list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support - batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts. + For diffusion request-level batching controls such as `max_num_seqs` and + `request_batch_max_wait_ms`, see + [Request-Level Batching](../user_guide/diffusion/request_batching.md). For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md) diff --git a/docs/mkdocs/hooks/generate_api_readme.py b/docs/mkdocs/hooks/generate_api_readme.py index 3214abd56ea..e09c6e1ea07 100644 --- a/docs/mkdocs/hooks/generate_api_readme.py +++ b/docs/mkdocs/hooks/generate_api_readme.py @@ -144,6 +144,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]: "vllm_omni.entrypoints.async_diffusion", "vllm_omni.entrypoints.openai", "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts", + "vllm_omni.experimental", # optional serving deps not installed in docs build ] if any(module_path.startswith(prefix) for prefix in excluded_prefixes): continue diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 81f2bac24a9..6cb601b962e 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -16,8 +16,35 @@ EXAMPLE_DIR = ROOT_DIR / "examples" EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples" NAV_FILE = ROOT_DIR / "docs/.nav.yml" +MODEL_DISPLAY_NAMES_FILE = EXAMPLE_DIR / "model_display_names.yml" MAX_INLINE_MATERIAL_SIZE = 8 * 1024 +SERVING_MODE_TITLES = { + "offline_inference": "Offline inference", + "online_serving": "Online serving", +} + + +def load_model_display_names() -> dict[str, str]: + try: + with open(MODEL_DISPLAY_NAMES_FILE, encoding="utf-8") as f: + model_display_names = yaml.safe_load(f) + except (OSError, yaml.YAMLError) as exc: + raise ValueError(f"Failed to load MODEL_DISPLAY_NAMES_FILE at {MODEL_DISPLAY_NAMES_FILE}: {exc}") from exc + + if not isinstance(model_display_names, dict) or not all( + isinstance(key, str) and isinstance(value, str) for key, value in model_display_names.items() + ): + raise ValueError( + "MODEL_DISPLAY_NAMES_FILE " + f"at {MODEL_DISPLAY_NAMES_FILE} must contain a YAML mapping " + "of string model IDs to string display names" + ) + return model_display_names + + +MODEL_DISPLAY_NAMES = load_model_display_names() + def fix_case(text: str) -> str: subs = { @@ -54,6 +81,7 @@ class Example: main_file (Path): The main file in the directory. other_files (list[Path]): list of other files in the directory. title (str): The title of the document. + nav_title (str): The concise title used in navigation. Methods: __post_init__(): Initializes the main_file, other_files, and title attributes. @@ -78,6 +106,16 @@ def __post_init__(self): def is_code(self) -> bool: return self.main_file.suffix != ".md" + @property + def model_display_name(self) -> str | None: + if self.category not in SERVING_MODE_TITLES: + return None + return MODEL_DISPLAY_NAMES.get(self.path.stem) + + @property + def nav_title(self) -> str: + return self.model_display_name or self.title + def determine_main_file(self) -> Path: """ Determines the main file in the given path. @@ -151,13 +189,25 @@ def is_other_file(file: Path) -> bool: return [file for file in self.path.rglob("*") if is_other_file(file)] def determine_title(self) -> str: + model_display_name = self.model_display_name if not self.is_code: # Specify encoding for building on Windows with open(self.main_file, encoding="utf-8") as f: first_line = f.readline().strip() match = re.match(r"^#\s+(?P.+)$", first_line) + if model_display_name: + expected_title = f"{model_display_name}: {SERVING_MODE_TITLES[self.category]}" + actual_title = match.group("title") if match else first_line + if actual_title != expected_title: + raise ValueError( + f"Model example title mismatch in {self.main_file}: " + f"expected '# {expected_title}', got {first_line!r}" + ) + return expected_title if match: return match.group("title") + elif model_display_name: + raise ValueError(f"Mapped model example must use a Markdown README: {self.path}") return fix_case(self.path.stem.replace("_", " ").title()) def fix_relative_links(self, content: str) -> str: @@ -300,7 +350,7 @@ def update_nav_file(examples: list[Example]): for example in category_examples: doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md" rel_path = doc_path.relative_to(ROOT_DIR / "docs") - category_items.append({example.title: rel_path.as_posix()}) + category_items.append({example.nav_title: rel_path.as_posix()}) if category_items: # Format category name (e.g., "offline_inference" -> "Offline Inference") diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 023ad193293..c1dac6ca72e 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -6,6 +6,8 @@ vLLM-Omni supports unified multimodal comprehension and generation models across If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>. +For deployment recipes (TTS and multimodal), see [recipes.vllm.ai](https://recipes.vllm.ai) and the in-repo [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes) directory. + ## List of Supported Models <style> @@ -22,6 +24,7 @@ th { | `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | | | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | | `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | +| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | | | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | @@ -31,23 +34,26 @@ th { | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | | | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Krea2Pipeline` | Krea 2 (Raw + Turbo) | `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | ✅︎ | | | | | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | -| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | -| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | | | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | | | `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | | `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | | -| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | -| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | +| `LTX23Pipeline` | LTX-2.3-T2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | +| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | | | `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | | | `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | | | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ | | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | @@ -59,21 +65,25 @@ th { | `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | +| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | | +| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | | | `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | -| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | +| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | | | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | | | `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | | | `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | | | `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | | | `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | | +| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | | +| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | | | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ | | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | | | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | | | `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | | -| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | | +| `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` | ✅︎ | | | | diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index f64d37f44eb..fef28142e46 100644 --- a/docs/serving/image_edit_api.md +++ b/docs/serving/image_edit_api.md @@ -252,7 +252,7 @@ Validation errors (missing required fields): ```bash # Check if server is responding -curl -X http://localhost:8000/v1/images/edit \ +curl http://localhost:8000/v1/images/edits \ -F "prompt='test'" ``` diff --git a/docs/serving/image_generation_api.md b/docs/serving/image_generation_api.md index 747cef99567..12e2935ea91 100644 --- a/docs/serving/image_generation_api.md +++ b/docs/serving/image_generation_api.md @@ -32,6 +32,20 @@ curl -X POST http://localhost:8000/v1/images/generations \ }' | jq -r '.data[0].b64_json' | base64 -d > dragon.png ``` +**Using curl save to file:** + +```bash +curl -o dragon.png -X POST http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "a dragon laying over the spine of the Green Mountains of Vermont", + "size": "1024x1024", + "seed": 42, + "response_format":"file" + }' +``` + + **Using Python:** ```python @@ -57,6 +71,36 @@ img = Image.open(io.BytesIO(img_bytes)) img.save("cat.png") ``` +**Using Python save to file:** + +```python +import requests +import base64 +from PIL import Image +import io +import re + +response = requests.post( + "http://localhost:8000/v1/images/generations", + json={ + "prompt": "a black and white cat wearing a princess tiara", + "size": "1024x1024", + "num_inference_steps": 50, + "seed": 42, + "response_format":"file" + } +) + +# save to file +content_disposition = response.headers.get("Content-Disposition", "") +match = re.search(r'filename="?(.+)"?', content_disposition) +filename = match.group(1) if match else "save.png" +with open(filename, "wb") as f: + for chunk in response.iter_content(8192): + f.write(chunk) +print("saved:", filename) +``` + **Using OpenAI SDK:** ```python @@ -94,7 +138,7 @@ Content-Type: application/json | `model` | string | server's model | Model to use (optional, should match server if specified) | | `n` | integer | 1 | Number of images to generate (1-10) | | `size` | string | model defaults | Image dimensions in WxH format (e.g., "1024x1024", "512x512") | -| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) | +| `response_format` | string | "b64_json" | Response format (`"b64_json"` or `"file"`) | | `user` | string | null | User identifier for tracking | #### vllm-omni Extension Parameters diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index d8dd092b13e..2bf630476be 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -9,6 +9,13 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano. +!!! tip "Deployment recipes" + TTS deployment recipes are published at + [recipes.vllm.ai](https://recipes.vllm.ai) (e.g. + [Qwen3-TTS](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice), + [Higgs-Audio v3](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b)). + The in-repo runbooks live under [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes). + Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`). ## Quick Start @@ -115,7 +122,8 @@ Content-Type: application/json | `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"`) | +| `stream` | bool | false | When true, stream OpenAI `speech.audio.*` SSE events (requires `response_format="pcm"` or `"wav"`). For raw PCM/WAV byte streaming, set `stream_format="audio"`. | +| `stream_format` | string | null | Streaming output format. `"audio"` streams raw audio bytes as they are decoded; `"sse"` streams OpenAI `speech.audio.*` Server-Sent Events. If omitted, `stream=true` selects SSE and `stream=false` remains non-streaming. See [Response Format](#response-format). | **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. @@ -129,7 +137,57 @@ Content-Type: application/json ### Response Format -Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`). +The response shape depends on the streaming parameters: + +**Non-streaming (default).** With `stream=false` and no `stream_format`, returns the +complete clip as binary audio data with an appropriate `Content-Type` header (e.g. +`audio/wav`). The raw-bytes body has no JSON carrier, so no `usage` is reported. + +**Raw audio stream** (`stream_format="audio"`). Streams raw audio bytes (PCM or +WAV) as they are decoded. + +Both streaming modes (`stream_format="audio"` and `"sse"`) require +`response_format="pcm"` or `"wav"`, and `speed` must be `1.0` (or omitted). + +**SSE stream** (`stream=true` or `stream_format="sse"`). Streams [OpenAI +`speech.audio.*` Server-Sent Events](https://platform.openai.com/docs/api-reference/audio-streaming). +Each event has an `event:` line and a JSON `data:` line: + +- `speech.audio.delta` — a base64 audio chunk: + + ```json + { "type": "speech.audio.delta", "audio": "<base64>", "response_format": "pcm" } + ``` + +- `speech.audio.done` — terminal event, carrying token `usage`: + + ```json + { + "type": "speech.audio.done", + "usage": { + "input_tokens": 119, + "output_tokens": 77, + "total_tokens": 196, + "input_token_details": { "text_tokens": 18, "audio_tokens": 101 } + } + } + ``` + +- `speech.audio.error` — emitted instead of `speech.audio.done` if generation fails: + + ```json + { "type": "speech.audio.error", "error": { "message": "...", "type": "server_error", "param": null, "code": 500 } } + ``` + +The `usage` object on `speech.audio.done` is the same shape returned per item by the +[batch endpoint](#batch-speech-generation): + +- `input_tokens` = `text_tokens` + `audio_tokens` + - `text_tokens`: tokens of the synthesized text (`input` plus `instructions`) + - `audio_tokens`: reference-audio codec frames, non-zero only for in-context + voice cloning (Base task); `0` for CustomVoice/VoiceDesign or x-vector-only +- `output_tokens`: generated codec tokens +- `total_tokens` = `input_tokens` + `output_tokens` ### Voices Endpoint @@ -226,9 +284,9 @@ Server -> Client: | Message | Description | |---------|-------------| -| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence | +| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) | -| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence | +| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | | `{"type": "session.done", "total_sentences": N}` | Session complete | | `{"type": "error", "message": "..."}` | Non-fatal error | @@ -238,8 +296,7 @@ All REST API parameters are supported, plus: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket | -| `split_granularity` | string | "sentence" | Text splitting granularity | +| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket | ```bash @@ -484,7 +541,13 @@ Each item in the `items` array requires only `input` (the text). All other field "index": 0, "status": "success", "audio_data": "<base64-encoded audio>", - "media_type": "audio/wav" + "media_type": "audio/wav", + "usage": { + "input_tokens": 119, + "output_tokens": 77, + "total_tokens": 196, + "input_token_details": { "text_tokens": 18, "audio_tokens": 101 } + } }, { "index": 1, @@ -498,6 +561,18 @@ Each item in the `items` array requires only `input` (the text). All other field } ``` +Each successful item carries a `usage` object (errored items omit it): + +- `input_tokens` = `text_tokens` + `audio_tokens` + - `text_tokens`: tokens of the synthesized text (`input` plus `instructions`) + - `audio_tokens`: reference-audio codec frames, non-zero only for in-context + voice cloning (Base task); `0` for CustomVoice/VoiceDesign or x-vector-only +- `output_tokens`: generated codec tokens +- `total_tokens` = `input_tokens` + `output_tokens` + +This is the same `usage` object emitted on the terminal `speech.audio.done` event +of the single endpoint's [SSE stream](#response-format) (`stream_format="sse"`). + ### Examples **Basic batch with shared defaults:** diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 4d156bdf0f0..72d79c30427 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -160,7 +160,12 @@ JSON references currently support `image_url`/`video_url`; `file_id` references are not implemented yet. Models may expose additional V2V controls through `extra_params`. For example, Cosmos3 supports `condition_frame_indexes_vision` and `condition_video_keep` to select which -decoded reference frames are used as clean conditioning. +decoded reference frames are used as clean conditioning. Cosmos3 transfer mode +also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus +transfer options such as `control_path`, `control_guidance`, +`control_guidance_interval`, `num_video_frames_per_chunk`, +`num_conditional_frames`, `show_control_condition`, and `show_input`; see the +Cosmos3 recipe for complete examples. ### Speech-to-Video @@ -195,13 +200,16 @@ curl -X POST http://localhost:8091/v1/videos/sync \ ## Storage -Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are +Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are stored: ```bash -export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos +export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos ``` +> `VLLM_OMNI_STORAGE_PATH` is deprecated and will be removed in a future release; +> use `VLLM_OMNI_SERVER_STORAGE__PATH` instead. + ## Model-Specific Examples For complete text-to-video, image-to-video, and model-specific video-to-video diff --git a/docs/source/architecture/omni-modality-model-architecture.png b/docs/source/architecture/omni-modality-model-architecture.png index 53978b32366..7bbd5b9b3ec 100644 Binary files a/docs/source/architecture/omni-modality-model-architecture.png and b/docs/source/architecture/omni-modality-model-architecture.png differ diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md index abc63ebed72..50e02971f6f 100644 --- a/docs/user_guide/diffusion/attention_backends.md +++ b/docs/user_guide/diffusion/attention_backends.md @@ -20,6 +20,8 @@ The full set of backends and their platform defaults is in the **Backend Options | `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. | | `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. | | `SAGE_ATTN_3` | Requires `sageattn3` from `SageAttention/sageattention3_blackwell`. CUDA only, intended for Blackwell GPUs, with GQA/MQA requests falling back to PyTorch SDPA. | +| `FLASH_ATTN_HUB` | FlashAttention 2 from HuggingFace `kernels` library. Useful for train/rollout alignment. | +| `FLASH_ATTN_3_HUB` | FlashAttention 3 from HuggingFace `kernels` library. CUDA Hopper (sm_90+) only; falls back to `FLASH_ATTN_HUB` on older GPUs. | ## Configuration @@ -81,7 +83,7 @@ When constructing `OmniDiffusionConfig` directly: from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, OmniDiffusionConfig config = OmniDiffusionConfig( - attention_config=AttentionConfig( + diffusion_attention_config=AttentionConfig( default=AttentionSpec(backend="FLASH_ATTN"), per_role={ "cross": AttentionSpec(backend="TORCH_SDPA"), @@ -210,6 +212,36 @@ Notes: - `SAGE_ATTN_3` is only selected on CUDA when `sageattn3` is importable and the GPU is Blackwell-class. - SageAttention3's Blackwell kernel assumes `Hq == Hkv`. In vLLM-Omni, GQA/MQA diffusion requests fall back to PyTorch SDPA for correctness. +## HuggingFace Kernels Hub Backends + +To achieve perfect numerical consistency between **training** (typically using HuggingFace Diffusers with Hub kernels) and **serving/rollout** (using vLLM-Omni), you can use the Hub-based attention backends. This eliminates numerical drift / sampling divergence caused by executing different local kernel versions during rollout. + +The following backend options are supported: +- `FLASH_ATTN_HUB` (HuggingFace `kernels-community/flash-attn2`) +- `FLASH_ATTN_3_HUB` (HuggingFace `kernels-community/flash-attn3`, Hopper sm_90+ only) + +### Installation + +To use these backends, you must install the `kernels` library: + +```bash +pip install kernels==0.14.1 +``` + +If the `kernels` library is not available in the environment, vLLM-Omni will log a warning and fall back gracefully to the corresponding local backend implementations (`FLASH_ATTN`). On CUDA GPUs below Hopper (compute capability < 9.0), `FLASH_ATTN_3_HUB` falls back to `FLASH_ATTN_HUB`. + +### Usage + +Select a Hub backend using the global CLI flag or environment variables: + +```bash +# Environment variable +export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN_HUB + +# CLI flag +vllm-omni serve <model> --diffusion-attention-backend FLASH_ATTN_HUB +``` + ## Usage Examples ### Default (auto-route) diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index 39dc366485e..d6fcd894424 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -33,7 +33,7 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True) **CLI:** ```bash -vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-cpu-offload ``` ### To Support a Model @@ -82,6 +82,33 @@ mutual exclusion: when one group runs, the other moves to CPU. - Support single GPU only for now +### Component offloading for split models (e.g. Cosmos3) + +Some models split their transformer into mutually-exclusive *components* that run +in different phases of a single forward pass rather than as separate pipeline +components -- e.g. Cosmos3's understanding (reasoner) component runs once per +generation while the generation (generator) component runs every denoising step. +Such models have no separate text encoder to swap against, so the transformer +owns a small model-local offload path and wraps each phase with +`with self._offload_context(name):` + +```python +class Cosmos3VFMTransformer(nn.Module): + def forward(self, ...): + with self._offload_context("reasoner"): + ... # understanding pass, runs once + with self._offload_context("generator"): + ... # denoising pass, runs every step +``` + +Model-level offloading then keeps exactly one component GPU-resident at a time +(the other on CPU), reusing the same `SequentialOffloadHook` `.to()` movers. The +pipeline opts in by exposing `enable_omni_model_cpu_offload` (which drives the +transformer's `enable_model_cpu_offload` and pins the VAE). Layerwise offloading +works for these models too -- each component declares its own block container via +`_layerwise_offload_blocks_attrs`. + + ## Layerwise (Blockwise) Offloading ### How It Works @@ -119,10 +146,10 @@ m = Omni(model="Wan-AI/Wan2.2-I2V-A14B-Diffusers", enable_layerwise_offload=True **CLI:** ```bash # Text-to-video -vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-layerwise-offload +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --enable-layerwise-offload # Or image-to-video -vllm-omni serve diffusion Wan-AI/Wan2.2-I2V-A14B-Diffusers --enable-layerwise-offload +vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --enable-layerwise-offload ``` ### To Support a Model @@ -145,9 +172,9 @@ class Flux2Transformer2DModel(nn.Module): ``` ### Limitations -- Cold start latency increases because of - 1) components are loaded to CPU first at the very first during initialization, - 2) weight consolidation and pinning +- Cold start latency increases because offloaded components must be moved to CPU + during setup; layerwise offload may add extra weight consolidation and pinning + work. - Performance depends on compute cost and H2D bandwidth as well - Support single GPU only for now @@ -178,11 +205,18 @@ Both strategies use vLLM-Omni's hook registry system (`HookRegistry` and `ModelH ``` OffloadBackend (base class) -├── ModelLevelOffloadBackend → uses SequentialOffloadHook +├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap) +│ (delegates to a pipeline's enable_omni_model_cpu_offload +│ for split models like Cosmos3) └── LayerWiseOffloadBackend → uses LayerwiseOffloadHook ``` -Factory function `get_offload_backend()` selects the appropriate backend based on configuration. +Factory function `get_offload_backend()` selects the appropriate backend based on +configuration. + +For split models, `ModelLevelOffloadBackend.enable()` detects a pipeline's +`enable_omni_model_cpu_offload` hook and delegates to it; Cosmos3 then swaps its +reasoner/generator components inside the model forward pass. ## Supported Models @@ -193,9 +227,11 @@ Factory function `get_offload_backend()` selects the appropriate backend based o | NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` | | OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` | | QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` | +| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` | | StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` | | Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` | +| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` | | BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` | **Notes:** diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md index ab03363cf13..817ae67f74e 100644 --- a/docs/user_guide/diffusion/frame_interpolation.md +++ b/docs/user_guide/diffusion/frame_interpolation.md @@ -47,7 +47,8 @@ For supported Wan2.2 pipelines, the execution order is: 1. Diffusion worker finishes denoising and decodes the raw video tensor. 2. Worker-side model-specific post-processing runs. 3. If frame interpolation is enabled, RIFE interpolates the decoded video - tensor on the worker side and records a FPS multiplier in `custom_output`. + tensor on the worker side and records a FPS multiplier in + `metadata.video.video_fps_multiplier`. 4. The API server receives the already-interpolated video and only performs MP4 export. diff --git a/docs/user_guide/diffusion/mot_config.md b/docs/user_guide/diffusion/mot_config.md new file mode 100644 index 00000000000..6661c6c9aaa --- /dev/null +++ b/docs/user_guide/diffusion/mot_config.md @@ -0,0 +1,50 @@ +This directory contains auto-tuned Triton kernel configurations for the +MoT (Mixture-of-Tokens) GEMM and RMSNorm operators used by BAGEL and other +MoT-architecture diffusion models. + +File naming convention: + device_name=<GPU_NAME>,dtype=<DTYPE>.json + +For example: + device_name=NVIDIA_A100-SXM4-80GB,dtype=w16a16.json + +Each JSON file maps (K, N) matrix shapes to a dictionary of batch sizes (M) +and their optimal Triton tile configurations: + + { + "3584_4608": { // K=3584, N=4608 (QKV projection) + "1024": { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 3 + }, + ... + }, + "3584_3584": { ... } // K=3584, N=3584 (output projection) + } + +Config loading order (3-tier, see ops/mot_gemm.py): + 1. $VLLM_TUNED_CONFIG_FOLDER/<filename> (env override) + 2. This directory: vllm_omni/diffusion/layers/mot/configs/<filename> + 3. Conservative default config (compiles everywhere, sub-optimal perf) + +If no config file matches the current device, a warning is printed with +instructions to run the auto-tuning benchmark. + +To generate configs for your hardware: + + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --tp-size 1 --dtype w16a16 --tune \ + --save-dir vllm_omni/diffusion/layers/mot/configs/ + +For multi-GPU tuning (uses Ray for parallel search): + + python benchmarks/kernels/mot_linear_benchmarks.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --tp-size 2 --tune + +See benchmarks/kernels/mot_linear_benchmarks.py for full options. diff --git a/docs/user_guide/diffusion/parallelism/overview.md b/docs/user_guide/diffusion/parallelism/overview.md index d482b5ea2fd..8419b480076 100644 --- a/docs/user_guide/diffusion/parallelism/overview.md +++ b/docs/user_guide/diffusion/parallelism/overview.md @@ -10,7 +10,7 @@ This guide covers the parallelism methods in vLLM-Omni for speeding up diffusion | **[Sequence Parallelism](sequence_parallel.md)** | Splits sequence dimension across GPUs (Ulysses-SP, Ring-Attention, or hybrid) for high-resolution images and videos | | **[CFG-Parallel](cfg_parallel.md)** | Runs CFG positive/negative branches on separate GPUs for ~1.8x speedup on guided generation | | **[Pipeline Parallelism](pipeline_parallel.md)** | Splits the denoising transformer block-wise across sequential GPU stages to reduce per-GPU model memory | -| **[VAE Patch Parallelism](vae_patch_parallel.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory | +| **[VAE Parallelism](vae_parallelism.md)** | Distributes VAE decode spatially across GPUs to reduce peak VAE memory | | **[HSDP](hsdp.md)** | Shards full model weights via PyTorch FSDP2 to enable large-model inference on memory-constrained GPUs | | **[Expert Parallelism](expert_parallel.md)** | Shards MoE expert blocks across GPUs for MoE models (e.g. HunyuanImage3.0) | diff --git a/docs/user_guide/diffusion/parallelism/vae_patch_parallel.md b/docs/user_guide/diffusion/parallelism/vae_parallelism.md similarity index 71% rename from docs/user_guide/diffusion/parallelism/vae_patch_parallel.md rename to docs/user_guide/diffusion/parallelism/vae_parallelism.md index 4e8513eabf4..ca0c428e181 100644 --- a/docs/user_guide/diffusion/parallelism/vae_patch_parallel.md +++ b/docs/user_guide/diffusion/parallelism/vae_parallelism.md @@ -1,4 +1,4 @@ -# VAE Patch Parallelism Guide +# VAE Parallelism Guide ## Table of Content @@ -15,7 +15,7 @@ ## Overview -VAE Patch Parallelism distributes the VAE (Variational AutoEncoder) decode/encode computation across multiple GPUs by splitting the latent space into spatial tiles or patches. Each GPU processes a subset of tiles in parallel, significantly reducing peak memory consumption during the VAE decode stage while maintaining output quality. +VAE parallelism distributes VAE (Variational AutoEncoder) decode/encode work across multiple GPUs. This guide covers VAE patch/tile parallelism, which splits latent space into spatial tiles or patches, and Wan spatial-shard decode, which shards decoder feature maps along height or width. This is particularly useful for: - **High-resolution image generation** where VAE decode can become a memory bottleneck @@ -25,7 +25,7 @@ This is particularly useful for: See supported models list in [Supported Models](../../diffusion_features.md#supported-models). -VAE Patch Parallelism uses two strategies based on image size: +VAE patch parallelism uses two strategies based on image size: | Strategy | Use Case | How It Works | Overlap Handling | Output Quality | |----------|----------|--------------|------------------|----------------| @@ -113,6 +113,7 @@ In `DiffusionParallelConfig`: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `vae_patch_parallel_size` | int | 1 | Number of GPUs for VAE patch/tile parallelism. Set to 2 or higher to enable. Should typically match `tensor_parallel_size` as they share the same process group. | +| `vae_parallel_mode` | str | `"tile"` | VAE parallel decode strategy: `"tile"` (default tile/patch parallel decode), `"spatial_shard_height"`, or `"spatial_shard_width"` (spatially-sharded decode, Wan only). See [Spatially-Sharded Decode](#spatially-sharded-decode-wan). | Additional requirements: @@ -125,6 +126,51 @@ Additional requirements: --- +## Spatially-Sharded Decode (Wan) + +The default `vae_parallel_mode="tile"` distributes whole tiles across ranks. For the **Wan** VAE there is an alternative decode strategy, **spatially-sharded decode**, selected via `vae_parallel_mode="spatial_shard_height"` or `vae_parallel_mode="spatial_shard_width"`. + +Instead of assigning independent tiles to ranks, spatial-shard decode shards the decoder feature maps along the height (`spatial_shard_height`) or width (`spatial_shard_width`) dimension and exchanges halo rows/columns between neighboring ranks around the spatial convolutions. This keeps the receptive field correct across shard boundaries, so the result matches the single-GPU decode within numerical tolerance. + +```python +from vllm_omni import Omni +from vllm_omni.diffusion.data import DiffusionParallelConfig + +omni = Omni( + model="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + parallel_config=DiffusionParallelConfig( + tensor_parallel_size=2, + vae_patch_parallel_size=2, # must match the DiT group size + vae_parallel_mode="spatial_shard_width", # or "spatial_shard_height" + ), +) +``` + +Or from the CLI / serving entrypoint: + +```bash +vllm serve Wan-AI/Wan2.1-T2V-1.3B-Diffusers --omni \ + --tensor-parallel-size 2 \ + --vae-patch-parallel-size 2 \ + --vae-parallel-mode spatial_shard_width +``` + +**Constraints and behavior:** + +- Spatial-shard decode is **decode-only** and currently implemented for the **Wan** VAE. Other models ignore `spatial_shard_*` modes. +- It requires `vae_patch_parallel_size` to **match the DiT process group size**. If it does not, the VAE logs a warning and **falls back to tile-parallel decode** at runtime. +- `spatial_shard_height` and `spatial_shard_width` are mutually exclusive for a given VAE instance (the decoder is patched in place for a single split dimension). + +For end-to-end latency/throughput, launch serving with the desired `vae_parallel_mode` and use the existing diffusion serving benchmark: + +```bash +python3 benchmarks/diffusion/diffusion_benchmark_serving.py \ + --endpoint /v1/videos --dataset random --task t2v --num-prompts 1 \ + --height 480 --width 832 --num-frames 17 --max-concurrency 1 +``` + +--- + ## Best Practices ### When to Use diff --git a/docs/user_guide/diffusion/request_batching.md b/docs/user_guide/diffusion/request_batching.md new file mode 100644 index 00000000000..5cd24802bcc --- /dev/null +++ b/docs/user_guide/diffusion/request_batching.md @@ -0,0 +1,110 @@ +# Request-Level Batching + +Request-level batching lets diffusion serving combine multiple compatible +logical requests into one pipeline `forward()` call. Each prompt remains a +separate request with its own `request_id`, sampling parameters, seed, output, +error, and abort state. The scheduler decides which requests can run together. + +!!! warning "Prompt List Semantics" + + Diffusion request-level batching does not support a top-level packed + list-prompt request. Submit multiple prompts as independent requests and let + the scheduler batch compatible in-flight requests. Multimodal payloads stay + inside a single prompt dict, for example + `{"prompt": "...", "multi_modal_data": {"image": image}}`. + +## Enablement + +Increase `max_num_seqs` above `1` to allow the request scheduler to keep more +than one compatible request active: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --max-num-seqs 4 +``` + +For bursty online traffic, you can also set a small admission wait window. This +lets the engine wait briefly before the first `schedule()` of a new wave so +nearby compatible requests can arrive and share the same fused forward pass: + +```bash +vllm serve Qwen/Qwen-Image --omni \ + --port 8091 \ + --max-num-seqs 4 \ + --request-batch-max-wait-ms 20 +``` + +`--request-batch-max-wait-ms 0` is the default and disables admission waiting, +so there is no added wait latency. + +For deploy YAMLs, configure the diffusion stage engine args: + +```yaml +stage_args: + - stage_id: 0 + stage_type: diffusion + engine_args: + max_num_seqs: 4 + request_batch_max_wait_ms: 20 +``` + +## Compatibility + +Only pipelines that declare request-batch support use the fused request-batch +path. The engine validates that the pipeline `forward()` uses the request-batch +contract and returns `list[DiffusionOutput]`. Pipelines that do not support this +contract do not use fused `pipeline.forward(batch)`; scheduled requests are +executed through per-request worker calls. + +The scheduler batches only compatible requests. Compatibility is based on +shape-sensitive and guidance-sensitive sampling fields, including resolution, +frame count, CFG settings, output count, and LoRA identity. Requests with +different LoRA adapters or scales are kept in separate batches so the worker +activates one adapter per batch. + +Request-level batching applies when `step_execution=False`. For the separate +step-wise runtime, see [Step Execution](step_execution.md). + +## Tuning + +- `max_num_seqs` caps the number of active compatible requests in one scheduler + wave. +- `request_batch_max_wait_ms` is an upper bound on extra admission wait before a + new wave starts. Keep it small for latency-sensitive serving; values such as + `10` to `50` ms are a practical starting range for bursty HTTP ingress. +- `0` disables admission waiting and preserves the lowest first-request latency. +- FIFO ordering is conservative: an incompatible request at the front of the + waiting queue can block later compatible requests from joining the current + batch. + +## Python API + +When constructing `Omni`, pass the same engine arguments: + +```python +from vllm_omni.entrypoints.omni import Omni + +omni = Omni( + model="Qwen/Qwen-Image", + max_num_seqs=4, + request_batch_max_wait_ms=20.0, +) + +outputs = omni.generate( + [ + "a cup of coffee on a table", + "a toy dinosaur on a sandy beach", + "a fox waking up in bed and yawning", + ] +) +``` + +`Omni.generate([...])` submits each list item as its own logical diffusion +request. The runtime may batch those requests internally when their sampling +parameters are compatible. + +## For Contributors + +For implementation details and model-author guidance, see +[Request-Level Batching for Diffusion](../../design/feature/diffusion_request_level_batching.md). diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 234446c0128..c34c1338cda 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -61,7 +61,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on |--------|-------------|----------| | **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs | | **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss | -| **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak | +| **[VAE Parallelism](diffusion/parallelism/vae_parallelism.md)** | Distributes VAE decode work across GPUs | High-resolution generation with reduced VAE memory peak | ### Extensions @@ -75,13 +75,17 @@ Extension methods add specialized capabilities to diffusion models beyond standa ### Execution Modes -Execution modes control how the diffusion pipeline processes denoise steps. +Execution modes control how the diffusion pipeline processes requests and +denoise steps. | Method | Description | Best For | |--------|-------------|----------| +| **[Request-Level Batching](diffusion/request_batching.md)** | Scheduler batches compatible independent diffusion requests into one pipeline forward pass | Bursty online serving and multi-request throughput | | **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control | -**Note:** Step execution is currently supported by QwenImagePipeline only. See [Supported Models](#supported-models) for details. +**Note:** Request-level batching is available for pipelines that declare the +request-batch forward contract. Step execution is currently supported by +QwenImagePipeline only. See [Supported Models](#supported-models) for details. ### Quantization Methods @@ -101,7 +105,7 @@ The following tables show which models support each feature: > Notes: -> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. +> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only. Split models like Cosmos3 (no separate text encoder) swap their reasoner/generator components for module-wise offload; see the [CPU Offload Guide](diffusion/cpu_offload_diffusion.md). > 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details. ### ImageGen @@ -117,19 +121,21 @@ The following tables show which models support each feature: | **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | | **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | +| **Krea 2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **LongCat-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **LongCat-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | -| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **OmniGen2** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | | **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | | **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ | | **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ | | **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **SenseNova-U1** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **SenseNova-U1** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | +| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | | **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ | | **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | @@ -138,17 +144,18 @@ The following tables show which models support each feature: > Notes: > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT. > 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary. -> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload. +> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload swaps the nested UND reasoner and GEN generator pathways; layerwise offload remains available for blockwise GEN/UND offload. +> 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. ### VideoGen | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | |------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:| | **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | -| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ | +| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ | | **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | -| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | -| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | +| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ | | **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | @@ -269,7 +276,7 @@ Measured on NVIDIA H800: **Memory Optimization:** - **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage -- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images +- **[VAE Parallelism Guide](diffusion/parallelism/vae_parallelism.md)** - Distribute VAE decode work across GPUs for high-resolution images and videos - **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models **Extensions:** diff --git a/docs/user_guide/examples/offline_inference/bagel.md b/docs/user_guide/examples/offline_inference/bagel.md index 0d3498b28d9..391c41b35c8 100644 --- a/docs/user_guide/examples/offline_inference/bagel.md +++ b/docs/user_guide/examples/offline_inference/bagel.md @@ -257,7 +257,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ | :------- | :--- | :------ | :---------- | | `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) | | `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` | -| `--worker-backend` | choice | `process` | `process` or `ray` | +| `--worker-backend` | choice | `multi_process` | `multi_process` or `ray` | | `--ray-address` | string | `None` | Ray cluster address | | `--quantization` | string | `None` | Quantization method (e.g. `fp8`) | | `--log-stats` | flag | `False` | Enable statistics logging | diff --git a/docs/user_guide/examples/offline_inference/glm_image.md b/docs/user_guide/examples/offline_inference/glm_image.md index c6ac6e33ffd..09d927ff75a 100644 --- a/docs/user_guide/examples/offline_inference/glm_image.md +++ b/docs/user_guide/examples/offline_inference/glm_image.md @@ -26,18 +26,22 @@ Stage 0 (AR Model) Stage 1 (Diffusion) ```python from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams if __name__ == "__main__": omni = Omni(model="zai-org/GLM-Image") outputs = omni.generate( "A photorealistic mountain landscape at sunset", - sampling_params={ - "height": 1024, - "width": 1024, - "num_inference_steps": 50, - "guidance_scale": 1.5, - "seed": 42, - }, + sampling_params_list=[ + omni.default_sampling_params_list[0], + OmniDiffusionSamplingParams( + height=1024, + width=1024, + num_inference_steps=50, + guidance_scale=1.5, + seed=42, + ), + ], ) outputs[0].request_output.images[0].save("output.png") ``` @@ -46,6 +50,7 @@ if __name__ == "__main__": ```python from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams if __name__ == "__main__": omni = Omni(model="zai-org/GLM-Image") @@ -56,13 +61,16 @@ if __name__ == "__main__": "image": "input.png", }, }, - sampling_params={ - "height": 1024, - "width": 1024, - "num_inference_steps": 50, - "guidance_scale": 1.5, - "seed": 42, - }, + sampling_params_list=[ + omni.default_sampling_params_list[0], + OmniDiffusionSamplingParams( + height=1024, + width=1024, + num_inference_steps=50, + guidance_scale=1.5, + seed=42, + ), + ], ) outputs[0].request_output.images[0].save("output.png") ``` diff --git a/docs/user_guide/examples/offline_inference/helios.md b/docs/user_guide/examples/offline_inference/helios.md deleted file mode 100644 index f83ebc2de25..00000000000 --- a/docs/user_guide/examples/offline_inference/helios.md +++ /dev/null @@ -1,203 +0,0 @@ -# Helios Video Generation - -Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/helios>. - - -Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants: - -| Variant | Description | Key Features | -|---------|-------------|--------------| -| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` | -| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support | -| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` | - -## Setup - -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. - -## Run Examples - -Get into the example folder: -```bash -cd examples/offline_inference/helios -``` - -### Text-to-Video (T2V) - -**Helios-Base** (Stage 1 only): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Base \ - --sample-type t2v \ - --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ - --guidance-scale 5.0 \ - --output helios_t2v_base.mp4 -``` - -**Helios-Mid** (Stage 2 + CFG-Zero*): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Mid \ - --sample-type t2v \ - --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 \ - --output helios_t2v_mid.mp4 -``` - -**Helios-Distilled** (Stage 2 pyramid + DMD): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Distilled \ - --sample-type t2v \ - --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ - --num-frames 240 \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk \ - --output helios_t2v_distilled.mp4 -``` - -### Image-to-Video (I2V) - -**Helios-Base**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Base \ - --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ - --guidance-scale 5.0 \ - --output helios_i2v_base.mp4 -``` - -**Helios-Mid** (Stage 2 + CFG-Zero*): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Mid \ - --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 \ - --output helios_i2v_mid.mp4 -``` - -**Helios-Distilled**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Distilled \ - --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ - --num-frames 240 \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk \ - --output helios_i2v_distilled.mp4 -``` - -### Video-to-Video (V2V) - -**Helios-Base**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Base \ - --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ - --guidance-scale 5.0 \ - --output helios_v2v_base.mp4 -``` - -**Helios-Mid** (Stage 2 + CFG-Zero*): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Mid \ - --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 \ - --output helios_v2v_mid.mp4 -``` - -**Helios-Distilled**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Distilled \ - --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ - --num-frames 240 \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk \ - --output helios_v2v_distilled.mp4 -``` - -## Common Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path | -| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` | -| `--prompt` | — | Text prompt describing the video | -| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) | -| `--image-path` | — | Input image (required for `i2v`) | -| `--video-path` | — | Input video (required for `v2v`) | -| `--height` | `384` | Video height in pixels | -| `--width` | `640` | Video width in pixels | -| `--num-frames` | `99` | Number of output frames | -| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) | -| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) | -| `--seed` | `42` | Random seed | -| `--fps` | `16` | Output video frame rate | -| `--output` | `helios_output.mp4` | Output file path | - -### Stage 2 / Pyramid Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising | -| `--pyramid-num-stages` | `3` | Number of pyramid stages | -| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage | -| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) | - -### CFG-Zero* Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) | -| `--use-zero-init` | off | Zero init for first steps | -| `--zero-steps` | `1` | Number of zero-init steps | - -### Memory & Parallelism - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--vae-use-slicing` | off | Enable VAE slicing | -| `--vae-use-tiling` | off | Enable VAE tiling | -| `--enforce-eager` | off | Disable torch.compile | -| `--enable-cpu-offload` | off | CPU offloading | -| `--enable-layerwise-offload` | off | Layerwise offloading | -| `--ulysses-degree` | `1` | Ulysses SP degree | -| `--ring-degree` | `1` | Ring SP degree | -| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) | -| `--tensor-parallel-size` | `1` | Tensor parallelism size | - -## Example materials - -??? abstract "end2end.py" - ``````py - --8<-- "examples/offline_inference/helios/end2end.py" - `````` diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index 6e105741a7e..b3d76ac46db 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -3,7 +3,9 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video>. -This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API. +This example demonstrates how to generate videos from images using Wan2.2, +LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's +offline inference API. ## Local CLI Usage @@ -51,13 +53,41 @@ python image_to_video.py \ --output i2v_output.mp4 ``` +### LTX-2.3 + +```bash +python image_to_video.py \ + --model diffusers/LTX-2.3-Diffusers \ + --model-class-name LTX23ImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --guidance-scale 4.0 \ + --num-inference-steps 20 \ + --frame-rate 24 \ + --fps 24 \ + --output ltx23_i2v_output.mp4 +``` + +Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select +the LTX-2.3 image-to-video pipeline. + Key arguments: -- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V). +- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or + LTX-2/LTX-2.3). - `--image`: Path to input image (required). - `--prompt`: Text description of desired motion/animation. -- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16. -- `--num-frames`: Number of frames (default 81). +- `--height/--width`: Output resolution (auto-calculated from image if not set). + Wan dimensions should be multiples of 16; LTX dimensions should be multiples + of 32. +- `--num-frames`: Number of frames (model-specific default; LTX-style models + work best with `8k + 1`). - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE). - `--negative-prompt`: Optional list of artifacts to suppress. - `--boundary-ratio`: Boundary split ratio for two-stage MoE models. @@ -65,6 +95,7 @@ Key arguments: - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints. - `--num-inference-steps`: Number of denoising steps (default 50). - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video). +- `--audio-sample-rate`: fallback audio sample rate for embedded audio. - `--output`: Path to save the generated video. - `--vae-use-slicing`: Enable VAE slicing for memory optimization. - `--vae-use-tiling`: Enable VAE tiling for memory optimization. diff --git a/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md b/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md index 85227b74809..60f1cbe5f5a 100644 --- a/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md +++ b/docs/user_guide/examples/offline_inference/mammothmodal2_preview.md @@ -7,29 +7,22 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf Download model ```bash -hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview +hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview ``` ### Text-to-Image (T2I) -```bash -python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \ - --model ./MammothModa2-Preview \ - --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --height 1024 \ - --width 1024 \ - --num-inference-steps 50 \ - --text-guidance-scale 4.0 \ - --out output.png -``` +Text-to-image now runs through the shared offline image example +(`examples/offline_inference/text_to_image/text_to_image.py`). See the recipe +`recipes/MammothModa2/MammothModa2-Preview.md` for the full command and the +`extra_body` knobs (`text_guidance_scale`, `cfg_range`, `num_inference_steps`). ### Image Summary ```bash python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ --model ./MammothModa2-Preview \ - --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \ + --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \ --question "Summarize this image." \ --image ./image.png ``` @@ -40,7 +33,3 @@ python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_s ``````py --8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py" `````` -??? abstract "run_mammothmoda2_t2i.py" - ``````py - --8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py" - `````` diff --git a/docs/user_guide/examples/offline_inference/mimo_audio.md b/docs/user_guide/examples/offline_inference/mimo_audio.md index 1cba2f77dcf..a43d8e4082e 100644 --- a/docs/user_guide/examples/offline_inference/mimo_audio.md +++ b/docs/user_guide/examples/offline_inference/mimo_audio.md @@ -1,4 +1,4 @@ -# MiMo-Audio Offline Inference +# MiMo-Audio: Offline inference Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mimo_audio>. diff --git a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md index c54976b540d..efd14bc1ce4 100644 --- a/docs/user_guide/examples/offline_inference/qwen2_5_omni.md +++ b/docs/user_guide/examples/offline_inference/qwen2_5_omni.md @@ -1,4 +1,4 @@ -# Qwen2.5-Omni +# Qwen2.5-Omni: Offline inference Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen2_5_omni>. diff --git a/docs/user_guide/examples/offline_inference/qwen3_omni.md b/docs/user_guide/examples/offline_inference/qwen3_omni.md index 2d856f7380a..3781c1cc87c 100644 --- a/docs/user_guide/examples/offline_inference/qwen3_omni.md +++ b/docs/user_guide/examples/offline_inference/qwen3_omni.md @@ -1,4 +1,4 @@ -# Qwen3-Omni +# Qwen3-Omni: Offline inference Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_omni>. diff --git a/docs/user_guide/examples/offline_inference/text_to_image.md b/docs/user_guide/examples/offline_inference/text_to_image.md index 3a97ffbf74b..7cd508d297a 100644 --- a/docs/user_guide/examples/offline_inference/text_to_image.md +++ b/docs/user_guide/examples/offline_inference/text_to_image.md @@ -160,9 +160,11 @@ python examples/offline_inference/text_to_image/text_to_image.py \ --output flux2-dev.png ``` -### Batch Requests (Multiple Prompts) +### Multiple Prompts -You can pass multiple prompts in a single `generate` call. +You can pass multiple prompts in a single `generate` call. For diffusion +pipelines, each prompt is submitted as a separate logical request; compatible +requests may be automatically batched by the scheduler and runner. ```python from vllm_omni.entrypoints.omni import Omni @@ -181,17 +183,8 @@ if __name__ == "__main__": !!! info - Not all models support batch inference, and batch requesting mostly does not provide significant - performance improvement. This feature is primarily for interface compatibility with vLLM and to - allow for future improvements. - -!!! info - - For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by - default, and the input list is sliced into single-item requests before feeding into the diffusion - pipeline. For models that do internally support batched inputs, you can - [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a - longer batch of prompts. + For diffusion request-level batching controls such as `max_num_seqs`, see + [Request-Level Batching](../../diffusion/request_batching.md). ### Negative Prompts diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index a09dbfc979f..61c498245f5 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | |---|---|---|---|---|---| | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | @@ -50,6 +51,30 @@ python text_to_video.py \ --output ltx2_out.mp4 ``` +### LTX-2.3 + +```bash +python text_to_video.py \ + --model diffusers/LTX-2.3-Diffusers \ + --model-class-name LTX23Pipeline \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --num-inference-steps 20 \ + --guidance-scale 4.0 \ + --frame-rate 24 \ + --fps 24 \ + --audio-sample-rate 48000 \ + --output ltx23_t2v_output.mp4 +``` + +Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3 +text-to-video pipeline explicitly. + ### HunyuanVideo-1.5 (480p) ```bash @@ -110,6 +135,8 @@ python text_to_video.py \ ### Common - `--model`: Diffusers model ID or local path. +- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline` + for LTX-2.3 text-to-video. - `--prompt`: text description (string). - `--height/--width`: output resolution. Default depends on model. - `--num-frames`: number of frames. Default depends on model. @@ -124,7 +151,8 @@ python text_to_video.py \ - `--enable-cpu-offload`: enable CPU offloading for diffusion models. - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). -- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio). +- `--audio-sample-rate`: audio sample rate for embedded audio (when the + pipeline returns audio; LTX-2.3 outputs 48kHz audio). - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). - `--flow-shift`: scheduler flow_shift parameter. diff --git a/docs/user_guide/examples/online_serving/aura_omni.md b/docs/user_guide/examples/online_serving/aura_omni.md new file mode 100644 index 00000000000..8b7f56ac5e1 --- /dev/null +++ b/docs/user_guide/examples/online_serving/aura_omni.md @@ -0,0 +1,65 @@ +# AURA Omni: Online serving + +`aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline: + +```text +ASR -> AURA -> Qwen3-TTS Talker -> Code2Wav +``` + +Qwen3-TTS remains two engine stages so the pipeline reuses the existing native +Talker and Code2Wav implementation. + +```bash +vllm serve aurateam/AURA \ + --omni \ + --deploy-config vllm_omni/deploy/aura_omni.yaml \ + --served-model-name aurateam/AURA \ + --trust-remote-code +``` + +Configure local checkpoints by editing per-stage `model` values in +`vllm_omni/deploy/aura_omni.yaml`. The deploy file sets +`pipeline: aura_omni`, so the four-stage topology is used even if the +command-line `--model` points at one of the component checkpoints. + +Send requests with `"model": "aurateam/AURA"`. The ASR, AURA, and Qwen3-TTS +checkpoint paths are internal stage models from the deploy YAML, not the +OpenAI-facing served model name. + +The AURA stage can emit `<|silent|>`. Silent outputs are treated as a gate: +they produce no Qwen3-TTS Talker input, so no audio is synthesized for that +turn. + +## GPU Utilization Recommendation + +`gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` controls how much +VRAM each stage can reserve. Start with this split for a single GPU: + +- Stage 0 (ASR): `0.10` +- Stage 1 (AURA): `0.40` +- Stage 2 (Qwen3-TTS Talker): `0.20` +- Stage 3 (Qwen3-TTS Code2Wav): `0.20` + +## TTS Modes + +`aura_omni` can pass AURA text to Qwen3-TTS in two task modes: + +- `Base`: voice clone from `tts_ref_audio` with ICL enabled in the AURA + pipeline. Provide both `tts_ref_audio` and `tts_ref_text`. Set + `tts_x_vector_only_mode=true` to disable ICL and use speaker embedding only. +- `CustomVoice`: predefined speaker mode. Use a Qwen3-TTS CustomVoice + checkpoint for stages 2 and 3 in `aura_omni.yaml`, then pass + `tts_task_type=CustomVoice` and `tts_speaker`. + +By default, AURA responses are passed to Qwen3-TTS as text. Set +`additional_information.tts_pass_token_ids=true` to pass AURA-generated +assistant token ids directly instead. Even when token passthrough is disabled, +the stage processor uses AURA token ids when available to estimate the Talker +prefill length, so scheduling does not rely on raw character length. + +The example client exposes this as: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-pass-token-ids +``` diff --git a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md index 1adc6119363..c6be3c014a0 100644 --- a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md +++ b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md @@ -9,7 +9,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online. Currently, the following features are NOT yet supported. -It is not guaranteed whether they will be supported in the future. +Community contributions are welcome to add these features when they can be +implemented with clear Diffusers or vLLM-Omni behavior. - CFG parallel execution - Sequence parallel execution @@ -66,6 +67,35 @@ This is suitable for sampling parameters not available through the vLLM-Omni int When a parameter is available in the vLLM-Omni interface, it will be adapted here. But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time). +### Quantization + +Use `diffusers_load_kwargs` for Diffusers-native quantization options. If +`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the +diffusers backend builds a Diffusers quantization config and lets Diffusers +validate it before calling `DiffusionPipeline.from_pretrained()`. + +The diffusers backend also provides a small compatibility shortcut for +vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already +contain `quantization_config`. Currently this maps online/dynamic `fp8` and +online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer +components such as `transformer` or `transformer_2`. This path requires +`torchao` to be installed. + +For example, the CLI can request dynamic FP8 through the vLLM-Omni interface: + +```bash +vllm serve "Qwen/Qwen-Image" \ + --omni \ + --diffusion-load-format diffusers \ + --quantization-config '{"method": "fp8"}' +``` + +Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`, +`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists +such as `ignored_layers`, are intentionally not translated by this compatibility +shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or +a native vLLM-Omni pipeline for those cases. + ### Attention Backends The diffusers backend converts diff --git a/docs/user_guide/examples/online_serving/glm_image.md b/docs/user_guide/examples/online_serving/glm_image.md index 410b1180a4f..ea404faa913 100644 --- a/docs/user_guide/examples/online_serving/glm_image.md +++ b/docs/user_guide/examples/online_serving/glm_image.md @@ -19,7 +19,7 @@ both stages on a single GPU, override per stage: ```bash vllm serve zai-org/GLM-Image --omni --port 8091 \ - --stage-0-devices 0 --stage-1-devices 0 + --stage-overrides '{"0": {"devices": "0"}, "1": {"devices": "0"}}' ``` ## API Calls diff --git a/docs/user_guide/examples/online_serving/image_to_video.md b/docs/user_guide/examples/online_serving/image_to_video.md index 781f0c2a5ed..2dbc57633b8 100644 --- a/docs/user_guide/examples/online_serving/image_to_video.md +++ b/docs/user_guide/examples/online_serving/image_to_video.md @@ -84,16 +84,26 @@ curl -X POST http://localhost:8091/v1/videos/sync \ Generated video files are stored on local disk by the async video API. Local file storage behavior can be controlled via the following environment variables: -- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`) -- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`) +- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files + (default: `/tmp/storage`) +- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open + operations (default: `4`) +- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in + seconds +- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in + seconds for enforcing file TTL; defaults to `300` when TTL is set Example: ```bash -export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos -export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8 +export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos +export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8 +export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400 +export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300 ``` +See also: [Configuration Options](../../../configuration/README.md) + ## API Calls ### Method 1: Using curl diff --git a/docs/user_guide/examples/online_serving/mimo_audio.md b/docs/user_guide/examples/online_serving/mimo_audio.md index c8752f5782e..9db59f9e8e5 100644 --- a/docs/user_guide/examples/online_serving/mimo_audio.md +++ b/docs/user_guide/examples/online_serving/mimo_audio.md @@ -1,4 +1,4 @@ -# Online serving Example of vLLM-Omni for MiMo-Audio +# MiMo-Audio: Online serving Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/mimo_audio>. diff --git a/docs/user_guide/examples/online_serving/qwen2_5_omni.md b/docs/user_guide/examples/online_serving/qwen2_5_omni.md index b3a2c9f2ac9..d8c48a403b5 100644 --- a/docs/user_guide/examples/online_serving/qwen2_5_omni.md +++ b/docs/user_guide/examples/online_serving/qwen2_5_omni.md @@ -1,4 +1,4 @@ -# Qwen2.5-Omni +# Qwen2.5-Omni: Online serving Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen2_5_omni>. diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 1973b198bb4..e372692a85c 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -1,4 +1,4 @@ -# Qwen3-Omni +# Qwen3-Omni: Online serving Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_omni>. @@ -140,7 +140,7 @@ You can control output modalities to specify which types of output the model sho | Modalities | Output | |------------|--------| | `["text"]` | Text only | -| `["audio"]` | Text + Audio | +| `["audio"]` | Audio only | | `["text", "audio"]` | Text + Audio | | Not specified | Text + Audio (default) | @@ -161,13 +161,16 @@ curl http://localhost:8091/v1/chat/completions \ #### Text + Audio ```bash -curl http://localhost:8091/v1/chat/completions \ +response=$(curl -s http://localhost:8091/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "messages": [{"role": "user", "content": "Describe vLLM in brief."}], - "modalities": ["audio"] - }' + "modalities": ["text", "audio"] + }') + +echo "$response" | jq -r '.choices[0].message.content' +echo "$response" | jq -r '.choices[1].message.audio.data' | base64 -d > output.wav ``` ### Using Python client @@ -205,7 +208,7 @@ client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") response = client.chat.completions.create( model="Qwen/Qwen3-Omni-30B-A3B-Instruct", messages=[{"role": "user", "content": "Describe vLLM in brief."}], - modalities=["audio"] + modalities=["text", "audio"] ) # Response contains two choices: one with text, one with audio print(response.choices[0].message.content) # Text response diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 9ee6b33fbf7..5640845d81b 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -79,7 +79,7 @@ response = client.audio.speech.create( response.stream_to_file("output.wav") ``` -Streaming PCM output (where supported) — set `stream=true` with `response_format="pcm"`: +Streaming PCM output (where supported) — set `stream=true`, `stream_format="audio"`, and `response_format="pcm"`: ```bash curl -X POST http://localhost:8091/v1/audio/speech \ @@ -88,6 +88,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "input": "Hello, how are you?", "voice": "default", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` @@ -414,10 +415,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "voice": "vivian", "language": "English", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Streaming requires `response_format="pcm"` and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. +Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. ### Streaming WebSocket The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence: @@ -496,7 +498,7 @@ After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use ```bash python voxcpm2/gradio_demo.py ``` -Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Audio is streamed from the OpenAI Speech endpoint with `stream=true`. +Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`. --- diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index b918aac19d0..0f1f33a25c3 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -13,6 +13,7 @@ This example demonstrates how to deploy text-to-video models for online video ge | Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | | Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | | LTX-2 | `Lightricks/LTX-2` | +| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | ## Wan2.2 T2V @@ -90,16 +91,26 @@ curl -X POST http://localhost:8091/v1/videos/sync \ Generated video files are stored on local disk by the async video API. Local file storage behavior can be controlled via the following environment variables: -- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`) -- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`) +- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files + (default: `/tmp/storage`) +- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open + operations (default: `4`) +- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in + seconds +- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in + seconds for enforcing file TTL; defaults to `300` when TTL is set Example: ```bash -export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos -export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8 +export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos +export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8 +export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400 +export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300 ``` +See also: [Configuration Options](../../../configuration/README.md) + ## API Calls ### Method 1: Using curl @@ -364,6 +375,17 @@ curl -sS -X POST http://localhost:8098/v1/videos \ -F "seed=42" ``` +## LTX-2.3 + +LTX-2.3 uses a Diffusers-format checkpoint and a dedicated pipeline class: + +```bash +vllm serve diffusers/LTX-2.3-Diffusers --omni --port 8098 \ + --model-class-name LTX23Pipeline +``` + +Send requests through the same `/v1/videos` API shown in the LTX-2 example. + ## Example materials ??? abstract "response.json" diff --git a/docs/user_guide/fault_injection_reliability_matrix.md b/docs/user_guide/fault_injection_reliability_matrix.md index 2b5fd240168..5e5122bfce9 100644 --- a/docs/user_guide/fault_injection_reliability_matrix.md +++ b/docs/user_guide/fault_injection_reliability_matrix.md @@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b ## Source of Conclusions -The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`. +The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/cuda/test-weekly.yml`. ### Example: `SIGTERM` fault injection log (Qwen3-Omni) diff --git a/docs/user_guide/quantization/autoround.md b/docs/user_guide/quantization/autoround.md index 3afc8b06233..bafadec9c11 100644 --- a/docs/user_guide/quantization/autoround.md +++ b/docs/user_guide/quantization/autoround.md @@ -3,7 +3,7 @@ ## Overview [AutoRound](https://github.com/intel/auto-round) produces pre-quantized -checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the +checkpoints for LLMs, VLMs, diffusion models, and world models. vLLM-Omni reads the checkpoint's `config.json` and auto-detects `quantization_config.quant_method = "auto-round"`. @@ -36,6 +36,13 @@ guide. AutoRound is Intel-supported. | Wan2.2-T2V | `Intel/Wan2.2-T2V-A14B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | | Wan2.2-TI2V | `Intel/Wan2.2-TI2V-5B-Diffusers-int4-AutoRound` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | +### World Model (Cosmos3) + +| Model | Checkpoint | Scope | Scheme | Backend | +|-------|------------|-------|--------|---------| +| Cosmos3-Nano | `Intel/Cosmos3-Nano-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | +| Cosmos3-Super | `Intel/Cosmos3-Super-int4-AutoRound` | World-model transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend | + ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) | Model | Checkpoint | Scope | Scheme | Backend | diff --git a/docs/user_guide/quantization/bitsandbytes.md b/docs/user_guide/quantization/bitsandbytes.md new file mode 100644 index 00000000000..37ccd2471f2 --- /dev/null +++ b/docs/user_guide/quantization/bitsandbytes.md @@ -0,0 +1,120 @@ +# BitsAndBytes W4 Quantization + +## Overview + +BitsAndBytes 4-bit quantization supports weight-only NF4/FP4 diffusion transformer +inference on CUDA GPUs. It quantizes BF16/FP16 weights at load time; activations +stay in BF16/FP16. No pre-quantized checkpoint is required. + +This is an online (dynamic) path: load the normal HuggingFace checkpoint and +quantize during model loading via the `bitsandbytes` CUDA kernels. + +## Hardware Support + +| Device | Support | +|--------|---------| +| NVIDIA CUDA GPU (SM 75+) | ✅ | +| NVIDIA Blackwell GPU (SM 100+) | ✅ | +| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | +| NVIDIA Ampere GPU (SM 80+) | ✅ | +| AMD ROCm | ❌ | +| Intel XPU | ❌ | +| Ascend NPU | ❌ | + +Legend: `✅` supported, `❌` unsupported. Non-CUDA platforms raise in +`get_quant_method()`; CUDA requires compute capability 7.5+ (`SM 75+`). + +Requires the optional `bitsandbytes` package (`pip install bitsandbytes`). + +## Model Type Support + +### Diffusion Model (Qwen-Image, Wan2.2) + +| Model | HF models | CUDA | Mode | Recommendation | +|-------|-----------|:----:|------|----------------| +| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Online W4 weight-only | All heavy linear layers; sensitive embedders stay BF16 | +| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Not validated | Online W4 weight-only | Compare vs BF16 before enabling | +| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Online W4 weight-only | Validate before enabling in docs | + +Other diffusion models may work if their transformer uses supported linear +layers, but they are not validated in this guide. + +### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) + +| Model | Scope | Status | Notes | +|-------|-------|--------|-------| +| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths | +| Qwen3-TTS | TTS language-model stage | Not validated | No BitsAndBytes TTS stage support is documented | + +### Multi-Stage Diffusion Model (BAGEL, GLM-Image) + +| Model | Scope | Status | Notes | +|-------|-------|--------|-------| +| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing | +| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 | + +## Configuration + +Python API: + +```python +from vllm_omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +omni = Omni(model="Tongyi-MAI/Z-Image-Turbo", quantization="bitsandbytes") + +omni_with_skips = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + quantization_config={ + "method": "bitsandbytes", + "ignored_layers": ["to_out"], + }, +) + +outputs = omni.generate( + "A cup of coffee on the table", + OmniDiffusionSamplingParams(num_inference_steps=50), +) +``` + +CLI: + +```bash +python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes +python text_to_image.py --model Tongyi-MAI/Z-Image-Turbo --quantization bitsandbytes --ignored-layers "to_out" +vllm serve Tongyi-MAI/Z-Image-Turbo --omni --quantization bitsandbytes +``` + +## Parameters + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `method` | str | - | Quantization method (`"bitsandbytes"`) | +| `quant_type` | str | `"nf4"` | 4-bit data type: `"nf4"` (recommended) or `"fp4"` | +| `compress_statistics` | bool | `True` | Double-quantize block scaling statistics for better accuracy | +| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 | + +## Validation and Notes + +On Z-Image-Turbo (single GPU, 1024×1024, 50 steps), BitsAndBytes W4 typically +reduces peak VRAM from roughly 24.5 GiB (BF16) to roughly 17 GiB. Compare output +quality against a BF16 baseline before enabling on new models. + +Multi-GPU tensor parallelism (`tensor_parallel_size` > 1) is not validated for +BitsAndBytes in diffusion models; each rank quantizes its own weight shard +independently. + +If quality regresses, use `ignored_layers` to keep sensitive projections in BF16 +(for example `to_out` or `w2`). + +```python +omni = Omni( + model="Tongyi-MAI/Z-Image-Turbo", + quantization="bitsandbytes", + cache_backend="tea_cache", + cache_config={"rel_l1_thresh": 0.2}, +) +``` + +Only add a new model to the supported table after comparing the BitsAndBytes +output against a BF16 baseline and documenting any required `ignored_layers`. diff --git a/docs/user_guide/quantization/fp8.md b/docs/user_guide/quantization/fp8.md index 7981c7ca21a..c76cfd32e52 100644 --- a/docs/user_guide/quantization/fp8.md +++ b/docs/user_guide/quantization/fp8.md @@ -44,10 +44,16 @@ pip install vllm-omni[quack] 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. +It is enabled automatically once installed (no flag needed) and is **datacenter +Blackwell only** (`sm_100` / `sm_101` / `sm_103`, compute capability `10.x`, e.g. +B200): quack's CuteDSL GEMM uses the 5th-gen `tcgen05` tensor-core MMA, which exists +only on those parts. On Hopper/Ada the CUTLASS FP8 kernel already fuses bias, and on +workstation/consumer Blackwell (`sm_120` / `sm_121`, compute capability `12.x`, e.g. +RTX PRO 6000 / RTX 50-series) `tcgen05` is absent — so quack is **not** auto-enabled +there and FlashInfer's native FP8 path is used instead. Set +`VLLM_OMNI_USE_QUACK_FP8=1` to force quack on, or `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 @@ -82,10 +88,11 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)]) | Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | | | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ | | FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | | +| FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | Yes | Not validated | All layers | None | ✅︎ | | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | | | HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | | | HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | | -| Cosmos3 | `nvidia/Cosmos3-Nano` | Yes | Not validated | All layers | None | | +| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | | ### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS) diff --git a/docs/user_guide/quantization/gguf.md b/docs/user_guide/quantization/gguf.md index e70d2d26ecc..206e25319e8 100644 --- a/docs/user_guide/quantization/gguf.md +++ b/docs/user_guide/quantization/gguf.md @@ -9,6 +9,15 @@ for the transformer. GGUF is static quantization: the quantized weights are produced before serving. +!!! note + GGUF support has migrated to OOT [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin). Make sure you have GGUF plugin installed before serving a GGUF model. + +Before serving a GGUF model, make sure to install the [vllm-gguf-plugin](https://github.com/vllm-project/vllm-gguf-plugin): + +```bash +uv pip install vllm-gguf-plugin>=0.0.3 +``` + ## Hardware Support | Device | Support | diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md index 66e8427bb82..d8da4e1a92f 100644 --- a/docs/user_guide/quantization/modelopt.md +++ b/docs/user_guide/quantization/modelopt.md @@ -89,11 +89,54 @@ and | Model | Scope | Status | |-------|-------|--------| | Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path | +| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | Validated; see [Qwen3-Omni NVFP4 W4A4](#qwen3-omni-nvfp4-w4a4-thinker) below | | Qwen3-TTS | TTS language-model stage | Not validated | Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless a model-specific guide documents otherwise. +#### Qwen3-Omni NVFP4 W4A4 (thinker) + +vLLM-Omni serves ModelOpt NVFP4 W4A4 quantizations of the +Qwen3-Omni-30B-A3B-Instruct thinker language model. The thinker text body +(attention + MoE experts) is quantized to NVFP4 with FP8 per-tensor input +scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16. + +| Variant | HF checkpoint | Hardware | +|---------|---------------|----------| +| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ (Blackwell, FlashInfer FP4 GEMM) | + +Calibration uses ModelOpt `mtq.NVFP4_DEFAULT_CFG` with the `awq_clip` algorithm +on 1024 ultrachat samples chat-templated through the Qwen3-Omni tokenizer. +Excluded modules: `*audio_tower*`, `*visual*`, `*talker*`, `*code2wav*`, +`*lm_head*`, `*mlp.gate*`. See `scripts/nvfp4/calibrate.py` for the reference +recipe. + +!!! note "ModelOpt 0.44 NaN regression workaround" + ModelOpt 0.44's float32 -> FP8 E4M3 cast of per-block weight scales + occasionally emits literal NaN bytes (E4M3 encoding 0x7F / 0xFF) for + blocks whose pre-cast scale rounds above the FP8 max of 448 after the + global-scale division. A single NaN byte in any `weight_scale` + propagates through the FP4 GEMM and collapses the served model output + to `!!!!`. vLLM-Omni's `vllm_omni.patch` installs a defensive override + of `ModelOptNvFp4LinearMethod.process_weights_after_loading` that + clamps these bytes to the FP8 E4M3 max at load time. The override + self-extinguishes once vllm-omni's vllm pin moves to a release + containing the upstream fix. Set `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP=1` + to disable the override for diagnostics. + +Serving: + +```bash +vllm serve YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip \ + --omni --port 8000 +``` + +> Do **not** pass `--enforce-eager` for production / benchmarks. CUDA graphs +> amortize launch overhead and unlock the FP4 throughput wins; with +> `--enforce-eager` set, W4A4 TPOT degrades ~10x relative to the CUDA-graph +> configuration. + ### Multi-Stage Diffusion Model ModelOpt checkpoints must be routed to the stage whose checkpoint contains the diff --git a/docs/user_guide/quantization/mxfp4.md b/docs/user_guide/quantization/mxfp4.md index 401a55ad4d8..b18b3e7b39f 100644 --- a/docs/user_guide/quantization/mxfp4.md +++ b/docs/user_guide/quantization/mxfp4.md @@ -39,7 +39,7 @@ vLLM-Omni provides two quantization methods with different scale structures: | NVIDIA Blackwell GPU (SM 100+) | ⭕ | | NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | | NVIDIA Ampere GPU (SM 80+) | ⭕ | -| AMD ROCm | ⭕ | +| AMD ROCm (gfx950 / MI355X) | ✅ | | Intel XPU | ⭕ | | Ascend NPU (Atlas 950 A5) | ✅ | diff --git a/docs/user_guide/quantization/overview.md b/docs/user_guide/quantization/overview.md index 38d027f12b8..3a4a45573ee 100644 --- a/docs/user_guide/quantization/overview.md +++ b/docs/user_guide/quantization/overview.md @@ -20,7 +20,7 @@ type has a different quantization scope. | NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | | NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ | | NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ | -| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ❌ | +| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ | | Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ | | Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | diff --git a/examples/model_display_names.yml b/examples/model_display_names.yml new file mode 100644 index 00000000000..be3f8e3810d --- /dev/null +++ b/examples/model_display_names.yml @@ -0,0 +1,12 @@ +# Canonical display names for model examples that have both offline and online pages. +# The docs generator combines these names with the serving mode for each page H1. +aura_omni: AURA Omni +covo_audio: Covo-Audio-Chat +dynin_omni: Dynin-Omni +lance: Lance +mimo_audio: MiMo-Audio +ming_flash_omni: Ming-flash-omni 2.0 +qwen2_5_omni: Qwen2.5-Omni +qwen3_omni: Qwen3-Omni +sensenova_u1: SenseNova-U1 +step_audio2: Step-Audio2 diff --git a/examples/offline_inference/audiox/README.md b/examples/offline_inference/audiox/README.md deleted file mode 100644 index a0b46f9196f..00000000000 --- a/examples/offline_inference/audiox/README.md +++ /dev/null @@ -1,40 +0,0 @@ -# AudioX offline inference - -Generate audio with the [AudioX](https://zeyuet.github.io/AudioX/) MMDiT diffusion -pipeline (`AudioXPipeline`). Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. - -## Prerequisites - -Download a vLLM-Omni weight bundle (component-sharded safetensors): - -```bash -huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights -``` - -The Hugging Face id `zhangj1an/AudioX` also works directly without prefetching. - -## Usage - -```bash -# Text-to-audio only (default uses zhangj1an/AudioX from the Hub): -python end2end.py --tasks t2a - -# All six tasks against a local bundle and a sample video for v2*/tv2*: -python end2end.py \ - --model ./audiox_weights \ - --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 - -# Subset of tasks, custom seed and steps: -python end2end.py --tasks t2a tv2a --num-inference-steps 100 --seed 0 -``` - -## Arguments - -- `--model`: HF id or local bundle path (default: `zhangj1an/AudioX`). -- `--tasks`: any subset of `t2a t2m v2a v2m tv2a tv2m` (default: all). -- `--video`: video file/URL — required for `v2*` and `tv2*`. -- `--reference-audio`: optional audio prompt (audio-conditioned generation). -- `--num-inference-steps`, `--guidance-scale`, `--seed`, `--seconds-total`, - `--sample-rate`, `--output-dir`: generation knobs. - -Outputs land in `<output-dir>/<task>.wav` as 16-bit stereo WAV. diff --git a/examples/offline_inference/audiox/end2end.py b/examples/offline_inference/audiox/end2end.py deleted file mode 100644 index 34b5425724e..00000000000 --- a/examples/offline_inference/audiox/end2end.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""End-to-end AudioX offline example covering the 6 t2*/v2*/tv2* tasks. - -Provide a directory with the **vLLM-Omni AudioX safetensors bundle** (e.g. from -``zhangj1an/AudioX`` on Hugging Face):: - - huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights - python end2end.py --model ./audiox_weights - python end2end.py --model ./audiox_weights --tasks t2a tv2a -""" - -from __future__ import annotations - -import argparse -import time -from pathlib import Path - -import soundfile -import torch -import torchaudio.functional as TF - -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.platforms import current_omni_platform - -ROOT = Path(__file__).resolve().parent - -SAMPLE_PROMPTS: dict[str, str] = { - "t2a": "Fireworks burst twice, followed by a period of silence before a clock begins ticking.", - "t2m": "Uplifting ukulele tune for a travel vlog", - "v2a": "", - "v2m": "", - "tv2a": "drum beating sound and human talking", - "tv2m": "uplifting music matching the scene", -} - -ALL_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m") -VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"}) -TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"}) - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description="AudioX offline end-to-end (6 t2*/v2*/tv2* tasks).") - p.add_argument("--model", default="zhangj1an/AudioX", help="HF id or local AudioX bundle path.") - p.add_argument("--tasks", nargs="+", default=list(ALL_TASKS), choices=ALL_TASKS) - p.add_argument("--video", default="", help="Video path / URL (required for v2*/tv2*).") - p.add_argument("--reference-audio", default="", help="Optional audio prompt for audio-conditioned generation.") - p.add_argument("--output-dir", default=str(ROOT / "audiox_task_outputs")) - p.add_argument("--num-inference-steps", type=int, default=250) - p.add_argument("--seconds-total", type=float, default=10.0) - p.add_argument("--guidance-scale", type=float, default=6.0) - p.add_argument("--seed", type=int, default=42) - p.add_argument("--sample-rate", type=int, default=48000, help="Output WAV rate (resampled if != model rate).") - return p.parse_args() - - -def save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None: - """Write 16-bit PCM WAV. ``audio`` is ``[channels, samples]`` float in [-1, 1].""" - path.parent.mkdir(parents=True, exist_ok=True) - soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16") - - -def main() -> None: - args = parse_args() - - omni = Omni(model=args.model, model_class_name="AudioXPipeline") - - for task in args.tasks: - if task in VIDEO_TASKS and not args.video: - raise SystemExit(f"task={task!r} requires --video") - prompt = SAMPLE_PROMPTS[task] if task in TEXT_TASKS else "" - extra: dict = {"audiox_task": task, "seconds_start": 0.0, "seconds_total": float(args.seconds_total)} - if task in VIDEO_TASKS: - extra["video_path"] = args.video - if args.reference_audio: - extra["audio_path"] = args.reference_audio - - generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - t0 = time.perf_counter() - outputs = omni.generate( - prompt, - OmniDiffusionSamplingParams( - generator=generator, - guidance_scale=args.guidance_scale, - num_inference_steps=args.num_inference_steps, - seed=args.seed, - extra_args=extra, - ), - ) - audio = outputs[0].request_output.multimodal_output.get("audio") - if audio is None: - raise RuntimeError(f"No audio produced for task {task!r}") - audio = torch.as_tensor(audio).detach().cpu().float() - if audio.ndim == 3: - audio = audio[0] - - model_sr = int(outputs[0].request_output.multimodal_output.get("audio_sample_rate") or 44100) - if model_sr != args.sample_rate: - audio = TF.resample(audio, model_sr, args.sample_rate) - - out_path = Path(args.output_dir) / f"{task}.wav" - save_wav(audio, out_path, args.sample_rate) - print(f"[{task}] saved {out_path} ({time.perf_counter() - t0:.2f}s)") - - omni.close() - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/aura_omni/README.md b/examples/offline_inference/aura_omni/README.md new file mode 100644 index 00000000000..b73976a15c8 --- /dev/null +++ b/examples/offline_inference/aura_omni/README.md @@ -0,0 +1,61 @@ +# AURA Omni: Offline inference + +This example runs the native AURA Omni pipeline offline: + +```text +Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav +``` + +The first stage consumes speech and produces a transcript. The AURA stage then combines the transcript with the video frames from the original request and returns text or `<|silent|>`. Non-silent responses are passed to Qwen3-TTS as AURA-generated token ids. + +## Run + +```bash +cd examples/offline_inference/aura_omni +bash run_single_prompt.sh +``` + +Use local media: + +```bash +python end2end.py \ + --audio-path /path/to/input.wav \ + --video-path /path/to/video.mp4 \ + --modalities text,audio +``` + +Base voice clone mode is the default. It uses the AURA reference audio unless +you override it: + +```bash +python end2end.py \ + --tts-task-type Base \ + --tts-ref-audio /data/yrr/rein_test/shuhan.mp3 \ + --tts-x-vector-only-mode +``` + +CustomVoice mode requires stages 2 and 3 in `vllm_omni/deploy/aura_omni.yaml` +to point at a Qwen3-TTS CustomVoice checkpoint: + +```bash +python end2end.py \ + --tts-task-type CustomVoice \ + --tts-speaker Vivian +``` + +For local checkpoints, edit the stage `model` entries in +`vllm_omni/deploy/aura_omni.yaml` or pass a copied deploy config with +`--deploy-config`. + +## GPU Utilization Recommendation + +Set `gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` by stage. +Suggested starting point: + +- Stage 0 (ASR): `0.10` +- Stage 1 (AURA): `0.4` +- Stage 2 (Qwen3-TTS Talker): `0.20` +- Stage 3 (Qwen3-TTS Code2Wav): `0.20` + +Generated text and audio are written to `--output-dir` +(default: `output_aura_omni`). diff --git a/examples/offline_inference/aura_omni/end2end.py b/examples/offline_inference/aura_omni/end2end.py new file mode 100644 index 00000000000..c082875158e --- /dev/null +++ b/examples/offline_inference/aura_omni/end2end.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline AURA Omni example. + +The native AURA pipeline is: + + Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav + +Stage 0 consumes audio and emits a transcript. The stage input processor then +combines that transcript with the original visual payload for AURA. +""" + +from __future__ import annotations + +import os +from typing import NamedTuple + +import numpy as np +import soundfile as sf +import vllm +from vllm import SamplingParams +from vllm.assets.audio import AudioAsset +from vllm.assets.video import VideoAsset, video_to_ndarrays +from vllm.multimodal.media.audio import load_audio +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + DEFAULT_QWEN3_TTS_REF_TEXT, + default_qwen3_tts_ref_audio_path, +) + +SEED = 42 +DEFAULT_MODEL = "aurateam/AURA" +DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/aura_omni.yaml" + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +def _load_audio(audio_path: str | None, sampling_rate: int) -> tuple[np.ndarray, int]: + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + audio_signal, sr = load_audio(audio_path, sr=sampling_rate) + return audio_signal.astype(np.float32), int(sr) + audio_signal, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate + return audio_signal.astype(np.float32), int(sr) + + +def _load_video(video_path: str | None, num_frames: int) -> list[np.ndarray]: + if video_path: + if not os.path.exists(video_path): + raise FileNotFoundError(f"Video file not found: {video_path}") + return video_to_ndarrays(video_path, num_frames=num_frames) + return VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays + + +def build_query(args) -> QueryResult: + prompt = ( + "<|im_start|>user\n" + "<|audio_start|><|audio_pad|><|audio_end|>" + "Please transcribe this speech and use the video context to decide whether a reply is needed.<|im_end|>\n" + "<|im_start|>assistant\n" + ) + audio_data = _load_audio(args.audio_path, args.sampling_rate) + video_frames = _load_video(args.video_path, args.num_frames) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "audio": audio_data, + "video": video_frames, + }, + "additional_information": { + "aura_system_prompt": args.aura_system_prompt, + "tts_task_type": args.tts_task_type, + "tts_language": args.tts_language, + "tts_speaker": args.tts_speaker, + "tts_instruct": args.tts_instruct, + "tts_ref_audio": args.tts_ref_audio, + "tts_ref_text": args.tts_ref_text, + "tts_x_vector_only_mode": args.tts_x_vector_only_mode, + "tts_pass_token_ids": args.tts_pass_token_ids, + }, + }, + limit_mm_per_prompt={"audio": 1, "video": 1}, + ) + + +def build_sampling_params() -> list[SamplingParams]: + asr_sampling = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=256, + detokenize=True, + seed=SEED, + ) + aura_sampling = SamplingParams( + temperature=0.5, + top_p=1.0, + top_k=-1, + max_tokens=256, + detokenize=True, + repetition_penalty=1.0, + seed=SEED, + ) + talker_sampling = SamplingParams( + temperature=0.9, + top_k=50, + max_tokens=4096, + detokenize=False, + repetition_penalty=1.05, + stop_token_ids=[2150], + seed=SEED, + ) + code2wav_sampling = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + detokenize=True, + repetition_penalty=1.0, + seed=SEED, + ) + return [asr_sampling, aura_sampling, talker_sampling, code2wav_sampling] + + +def save_stage_output(stage_outputs, output_dir: str) -> None: + output = stage_outputs.request_output + request_id = output.request_id + if stage_outputs.final_output_type == "text": + text_output = output.outputs[0].text + out_txt = os.path.join(output_dir, f"{request_id}.txt") + with open(out_txt, "w", encoding="utf-8") as f: + f.write(str(text_output).strip() + "\n") + print(f"Request ID: {request_id}, text saved to {out_txt}") + print(text_output) + return + + if stage_outputs.final_output_type == "audio": + mm = output.outputs[0].multimodal_output + if not mm or "audio" not in mm: + print(f"Request ID: {request_id}, no audio output") + return + audio = mm["audio"] + if isinstance(audio, list): + import torch + + audio = torch.cat([(item if hasattr(item, "detach") else torch.tensor(item)).flatten() for item in audio]) + sample_rate = int(mm.get("sr", 24000)) + audio_numpy = audio.float().detach().cpu().numpy().flatten() + out_wav = os.path.join(output_dir, f"{request_id}.wav") + sf.write(out_wav, audio_numpy, samplerate=sample_rate, format="WAV") + print(f"Request ID: {request_id}, audio saved to {out_wav}") + + +def main(args) -> None: + print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20) + os.makedirs(args.output_dir, exist_ok=True) + + query = build_query(args) + prompts = [query.inputs for _ in range(args.num_prompts)] + if args.modalities: + modalities = [item.strip() for item in args.modalities.split(",") if item.strip()] + for prompt in prompts: + prompt["modalities"] = modalities + + omni_kwargs = vars(args).copy() + omni_kwargs["model"] = args.model + omni = Omni(**omni_kwargs) + sampling_params_list = build_sampling_params()[: omni.num_stages] + + try: + for stage_outputs in omni.generate(prompts, sampling_params_list): + save_stage_output(stage_outputs, args.output_dir) + finally: + omni.close() + + +def parse_args(): + parser = FlexibleArgumentParser(description="AURA Omni offline inference example") + parser.add_argument("--model", default=DEFAULT_MODEL, help="Top-level model name/path.") + parser.add_argument("--deploy-config", default=DEFAULT_DEPLOY_CONFIG, help="AURA Omni deploy config.") + parser.add_argument("--audio-path", "-a", default=None, help="Local audio file. Uses a built-in asset by default.") + parser.add_argument("--video-path", "-v", default=None, help="Local video file. Uses a built-in asset by default.") + parser.add_argument("--num-frames", type=int, default=16, help="Number of video frames to sample.") + parser.add_argument("--sampling-rate", type=int, default=16000, help="Audio sampling rate for ASR input.") + parser.add_argument("--num-prompts", type=int, default=1, help="Number of repeated prompts.") + parser.add_argument("--modalities", default="text,audio", help="Comma-separated final modalities.") + parser.add_argument("--output-dir", default="output_aura_omni", help="Directory for generated files.") + parser.add_argument( + "--aura-system-prompt", + default=( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in Chinese." + ), + help="System prompt for the AURA stage.", + ) + parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"], help="Qwen3-TTS task.") + parser.add_argument("--tts-language", default="Chinese", help="Qwen3-TTS language.") + parser.add_argument("--tts-speaker", default="Vivian", help="Qwen3-TTS speaker.") + parser.add_argument("--tts-instruct", default="", help="Optional Qwen3-TTS style instruction.") + parser.add_argument( + "--tts-ref-audio", + default=default_qwen3_tts_ref_audio_path(), + help="Base-mode reference audio path/URL.", + ) + parser.add_argument( + "--tts-ref-text", + default=DEFAULT_QWEN3_TTS_REF_TEXT, + help="Base-mode reference audio transcript.", + ) + parser.add_argument( + "--tts-x-vector-only-mode", + action="store_true", + help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).", + ) + parser.add_argument( + "--tts-pass-token-ids", + action="store_true", + help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.", + ) + parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.") + parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.") + parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.") + nullify_stage_engine_defaults(parser) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/aura_omni/run_single_prompt.sh b/examples/offline_inference/aura_omni/run_single_prompt.sh new file mode 100755 index 00000000000..f4db0ac6a02 --- /dev/null +++ b/examples/offline_inference/aura_omni/run_single_prompt.sh @@ -0,0 +1,8 @@ +#!/usr/bin/env bash +set -euo pipefail + +python end2end.py \ + --model aurateam/AURA \ + --deploy-config vllm_omni/deploy/aura_omni.yaml \ + --modalities text,audio \ + "$@" diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md deleted file mode 100644 index f2685d56fd3..00000000000 --- a/examples/offline_inference/bagel/README.md +++ /dev/null @@ -1,391 +0,0 @@ -# BAGEL-7B-MoT - -## Setup - -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. - -## Architecture - -BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies: - -| Topology | Stages | Description | -| :------- | :----- | :---------- | -| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. | -| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. | - -Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`. - -## Quick Start - -```bash -cd examples/offline_inference/bagel - -# Default two-stage mode (auto-detected) -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" - -# Single-stage mode -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml -``` - -> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices. - -## Modality Control - -Control the mode using the `--modality` argument: - -| Modality | Input | Output | Description | -| :------- | :---- | :----- | :---------- | -| `text2img` | Text | Image | Generate images from text prompts | -| `img2img` | Image + Text | Image | Transform images using text guidance | -| `img2text` | Image + Text | Text | Generate text descriptions from images | -| `text2text` | Text | Text | Pure text generation (language model mode) | - -### Text to Image (text2img) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --steps 50 -``` - -### Image to Image (img2img) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2img \ - --image-path /path/to/image.jpg \ - --prompts "Let the woman wear a blue dress" -``` - -### Image to Text (img2text) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "Describe this image in detail" -``` - -### Text to Text (text2text) - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2text \ - --prompts "What is the capital of France?" - -# Load prompts from a text file (one prompt per line): -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2text \ - --txt-prompts /path/to/prompts.txt -``` - -## Think Mode - -Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts. - -- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation. -- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise. - -```bash -# Think + text2img: plan before generating -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A futuristic city with flying cars" \ - --think \ - --max-think-tokens 1000 - -# Think + img2img: reason about the edit -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2img \ - --image-path /path/to/image.jpg \ - --prompts "Make it look like a watercolor painting" \ - --think - -# Think + img2text: reason before describing -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality img2text \ - --image-path /path/to/image.jpg \ - --prompts "What is happening in this image?" \ - --think - -# Think + text2text: chain-of-thought reasoning -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2text \ - --prompts "Solve: 23 * 47" \ - --think -``` - -Think mode parameters: - -| Argument | Default | Description | -| :------- | :------ | :---------- | -| `--think` | `False` | Enable thinking mode | -| `--max-think-tokens` | `1000` | Maximum tokens for think generation | -| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation | -| `--text-temperature` | `0.3` | Temperature for text generation sampling | - -## Classifier-Free Guidance (CFG) - -CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`). - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A photorealistic portrait" \ - --cfg-text-scale 6.0 \ - --cfg-img-scale 2.0 \ - --negative-prompt "blurry, low quality, distorted" \ - --cfg-interval 0.4 1.0 \ - --cfg-renorm-type global \ - --cfg-renorm-min 0.0 -``` - -| Argument | Default | Description | -| :------- | :------ | :---------- | -| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) | -| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) | -| `--negative-prompt` | `None` | Negative prompt for CFG conditioning | -| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps | -| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` | -| `--cfg-renorm-min` | `None` | Minimum renormalization value | -| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel | - -## Deployment Topologies - -### Two-Stage (Default) - -The default topology auto-detected from the model. No extra flags needed. - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" -``` - -The pipeline is defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1. - -### Single-Stage - -Pass the single-stage deploy config via `--deploy-config`: - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml -``` - -See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry. - -### Tensor Parallelism (TP) - -For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`: - -```yaml -# Example: TP=2 on GPUs 0,1 for the Thinker stage -stages: - - stage_id: 0 - tensor_parallel_size: 2 - devices: "0,1" -``` - -### VAE Patch Parallelism - -[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes. - -**Bagel-specific notes:** - -- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only). -- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel. -- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel. -- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker. - -**Single-stage via deploy YAML** (recommended for `end2end.py`): - -```yaml -pipeline: bagel_single_stage -async_chunk: false - -stages: - - stage_id: 0 - max_num_batched_tokens: 32768 - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - devices: "0,1" - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 2 - vae_patch_parallel_size: 2 - default_sampling_params: - seed: 52 -``` - -```bash -cd examples/offline_inference/bagel - -CUDA_VISIBLE_DEVICES=0,1 python end2end.py \ - --model /path/to/BAGEL-7B-MoT \ - --deploy-config /path/to/bagel_single_stage_vae_pp.yaml \ - --modality text2img \ - --prompts "A cute cat" \ - --steps 10 \ - --output ./out_vae_pp -``` - -**Single-stage via `Omni` kwargs** (same flags as online serving): - -```python -from vllm_omni.entrypoints.omni import Omni - -omni = Omni( - model="ByteDance-Seed/BAGEL-7B-MoT", - deploy_config="vllm_omni/deploy/bagel_single_stage.yaml", - tensor_parallel_size=2, - vae_patch_parallel_size=2, - vae_use_tiling=True, -) -# Then call omni.generate(...) as in end2end.py -``` - -**Two-stage (VAE PP on DiT only):** - -```yaml -stages: - - stage_id: 0 - devices: "0" - # AR Thinker — no vae_patch_parallel here - - - stage_id: 1 - devices: "0,1" - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 2 - vae_patch_parallel_size: 2 -``` - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --deploy-config /path/to/bagel_vae_pp.yaml \ - --modality text2img \ - --prompts "A cute cat" -``` - -**Startup log checks:** - -```text -INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it. -``` - -| Setting | Role | -| :------ | :--- | -| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) | -| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) | -| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) | - -#### Hybrid Sharded Data Parallel (HSDP) - -For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. - -1. **Enable HSDP**: Set `use_hsdp: true`. -2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`). -3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups. -4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`). - -Example configuration for HSDP across 4 GPUs: -```yaml - - stage_id: 1 - devices: "0,1,2,3" - parallel_config: - use_hsdp: true - hsdp_shard_size: 4 - hsdp_replicate_size: 1 -``` - -Then pass the custom deploy YAML: - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --deploy-config /path/to/custom_bagel.yaml -``` - -### FP8 Quantization - -```bash -python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \ - --modality text2img \ - --prompts "A cute cat" \ - --quantization fp8 -``` - -## Command Line Reference - -### Core Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name | -| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` | -| `--prompts` | list | `None` | Input text prompts | -| `--txt-prompts` | string | `None` | Path to text file with one prompt per line | -| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) | -| `--output` | string | `.` | Output directory for saved images | -| `--steps` | int | `50` | Number of diffusion inference steps | -| `--seed` | int | `None` | Random seed for reproducibility | - -### Think Mode Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning | -| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation | -| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding | -| `--text-temperature` | float | `0.3` | Sampling temperature for text generation | - -### CFG Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale | -| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale | -| `--negative-prompt` | string | `None` | Negative prompt for CFG | -| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) | -| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` | -| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` | -| `--cfg-renorm-min` | float | `None` | Minimum renormalization value | - -### Engine Arguments - -| Argument | Type | Default | Description | -| :------- | :--- | :------ | :---------- | -| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) | -| `--worker-backend` | choice | `process` | `process` or `ray` | -| `--ray-address` | string | `None` | Ray cluster address | -| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) | -| `--log-stats` | flag | `False` | Enable statistics logging | -| `--init-timeout` | int | `300` | Initialization timeout (seconds) | -| `--batch-timeout` | int | `5` | Batch timeout (seconds) | -| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations | - -## FAQ - -- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML. - -**Two-stage VRAM usage:** - -| Stage | VRAM | -| :---- | :--- | -| Stage 0 (Thinker) | **15.04 GiB + KV Cache** | -| Stage 1 (DiT) | **26.50 GiB** | -| Total | **~42 GiB + KV Cache** | - -**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process. diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py deleted file mode 100644 index 08868725315..00000000000 --- a/examples/offline_inference/bagel/end2end.py +++ /dev/null @@ -1,278 +0,0 @@ -import argparse -import os - -from vllm_omni.inputs.data import OmniPromptType -from vllm_omni.model_executor.stage_input_processors.bagel import ( - GEN_THINK_SYSTEM_PROMPT, - VLM_THINK_SYSTEM_PROMPT, -) - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--model", - default="ByteDance-Seed/BAGEL-7B-MoT", - help="Path to merged model directory.", - ) - parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.") - parser.add_argument( - "--txt-prompts", - type=str, - default=None, - help="Path to a .txt file with one prompt per line (preferred).", - ) - parser.add_argument("--prompt-type", default="text", choices=["text"]) - - parser.add_argument( - "--modality", - default="text2img", - choices=["text2img", "img2img", "img2text", "text2text"], - help="Modality mode to control stage execution.", - ) - - parser.add_argument( - "--image-path", - type=str, - default=None, - help="Path to input image for img2img.", - ) - - parser.add_argument( - "--output", - type=str, - default=".", - help="Output directory to save images.", - ) - - # OmniLLM init args - parser.add_argument("--log-stats", action="store_true", default=False) - parser.add_argument("--init-sleep-seconds", type=int, default=20) - parser.add_argument("--batch-timeout", type=int, default=5) - parser.add_argument("--init-timeout", type=int, default=300) - parser.add_argument("--shm-threshold-bytes", type=int, default=65536) - parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"]) - parser.add_argument("--ray-address", type=str, default=None) - parser.add_argument( - "--deploy-config", - type=str, - default=None, - help="Path to deploy YAML. If unset, auto-loads vllm_omni/deploy/bagel.yaml based on the HF model_type.", - ) - parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.") - - parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)") - parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)") - parser.add_argument( - "--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)" - ) - parser.add_argument( - "--cfg-parallel-size", - type=int, - default=1, - choices=[1, 2, 3], - help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).", - ) - parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.") - parser.add_argument( - "--cfg-interval", - type=float, - nargs=2, - default=None, - help="CFG interval [start, end] (default: pipeline default)", - ) - parser.add_argument( - "--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel" - ) - parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min") - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - parser.add_argument( - "--quantization", - type=str, - default=None, - help="Quantization method (e.g. 'fp8').", - ) - parser.add_argument( - "--think", - action="store_true", - default=False, - help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.", - ) - parser.add_argument( - "--max-think-tokens", - type=int, - default=1000, - help="Maximum number of tokens for thinking text generation (default: 1000).", - ) - parser.add_argument( - "--do-sample", - action="store_true", - default=False, - help="Enable sampling for text generation (default: greedy).", - ) - parser.add_argument( - "--text-temperature", - type=float, - default=0.3, - help="Temperature for text generation sampling (default: 0.3).", - ) - - args = parser.parse_args() - return args - - -def main(): - args = parse_args() - os.makedirs(args.output, exist_ok=True) - model_name = args.model - prompts: list[OmniPromptType] = [] - try: - if getattr(args, "txt_prompts", None) and args.prompt_type == "text": - with open(args.txt_prompts, encoding="utf-8") as f: - lines = [ln.strip() for ln in f.readlines()] - prompts = [ln for ln in lines if ln != ""] - print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}") - else: - prompts = args.prompts - except Exception as e: - print(f"[Error] Failed to load prompts: {e}") - raise - - if not prompts: - prompts = ["A cute cat"] - print(f"[Info] No prompts provided, using default: {prompts}") - - from PIL import Image - - from vllm_omni.entrypoints.omni import Omni - - omni_kwargs = vars(args).copy() - omni_kwargs.pop("model", None) - deploy_config = args.deploy_config - if args.think and deploy_config is None: - deploy_config = "vllm_omni/deploy/bagel_think.yaml" - print(f"[Info] Think mode enabled, using deploy config: {deploy_config}") - if deploy_config: - omni_kwargs["deploy_config"] = deploy_config - - if args.quantization: - omni_kwargs["quantization_config"] = args.quantization - - omni = Omni(model=model_name, **omni_kwargs) - - formatted_prompts = [] - for p in prompts: - if args.modality == "img2img": - if not args.image_path or not os.path.exists(args.image_path): - raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}") - loaded_image = Image.open(args.image_path).convert("RGB") - think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else "" - final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>" - prompt_dict = { - "prompt": final_prompt_text, - "multi_modal_data": {"img2img": loaded_image}, - "modalities": ["img2img"], - } - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - formatted_prompts.append(prompt_dict) - elif args.modality == "img2text": - if args.image_path: - loaded_image = Image.open(args.image_path).convert("RGB") - think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else "" - final_prompt_text = ( - f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n" - ) - prompt_dict = { - "prompt": final_prompt_text, - "multi_modal_data": {"image": loaded_image}, - "modalities": ["text"], - } - formatted_prompts.append(prompt_dict) - elif args.modality == "text2text": - think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else "" - final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>" - prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]} - formatted_prompts.append(prompt_dict) - else: - think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else "" - final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>" - prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]} - if args.negative_prompt is not None: - prompt_dict["negative_prompt"] = args.negative_prompt - formatted_prompts.append(prompt_dict) - - params_list = omni.default_sampling_params_list - # Bagel exposes 1 sampling param set for single-stage (DiT-only) and - # 2 for two-stage (Thinker + DiT). This heuristic may need updating - # if future pipelines break that 1:1 mapping. - is_single_stage = len(params_list) == 1 - - diffusion_params_idx = 0 if is_single_stage else (1 if len(params_list) > 1 else 0) - diffusion_params = params_list[diffusion_params_idx] - - if args.modality in ("text2img", "img2img"): - diffusion_params.num_inference_steps = args.steps # type: ignore - diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore - if args.seed is not None: - diffusion_params.seed = args.seed # type: ignore - - extra = getattr(diffusion_params, "extra_args", {}) or {} - extra["cfg_text_scale"] = args.cfg_text_scale - extra["cfg_img_scale"] = args.cfg_img_scale - if args.cfg_interval is not None: - extra["cfg_interval"] = tuple(args.cfg_interval) - if args.cfg_renorm_type is not None: - extra["cfg_renorm_type"] = args.cfg_renorm_type - if args.cfg_renorm_min is not None: - extra["cfg_renorm_min"] = args.cfg_renorm_min - if args.negative_prompt is not None: - extra["negative_prompt"] = args.negative_prompt - - needs_text_gen = is_single_stage and (args.think or args.modality in ("text2text", "img2text")) - if needs_text_gen: - if args.think: - extra["think"] = True - extra["max_think_tokens"] = args.max_think_tokens - extra["do_sample"] = args.do_sample - extra["text_temperature"] = args.text_temperature - diffusion_params.extra_args = extra # type: ignore - - omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) - - img_idx = 0 - for req_output in omni_outputs: - # 2-stage think mode: text output from thinker stage - ro = getattr(req_output, "request_output", None) - if ro and getattr(ro, "outputs", None): - txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs) - if txt: - if args.think: - print(f"[Think]\n{txt}") - else: - print(f"[Output] Text:\n{txt}") - - # Single-stage DiT: text from custom_output - custom = getattr(req_output, "_custom_output", {}) or {} - if custom.get("think_text"): - print(f"[Think]\n{custom['think_text']}") - if custom.get("text_output"): - print(f"[Output] Text:\n{custom['text_output']}") - - images = getattr(req_output, "images", None) - if not images: - continue - - for j, img in enumerate(images): - save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") - img.save(save_path) - print(f"[Output] Saved image to {save_path}") - img_idx += 1 - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/cosmos3/inputs/v2v.json b/examples/offline_inference/cosmos3/inputs/v2v.json deleted file mode 100644 index c16c4bbb719..00000000000 --- a/examples/offline_inference/cosmos3/inputs/v2v.json +++ /dev/null @@ -1,14 +0,0 @@ -{ - "prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.", - "negative_prompt": "", - "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4", - "height": 720, - "width": 1280, - "num_frames": 189, - "num_inference_steps": 35, - "guidance_scale": 6.0, - "flow_shift": 10.0, - "fps": 24, - "condition_frame_indexes_vision": [0, 1], - "condition_video_keep": "first" -} diff --git a/examples/offline_inference/covo_audio/README.md b/examples/offline_inference/covo_audio/README.md index d13783d392f..80a822fa290 100644 --- a/examples/offline_inference/covo_audio/README.md +++ b/examples/offline_inference/covo_audio/README.md @@ -1,4 +1,4 @@ -# Covo-Audio-Chat (Offline Inference) +# Covo-Audio-Chat: Offline inference ## Setup diff --git a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py index b1c74e2b345..fb14d9d33c7 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/custom_pipeline.py @@ -2,14 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import logging -from typing import Any -import PIL.Image import torch -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import QwenImageEditPipeline -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch logger = logging.getLogger(__name__) @@ -18,58 +16,14 @@ class CustomPipeline(QwenImageEditPipeline): def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 1.0, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch): """Forward pass for image editing with dummy trajectory data.""" - # Call parent's forward to get the normal output - output = super().forward( - req=req, - prompt=prompt, - negative_prompt=negative_prompt, - image=image, - true_cfg_scale=true_cfg_scale, - height=height, - width=width, - num_inference_steps=num_inference_steps, - sigmas=sigmas, - guidance_scale=guidance_scale, - num_images_per_prompt=num_images_per_prompt, - generator=generator, - latents=latents, - prompt_embeds=prompt_embeds, - prompt_embeds_mask=prompt_embeds_mask, - negative_prompt_embeds=negative_prompt_embeds, - negative_prompt_embeds_mask=negative_prompt_embeds_mask, - output_type=output_type, - attention_kwargs=attention_kwargs, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - max_sequence_length=max_sequence_length, - ) + # Customize num_inference_steps that best suits your model + actual_num_steps = req.sampling_params.num_inference_steps or 50 + req.sampling_params.num_inference_steps = actual_num_steps - # Get actual num_inference_steps used - actual_num_steps = req.sampling_params.num_inference_steps or num_inference_steps + # Call parent's forward to get the normal output + output = super().forward(req=req) # Create dummy trajectory data dummy_trajectory_latents = torch.randn(actual_num_steps, 1, 16, 64, 64, dtype=torch.float32) diff --git a/examples/offline_inference/dreamzero/client_schedule.py b/examples/offline_inference/dreamzero/client_schedule.py new file mode 100644 index 00000000000..e4dbad55a03 --- /dev/null +++ b/examples/offline_inference/dreamzero/client_schedule.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``.""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np + +# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py). +ACTION_HORIZON = 24 +DEFAULT_ACTION_DIM = 8 +RELATIVE_OFFSETS = [-23, -16, -8, 0] +DEFAULT_NUM_AR_CHUNKS = 15 + +CAMERA_FILES = { + "observation/exterior_image_0_left": "exterior_image_1_left.mp4", + "observation/exterior_image_1_left": "exterior_image_2_left.mp4", + "observation/wrist_image_left": "wrist_image_left.mp4", +} + + +def build_frame_schedule( + total_frames: int, + num_chunks: int, + *, + logger: logging.Logger | None = None, +) -> list[list[int]]: + """Build the frame index schedule for multi-frame chunks. + + Args: + total_frames: Number of frames available in each camera video. + num_chunks: Number of 4-frame chunks to schedule after the initial frame. + + Returns: + A list of frame-index lists. Each inner list has four indices. The + returned list may be shorter than ``num_chunks`` when the videos run + out of frames (upstream stops early instead of erroring). + """ + chunks: list[list[int]] = [] + current_frame = 23 + for _ in range(num_chunks): + indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + if logger is not None: + logger.info( + "Frame %s >= %s, stopping at %s chunks", + indices[-1], + total_frames, + len(chunks), + ) + break + chunks.append(indices) + current_frame += ACTION_HORIZON + return chunks + + +def make_obs_from_video( + camera_frames: dict[str, np.ndarray], + frame_indices: list[int], + *, + prompt: str, + session_id: str, +) -> dict[str, Any]: + """Build an observation dict from real video frames. + + For one frame each image key is ``(H, W, 3)``. For four frames each key is + ``(4, H, W, 3)``. + """ + obs: dict[str, Any] = {} + for camera_key, all_frames in camera_frames.items(): + selected = all_frames[frame_indices] + obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected + + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def build_ar_observations( + camera_frames: dict[str, np.ndarray], + *, + prompt: str, + session_id: str, + num_chunks: int = DEFAULT_NUM_AR_CHUNKS, + repeat_chunk_observations: bool = False, + logger: logging.Logger | None = None, +) -> list[dict[str, Any]]: + """Build the AR observation sequence used by upstream ``test_client_AR.py``. + + Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk. + + ``num_chunks`` counts only the 4-frame chunks after the initial frame, not + the total number of inferences. + """ + if num_chunks < 0: + raise ValueError("num_chunks must be non-negative") + + total_frames = min(frames.shape[0] for frames in camera_frames.values()) + observations = [ + make_obs_from_video( + camera_frames, + [0], + prompt=prompt, + session_id=session_id, + ) + ] + + chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger) + if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks: + while len(chunk_schedule) < num_chunks: + chunk_schedule.append(chunk_schedule[-1]) + + for frame_indices in chunk_schedule: + observations.append( + make_obs_from_video( + camera_frames, + frame_indices, + prompt=prompt, + session_id=session_id, + ) + ) + return observations diff --git a/examples/offline_inference/dreamzero/export_prediction_video.py b/examples/offline_inference/dreamzero/export_prediction_video.py index 82f56ff26e5..4c560418c2b 100644 --- a/examples/offline_inference/dreamzero/export_prediction_video.py +++ b/examples/offline_inference/dreamzero/export_prediction_video.py @@ -11,6 +11,11 @@ import cv2 import numpy as np import torch +from client_schedule import ( + CAMERA_FILES, + DEFAULT_NUM_AR_CHUNKS, + build_ar_observations, +) from PIL import Image from vllm_omni import Omni @@ -20,13 +25,6 @@ WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension" ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets" DEFAULT_SESSION_PREFIX = "dreamzero-export" -RELATIVE_OFFSETS = [-23, -16, -8, 0] -ACTION_HORIZON = 24 -CAMERA_FILES = { - "observation/exterior_image_0_left": "exterior_image_1_left.mp4", - "observation/exterior_image_1_left": "exterior_image_2_left.mp4", - "observation/wrist_image_left": "wrist_image_left.mp4", -} def _parse_args() -> argparse.Namespace: @@ -49,6 +47,17 @@ def _parse_args() -> argparse.Namespace: default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", ) parser.add_argument("--session-id", default=None) + parser.add_argument( + "--num-chunks", + type=int, + default=DEFAULT_NUM_AR_CHUNKS, + help="Number of 4-frame chunks to send after the initial frame (default: 15).", + ) + parser.add_argument( + "--repeat-chunk-observations", + action="store_true", + help="Repeat the last valid chunk observation when bundled assets run out of frames.", + ) parser.add_argument("--save-input-video", action="store_true") parser.add_argument("--save-gif", action="store_true") parser.add_argument("--save-actions", action="store_true") @@ -84,57 +93,25 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]: return camera_frames -def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: - chunks: list[list[int]] = [] - current_frame = 23 - for _ in range(num_chunks): - indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] - if indices[-1] >= total_frames: - break - chunks.append(indices) - current_frame += ACTION_HORIZON - return chunks - - -def _make_obs_from_video( - camera_frames: dict[str, np.ndarray], - frame_indices: list[int], - *, +def _build_observations( + video_dir: Path, prompt: str, session_id: str, -) -> dict: - obs: dict = {} - for camera_key, all_frames in camera_frames.items(): - selected = all_frames[frame_indices] - obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected - - obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) - obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) - obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) - obs["prompt"] = prompt - obs["session_id"] = session_id - return obs - - -def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]: + *, + num_chunks: int, + repeat_chunk_observations: bool, +) -> tuple[dict[str, np.ndarray], list[dict]]: camera_frames = _load_camera_frames(video_dir) - total_frames = min(frames.shape[0] for frames in camera_frames.values()) - chunks = _build_frame_schedule(total_frames, 1) - observations = [ - _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id), - ] - if chunks: - observations.append( - _make_obs_from_video( - camera_frames, - chunks[0], - prompt=prompt, - session_id=session_id, - ) - ) + observations = build_ar_observations( + camera_frames, + prompt=prompt, + session_id=session_id, + num_chunks=num_chunks, + repeat_chunk_observations=repeat_chunk_observations, + ) if len(observations) < 2: - raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.") - return camera_frames, observations[:2] + raise RuntimeError("Need at least two DreamZero observations to export a prediction video.") + return camera_frames, observations def _extract_latents(output: OmniRequestOutput) -> torch.Tensor: @@ -229,6 +206,26 @@ def _run_generation( return omni, outputs +def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray: + """Decode server-side accumulated latents when the pipeline supports it.""" + stage_client = omni.engine.stage_clients[0] + engine = getattr(stage_client, "_engine", None) + if engine is None: + raise RuntimeError("DreamZero export requires inline diffusion stage access.") + + decoded = engine.executor.collective_rpc( + "decode_accumulated_video_latents_to_uint8", + args=(session_id,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if isinstance(decoded, torch.Tensor): + decoded = decoded.numpy() + if not isinstance(decoded, np.ndarray): + raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}") + return decoded + + def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray: stage_client = omni.engine.stage_clients[0] engine = getattr(stage_client, "_engine", None) @@ -248,6 +245,15 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray: return decoded +def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray: + """Decode full AR rollout by concatenating per-step latents from all outputs.""" + latent_steps = [_extract_latents(output) for output in outputs] + latent_t = [step.shape[2] for step in latent_steps] + print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}") + full_latents = torch.cat(latent_steps, dim=2) + return _decode_with_worker(omni, full_latents) + + def main() -> None: args = _parse_args() session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}" @@ -256,6 +262,12 @@ def main() -> None: video_dir=args.video_dir, prompt=args.prompt, session_id=session_id, + num_chunks=args.num_chunks, + repeat_chunk_observations=args.repeat_chunk_observations, + ) + print( + f"Running {len(observations)} inferences " + f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)" ) args.output_dir.mkdir(parents=True, exist_ok=True) @@ -273,17 +285,19 @@ def main() -> None: deploy_config_path=args.deploy_config, observations=observations, ) - latent_steps = [_extract_latents(output) for output in outputs] - full_latents = torch.cat(latent_steps, dim=2) - frames = _decode_with_worker(omni, full_latents) + frames = _decode_video(omni, session_id, outputs) finally: if omni is not None: omni.close() + expected_frames = 8 * len(observations) + 1 mp4_path = args.output_dir / f"{args.output_stem}.mp4" _write_mp4(mp4_path, frames, fps=args.fps) - print(f"SAVED_MP4={mp4_path}") + print( + f"SAVED_MP4={mp4_path} frames={len(frames)} " + f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s" + ) if args.save_gif: gif_path = args.output_dir / f"{args.output_stem}.gif" @@ -294,8 +308,10 @@ def main() -> None: npz_path = args.output_dir / f"{args.output_stem}_actions.npz" np.savez( npz_path, - step0=np.asarray(outputs[0].multimodal_output.get("actions")), - step1=np.asarray(outputs[1].multimodal_output.get("actions")), + **{ + f"step{index}": np.asarray(output.multimodal_output.get("actions")) + for index, output in enumerate(outputs) + }, ) print(f"SAVED_ACTIONS={npz_path}") diff --git a/examples/offline_inference/dynin_omni/README.md b/examples/offline_inference/dynin_omni/README.md index 0b1cc22f216..19fa8c69ea5 100644 --- a/examples/offline_inference/dynin_omni/README.md +++ b/examples/offline_inference/dynin_omni/README.md @@ -1,4 +1,4 @@ -# Dynin-Omni Offline End2End Example +# Dynin-Omni: Offline inference This folder contains a unified offline inference entrypoint: diff --git a/examples/offline_inference/helios/README.md b/examples/offline_inference/helios/README.md deleted file mode 100644 index 2c7ba525ef6..00000000000 --- a/examples/offline_inference/helios/README.md +++ /dev/null @@ -1,193 +0,0 @@ -# Helios Video Generation - -Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants: - -| Variant | Description | Key Features | -|---------|-------------|--------------| -| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` | -| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support | -| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` | - -## Setup - -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. - -## Run Examples - -Get into the example folder: -```bash -cd examples/offline_inference/helios -``` - -### Text-to-Video (T2V) - -**Helios-Base** (Stage 1 only): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Base \ - --sample-type t2v \ - --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ - --guidance-scale 5.0 \ - --output helios_t2v_base.mp4 -``` - -**Helios-Mid** (Stage 2 + CFG-Zero*): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Mid \ - --sample-type t2v \ - --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 \ - --output helios_t2v_mid.mp4 -``` - -**Helios-Distilled** (Stage 2 pyramid + DMD): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Distilled \ - --sample-type t2v \ - --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ - --num-frames 240 \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk \ - --output helios_t2v_distilled.mp4 -``` - -### Image-to-Video (I2V) - -**Helios-Base**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Base \ - --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ - --guidance-scale 5.0 \ - --output helios_i2v_base.mp4 -``` - -**Helios-Mid** (Stage 2 + CFG-Zero*): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Mid \ - --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 \ - --output helios_i2v_mid.mp4 -``` - -**Helios-Distilled**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Distilled \ - --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \ - --num-frames 240 \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk \ - --output helios_i2v_distilled.mp4 -``` - -### Video-to-Video (V2V) - -**Helios-Base**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Base \ - --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ - --guidance-scale 5.0 \ - --output helios_v2v_base.mp4 -``` - -**Helios-Mid** (Stage 2 + CFG-Zero*): -```bash -python end2end.py \ - --model BestWishYsh/Helios-Mid \ - --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 \ - --output helios_v2v_mid.mp4 -``` - -**Helios-Distilled**: -```bash -python end2end.py \ - --model BestWishYsh/Helios-Distilled \ - --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "A bright yellow Lamborghini speeds along a curving mountain road." \ - --num-frames 240 \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk \ - --output helios_v2v_distilled.mp4 -``` - -## Common Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path | -| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` | -| `--prompt` | — | Text prompt describing the video | -| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) | -| `--image-path` | — | Input image (required for `i2v`) | -| `--video-path` | — | Input video (required for `v2v`) | -| `--height` | `384` | Video height in pixels | -| `--width` | `640` | Video width in pixels | -| `--num-frames` | `99` | Number of output frames | -| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) | -| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) | -| `--seed` | `42` | Random seed | -| `--fps` | `16` | Output video frame rate | -| `--output` | `helios_output.mp4` | Output file path | - -### Stage 2 / Pyramid Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising | -| `--pyramid-num-stages` | `3` | Number of pyramid stages | -| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage | -| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) | - -### CFG-Zero* Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) | -| `--use-zero-init` | off | Zero init for first steps | -| `--zero-steps` | `1` | Number of zero-init steps | - -### Memory & Parallelism - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--vae-use-slicing` | off | Enable VAE slicing | -| `--vae-use-tiling` | off | Enable VAE tiling | -| `--enforce-eager` | off | Disable torch.compile | -| `--enable-cpu-offload` | off | CPU offloading | -| `--enable-layerwise-offload` | off | Layerwise offloading | -| `--ulysses-degree` | `1` | Ulysses SP degree | -| `--ring-degree` | `1` | Ring SP degree | -| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) | -| `--tensor-parallel-size` | `1` | Tensor parallelism size | diff --git a/examples/offline_inference/helios/end2end.py b/examples/offline_inference/helios/end2end.py deleted file mode 100644 index 88c3b865d42..00000000000 --- a/examples/offline_inference/helios/end2end.py +++ /dev/null @@ -1,355 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Helios video generation example supporting T2V, I2V, and V2V. - -Usage (T2V, Helios-Base, Stage 1 only): - python end2end.py \ - --model /path/to/Helios-Base --sample-type t2v \ - --prompt "A serene lakeside sunrise with mist over the water." \ - --height 384 --width 640 --num-frames 99 \ - --num-inference-steps 50 --guidance-scale 5.0 - -Usage (I2V, Helios-Base): - python end2end.py \ - --model /path/to/Helios-Base --sample-type i2v \ - --image-path /path/to/image.jpg \ - --prompt "Description of desired animation." \ - --guidance-scale 5.0 - -Usage (V2V, Helios-Base): - python end2end.py \ - --model /path/to/Helios-Base --sample-type v2v \ - --video-path /path/to/video.mp4 \ - --prompt "Description of desired transformation." \ - --guidance-scale 5.0 - -Usage (Helios-Mid, Stage 2 + CFG-Zero*): - python end2end.py \ - --model /path/to/Helios-Mid --sample-type t2v \ - --prompt "A serene lakeside sunrise with mist over the water." \ - --guidance-scale 5.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 20 20 20 \ - --use-cfg-zero-star --use-zero-init --zero-steps 1 - -Usage (Helios-Distilled, Stage 2 pyramid + DMD): - python end2end.py \ - --model /path/to/Helios-Distilled --sample-type t2v \ - --prompt "A serene lakeside sunrise with mist over the water." \ - --guidance-scale 1.0 \ - --is-enable-stage2 \ - --pyramid-num-inference-steps-list 2 2 2 \ - --is-amplify-first-chunk -""" - -import argparse -import time -from pathlib import Path - -import numpy as np -import torch - -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.platforms import current_omni_platform - - -def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor: - """Load an image and return a (1, C, H, W) tensor normalized to [-1, 1].""" - from PIL import Image - from torchvision import transforms - - image = Image.open(image_path).convert("RGB").resize((width, height)) - transform = transforms.Compose( - [ - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ] - ) - return transform(image).unsqueeze(0) - - -def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor: - """Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1].""" - from torchvision import transforms - from torchvision.io import read_video - - video_frames, _, _ = read_video(video_path, output_format="TCHW") - video_frames = video_frames.float() / 255.0 - - transform = transforms.Compose( - [ - transforms.Resize((height, width), antialias=True), - transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]), - ] - ) - video_frames = torch.stack([transform(f) for f in video_frames]) - return video_frames.permute(1, 0, 2, 3).unsqueeze(0) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).") - parser.add_argument( - "--model", - default="BestWishYsh/Helios-Base", - help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).", - ) - parser.add_argument( - "--sample-type", - choices=["t2v", "i2v", "v2v"], - default="t2v", - help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).", - ) - parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.") - parser.add_argument( - "--negative-prompt", - default=( - "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, " - "images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, " - "incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, " - "misshapen limbs, fused fingers, still picture, messy background, three legs, many people " - "in the background, walking backwards" - ), - help="Negative prompt.", - ) - parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.") - parser.add_argument("--height", type=int, default=384, help="Video height.") - parser.add_argument("--width", type=int, default=640, help="Video width.") - parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.") - parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).") - parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.") - parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.") - - # I2V / V2V - parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.") - parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.") - - # Stage 2 (pyramid multi-stage denoising) - parser.add_argument( - "--is-enable-stage2", - action="store_true", - help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.", - ) - parser.add_argument( - "--pyramid-num-stages", - type=int, - default=3, - help="Number of pyramid stages for Stage 2.", - ) - parser.add_argument( - "--pyramid-num-inference-steps-list", - type=int, - nargs="+", - default=[10, 10, 10], - help="Inference steps per pyramid stage.", - ) - - # DMD - parser.add_argument( - "--is-amplify-first-chunk", - action="store_true", - help="Enable DMD amplification for the first chunk (Helios-Distilled).", - ) - - # CFG Zero Star - parser.add_argument( - "--use-cfg-zero-star", - action="store_true", - help="Enable CFG Zero Star guidance (recommended for Helios-Mid).", - ) - parser.add_argument( - "--use-zero-init", - action="store_true", - help="Use zero initialization for the first denoising steps with CFG-Zero*.", - ) - parser.add_argument( - "--zero-steps", - type=int, - default=1, - help="Number of initial denoising steps using zero prediction (default: 1).", - ) - parser.add_argument( - "--enable-diffusion-pipeline-profiler", - action="store_true", - help="Enable diffusion pipeline profiler to display stage durations.", - ) - parser.add_argument( - "--quantization", - type=str, - default=None, - help="Quantization method (e.g. 'fp8').", - ) - - # Memory & parallelism - parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.") - parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.") - parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") - parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.") - parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.") - parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.") - parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.") - parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.") - parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.") - - return parser.parse_args() - - -def main(): - args = parse_args() - generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - cfg_parallel_size=args.cfg_parallel_size, - tensor_parallel_size=args.tensor_parallel_size, - ) - - omni = Omni( - model=args.model, - enable_layerwise_offload=args.enable_layerwise_offload, - vae_use_slicing=args.vae_use_slicing, - vae_use_tiling=args.vae_use_tiling, - enable_cpu_offload=args.enable_cpu_offload, - parallel_config=parallel_config, - enforce_eager=args.enforce_eager, - enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, - quantization=args.quantization, - ) - - # Validate I2V / V2V arguments - if args.sample_type == "i2v" and not args.image_path: - raise ValueError("--image-path is required for I2V mode.") - if args.sample_type == "v2v" and not args.video_path: - raise ValueError("--video-path is required for V2V mode.") - - # Build extra_args for Helios-specific parameters - extra_args = {} - if args.is_enable_stage2: - extra_args["is_enable_stage2"] = True - extra_args["pyramid_num_stages"] = args.pyramid_num_stages - extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list - if args.is_amplify_first_chunk: - extra_args["is_amplify_first_chunk"] = True - if args.use_cfg_zero_star: - extra_args["use_cfg_zero_star"] = True - if args.use_zero_init: - extra_args["use_zero_init"] = True - extra_args["zero_steps"] = args.zero_steps - - if args.sample_type == "i2v": - image_tensor = load_image_as_tensor(args.image_path, args.height, args.width) - extra_args["image"] = image_tensor - elif args.sample_type == "v2v": - video_tensor = load_video_as_tensor(args.video_path, args.height, args.width) - extra_args["video"] = video_tensor - - # Print generation configuration - print(f"\n{'=' * 60}") - print("Helios Generation Configuration:") - print(f" Model: {args.model}") - print(f" Sample type: {args.sample_type.upper()}") - if args.sample_type == "i2v": - print(f" Image: {args.image_path}") - elif args.sample_type == "v2v": - print(f" Video: {args.video_path}") - print(f" Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}") - print(f" Video size: {args.width}x{args.height}, {args.num_frames} frames") - print(f" Inference steps: {args.num_inference_steps}") - print(f" Guidance scale: {args.guidance_scale}") - if args.is_enable_stage2: - print(f" Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})") - if args.is_amplify_first_chunk: - print(" DMD amplify first chunk: enabled") - if args.use_cfg_zero_star: - print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})") - else: - if args.use_cfg_zero_star: - print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})") - print(" Stage 2: disabled (Stage 1 only)") - print(f"{'=' * 60}\n") - - generation_start = time.perf_counter() - frames = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - }, - OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - generator=generator, - guidance_scale=args.guidance_scale, - num_inference_steps=args.num_inference_steps, - num_frames=args.num_frames, - extra_args=extra_args, - ), - ) - generation_end = time.perf_counter() - generation_time = generation_end - generation_start - - print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") - - # Extract video frames from OmniRequestOutput - if isinstance(frames, list) and len(frames) > 0: - first_item = frames[0] - - if hasattr(first_item, "final_output_type"): - if first_item.final_output_type != "image": - raise ValueError( - f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation." - ) - - if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output: - inner_output = first_item.request_output - if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"): - frames = inner_output.images[0] if inner_output.images else None - if frames is None: - raise ValueError("No video frames found in output.") - elif hasattr(first_item, "images") and first_item.images: - frames = first_item.images - else: - raise ValueError("No video frames found in OmniRequestOutput.") - - # Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C) - if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray): - frames = frames[0] - - output_path = Path(args.output) - output_path.parent.mkdir(parents=True, exist_ok=True) - try: - from diffusers.utils import export_to_video - except ImportError: - raise ImportError("diffusers is required for export_to_video.") - - if isinstance(frames, torch.Tensor): - video_tensor = frames.detach().cpu() - if video_tensor.dim() == 5: - if video_tensor.shape[1] in (3, 4): - video_tensor = video_tensor[0].permute(1, 2, 3, 0) - else: - video_tensor = video_tensor[0] - elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): - video_tensor = video_tensor.permute(1, 2, 3, 0) - if video_tensor.is_floating_point(): - video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5 - video_array = video_tensor.float().numpy() - elif isinstance(frames, np.ndarray): - video_array = frames - else: - video_array = frames - - if isinstance(video_array, np.ndarray) and video_array.ndim == 5: - video_array = video_array[0] - - export_to_video(video_array, str(output_path), fps=args.fps) - print(f"Saved generated video to {output_path}") - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 3ba69c2e216..62836f9633a 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -74,6 +74,12 @@ def parse_args(): parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.") parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.") parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.") + parser.add_argument( + "--stream", + action="store_true", + default=False, + help="Enable streaming CoT display; print AR text token-by-token in real time.", + ) parser.add_argument("--log-stats", action="store_true", default=False) parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") @@ -222,6 +228,8 @@ def main(): prompt_dict["modalities"] = ["text"] formatted_prompts.append(prompt_dict) + from vllm.sampling_params import RequestOutputKind + params_list = list(omni.default_sampling_params_list) from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -253,6 +261,10 @@ def main(): sp.width = args.width elif hasattr(sp, "stop_token_ids"): sp.stop_token_ids = ar_stop_token_ids + # When --stream is set, request DELTA output from the AR stage + # so we can display CoT text token-by-token in real time. + if args.stream and hasattr(sp, "output_kind"): + sp.output_kind = RequestOutputKind.DELTA print(f"\n{'=' * 60}") print("HunyuanImage-3.0 Generation Configuration:") @@ -281,22 +293,32 @@ def main(): print(f" Prompts: {prompts}") print(f"{'=' * 60}\n") - omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list)) + # When --stream is set, print AR CoT text token-by-token in real time. + # Otherwise, collect and print the full AR text once when stage 0 finishes. + omni_outputs = omni.generate( + prompts=formatted_prompts, + sampling_params_list=params_list, + py_generator=True, + use_tqdm=False, + ) img_idx = 0 for req_output in omni_outputs: ro = getattr(req_output, "request_output", None) - txt = "" - if ro and getattr(ro, "outputs", None): - txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs) - if not txt: - ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text") - if isinstance(ar_text, list): - txt = "\n".join(text for text in ar_text if text) - else: - txt = ar_text or "" - if txt: - print(f"[Output] Text:\n{txt}") - + stage_id = getattr(req_output, "stage_id", None) + + # AR stage text — each CompletionOutput.text is already a delta when + # output_kind=DELTA, so we can print it directly (matching the pattern + # in serving_chat.py). + if stage_id == 0 and ro and getattr(ro, "outputs", None): + for o in ro.outputs: + text = getattr(o, "text", "") or "" + if text: + print(text, end="", flush=True) + # Non-streaming: one shot with full text — emit a trailing newline. + if not args.stream: + print(flush=True) + + # Collect images from diffusion stage images = getattr(req_output, "images", None) if not images and ro and hasattr(ro, "images"): images = ro.images @@ -304,7 +326,7 @@ def main(): for j, img in enumerate(images): save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png") img.save(save_path) - print(f"[Output] Saved image to {save_path}") + print(f"\n[Output] Saved image to {save_path}") img_idx += 1 diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 0f2445673b1..12bd978a3aa 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -97,9 +97,16 @@ from PIL import Image from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.model_extras import ( + build_image_to_image_prompt, + get_extra_body_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) from vllm_omni.platforms import current_omni_platform @@ -131,6 +138,15 @@ def parse_args() -> argparse.Namespace: required=True, help="Path(s) to input image file(s) (PNG, JPG, etc.). Can specify multiple images.", ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help=( + "Path to a deploy YAML. Required for multi-stage image-edit pipelines " + "whose deploy config is not auto-loaded." + ), + ) parser.add_argument( "--prompt", type=str, @@ -187,6 +203,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--guidance-scale-2", type=float, default=None, help="image guidance scale for image-to-image generation." ) + parser.add_argument( + "--extra-args", + type=parse_profiler_config, + default=None, + help="JSON object copied to OmniDiffusionSamplingParams.extra_args, e.g. '{\"cfg_text_scale\": 4.0}'.", + ) parser.add_argument( "--output", type=str, @@ -336,7 +358,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--enforce-eager", action="store_true", - help="Disable torch.compile and force eager execution.", + default=None, + help=( + "Disable torch.compile and force eager execution. Left unset (None) " + "so it is only forwarded when explicitly given; " + "otherwise the per-stage deploy YAML value wins." + ), ) parser.add_argument( "--vae-use-slicing", @@ -455,7 +482,7 @@ def main(): } # Initialize Omni with appropriate pipeline - omni = Omni( + omni_kwargs: dict[str, Any] = dict( model=args.model, enable_layerwise_offload=args.enable_layerwise_offload, vae_use_slicing=args.vae_use_slicing, @@ -463,11 +490,17 @@ def main(): cache_backend=args.cache_backend, cache_config=cache_config, parallel_config=parallel_config, - enforce_eager=args.enforce_eager, enable_cpu_offload=args.enable_cpu_offload, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, profiler_config=args.profiler_config, ) + if args.enforce_eager is not None: + omni_kwargs["enforce_eager"] = args.enforce_eager + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config + omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) + declared_extra_body_params = get_extra_body_params(model_class_name) print("Pipeline loaded") profiler_enabled = args.profiler_config is not None @@ -478,6 +511,8 @@ def main(): print(f" Model: {args.model}") print(f" Inference steps: {args.num_inference_steps}") print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") + if args.height is not None or args.width is not None: + print(f" Output size: {args.width or 'auto'}x{args.height or 'auto'}") if isinstance(input_image, list): print(f" Number of input images: {len(input_image)}") for idx, img in enumerate(input_image): @@ -487,7 +522,6 @@ def main(): print( f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}" ) - print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}") print(f"{'=' * 60}\n") generation_start = time.perf_counter() @@ -496,10 +530,20 @@ def main(): print("[Profiler] Starting profiling...") omni.start_profile() - # Generate edited image - sampling_params = OmniDiffusionSamplingParams( - width=args.width, + prompt_dict = build_image_to_image_prompt( + model_class_name=model_class_name, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + input_image=input_image, height=args.height, + width=args.width, + ) + + extra_args_from_cli = dict(args.extra_args or {}) + if args.negative_prompt is not None: + extra_args_from_cli.setdefault("negative_prompt", args.negative_prompt) + + diffusion_params = OmniDiffusionSamplingParams( generator=generator, true_cfg_scale=args.cfg_scale, guidance_scale=args.guidance_scale, @@ -507,17 +551,44 @@ def main(): num_inference_steps=args.num_inference_steps, num_outputs_per_prompt=args.num_outputs_per_prompt, layers=args.layers, + resolution=args.resolution, + height=args.height, + width=args.width, ) - if args.resolution: - sampling_params.resolution = args.resolution - outputs = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - "multi_modal_data": {"image": input_image}, - }, - sampling_params, - ) + if declared_extra_body_params: + apply_declared_extra_args( + diffusion_params, + declared_extra_body_params, + extra_args_from_cli, + ) + else: + diffusion_params.extra_args.update({k: v for k, v in extra_args_from_cli.items() if v is not None}) + + # Build per-stage sampling params for multi-stage models + init_non_diffusion = should_init_extra_args_for_non_diffusion_stages( + model_class_name, + ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + if not sampling_params_list: + sampling_params_list = [diffusion_params] + + diffusion_replaced = False + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + merged_extra = dict(getattr(params, "extra_args", {}) or {}) + merged_extra.update(diffusion_params.extra_args) + diffusion_params.extra_args = merged_extra + sampling_params_list[idx] = diffusion_params + diffusion_replaced = True + elif init_non_diffusion and hasattr(params, "extra_args"): + if params.extra_args is None: + params.extra_args = {} + + if not diffusion_replaced and len(sampling_params_list) == 1: + sampling_params_list = [diffusion_params] + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start @@ -544,17 +615,16 @@ def main(): if not outputs: raise ValueError("No output generated from omni.generate()") - # Extract images from OmniRequestOutput - # omni.generate() returns list[OmniRequestOutput], extract images from request_output.images - first_output = outputs[0] - if not hasattr(first_output, "request_output") or not first_output.request_output: - raise ValueError("No request_output found in OmniRequestOutput") - - req_out = first_output.request_output - if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): - raise ValueError("Invalid request_output structure or missing 'images' key") + images = None + for output in outputs: + images = getattr(output, "images", None) + if images: + break + req_out = getattr(output, "request_output", None) + images = getattr(req_out, "images", None) if req_out is not None else None + if images: + break - images = req_out.images if not images: raise ValueError("No images found in request_output") diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index a458850a02b..4604a64bd43 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -1,15 +1,142 @@ # Image-To-Video -This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API. +This shared example generates videos from images with VACE, Wan2.2, +LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines. -## Local CLI Usage +- `image_to_video.py`: command-line script for single video generation with advanced options. -Download the example image: +## Table of Contents + +- [Overview](#overview) +- [Prerequisites](#prerequisites) +- [Quick Start](#quick-start) +- [Key Arguments](#key-arguments) +- [More CLI Examples](#more-cli-examples) +- [Advanced Features](#advanced-features) +- [FAQ](#faq) + +## Overview + +This folder provides a unified CLI script for image-to-video generation using vLLM-Omni diffusion/video pipelines. The script selects practical defaults for supported model families while still exposing common sampling, memory, and parallelism options. + +### Supported Models + +| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM Notes | +| ----- | ------------------ | -------------- | ------------- | -------- | ---------- | +| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage | +| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model | +| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card | +| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism | + +!!! info + Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading. + +Default model: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`. + +## Prerequisites + +Download the example image used in the snippets below: ```bash -wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg +curl -L -o cherry_blossom.jpg https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg +``` + +## Quick Start + +### Python API + +Single-prompt generation using TI2V-5B (lightweight): + +```python +import PIL.Image +import torch + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +if __name__ == "__main__": + image = PIL.Image.open("cherry_blossom.jpg").convert("RGB") + image = image.resize((576, 320)) + + omni = Omni( + model="Wan-AI/Wan2.2-TI2V-5B-Diffusers", + flow_shift=12.0, + ) + outputs = omni.generate( + { + "prompt": "Cherry blossoms swaying gently in the breeze, petals falling", + "multi_modal_data": {"image": image}, + }, + OmniDiffusionSamplingParams( + height=320, + width=576, + num_frames=17, + num_inference_steps=20, + guidance_scale=4.0, + generator=torch.Generator(device="cuda").manual_seed(42), + ), + ) + from diffusers.utils import export_to_video + frames = outputs[0].request_output.images + export_to_video(frames, "quick_test_i2v.mp4", fps=16) +``` + +### Local CLI Usage + +Quick test using TI2V-5B with a small resolution and few frames: + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ + --height 320 \ + --width 576 \ + --num-frames 17 \ + --guidance-scale 4.0 \ + --num-inference-steps 20 \ + --flow-shift 12.0 \ + --fps 16 \ + --output quick_test_i2v.mp4 ``` +## Key Arguments + +| Argument | Type | Default | Description | +| -------- | ---- | ------- | ----------- | +| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path | +| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) | +| `--image` | str | (required) | Path to input image | +| `--prompt` | str | `""` | Text description of desired motion/animation | +| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress | +| `--seed` | int | `42` | Random seed for deterministic sampling | +| `--guidance-scale` | float | `5.0` | CFG scale | +| `--guidance-scale-high` | float | `None` | Separate CFG for high-noise stage (MoE only) | +| `--height` | int | auto | Video height (auto-calculated from image if not set). Multiples of 16 | +| `--width` | int | auto | Video width (auto-calculated from image if not set). Multiples of 16 | +| `--num-frames` | int | `81` | Number of frames | +| `--num-inference-steps` | int | `50` | Number of denoising steps | +| `--boundary-ratio` | float | `0.875` | Boundary split ratio for two-stage MoE models | +| `--flow-shift` | float | `5.0` | Scheduler flow shift (5.0 for 720p, 12.0 for 480p) | +| `--sample-solver` | str | `unipc` | Wan2.2 sampling solver (`unipc` or `euler` for Lightning/Distill) | +| `--fps` | int | `None` | Frames per second for the saved MP4 | +| `--frame-rate` | float | `None` | Generation frame rate for pipelines that require it (e.g., LTX2) | +| `--output` | str | `i2v_output.mp4` | Path to save the generated video | +| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization | +| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization | +| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models | +| `--enable-layerwise-offload` | flag | off | Enable layerwise offloading on DiT modules | +| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel | +| `--tensor-parallel-size` | int | `1` | Tensor parallel size (effective for models that support TP, e.g. LTX2) | +| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree | +| `--ring-degree` | int | `1` | Ring sequence parallel degree | +| `--cache-backend` | str | `None` | Cache backend: `cache_dit` or `tea_cache` | +| `--use-hsdp` | flag | off | Enable Hybrid Sharded Data Parallel | +| `--hsdp-shard-size` | int | `-1` | GPUs per shard group (-1 auto-calculates) | +| `--hsdp-replicate-size` | int | `1` | Number of replica groups for HSDP | + +## More CLI Examples + ### Wan2.2-I2V-A14B-Diffusers (MoE) ```bash @@ -17,7 +144,7 @@ python image_to_video.py \ --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \ --image cherry_blossom.jpg \ --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ - --negative-prompt "<optional quality filter>" \ + --negative-prompt "low quality, blurry" \ --height 480 \ --width 832 \ --num-frames 48 \ @@ -27,7 +154,7 @@ python image_to_video.py \ --boundary-ratio 0.875 \ --flow-shift 12.0 \ --fps 16 \ - --output i2v_output.mp4 + --output i2v_wan_moe.mp4 ``` ### Wan2.2-TI2V-5B-Diffusers (Unified) @@ -37,7 +164,7 @@ python image_to_video.py \ --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \ --image cherry_blossom.jpg \ --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ - --negative-prompt "<optional quality filter>" \ + --negative-prompt "low quality, blurry" \ --height 480 \ --width 832 \ --num-frames 48 \ @@ -45,23 +172,129 @@ python image_to_video.py \ --num-inference-steps 40 \ --flow-shift 12.0 \ --fps 16 \ - --output i2v_output.mp4 + --output i2v_wan_ti2v.mp4 +``` + +### HunyuanVideo-1.5 I2V (480p) + +```bash +python image_to_video.py \ + --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \ + --height 480 \ + --width 832 \ + --num-frames 121 \ + --guidance-scale 6.0 \ + --flow-shift 5.0 \ + --num-inference-steps 50 \ + --fps 24 \ + --enable-cpu-offload \ + --vae-use-tiling \ + --vae-use-slicing \ + --output hunyuan_i2v.mp4 +``` + +### LTX2 Image-to-Video + +```bash +python image_to_video.py \ + --model /path/to/LTX-2 \ + --model-class-name LTX2ImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "A cinematic dolly shot of cherry blossoms" \ + --height 512 \ + --width 768 \ + --num-frames 121 \ + --num-inference-steps 40 \ + --guidance-scale 4.0 \ + --frame-rate 24 \ + --fps 24 \ + --output ltx2_i2v.mp4 +``` + +## Advanced Features + +### CFG Parallel + +Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups. +See more examples in the [cfg_parallel user guide](../../../docs/user_guide/diffusion/parallelism/cfg_parallel.md). + +### Cache Acceleration + +Use `--cache-backend cache_dit` for Cache-DiT acceleration or `--cache-backend tea_cache` for Timestep Embedding Aware Cache: + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze" \ + --cache-backend cache_dit \ + --output i2v_cached.mp4 +``` + +### LTX-2.3 + +```bash +python image_to_video.py \ + --model diffusers/LTX-2.3-Diffusers \ + --model-class-name LTX23ImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --guidance-scale 4.0 \ + --num-inference-steps 20 \ + --frame-rate 24 \ + --fps 24 \ + --output ltx23_i2v_output.mp4 +``` + +Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select +the LTX-2.3 image-to-video pipeline. + +### Cosmos3 + +```bash +# Cosmos3 bundles example frames under assets/ (any RGB image works too): +python image_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \ + --prompt "The scene comes to life with smooth, natural motion." \ + --negative-prompt "blurry, distorted, low quality" \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \ + --output cosmos3_i2v.mp4 ``` Key arguments: -- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V). -- `--image`: Path to input image (required). +- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, + LTX-2/LTX-2.3, Cosmos3, or VACE). +- `--image`: Path to the first-frame or source image. +- `--last-image`: Optional last-frame condition for models such as VACE. +- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved. +- `--reference-image`: Optional reference image; repeat it to provide multiple references. +- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3. - `--prompt`: Text description of desired motion/animation. -- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16. -- `--num-frames`: Number of frames (default 81). +- `--height/--width`: Output resolution (auto-calculated from image if not set). + Wan dimensions should be multiples of 16; LTX dimensions should be multiples + of 32. +- `--num-frames`: Number of frames (model-specific default; LTX-style models + work best with `8k + 1`; Cosmos3 defaults to 189). - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE). - `--negative-prompt`: Optional list of artifacts to suppress. - `--boundary-ratio`: Boundary split ratio for two-stage MoE models. -- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p). +- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0). - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints. -- `--num-inference-steps`: Number of denoising steps (default 50). +- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35). - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video). +- `--audio-sample-rate`: fallback audio sample rate for embedded audio. - `--output`: Path to save the generated video. - `--vae-use-slicing`: Enable VAE slicing for memory optimization. - `--vae-use-tiling`: Enable VAE tiling for memory optimization. @@ -76,5 +309,106 @@ Key arguments: > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. +## Wan2.1 VACE Conditional Tasks + +The shared script selects the VACE conditioning structure from the media inputs. +No explicit mode parameter is required: the script constructs the source video, +mask, or reference images consumed by the VACE pipeline. +Download the same Hugging Face assets used by the original VACE example: + +```bash +wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg +wget -O vace_first_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png +wget -O vace_last_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png +``` + +### Image-to-Video (I2V) + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.1-VACE-1.3B-diffusers \ + --image astronaut.jpg \ + --prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \ + --seed 42 --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \ + --vae-use-tiling --output vace_i2v_output.mp4 +``` + +### Video-to-Last-Frame (V2LF) + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.1-VACE-1.3B-diffusers \ + --last-image astronaut.jpg \ + --prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \ + --seed 42 --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \ + --vae-use-tiling --output vace_v2lf_output.mp4 +``` + +### First-Last-Frame-to-Video (FLF2V) + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.1-VACE-1.3B-diffusers \ + --image vace_first_frame.png \ + --last-image vace_last_frame.png \ + --prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \ + --seed 42 --height 512 --width 512 --num-frames 81 \ + --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \ + --vae-use-tiling --output vace_flf2v_output.mp4 +``` + +### Inpainting + +Create a mask matching the original VACE example: a 160-pixel-wide white +vertical stripe marks the region to regenerate. + +```bash +python - <<'PY' +from PIL import Image + +mask = Image.new("L", (832, 480), 0) +mask.paste(255, (336, 0, 496, 480)) +mask.save("vace_center_mask.png") +PY +``` + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.1-VACE-1.3B-diffusers \ + --image astronaut.jpg \ + --mask-image vace_center_mask.png \ + --prompt "Shrek, the ogre, walks out of a building in a happy mood" \ + --seed 42 --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \ + --vae-use-tiling --output vace_inpaint_output.mp4 +``` + +### Reference-to-Video (R2V) + +Repeat `--reference-image` to provide more than one reference image. + +```bash +python image_to_video.py \ + --model Wan-AI/Wan2.1-VACE-1.3B-diffusers \ + --reference-image astronaut.jpg \ + --prompt "Camera slowly zooms out from the character walking in a garden" \ + --seed 42 --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \ + --vae-use-tiling --output vace_r2v_output.mp4 +``` + +The VACE T2V command is documented in the shared +[`text_to_video.py`](../text_to_video/text_to_video.md#wan21-vace-t2v) example. + For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly). + +## FAQ + +**OOM errors**: Try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage. For very large models, add `--enable-cpu-offload` or `--enable-layerwise-offload`. + +**Auto-calculated resolution**: If `--height` and `--width` are not provided, the script calculates output dimensions from the input image while maintaining aspect ratio and targeting 480 x 832 area (or 512 x 768 for LTX2). + +**Wan2.2 MoE vs unified**: `I2V-A14B` is a larger Mixture-of-Experts model with separate low/high-noise DiT stages (use `--guidance-scale-high` and `--boundary-ratio`). `TI2V-5B` is a smaller unified T2V+I2V model that does not need these extra arguments. diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index e5bbacdca69..164290d432b 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5. +Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3, +HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE. Supports: - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B) - LTX2 image-to-video pipeline - HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning +- Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning Usage: # Wan I2V-A14B (MoE) @@ -26,13 +28,27 @@ --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \ --frame-rate 24 --fps 24 --output ltx2_i2v.mp4 + # LTX-2.3 image-to-video + python image_to_video.py --model diffusers/LTX-2.3-Diffusers \ + --model-class-name LTX23ImageToVideoPipeline \ + --image input.jpg --prompt "A cinematic dolly shot of a boat" \ + --height 384 --width 512 --num-frames 25 --num-inference-steps 20 \ + --frame-rate 24 --fps 24 --output ltx23_i2v.mp4 + # HunyuanVideo-1.5 I2V (480p) python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \ --image input.jpg --prompt "A cat playing with yarn" \ --flow-shift 5.0 --guidance-scale 6.0 + + # Cosmos3 I2V (image conditioning) + python image_to_video.py --model nvidia/Cosmos3-Nano \ + --image input.jpg --prompt "The scene comes to life with smooth, natural motion." \ + --num-frames 189 --num-inference-steps 35 --guidance-scale 6.0 --fps 24 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' """ import argparse +import functools import json import time from pathlib import Path @@ -43,39 +59,67 @@ import torch from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import ( + build_image_to_video_prompt, + get_extra_body_params, + get_model_class_name, +) from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform -def parse_profiler_config(value: str) -> dict[str, Any]: +def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]: + """Parse a CLI value as a JSON object, attributing errors to ``flag_name``.""" try: config = json.loads(value) except json.JSONDecodeError as e: - raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e if not isinstance(config, dict): - raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object") return config +parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config") + + def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).") + parser = argparse.ArgumentParser( + description=( + "Generate a video from one or more images " + "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)." + ) + ) parser.add_argument( "--model", default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", - help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).", + help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).", ) parser.add_argument( "--model-class-name", default=None, - help="Override model class name (e.g., LTX2ImageToVideoPipeline).", + help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).", + ) + parser.add_argument( + "--deploy-config", + default=None, + help="Optional deploy config YAML to use for pipeline-backed runs.", + ) + parser.add_argument("--image", help="Path to the first-frame or source image.") + parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).") + parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).") + parser.add_argument( + "--reference-image", + action="append", + default=None, + help="Path to a reference image. Repeat to provide multiple references.", ) - parser.add_argument("--image", required=True, help="Path to input image.") parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.") parser.add_argument("--negative-prompt", default="", help="Negative prompt.") parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.") + parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.") parser.add_argument( "--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)." ) @@ -83,8 +127,13 @@ def parse_args() -> argparse.Namespace: "--height", type=int, default=None, help="Video height (auto-calculated from image if not set)." ) parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).") - parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.") - parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.") + parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.") + parser.add_argument( + "--num-inference-steps", + type=int, + default=None, + help="Sampling steps. Default: model-specific.", + ) parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.") parser.add_argument( "--frame-rate", @@ -93,7 +142,10 @@ def parse_args() -> argparse.Namespace: help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.", ) parser.add_argument( - "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)." + "--flow-shift", + type=float, + default=None, + help="Scheduler flow_shift. Default: model-specific (Wan 5.0, Cosmos3 10.0).", ) parser.add_argument( "--sample-solver", @@ -173,7 +225,7 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"], help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) @@ -244,6 +296,19 @@ def parse_args() -> argparse.Namespace: default=None, help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.', ) + parser.add_argument( + "--extra-body", + type=functools.partial(parse_json_object, flag_name="--extra-body"), + default=None, + help=( + "Model-specific generation params as a JSON object. Keys are filtered " + "against the model's declared extra_body_params (see vllm_omni/model_extras), " + "so unknown keys for the chosen model are silently dropped. " + 'Cosmos3 V2V example: \'{"condition_frame_indexes_vision": [0, 1], ' + '"condition_video_keep": "first", "flow_shift": 10.0, ' + '"max_sequence_length": 4096, "guardrails": false}\'.' + ), + ) return parser.parse_args() @@ -266,32 +331,70 @@ def main(): generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) model_name = str(args.model).lower() if args.model is not None else "" model_class_name = args.model_class_name - is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower()) - if model_class_name is None and is_ltx2: - model_class_name = "LTX2ImageToVideoPipeline" - - # Load input image - image = PIL.Image.open(args.image).convert("RGB") + is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"} + is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower()) + + image = PIL.Image.open(args.image).convert("RGB") if args.image else None + last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None + mask_image = PIL.Image.open(args.mask_image).convert("L") if args.mask_image else None + reference_images = ( + [PIL.Image.open(path).convert("RGB") for path in args.reference_image] if args.reference_image else None + ) + dimension_image = image or last_image or (reference_images[0] if reference_images else None) + if dimension_image is None: + raise ValueError("Provide --image, --last-image, or at least one --reference-image.") + + # Per-model generation defaults, applied only when the matching flag is omitted. + # Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift). + if is_cosmos: + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 24, + 6.0, + 189, + 35, + 10.0, + 1280 * 720, + 16, + ) + elif is_ltx2: + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 24, + 4.0, + 121, + 40, + 5.0, + 512 * 768, + 32, + ) + else: # Wan2.2 / HunyuanVideo-1.5 + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16 - fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16) + fps = args.fps if args.fps is not None else d_fps frame_rate = args.frame_rate if args.frame_rate is not None else float(fps) - guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0) - num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81) - num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50) + guidance_scale = args.guidance_scale if args.guidance_scale is not None else d_guidance + num_frames = args.num_frames if args.num_frames is not None else d_num_frames + num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else d_steps + flow_shift = args.flow_shift if args.flow_shift is not None else d_flow_shift - # Calculate dimensions if not provided + # Calculate dimensions if not provided (model-aware max area). height = args.height width = args.width if height is None or width is None: - # Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2 - max_area = 512 * 768 if is_ltx2 else 480 * 832 - mod_value = 32 if is_ltx2 else 16 - calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value) + calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod) height = height or calc_height width = width or calc_width - # Resize image to target dimensions - image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + media_inputs: dict[str, Any] = {} + if image is not None: + media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + if last_image is not None: + media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS) + if mask_image is not None: + media_inputs["mask"] = mask_image.resize((width, height), PIL.Image.Resampling.NEAREST) + if reference_images is not None: + media_inputs["reference_images"] = [ + reference.resize((width, height), PIL.Image.Resampling.LANCZOS) for reference in reference_images + ] # Configure cache based on backend type cache_config = None @@ -342,7 +445,7 @@ def main(): vae_use_slicing=args.vae_use_slicing, vae_use_tiling=args.vae_use_tiling, boundary_ratio=args.boundary_ratio, - flow_shift=args.flow_shift, + flow_shift=flow_shift, diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype, diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps, diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers, @@ -355,9 +458,17 @@ def main(): enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, profiler_config=args.profiler_config, ) + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config if args.quantization is not None: omni_kwargs["quantization"] = args.quantization + # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails + # gate is an engine-level config (offline analog of the server's --no-guardrails). + if args.extra_body and "guardrails" in args.extra_body: + omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])} omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) or model_class_name + declared_extra_body_params = get_extra_body_params(model_class_name) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -367,8 +478,8 @@ def main(): print(f"\n{'=' * 60}") print("Generation Configuration:") print(f" Model: {args.model}") - print(f" Inference steps: {args.num_inference_steps}") - print(f" Frames: {args.num_frames}") + print(f" Inference steps: {num_inference_steps}") + print(f" Frames: {num_frames}") print(f" Solver: {args.sample_solver}") print(f" diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}") print(f" diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}") @@ -378,33 +489,46 @@ def main(): f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}," f" pipeline_parallel_size={args.pipeline_parallel_size}" ) - print(f" Video size: {args.width}x{args.height}") + print(f" Video size: {width}x{height}") print(f"{'=' * 60}\n") - generation_start = time.perf_counter() - # omni.generate() returns Generator[OmniRequestOutput, None, None] - frames = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - "multi_modal_data": {"image": image}, + prompt_dict = build_image_to_video_prompt( + model_class_name=model_class_name, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + media_inputs=media_inputs, + height=height, + width=width, + num_frames=num_frames, + ) + sampling_params = OmniDiffusionSamplingParams( + height=height, + width=width, + generator=generator, + guidance_scale=guidance_scale, + guidance_scale_2=args.guidance_scale_high, + boundary_ratio=args.boundary_ratio, + num_inference_steps=num_inference_steps, + num_frames=num_frames, + frame_rate=frame_rate, + extra_args={ + "sample_solver": args.sample_solver, + "flow_shift": flow_shift, }, - OmniDiffusionSamplingParams( - height=height, - width=width, - generator=generator, - guidance_scale=guidance_scale, - guidance_scale_2=args.guidance_scale_high, - boundary_ratio=args.boundary_ratio, - num_inference_steps=num_inference_steps, - num_frames=num_frames, - frame_rate=frame_rate, - extra_args={ - "sample_solver": args.sample_solver, - "flow_shift": args.flow_shift, - }, - ), ) + + # Route model-specific knobs through extra_body, filtered against the model's + # declared extra_body_params. Models without a declaration only forward explicit + # --extra-body JSON (e.g. Cosmos3 V2V's condition_frame_indexes_vision). + extra_body = dict(args.extra_body or {}) + if declared_extra_body_params: + apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body) + elif extra_body: + sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None}) + + generation_start = time.perf_counter() + # omni.generate() returns Generator[OmniRequestOutput, None, None] + frames = omni.generate(prompt_dict, sampling_params) generation_end = time.perf_counter() generation_time = generation_end - generation_start @@ -412,6 +536,7 @@ def main(): print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") audio = None + audio_sample_rate = args.audio_sample_rate if isinstance(frames, list): frames = frames[0] if frames else None @@ -422,11 +547,13 @@ def main(): ) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] + audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) if frames.is_pipeline_output and frames.request_output is not None: inner_output = frames.request_output if isinstance(inner_output, OmniRequestOutput): if inner_output.multimodal_output and "audio" in inner_output.multimodal_output: audio = inner_output.multimodal_output["audio"] + audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate) frames = inner_output if isinstance(frames, OmniRequestOutput): if frames.images: @@ -434,6 +561,7 @@ def main(): frames, audio = frames.images[0] elif len(frames.images) == 1 and isinstance(frames.images[0], dict): audio = frames.images[0].get("audio") + audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate) frames = frames.images[0].get("frames") or frames.images[0].get("video") else: frames = frames.images @@ -446,6 +574,7 @@ def main(): frames, audio = first_item elif isinstance(first_item, dict): audio = first_item.get("audio") + audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate) frames = first_item.get("frames") or first_item.get("video") elif isinstance(first_item, list): frames = first_item @@ -454,6 +583,7 @@ def main(): frames, audio = frames elif isinstance(frames, dict): audio = frames.get("audio") + audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate) frames = frames.get("frames") or frames.get("video") if frames is None: @@ -572,7 +702,7 @@ def _ensure_frame_list(video_array): frames_u8, audio_np, fps=float(fps), - audio_sample_rate=args.audio_sample_rate, + audio_sample_rate=audio_sample_rate, ) with open(str(output_path), "wb") as f: f.write(video_bytes) diff --git a/examples/offline_inference/lance/README.md b/examples/offline_inference/lance/README.md index d029abdeaac..f3a422a3b3b 100644 --- a/examples/offline_inference/lance/README.md +++ b/examples/offline_inference/lance/README.md @@ -1,4 +1,4 @@ -# Lance (ByteDance) — offline inference +# Lance: Offline inference [Lance](https://huggingface.co/bytedance-research/Lance) is a 3B unified autoregressive + diffusion multimodal model on a Qwen2.5-VL backbone. It is diff --git a/examples/offline_inference/lance/end2end.py b/examples/offline_inference/lance/end2end.py index 504dceec613..d2691016187 100644 --- a/examples/offline_inference/lance/end2end.py +++ b/examples/offline_inference/lance/end2end.py @@ -293,16 +293,19 @@ def main(): if args.modality in ("img2text", "video2text"): for i, req_output in enumerate(outputs): - text = getattr(req_output, "output", None) or getattr(req_output, "text", None) - if text is None and hasattr(req_output, "custom_output"): - text = (req_output.custom_output or {}).get("text_output") + multimodal_output = getattr(req_output, "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} + text = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or ( + text_metadata.get("text_output") if isinstance(text_metadata, dict) else None + ) + text = text or getattr(req_output, "output", None) or getattr(req_output, "text", None) print(f"[Output {i}] {text}") return if args.modality in ("text2video", "video2video", "image2video"): for i, req_output in enumerate(outputs): - custom = getattr(req_output, "custom_output", None) or {} - frames = custom.get("video_frames") + frames = getattr(req_output, "images", None) if not frames: print(f"[Output {i}] no video frames returned") continue diff --git a/examples/offline_inference/lance/gradio_demo.py b/examples/offline_inference/lance/gradio_demo.py index 7dfeb46c9e7..756b006cb4c 100644 --- a/examples/offline_inference/lance/gradio_demo.py +++ b/examples/offline_inference/lance/gradio_demo.py @@ -835,12 +835,15 @@ def run_task( print(f"[unified] task={task} {width}x{height} frames={num_frames} cfg={cfg} steps={steps} seed={seed}", flush=True) outputs = list(omni.generate(prompts=formatted, sampling_params_list=params)) - custom = getattr(outputs[0], "custom_output", None) or {} - # Lance/Bagel return text at outputs[0].output (string) or - # custom_output["text_output"]. + multimodal_output = getattr(outputs[0], "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} text_out = None if task in {"x2t_image", "x2t_video"}: - text_out = custom.get("text_output") or getattr(outputs[0], "output", None) + text_out = (multimodal_output.get("text") if isinstance(multimodal_output, dict) else None) or ( + text_metadata.get("text_output") if isinstance(text_metadata, dict) else None + ) + text_out = text_out or getattr(outputs[0], "output", None) if not isinstance(text_out, str): text_out = None @@ -856,13 +859,11 @@ def run_task( gr.update(value=str(text_out or "(empty)"), visible=True), ) if task in {"t2v", "i2v", "video_edit"}: - frames = custom.get("video_frames") + frames = getattr(outputs[0], "images", None) if frames is None: raise gr.Error("No video frames in output.") return (gr.update(visible=False), gr.update(value=_save_video(frames), visible=True), gr.update(visible=False)) - # image tasks — orchestrator strips ``outputs[0].output`` across the - # IPC boundary but preserves ``custom_output``, so prefer the latter. - img_out = custom.get("image") or getattr(outputs[0], "output", None) + img_out = (getattr(outputs[0], "images", None) or [None])[0] or getattr(outputs[0], "output", None) if img_out is None: raise gr.Error("No image in output.") return (gr.update(value=_save_image(img_out), visible=True), gr.update(visible=False), gr.update(visible=False)) diff --git a/examples/offline_inference/magi_human/README.md b/examples/offline_inference/magi_human/README.md deleted file mode 100644 index 2b89093d941..00000000000 --- a/examples/offline_inference/magi_human/README.md +++ /dev/null @@ -1,72 +0,0 @@ -# MagiHuman Generation - -MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt. - -Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks. - -## Setup - -### Install MagiCompiler (recommended) - -MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour. - -```bash -# Clone the repo -git clone https://github.com/SandAI-org/MagiCompiler.git -cd MagiCompiler - -# System dependencies (optional, for FX graph visualization; Debian/Ubuntu) -sudo apt update && sudo apt install -y graphviz - -# Python dependencies -pip install -r requirements.txt - -# Install MagiCompiler -pip install . # end users (recommended) -# pip install -e . # developers (editable install) -``` - -### Hardware requirements - -Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference). - -Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory. - -## Run Examples - -Get into the example folder: -```bash -cd examples/offline_inference/magi_human -``` - -### End-to-End Generation (Text to Video+Audio) - -Generate a video with synchronized speech natively generated by the model. - -```bash -python end2end.py \ - --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \ - --prompt "A young woman with long, wavy golden blonde hair..." \ - --tensor-parallel-size 4 \ - --output output_magihuman.mp4 -``` - -## Common Parameters - -| Parameter | Default | Description | -|-----------|---------|-------------| -| `--model` | *(Required)* | Local model path or HuggingFace ID | -| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text | -| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) | -| `--height` | `256` | Initial resolution height | -| `--width` | `448` | Initial resolution width | -| `--num-inference-steps` | `8` | Denoising steps | -| `--seed` | `52` | Random seed | -| `--output` | `output_magihuman.mp4` | Output video with audio path | - -## Example materials - -??? abstract "end2end.py" - ``````py - --8<-- "examples/offline_inference/magi_human/end2end.py" - `````` diff --git a/examples/offline_inference/magi_human/end2end.py b/examples/offline_inference/magi_human/end2end.py deleted file mode 100644 index 64f11c4658b..00000000000 --- a/examples/offline_inference/magi_human/end2end.py +++ /dev/null @@ -1,122 +0,0 @@ -import argparse - -from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - -def parse_args(): - parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.") - parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.") - parser.add_argument( - "--prompt", - type=str, - default="", - help="Text prompt containing visual description, dialogue, and background sound.", - ) - parser.add_argument( - "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)." - ) - parser.add_argument( - "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file." - ) - parser.add_argument("--height", type=int, default=256, help="Video height.") - parser.add_argument("--width", type=int, default=448, help="Video width.") - parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.") - parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.") - return parser.parse_args() - - -def main(): - args = parse_args() - - print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...") - omni = Omni( - model=args.model, - init_timeout=1200, - tensor_parallel_size=args.tensor_parallel_size, - devices=list(range(args.tensor_parallel_size)), - ) - - prompt = args.prompt - if not prompt: - prompt = ( - "A young woman with long, wavy golden blonde hair and bright blue eyes, " - "wearing a fitted ivory silk blouse with a delicate lace collar, sits " - "stationary in front of a softly lit, blurred warm-toned interior. Her " - "overall disposition is warm, composed, and gently confident. The camera " - "holds a static medium close-up, framing her from the shoulders up, " - "with shallow depth of field keeping her face in sharp focus. Soft " - "directional key light falls from the upper left, casting a gentle " - "highlight along her cheekbone and nose bridge. She draws a quiet breath, " - "the levator labii superiors relaxing as her lips part. She speaks in " - "clear, warm, unhurried American English: " - "\"The most beautiful things in life aren't things at all — " - "they're moments, feelings, and the people who make you feel truly alive.\" " - "Her jaw descends smoothly on each stressed syllable; the orbicularis oris " - "shapes each vowel with precision. A faint, genuine smile engages the " - "zygomaticus major, lifting her lip corners fractionally. Her brows rest " - "in a soft, neutral arch throughout. She maintains steady, forward-facing " - "eye contact. Head position remains level; no torso displacement occurs.\n\n" - "Dialogue:\n" - "<Young blonde woman, American English>: " - "\"The most beautiful things in life aren't things at all — " - "they're moments, feelings, and the people who make you feel truly alive.\"\n\n" - "Background Sound:\n" - "<Soft, warm indoor ambience with a faint distant piano melody>" - ) - - sampling_params = OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_inference_steps=args.num_inference_steps, - seed=args.seed, - extra_args={ - "seconds": 5, - "sr_height": 1080, - "sr_width": 1920, - "sr_num_inference_steps": 5, - }, - ) - - print(f"Generating with prompt: {prompt[:80]}...") - outputs = omni.generate( - prompts=[prompt], - sampling_params_list=[sampling_params], - ) - - print(f"Generation complete. Output type: {type(outputs)}") - if outputs: - first = outputs[0] - - if hasattr(first, "images") and first.images: - video_frames = first.images[0] - print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}") - - audio_waveform = None - mm = first.multimodal_output or {} - if mm: - audio_waveform = mm.get("audio") - if audio_waveform is not None: - print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}") - - output_fps = float(mm.get("fps", 25)) - output_sr = int(mm.get("audio_sample_rate", 24000)) - print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output") - - video_bytes = mux_video_audio_bytes( - video_frames, - audio_waveform, - fps=output_fps, - audio_sample_rate=output_sr, - ) - with open(args.output, "wb") as f: - f.write(video_bytes) - print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}") - print("SUCCESS: MagiHuman pipeline generation completed.") - else: - print("WARNING: No outputs returned.") - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/mammothmodal2_preview/README.md b/examples/offline_inference/mammothmodal2_preview/README.md index 44b20c76a14..bb68f384553 100644 --- a/examples/offline_inference/mammothmodal2_preview/README.md +++ b/examples/offline_inference/mammothmodal2_preview/README.md @@ -1,24 +1,17 @@ # MammothModa2-Preview -## Run examples (MammothModa2-Preview) +> Text-to-image (T2I) generation has moved to the shared offline image example. +> See the recipe: [`recipes/MammothModa2/MammothModa2-Preview.md`](../../../recipes/MammothModa2/MammothModa2-Preview.md). +> +> This directory retains the image-understanding (image → text) example below, +> which is outside the scope of the diffusion `model_extras` example migration +> (#4548 / RFC #4539: AR / Omni understanding examples keep dedicated scripts). -Download model -```bash -hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview -``` - -### Text-to-Image (T2I) +## Run example (MammothModa2-Preview) +Download model ```bash -python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \ - --model ./MammothModa2-Preview \ - --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --height 1024 \ - --width 1024 \ - --num-inference-steps 50 \ - --text-guidance-scale 4.0 \ - --out output.png +hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview ``` ### Image Summary @@ -26,7 +19,7 @@ python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py ```bash python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ --model ./MammothModa2-Preview \ - --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \ + --deploy-config ./vllm_omni/deploy/mammoth_moda2_ar.yaml \ --question "Summarize this image." \ --image ./image.png ``` diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py index ca87b9e9a94..7a22803c926 100644 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py @@ -4,7 +4,7 @@ Example: uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \ --model path/to/MammothModa2-Preview \ - --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \ + --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \ --image /path/to/input.jpg \ --question "Please summarize the content of this image." """ @@ -27,9 +27,7 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).") parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.") - parser.add_argument( - "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)." - ) + parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.") parser.add_argument("--image", type=str, required=True, help="Path to input image.") parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.") parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.") @@ -73,7 +71,7 @@ def main() -> None: omni = Omni( model=args.model, - stage_configs_path=args.stage_config, + deploy_config=args.deploy_config, trust_remote_code=args.trust_remote_code, enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, ) diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py deleted file mode 100644 index a4c41fee1f8..00000000000 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py +++ /dev/null @@ -1,248 +0,0 @@ -""" -Offline inference example for MammothModa2 Text-to-Image (T2I) generation. -This script uses the vllm_omni.Omni pipeline with a multi-stage configuration. - -Workflow: -1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states. -2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion - and VAE decoding to produce the final image. - -Example Usage: - uv run python examples/offline_inference/run_mammothmoda2_t2i.py \ - --model path/to/MammothModa2-Preview \ - --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \ - --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ - --out output.png -""" - -from __future__ import annotations - -import argparse -import json -import logging -import os -from pathlib import Path -from typing import NamedTuple - -import torch -from PIL import Image -from vllm.sampling_params import SamplingParams - -from vllm_omni import Omni - -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") -logger = logging.getLogger(__name__) - -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- -_PATCH_SIZE = 16 # AR image grid patch size (pixels per token) - - -class T2IGenConfig(NamedTuple): - eol_token_id: int - visual_token_start_id: int - visual_token_end_id: int - top_k: int # AR sampling top-k (covers the full visual generation vocabulary) - # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|> - visual_ids: list[int] - - -def load_t2i_generation_config(model_dir: str) -> T2IGenConfig: - """Load T2I token IDs from t2i_generation_config.json and config.json. - - Supports both local directory paths and HuggingFace Hub model IDs. - """ - model_path = Path(model_dir) - - def _read_json(filename: str) -> dict: - local = model_path / filename - if local.exists(): - with local.open(encoding="utf-8") as f: - return json.load(f) - # Fall back to HuggingFace Hub when model_dir is a repo ID. - try: - from huggingface_hub import hf_hub_download - except ImportError as exc: - raise ImportError( - "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub" - ) from exc - cached = hf_hub_download(repo_id=model_dir, filename=filename) - with open(cached, encoding="utf-8") as f: - return json.load(f) - - gen_cfg = _read_json("t2i_generation_config.json") - llm_cfg = _read_json("config.json").get("llm_config", {}) - - return T2IGenConfig( - eol_token_id=int(gen_cfg["eol_token_id"]), - visual_token_start_id=int(gen_cfg["visual_token_start_id"]), - visual_token_end_id=int(gen_cfg["visual_token_end_id"]), - top_k=int(gen_cfg["top_k"]), - visual_ids=[ - int(llm_cfg["image_token_id"]), - int(llm_cfg["video_token_id"]), - int(llm_cfg["vision_start_token_id"]), - int(llm_cfg["vision_end_token_id"]), - ], - ) - - -def parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.") - p.add_argument("--model", type=str, required=True, help="Path to the model directory.") - p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.") - p.add_argument( - "--prompt", - type=str, - action="append", - default=None, - help=( - "Text prompt for image generation. Can be provided multiple times " - "to generate multiple images with shared height/width/CFG settings." - ), - ) - p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).") - p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).") - p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.") - p.add_argument( - "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT." - ) - p.add_argument( - "--cfg-range", - type=float, - nargs=2, - default=(0.0, 1.0), - help="Relative step range [start, end] where CFG is active.", - ) - p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.") - p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.") - args = p.parse_args() - if not args.prompt: - args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."] - return args - - -def tensor_to_pil(image: torch.Tensor) -> Image.Image: - """Convert a normalized torch tensor [-1, 1] to a PIL Image.""" - if image.ndim == 4: - image = image[0] - image = image.detach().to("cpu") - image = (image / 2 + 0.5).clamp(0, 1) - image = (image * 255).to(torch.uint8) - image = image.permute(1, 2, 0).contiguous().numpy() - return Image.fromarray(image) - - -def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str: - """Build the AR-stage prompt string including the image grid header.""" - return ( - "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" - f"<|im_start|>user\n{user_prompt}<|im_end|>\n" - "<|im_start|>assistant\n" - f"<|image start|>{ar_width}*{ar_height}<|image token|>" - ) - - -def _collect_images(outputs: list) -> list[torch.Tensor]: - """Extract all image tensors produced by the final (DiT) stage.""" - images: list[torch.Tensor] = [] - for out in outputs: - ro_item = getattr(out, "request_output", out) - for completion in getattr(ro_item, "outputs", None) or []: - mm = getattr(completion, "multimodal_output", None) - if not isinstance(mm, dict) or "image" not in mm: - raise RuntimeError(f"Missing image in multimodal output: {mm}") - payload = mm["image"] - for tensor in payload if isinstance(payload, list) else [payload]: - if not isinstance(tensor, torch.Tensor): - raise TypeError(f"Expected image tensor, got {type(tensor)}") - images.append(tensor) - return images - - -def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]: - """Save image tensors to disk. - - Single image: written to *out_path* exactly. - Multiple images: suffixed as ``<base>_0<ext>``, ``<base>_1<ext>``, … - """ - if not images: - raise RuntimeError("No images to save.") - base, ext = os.path.splitext(out_path) - ext = ext or ".png" - paths = [] - for i, tensor in enumerate(images): - path = out_path if len(images) == 1 else f"{base}_{i}{ext}" - tensor_to_pil(tensor).save(path) - paths.append(path) - return paths - - -def main() -> None: - args = parse_args() - os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) - - if args.height <= 0 or args.width <= 0: - raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}") - if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0: - raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}") - - ar_height = args.height // _PATCH_SIZE - ar_width = args.width // _PATCH_SIZE - gen_cfg = load_t2i_generation_config(args.model) - expected_grid_tokens = ar_height * (ar_width + 1) - - logger.info("Initializing Omni pipeline...") - omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code) - try: - ar_sampling = SamplingParams( - temperature=1.0, - top_p=1.0, - top_k=gen_cfg.top_k, - max_tokens=max(1, expected_grid_tokens + 1), # +1 for hidden state of eoi - detokenize=False, - ) - dit_sampling = SamplingParams( - temperature=0.0, - top_p=1.0, - top_k=-1, - max_tokens=1, - detokenize=False, - ) - - additional_information = { - "omni_task": ["t2i"], - "ar_width": [ar_width], - "ar_height": [ar_height], - "eol_token_id": [gen_cfg.eol_token_id], - "visual_token_start_id": [gen_cfg.visual_token_start_id], - "visual_token_end_id": [gen_cfg.visual_token_end_id], - "image_height": [args.height], - "image_width": [args.width], - "num_inference_steps": [args.num_inference_steps], - "text_guidance_scale": [args.text_guidance_scale], - "cfg_range": [args.cfg_range[0], args.cfg_range[1]], - "visual_ids": gen_cfg.visual_ids, - } - inputs = [ - { - "prompt": _format_prompt(p, ar_width, ar_height), - "additional_information": dict(additional_information), - } - for p in args.prompt - ] - - logger.info("Starting generation...") - # omni.generate() returns a Generator; consume it to run the full pipeline. - outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling])) - - logger.info("Post-processing and saving image(s)...") - for path in _save_images(_collect_images(outputs), args.out): - logger.info(f"Saved: {path}") - finally: - omni.close() - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/mimo_audio/README.md b/examples/offline_inference/mimo_audio/README.md index 5615dea5176..60f433a4f30 100644 --- a/examples/offline_inference/mimo_audio/README.md +++ b/examples/offline_inference/mimo_audio/README.md @@ -1,4 +1,4 @@ -# MiMo-Audio Offline Inference +# MiMo-Audio: Offline inference This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally. diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index 2fe46c9721d..139351aab79 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -219,31 +219,31 @@ def main(args): # Build query result based on query type # Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples if args.query_type == "tts_sft": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft query_result = query_func(text=text, read_text_only=True) elif args.query_type == "tts_sft_with_instruct": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice" query_result = query_func(text=text, instruct=instruct, read_text_only=True) elif args.query_type == "tts_sft_with_audio": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_audio --audio-path "./spoken_dialogue_assistant_turn_1.wav" audio_list = [get_audio_data(audio_path)] query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list) elif args.query_type == "tts_sft_with_natural_instruction": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!" query_result = query_func(text=text, read_text_only=False) elif args.query_type == "audio_trancribing_sft": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_trancribing_sft --audio-path "./spoken_dialogue_assistant_turn_1.wav" audio_path = "spoken_dialogue_assistant_turn_1.wav" text = "Please transcribe this audio and repeat it once." query_result = query_func(text=text, audio_path=audio_path, use_sostm=True) elif args.query_type == "audio_understanding_sft": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio-path "./spoken_dialogue_assistant_turn_1.wav" query_result = query_func(text=text, audio_path=audio_path) elif args.query_type == "audio_understanding_sft_with_thinking": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio-path "./spoken_dialogue_assistant_turn_1.wav" query_result = query_func(text=text, audio_path=audio_path, thinking=True) elif args.query_type == "spoken_dialogue_sft_multiturn": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio_path "./prompt_speech_zh_m.wav" + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio-path "./prompt_speech_zh_m.wav" first_turn_text_response = "我没办法获取实时的天气信息。不过呢，你可以试试几个方法来查看今天的天气。首先，你可以用手机自带的天气功能，比如苹果手机的天气应用，或者直接在系统设置里查看。其次，你也可以用一些专业的天气服务，像是国外的AccuWeather、Weather.com，或者国内的中国天气网、墨迹天气等等。再有就是，你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市，我也可以帮你分析一下历史天气趋势，不过最新的数据还是需要你通过官方渠道去获取哦。" audio_list = [] s1_audio_path = "weather_of_today.mp3" @@ -261,7 +261,7 @@ def main(args): ] query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list) elif args.query_type == "speech2text_dialogue_sft_multiturn": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn s1_audio_path = "weather_of_today.mp3" s2_audio_path = "beijing.mp3" audio_list = [] @@ -277,7 +277,7 @@ def main(args): ] query_result = query_func(message_list, thinking=True, audio_list=audio_list) elif args.query_type == "text_dialogue_sft_multiturn": - # python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type text_dialogue_sft_multiturn + # python3 -u end2end.py --deploy-config ${config_file} --model-name ${MODEL_PATH} --query-type text_dialogue_sft_multiturn message_list = [ {"role": "user", "content": "Could you recommend some tourist attractions in China?"}, {"role": "assistant", "content": "Hello, which city would you like to travel to?"}, diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md index 7e7ab0a8f7a..ee7e93fd060 100644 --- a/examples/offline_inference/ming_flash_omni/README.md +++ b/examples/offline_inference/ming_flash_omni/README.md @@ -1,4 +1,4 @@ -# Ming-flash-omni 2.0 +# Ming-flash-omni 2.0: Offline inference [Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs. @@ -90,14 +90,40 @@ Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to ### Image generation (text-to-image / img2img) -The diffusion-stage image-generation path is currently only wired through -the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with -`"modalities": ["image"]`). For payloads, optional knobs -(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img -reference-image flow, see the -[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image). -`end2end.py` does not yet exercise the imagegen stage; that is tracked as -follow-up work. +Image generation is served through the standard task examples. +The diffusion-stage knobs are declared centrally in `vllm_omni/model_extras/ming_flash_omni.py` and routed via `--extra-body`. With deploy yaml assigned properly, the model can run through the shared example scripts like any other diffusion model: + +Text-to-image (offline): + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model Jonathan1909/Ming-flash-omni-2.0 \ + --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \ + --prompt "Please draw a cute cat." \ + --height 1024 \ + --width 1024 \ + --extra-body '{"steps": 30, "cfg": 2.0, "seed": 42}' \ + --output ming_flash_omni_t2i.png +``` + +Image-to-image (offline): +```bash +# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo +# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py +wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png + +python examples/offline_inference/image_to_image/image_edit.py \ + --model Jonathan1909/Ming-flash-omni-2.0 \ + --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \ + --image person_gen_05.png \ + --prompt "Put a pair of sunglasses on the person." \ + --extra-args '{"steps": 30, "cfg": 2.0, "seed": 42}' \ + --output ming_flash_omni_i2i.png +``` + +For the online path and the full knob list +(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`/`height`/`width`), see the +[image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img). ## Online serving diff --git a/examples/offline_inference/qwen2_5_omni/README.md b/examples/offline_inference/qwen2_5_omni/README.md index e2eae8a96b5..059c102261d 100644 --- a/examples/offline_inference/qwen2_5_omni/README.md +++ b/examples/offline_inference/qwen2_5_omni/README.md @@ -1,4 +1,4 @@ -# Qwen2.5-Omni +# Qwen2.5-Omni: Offline inference ## Setup Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. @@ -60,3 +60,24 @@ If media file paths are not provided, the script will use default assets. Suppor - `mixed_modalities`: Audio + image + video - `use_audio_in_video`: Extract audio from video - `text`: Text-only query + +### Composable parallelism (strategy configs) + +You can shard or replicate a stage with a small composable-parallel +`strategy.yaml` overlaid onto the **bundled default deploy config** -- no bespoke +deploy YAML required. Supply the strategy with `--strategy-config` and any matching +device layout via `--stage-overrides` (explicit kwargs forwarded to `Omni`). + +Both examples below require >= 3 GPUs (the thinker uses 2, talker + code2wav share a third): + +```bash +# Tensor-parallel the thinker (TP=2) +python end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_tp2.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' + +# Replicate the thinker across 2 engines (round-robin load balancing) +python end2end.py --query-type text --num-prompts 6 \ + --strategy-config strategy_stage_replica.yaml \ + --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +``` diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 8290d27b5ca..d8769b5f7e9 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -553,6 +553,39 @@ def parse_args(): default=False, help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.", ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Optional explicit deploy YAML (otherwise the bundled default is used).", + ) + parser.add_argument( + "--strategy-config", + type=str, + default=None, + help=( + "Optional composable-parallel strategy.yaml mapping role -> parallel axes " + "(e.g. tp / stage_replica). Overlaid onto the merged stage configs." + ), + ) + parser.add_argument( + "--stage-overrides", + type=str, + default=None, + help=( + "Optional JSON of per-stage overrides applied on top of the default " + 'deploy config, e.g. \'{"0": {"devices": "0,1"}}\' to give the ' + "thinker a 2-GPU pool. Lets you run a strategy on the bundled default " + "deploy config without writing a bespoke deploy YAML." + ), + ) + parser.add_argument( + "--omni-lb-policy", + type=str, + default=None, + choices=["random", "round-robin", "least-queue-length"], + help="StagePool load-balancer policy for replicated stages (orchestrator-level knob).", + ) return parser.parse_args() diff --git a/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml b/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml new file mode 100644 index 00000000000..f0e2d9bf5af --- /dev/null +++ b/examples/offline_inference/qwen2_5_omni/strategy_stage_replica.yaml @@ -0,0 +1,22 @@ +# Composable-parallel strategy for Qwen2.5-Omni. +# +# Replicate the thinker stage across 2 independent engines and let the omni +# StagePool balance requests round-robin across them. The thinker's `devices` +# entry acts as a per-replica template, so omni expands it to one GPU per replica. +# The `routing: round_robin` below auto-derives the StagePool load-balancer policy, +# so --omni-lb-policy is not needed. +# +# Rather than hand-write a bespoke deploy YAML, run on the bundled default deploy +# config and pass the matching 2-GPU device template via --stage-overrides (explicit +# kwargs): the thinker gets a two-GPU pool while talker + code2wav share a third +# card. Requires >= 3 GPUs. +# +# Usage: +# python3 end2end.py --query-type text --num-prompts 6 \ +# --strategy-config strategy_stage_replica.yaml \ +# --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +strategies: + thinker: + - axis: stage_replica + size: 2 + routing: round_robin diff --git a/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml b/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml new file mode 100644 index 00000000000..02d77395285 --- /dev/null +++ b/examples/offline_inference/qwen2_5_omni/strategy_tp2.yaml @@ -0,0 +1,19 @@ +# Composable-parallel strategy for Qwen2.5-Omni: tensor-parallel the thinker. +# +# Shard the thinker stage across 2 GPUs with tensor parallelism (TP=2). TP routing +# is fixed (broadcast), so no `routing` / load-balancer policy is involved -- this +# is the simplest axis and a good starting point for the mental model. +# +# The strategy sets tensor_parallel_size=2 on the thinker. Rather than hand-write a +# bespoke deploy YAML, run on the bundled default deploy config and pass a matching +# 2-GPU device layout via --stage-overrides (explicit kwargs), so the thinker gets a +# two-GPU pool while talker + code2wav share a third card. Requires >= 3 GPUs. +# +# Usage: +# python3 end2end.py --query-type text --num-prompts 6 \ +# --strategy-config strategy_tp2.yaml \ +# --stage-overrides '{"0": {"devices": "0,1"}, "1": {"devices": "2"}, "2": {"devices": "2"}}' +strategies: + thinker: + - axis: tp + size: 2 diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md index 0710faa133c..2f4ee11cc3f 100644 --- a/examples/offline_inference/qwen3_omni/README.md +++ b/examples/offline_inference/qwen3_omni/README.md @@ -1,4 +1,4 @@ -# Qwen3-Omni +# Qwen3-Omni: Offline inference ## Setup Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. diff --git a/examples/offline_inference/sensenova_u1/README.md b/examples/offline_inference/sensenova_u1/README.md index 3943b8d4e73..9fdc23a00ee 100644 --- a/examples/offline_inference/sensenova_u1/README.md +++ b/examples/offline_inference/sensenova_u1/README.md @@ -1,4 +1,4 @@ -# SenseNova-U1-8B-MoT +# SenseNova-U1: Offline inference ## Architecture @@ -38,9 +38,9 @@ python end2end.py --modality text2text \ python end2end.py --prompt "A futuristic cityscape at sunset" \ --width 2048 --height 1024 --think -# Text-to-image with Cache-DiT acceleration +# Text-to-image with cache acceleration (cache_dit or tea_cache) python end2end.py --prompt "A cute cat" \ - --cache-backend cache_dit \ + --cache-backend tea_cache ``` > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU. @@ -139,13 +139,13 @@ python end2end.py \ | `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism | | `--enforce-eager` | False | Disable torch.compile | | `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM | -| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration | +| `--cache-backend` | None | Set to `cache_dit` or `tea_cache` for image generation acceleration | | `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation | ## Reducing GPU Memory Usage For hardware with limited VRAM, enable **module-wise CPU offload** with -`--enable-cpu-offload`. The pipeline implements `SupportsModuleOffload`, so the +`--enable-cpu-offload`. The pipeline implements `SupportsComponentDiscovery`, so the vision encoder (`vision_model`) and the Qwen3 LLM (`language_model`) are swapped between CPU and GPU on demand: @@ -211,4 +211,4 @@ The img2img counterpart lives at ## Online Serving For OpenAI-compatible API serving, see -[`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/). +[`examples/online_serving/sensenova_u1/`](../../online_serving/sensenova_u1/README.md). diff --git a/examples/offline_inference/sensenova_u1/end2end.py b/examples/offline_inference/sensenova_u1/end2end.py index 0716c74d5bc..9547c40b954 100644 --- a/examples/offline_inference/sensenova_u1/end2end.py +++ b/examples/offline_inference/sensenova_u1/end2end.py @@ -139,6 +139,12 @@ def parse_args(): default=1, help="Number of GPUs for tensor parallelism.", ) + parser.add_argument( + "--cfg-parallel-size", + type=int, + default=1, + help="Number of GPUs for CFG parallelism.", + ) parser.add_argument( "--enforce-eager", action="store_true", @@ -158,7 +164,7 @@ def parse_args(): "--cache-backend", type=str, default=None, - choices=["cache_dit"], + choices=["cache_dit", "tea_cache"], help="Cache backend to use for image generation acceleration.", ) parser.add_argument( @@ -194,6 +200,7 @@ def main(): enable_cpu_offload=args.enable_cpu_offload, cache_backend=args.cache_backend, enable_cache_dit_summary=args.enable_cache_dit_summary, + cfg_parallel_size=args.cfg_parallel_size, ) extra_args = { @@ -237,6 +244,7 @@ def main(): print(f" Seed : {args.seed}") print(f" Think mode : {args.think}") print(f" TP size : {args.tensor_parallel_size}") + print(f" CFG size : {args.cfg_parallel_size}") print(f" Cache backend : {args.cache_backend or 'none'}") print(f"{'=' * 60}\n") @@ -270,12 +278,17 @@ def main(): ) for req_output in outputs: - custom = getattr(req_output, "_custom_output", {}) or {} - if args.print_think and custom.get("think_text"): - print(f"[Think]\n{custom['think_text']}\n") + multimodal_output = getattr(req_output, "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} + think_text = text_metadata.get("think_text") if isinstance(text_metadata, dict) else None + if args.print_think and think_text: + print(f"[Think]\n{think_text}\n") if is_text_output: - text = custom.get("text_output", "") + text = multimodal_output.get("text", "") if isinstance(multimodal_output, dict) else "" + if not text: + text = text_metadata.get("text_output", "") if isinstance(text_metadata, dict) else "" if not text: text = getattr(req_output, "text", "") or "" print(f"[Response]\n{text}") 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 d783ee2f775..bfa6d961ae5 100644 --- a/examples/offline_inference/speech_to_video/speech_to_video.py +++ b/examples/offline_inference/speech_to_video/speech_to_video.py @@ -172,6 +172,29 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of GPUs used for ulysses sequence parallelism.", ) + parser.add_argument( + "--use-hsdp", + action="store_true", + help="Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.", + ) + parser.add_argument( + "--hsdp-shard-size", + type=int, + default=-1, + help=( + "Number of GPUs to shard model weights across within each replica group. " + "-1 (default) auto-calculates as world_size / replicate_size." + ), + ) + parser.add_argument( + "--hsdp-replicate-size", + type=int, + default=1, + help=( + "Number of replica groups for HSDP. Each replica holds a full sharded copy. " + "Default 1 means pure sharding (no replication)." + ), + ) parser.add_argument( "--enforce-eager", action="store_true", @@ -234,6 +257,9 @@ def main(): vae_patch_parallel_size=args.vae_patch_parallel_size, ring_degree=args.ring_degree, ulysses_degree=args.ulysses_degree, + use_hsdp=args.use_hsdp, + hsdp_shard_size=args.hsdp_shard_size, + hsdp_replicate_size=args.hsdp_replicate_size, ) profiler_enabled = args.profiler_config is not None diff --git a/examples/offline_inference/step_audio2/README.md b/examples/offline_inference/step_audio2/README.md index 2ba69ceb527..abd21f371d9 100644 --- a/examples/offline_inference/step_audio2/README.md +++ b/examples/offline_inference/step_audio2/README.md @@ -1,4 +1,4 @@ -# Step-Audio2 Offline Inference Examples +# Step-Audio2: Offline inference This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni. @@ -51,7 +51,7 @@ Single request latency comparison between vLLM-Omni and official Step-Audio2 imp **Benchmark environment:** - GPU: NVIDIA H100 80GB (single card) -- Model: Step-Audio2-mini +- Model: Step-Audio-2-mini - Warmup: 1 run, Measured: 3 runs (averaged) ### Async Chunk Streaming Performance @@ -73,7 +73,7 @@ Comparison between sequential (non-async) and async chunk modes via `/v1/audio/s **Benchmark environment:** - GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav) -- Model: Step-Audio2-mini +- Model: Step-Audio-2-mini - Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1) - Measured via `bench_tts_serve.py` @@ -91,23 +91,22 @@ pip install step-audio2 ## Model Setup +Step-Audio2 ships a custom Transformers configuration, so this model-specific +example enables remote code loading when it initializes the model. + ### Option 1: Auto-download from HuggingFace (Recommended) -The script will **automatically download** the model on first run: +The script will **automatically download** the model on first run when a +HuggingFace model id is passed: ```bash -# Just run without specifying --model, it will auto-download stepfun-ai/Step-Audio2-mini -python end2end.py --query-type audio_to_text - -# Or explicitly specify the HuggingFace model -python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio2-mini +python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio-2-mini ``` Models will be cached in `~/.cache/huggingface/hub/` for future use. -**Available models**: -- `stepfun-ai/Step-Audio2-mini` (smaller, faster) -- `stepfun-ai/Step-Audio2-7B` (larger, better quality) +**Supported model**: +- `stepfun-ai/Step-Audio-2-mini` ### Option 2: Manual Download (for offline use) @@ -115,15 +114,15 @@ Download and use locally: ```bash # Download from HuggingFace -huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini +hf download stepfun-ai/Step-Audio-2-mini --local-dir ./models/Step-Audio-2-mini # Then use the local path -python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini +python end2end.py --query-type audio_to_text --model ./models/Step-Audio-2-mini ``` Ensure the model directory contains: ``` -Step-Audio2-mini/ +Step-Audio-2-mini/ ├── config.json ├── model.safetensors (or pytorch_model.bin) ├── tokenizer.json @@ -144,20 +143,18 @@ Transcribe audio to text: ```bash # Quick start - Using default model and test audio -python end2end.py --query-type audio_to_text - -# Using your own audio file (model will auto-download) python end2end.py --query-type audio_to_text \ - --audio-path /path/to/input.wav + --model stepfun-ai/Step-Audio-2-mini -# With specific model +# Using your own audio file (model will auto-download) python end2end.py --query-type audio_to_text \ --audio-path /path/to/input.wav \ - --model stepfun-ai/Step-Audio2-7B + --model stepfun-ai/Step-Audio-2-mini # With custom question python end2end.py --query-type audio_to_text \ --audio-path input.wav \ + --model stepfun-ai/Step-Audio-2-mini \ --question "What is the speaker saying?" ``` @@ -170,12 +167,8 @@ Convert text to speech: ```bash # Basic TTS (model auto-downloads) python end2end.py --query-type text_to_audio \ - --text "Hello, this is a test of Step Audio 2 synthesis." - -# With specific model -python end2end.py --query-type text_to_audio \ - --text "Hello, this is a test." \ - --model stepfun-ai/Step-Audio2-7B + --text "Hello, this is a test of Step Audio 2 synthesis." \ + --model stepfun-ai/Step-Audio-2-mini ``` **Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model. @@ -191,12 +184,8 @@ Process input audio and generate output audio: ```bash # Basic voice conversion (model auto-downloads) python end2end.py --query-type audio_to_audio \ - --audio-path /path/to/source_audio.wav - -# With specific model -python end2end.py --query-type audio_to_audio \ - --audio-path source.wav \ - --model stepfun-ai/Step-Audio2-7B + --audio-path /path/to/source_audio.wav \ + --model stepfun-ai/Step-Audio-2-mini ``` This mode: @@ -210,8 +199,14 @@ This mode: ```bash # Use custom stage configuration python end2end.py --query-type audio_to_text \ + --model stepfun-ai/Step-Audio-2-mini \ --stage-configs-path /path/to/custom_config.yaml +# Use custom deploy configuration +python end2end.py --query-type audio_to_text \ + --model stepfun-ai/Step-Audio-2-mini \ + --deploy-config /path/to/step_audio_2_asr.yaml + # Multiple prompts (for batch testing) python end2end.py --query-type audio_to_text \ --audio-path input.wav \ @@ -246,14 +241,14 @@ python end2end.py --query-type text_to_audio \ ## Configuration -### Stage Configuration +### Deploy Configuration -The default configuration (`step_audio_2.yaml`) uses: +The default configuration (`vllm_omni/deploy/step_audio_2.yaml`) uses: -- **Stage 0 (Thinker)**: GPU 0, 80% memory -- **Stage 1 (Token2Wav)**: GPU 1, 30% memory +- **Stage 0 (Thinker)**: GPUs 0-1 with tensor parallel size 2, 70% memory +- **Stage 1 (Token2Wav)**: GPU 1, 20% memory -For **single GPU** setup, edit the config to use `devices: "0"` for both stages. +For **single GPU** setup, edit a deploy config copy to use `devices: "0"` for both stages. ### Sampling Parameters @@ -351,7 +346,7 @@ Complete example from audio to final output: # 1. ASR: Transcribe audio python end2end.py --query-type audio_to_text \ --audio-path interview.wav \ - --model ./models/Step-Audio2-7B \ + --model ./models/Step-Audio-2-mini \ --output-dir ./outputs # 2. Check the transcription @@ -361,7 +356,7 @@ cat ./outputs/00000_text.txt STEP_AUDIO2_DEFAULT_PROMPT_WAV=./speaker_samples/female_voice.wav \ python end2end.py --query-type text_to_audio \ --text "The quick brown fox jumps over the lazy dog" \ - --model ./models/Step-Audio2-7B \ + --model ./models/Step-Audio-2-mini \ --output-dir ./outputs # 4. Listen to the result diff --git a/examples/offline_inference/step_audio2/end2end.py b/examples/offline_inference/step_audio2/end2end.py index c67f44926ed..8cc7e2a263f 100644 --- a/examples/offline_inference/step_audio2/end2end.py +++ b/examples/offline_inference/step_audio2/end2end.py @@ -232,32 +232,43 @@ def main(args): raise ValueError(f"Unknown query type: {args.query_type}") # Initialize vLLM-Omni with Step-Audio2 - # Resolve stage config path + # Resolve deploy config path. Keep explicit stage_configs_path for legacy + # custom configs, but use bundled deploy configs by default. + deploy_config_path = None + stage_config_path = None + if args.stage_configs_path and args.deploy_config: + raise ValueError("--stage-configs-path and --deploy-config are mutually exclusive") if args.stage_configs_path: stage_config_path = args.stage_configs_path + elif args.deploy_config: + deploy_config_path = args.deploy_config else: - configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs" + configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/deploy" if args.query_type == "audio_to_text": # ASR only needs Thinker (Stage 0), no Token2Wav - stage_config_path = str(configs_dir / "step_audio_2_asr.yaml") - if not os.path.exists(stage_config_path): - stage_config_path = str(configs_dir / "step_audio_2.yaml") + deploy_config_path = str(configs_dir / "step_audio_2_asr.yaml") else: # TTS/S2ST need both stages - stage_config_path = str(configs_dir / "step_audio_2.yaml") - - omni_llm = Omni( - model=model_name, - log_stats=args.enable_stats, - log_file=("step_audio2_pipeline.log" if args.enable_stats else None), - init_sleep_seconds=args.init_sleep_seconds, - batch_timeout=args.batch_timeout, - init_timeout=args.init_timeout, - shm_threshold_bytes=args.shm_threshold_bytes, - stage_configs_path=stage_config_path, - worker_backend=args.worker_backend, - ray_address=args.ray_address, - ) + deploy_config_path = str(configs_dir / "step_audio_2.yaml") + + omni_kwargs = { + "model": model_name, + "log_stats": args.enable_stats, + "log_file": ("step_audio2_pipeline.log" if args.enable_stats else None), + "init_sleep_seconds": args.init_sleep_seconds, + "batch_timeout": args.batch_timeout, + "init_timeout": args.init_timeout, + "shm_threshold_bytes": args.shm_threshold_bytes, + "worker_backend": args.worker_backend, + "ray_address": args.ray_address, + "trust_remote_code": True, + } + if stage_config_path is not None: + omni_kwargs["stage_configs_path"] = stage_config_path + else: + omni_kwargs["deploy_config"] = deploy_config_path + + omni_llm = Omni(**omni_kwargs) # Configure sampling parameters for each stage # Stage 0 (Thinker): Generate text and audio tokens @@ -405,16 +416,21 @@ def parse_args(): help=( "Model path (required). Can be:\n" " - Local path: /path/to/step-audio-2\n" - " - HuggingFace ID: stepfun-ai/Step-Audio2-mini" + " - HuggingFace ID: stepfun-ai/Step-Audio-2-mini" ), ) parser.add_argument( "--stage-configs-path", type=str, default=None, - help="Path to stage config YAML file (default: use built-in step_audio_2.yaml)", + help="Path to a legacy stage config YAML file (default: use bundled deploy config)", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Path to a deploy config YAML file (default: use bundled Step-Audio2 deploy config)", ) - # Query configuration parser.add_argument( "--query-type", diff --git a/examples/offline_inference/text_to_audio/README.md b/examples/offline_inference/text_to_audio/README.md index 10a8ae37ed1..89ffc8bd8a9 100644 --- a/examples/offline_inference/text_to_audio/README.md +++ b/examples/offline_inference/text_to_audio/README.md @@ -1,5 +1,12 @@ # Text-To-Audio +A unified script for text/video-to-audio generation. Supported models: + +| Model | Tasks | Notes | +|-------|-------|-------| +| `stabilityai/stable-audio-open-1.0` | text-to-audio | gated; uses `--audio-length` | +| `zhangj1an/AudioX` | `t2a` / `t2m` / `v2a` / `v2m` / `tv2a` / `tv2m` | pass `--task`; video tasks need `--video` | + The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts. ## Prerequisites @@ -43,13 +50,45 @@ python text_to_audio.py \ --output stable_audio_output.wav ``` +### AudioX + +AudioX supports six tasks. Sampler and reference knobs (declared in +`vllm_omni/model_extras/audiox.py`) are passed via the generic `--extra-body` +JSON flag, routed into sampling `extra_args`. + +Text tasks (`t2a` / `t2m`): + +```bash +python text_to_audio.py \ + --model zhangj1an/AudioX --task t2a \ + --prompt "Fireworks burst twice, followed by a clock ticking." \ + --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 --seed 42 \ + --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \ + --output t2a.wav +``` + +Video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) require `--video`: + +```bash +python text_to_audio.py \ + --model zhangj1an/AudioX --task tv2a \ + --prompt "drum beating sound and human talking" \ + --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \ + --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ + --output tv2a.wav +``` + Key arguments: - `--prompt`: text description (string). +- `--task`: [AudioX] one of `t2a`/`t2m`/`v2a`/`v2m`/`tv2a`/`tv2m`. +- `--video`: [AudioX `v2*`/`tv2*`] video file/URL for conditioning (→ `video_path`). +- `--audio-start`: audio start offset in seconds (→ `audio_start_in_s` for Stable Audio, `seconds_start` for AudioX). +- `--audio-length`: audio duration in seconds (audio length for Stable Audio, `seconds_total` for AudioX). +- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. For AudioX, sampler/reference knobs go here, e.g. `'{"sigma_min": 0.03, "sigma_max": 1000.0, "cfg_rescale": 0.0, "audio_path": "ref.wav"}'`. - `--negative-prompt`: negative prompt for classifier-free guidance. - `--seed`: integer seed for deterministic generation. - `--guidance-scale`: classifier-free guidance scale. -- `--audio-length`: audio duration in seconds. - `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower). - `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT. - `--hsdp-shard-size`: number of GPUs used for HSDP sharding. diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 3072cc06f47..3ad88a6f346 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -2,37 +2,70 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Example script for text-to-audio generation using Stable Audio Open. +Example script for text/video-to-audio generation with diffusion audio models. -This script demonstrates how to generate audio from text prompts using -the Stable Audio Open model with vLLM-Omni. +This script supports: + * Stable Audio Open (`stabilityai/stable-audio-open-1.0`) — text-to-audio. + * AudioX (`zhangj1an/AudioX`) — six tasks: t2a / t2m / v2a / v2m / tv2a / tv2m. + +Model-specific generation knobs are declared in `vllm_omni/model_extras` and +routed into `sampling_params.extra_args` via `apply_declared_extra_args`, so +the same script serves both models without per-model branching in the engine. Usage: + # Stable Audio Open python text_to_audio.py --prompt "The sound of a dog barking" python text_to_audio.py --prompt "A piano playing a gentle melody" --audio-length 10.0 python text_to_audio.py --prompt "Thunder and rain sounds" --negative-prompt "Low quality" python text_to_audio.py --prompt "A soft synth pad" --cache-backend tea_cache + + # AudioX (text-to-audio / text-to-music) + python text_to_audio.py --model zhangj1an/AudioX --task t2a \ + --prompt "Fireworks burst twice, then a clock ticking." \ + --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 + # AudioX video-conditioned (v2*/tv2*) — pass a video file/URL + python text_to_audio.py --model zhangj1an/AudioX --task tv2a \ + --prompt "drum beating sound and human talking" --video clip.mp4 \ + --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' """ import argparse +import json import time from pathlib import Path +from typing import Any import numpy as np import torch from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import get_extra_body_params, get_model_class_name from vllm_omni.platforms import current_omni_platform +AUDIOX_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m") +AUDIOX_VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"}) + + +def parse_extra_body(value: str) -> dict[str, Any]: + """Parse a JSON object string for --extra-body.""" + try: + obj = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e + if not isinstance(obj, dict): + raise argparse.ArgumentTypeError("--extra-body must be a JSON object") + return obj + def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Generate audio with Stable Audio Open.") + parser = argparse.ArgumentParser(description="Generate audio with supported diffusion audio models.") parser.add_argument( "--model", default="stabilityai/stable-audio-open-1.0", - help="Stable Audio model name or local path.", + help="Audio diffusion model name or local path. Supported: stabilityai/stable-audio-open-1.0, zhangj1an/AudioX.", ) parser.add_argument( "--prompt", @@ -41,8 +74,11 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--negative-prompt", - default="Low quality.", - help="Negative prompt for classifier-free guidance.", + default=None, + help=( + "Negative prompt for classifier-free guidance. Off by default; " + 'recommended for Stable Audio, e.g. --negative-prompt "Low quality".' + ), ) parser.add_argument( "--seed", @@ -60,13 +96,19 @@ def parse_args() -> argparse.Namespace: "--audio-start", type=float, default=0.0, - help="Audio start time in seconds.", + help=( + "Audio start offset in seconds. Applies to both models " + "(audio_start_in_s for Stable Audio, seconds_start for AudioX)." + ), ) parser.add_argument( "--audio-length", type=float, default=10.0, - help="Audio length in seconds (max ~47s for stable-audio-open-1.0).", + help=( + "Audio duration in seconds. Maps to audio length for Stable Audio " + "(max ~47s for stable-audio-open-1.0) and seconds_total for AudioX." + ), ) parser.add_argument( "--num-inference-steps", @@ -80,11 +122,41 @@ def parse_args() -> argparse.Namespace: default=1, help="Number of audio waveforms to generate for the given prompt.", ) + parser.add_argument( + "--task", + type=str, + default=None, + choices=list(AUDIOX_TASKS), + help=( + "[AudioX only] Generation task: t2a/t2m/v2a/v2m/tv2a/tv2m. " + "v2*/tv2* require --video. Ignored for Stable Audio Open." + ), + ) + parser.add_argument( + "--video", + type=str, + default=None, + help="[AudioX v2*/tv2*] Video file path or URL used for conditioning.", + ) + parser.add_argument( + "--extra-body", + type=parse_extra_body, + default=None, + help=( + "Model-specific generation params as a JSON object, e.g. " + '\'{"sigma_min": 0.03, "sigma_max": 1000.0}\'. Each key is filtered ' + "against the model's declared extra_body_params (see vllm_omni/model_extras), " + "so unknown keys are silently dropped. Values here override the equivalent flags." + ), + ) parser.add_argument( "--output", type=str, - default="stable_audio_output.wav", - help="Path to save the generated audio (WAV format).", + default=None, + help=( + "Path to save the generated audio (WAV format). Defaults to " + "audiox_output.wav for AudioX and stable_audio_output.wav otherwise." + ), ) parser.add_argument( "--sample-rate", @@ -203,6 +275,8 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410 def main(): args = parse_args() + model_name = str(args.model).lower() if args.model else "" + is_audiox = "audiox" in model_name generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) cache_config = None if args.cache_backend == "tea_cache": @@ -211,12 +285,14 @@ def main(): } print(f"\n{'=' * 60}") - print("Stable Audio Open - Text-to-Audio Generation") + print("Text/Video-to-Audio Generation") print(f"{'=' * 60}") print(f" Model: {args.model}") + if args.task is not None: + print(f" Task: {args.task}") print(f" Prompt: {args.prompt}") print(f" Negative prompt: {args.negative_prompt}") - print(f" Audio length: {args.audio_length}s") + print(f" Duration: {args.audio_length}s") print(f" Inference steps: {args.num_inference_steps}") print(f" Guidance scale: {args.guidance_scale}") print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}") @@ -235,8 +311,9 @@ def main(): hsdp_replicate_size=args.hsdp_replicate_size, ) - # Initialize Omni with Stable Audio model - omni = Omni( + # Initialize Omni. AudioX requires an explicit pipeline class; Stable Audio + # is resolved from its config. + omni_kwargs: dict[str, Any] = dict( model=args.model, parallel_config=parallel_config, cache_backend=args.cache_backend, @@ -245,43 +322,68 @@ def main(): enable_cpu_offload=args.enable_cpu_offload, enable_layerwise_offload=args.enable_layerwise_offload, ) + if is_audiox: + omni_kwargs["model_class_name"] = "AudioXPipeline" + omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) + declared_extra_body_params = get_extra_body_params(model_class_name) + + diffusion_params = OmniDiffusionSamplingParams( + generator=generator, + guidance_scale=args.guidance_scale, + num_inference_steps=args.num_inference_steps, + num_outputs_per_prompt=args.num_waveforms, + seed=args.seed, + ) - # Calculate audio end time - audio_end_in_s = args.audio_start + args.audio_length + # Negative prompt is opt-in for both models (off by default). Including a + # non-empty negative prompt enables the separate CFG conditioning branch. + prompt: dict[str, Any] = {"prompt": args.prompt} + if args.negative_prompt: + prompt["negative_prompt"] = args.negative_prompt + + if model_class_name == "AudioXPipeline": + if args.task is None: + raise ValueError("AudioX requires --task (one of t2a/t2m/v2a/v2m/tv2a/tv2m).") + if args.task in AUDIOX_VIDEO_TASKS and not args.video: + raise ValueError(f"AudioX task {args.task!r} requires --video.") + # AudioX core knobs are surfaced as shared flags; sampler/reference knobs + # (sigma_min/sigma_max/cfg_rescale/audio_path) arrive via --extra-body and + # are filtered against the declared extra_body contract. + user_extra: dict[str, Any] = { + "audiox_task": args.task, + "seconds_start": args.audio_start, + "seconds_total": args.audio_length, + "video_path": args.video, + } + if args.extra_body: + user_extra.update(args.extra_body) + apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra) + else: + # Stable Audio Open path (backward-compatible). + audio_end_in_s = args.audio_start + args.audio_length + diffusion_params.extra_args = { + "audio_start_in_s": args.audio_start, + "audio_end_in_s": audio_end_in_s, + } + if args.extra_body: + diffusion_params.extra_args.update(args.extra_body) - # Time profiling for generation generation_start = time.perf_counter() - # Generate audio - outputs = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - }, - OmniDiffusionSamplingParams( - generator=generator, - guidance_scale=args.guidance_scale, - num_inference_steps=args.num_inference_steps, - num_outputs_per_prompt=args.num_waveforms, - extra_args={ - "audio_start_in_s": args.audio_start, - "audio_end_in_s": audio_end_in_s, - }, - ), - ) + outputs = omni.generate(prompt, diffusion_params) generation_end = time.perf_counter() generation_time = generation_end - generation_start print(f"Total generation time: {generation_time:.2f} seconds") - # Process and save audio - output_path = Path(args.output) + output = args.output or ("audiox_output.wav" if model_class_name == "AudioXPipeline" else "stable_audio_output.wav") + output_path = Path(output) output_path.parent.mkdir(parents=True, exist_ok=True) suffix = output_path.suffix or ".wav" stem = output_path.stem or "stable_audio_output" - # Extract audio from omni.generate() outputs if not outputs: raise ValueError("No output generated from omni.generate()") @@ -296,7 +398,6 @@ def main(): if audio is None: raise ValueError("No audio output found in request_output") - # Handle different output formats if isinstance(audio, torch.Tensor): audio = audio.cpu().float().numpy() @@ -323,7 +424,7 @@ def main(): save_audio(audio, str(output_path), args.sample_rate) print(f"Saved generated audio to {output_path}") - print(f"\nGenerated {args.audio_length}s of audio at {args.sample_rate} Hz") + print(f"\nGenerated ~{args.audio_length}s of audio at {args.sample_rate} Hz") if __name__ == "__main__": diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index 02538b40fac..985c4ff7d2f 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -36,6 +36,7 @@ This folder provides several entrypoints for experimenting with text-to-image di | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) | | `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3) | 160 | | `HiDream-I1-Full` | 1024 x 1024 | 63.7 | 57.7 | +| `krea/Krea-2-Raw`, `krea/Krea-2-Turbo` | 1024 (Raw) / 2048 (Turbo) | — | ~30 | !!! info *Peak VRAM: based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU. @@ -179,7 +180,6 @@ python examples/offline_inference/text_to_image/text_to_image.py \ --output flux2-dev.png ``` - ### HiDream-I1-Full Models The `--auxiliary-text-encoder` parameter is required when running HiDream‑I1‑Full: @@ -197,9 +197,42 @@ python examples/offline_inference/text_to_image/text_to_image.py \ --output /output.png ``` +### Krea 2 + +Krea 2 is published as diffusers-format checkpoints: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw`. +The pipeline class (`Krea2Pipeline`) is auto-detected from `model_index.json`. The distilled Turbo checkpoint sets +`is_distilled=true` and is best run at 2048x2048 with few steps and guidance disabled: + +```bash +# Few-step distilled (Turbo) checkpoint +python examples/offline_inference/text_to_image/text_to_image.py \ + --model krea/Krea-2-Turbo \ + --prompt "a fox in the snow" \ + --num-inference-steps 8 \ + --guidance-scale 0.0 \ + --height 2048 --width 2048 \ + --output krea2.png +``` + +For the Raw checkpoint, use more steps with CFG enabled: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model krea/Krea-2-Raw \ + --prompt "a fox in the snow" \ + --num-inference-steps 28 \ + --guidance-scale 4.5 \ + --output krea2.png +``` + +`guidance-scale` follows the Krea 2 convention (`velocity = cond + guidance_scale * (cond - uncond)`); CFG is enabled +whenever `guidance-scale > 0`. + ### Batch Requests (Multiple Prompts) -You can pass multiple prompts in a single `generate` call. +You can pass multiple prompts in a single `Omni.generate` call. `Omni` +submits each prompt as an independent request and returns one output per +prompt. ```python from vllm_omni.entrypoints.omni import Omni @@ -224,11 +257,10 @@ if __name__ == "__main__": !!! info - For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by - default, and the input list is sliced into single-item requests before feeding into the diffusion - pipeline. For models that do internally support batched inputs, you can - [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a - longer batch of prompts. + For diffusion pipelines, the input list is sliced into single-item requests + before feeding into the diffusion pipeline. For request-level batching + controls such as `max_num_seqs`, see + [Request-Level Batching](../../../docs/user_guide/diffusion/request_batching.md). ### Negative Prompts diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index b65ffa0da56..61dd29261e0 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import functools import json import time from pathlib import Path @@ -10,10 +11,18 @@ import torch from vllm_omni.diffusion.data import logger +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni +from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id +from vllm_omni.model_extras import ( + build_text_to_image_prompt, + get_extra_body_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) from vllm_omni.platforms import current_omni_platform @@ -30,16 +39,20 @@ def is_nextstep_model(model_name: str) -> bool: return False -def parse_profiler_config(value: str) -> dict[str, Any]: +def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]: + """Parse a CLI value as a JSON object, attributing errors to ``flag_name``.""" try: config = json.loads(value) except json.JSONDecodeError as e: - raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e + raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e if not isinstance(config, dict): - raise argparse.ArgumentTypeError("--profiler-config must be a JSON object") + raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object") return config +parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config") + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.") parser.add_argument( @@ -56,7 +69,16 @@ def parse_args() -> argparse.Namespace: "--stage-configs-path", type=str, default=None, - help="Path to a YAML file containing stage configurations for Omni.", + help="[Deprecated] Path to a legacy stage_args-format YAML. Prefer --deploy-config.", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help=( + "Path to a deploy YAML (new pipeline/stages format). Required for multi-stage " + "text-to-image pipelines whose deploy config is not auto-loaded." + ), ) parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.") parser.add_argument( @@ -142,7 +164,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--enforce-eager", action="store_true", - help="Disable torch.compile and force eager execution.", + default=None, + help=( + "Disable torch.compile and force eager execution. Left unset (None) " + "so it is only forwarded when explicitly given; " + "otherwise the per-stage deploy YAML value wins." + ), ) parser.add_argument( "--enable-cpu-offload", @@ -175,17 +202,12 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "int8", "gguf"], + choices=["fp8", "int8", "bitsandbytes"], help="Quantization method for the transformer. " - "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). " + "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), " + "'int8' (Int8 W8A8), 'bitsandbytes' (NF4 4-bit weight-only, CUDA SM 75+). " "Default: None (no quantization, uses BF16).", ) - parser.add_argument( - "--gguf-model", - type=str, - default=None, - help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."), - ) parser.add_argument( "--ignored-layers", type=str, @@ -208,8 +230,13 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--tensor-parallel-size", type=int, - default=1, - help="Number of GPUs used for tensor parallelism (TP) inside the DiT.", + default=None, + help=( + "Number of GPUs used for tensor parallelism (TP) inside the DiT. " + "Left unset so it is only forwarded when explicitly given; " + "otherwise the per-stage deploy YAML (or engine default of 1) wins. " + "Passing it always overrides the deploy YAML." + ), ) parser.add_argument( "--enable-expert-parallel", @@ -259,6 +286,19 @@ def parse_args() -> argparse.Namespace: action="store_true", help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.", ) + parser.add_argument( + "--extra-body", + type=functools.partial(parse_json_object, flag_name="--extra-body"), + default=None, + help=( + "Model-specific generation params as a JSON object, e.g. " + '\'{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0]}\'. ' + "Each key is filtered against the model's declared extra_body_params " + "(see vllm_omni/model_extras), so unknown keys for the chosen model are " + "silently dropped. Values here take precedence over the equivalent " + "model-specific flags above." + ), + ) parser.add_argument( "--enable-diffusion-pipeline-profiler", action="store_true", @@ -357,14 +397,7 @@ def main(): # ignored_layers is specified so the list flows through OmniDiffusionConfig quant_kwargs: dict[str, Any] = {} ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None - if args.quantization == "gguf": - if not args.gguf_model: - raise ValueError("--gguf-model is required when --quantization gguf is set.") - quant_kwargs["quantization_config"] = { - "method": "gguf", - "gguf_model": args.gguf_model, - } - elif args.quantization and ignored_layers: + if args.quantization and ignored_layers: quant_kwargs["quantization_config"] = { "method": args.quantization, "ignored_layers": ignored_layers, @@ -384,10 +417,8 @@ def main(): "ring_degree": args.ring_degree, "ulysses_mode": args.ulysses_mode, "cfg_parallel_size": args.cfg_parallel_size, - "tensor_parallel_size": args.tensor_parallel_size, "vae_patch_parallel_size": args.vae_patch_parallel_size, "enable_expert_parallel": args.enable_expert_parallel, - "enforce_eager": args.enforce_eager, "enable_cpu_offload": args.enable_cpu_offload, "mode": "text-to-image", "log_stats": args.log_stats, @@ -399,12 +430,24 @@ def main(): **lora_args, **quant_kwargs, } + if args.tensor_parallel_size is not None: + omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size + if args.enforce_eager is not None: + omni_kwargs["enforce_eager"] = args.enforce_eager if args.stage_configs_path: omni_kwargs["stage_configs_path"] = args.stage_configs_path + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config if use_nextstep: # NextStep-1.1 requires explicit pipeline class omni_kwargs["model_class_name"] = "NextStep11Pipeline" + # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails + # gate is an engine-level config (offline analog of the server's --no-guardrails). + if args.extra_body and "guardrails" in args.extra_body: + omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])} omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) + declared_extra_body_params = get_extra_body_params(model_class_name) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -419,8 +462,9 @@ def main(): print(f" Quantization: {args.quantization if args.quantization else 'None (BF16)'}") if ignored_layers: print(f" Ignored layers: {ignored_layers}") + tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default" print( - f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, " + f" Parallel configuration: tensor_parallel_size={tp_display}, " f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, " f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, " f"vae_patch_parallel_size={args.vae_patch_parallel_size}, " @@ -445,34 +489,79 @@ def main(): ) generation_start = time.perf_counter() - extra_args = { + + prompt_dict = build_text_to_image_prompt( + model_class_name=model_class_name, + prompt=args.prompt, + negative_prompt=args.negative_prompt, + height=args.height, + width=args.width, + ) + + diffusion_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + seed=args.seed, + generator=generator, + true_cfg_scale=args.cfg_scale, + guidance_scale=args.guidance_scale, + guidance_scale_2=args.guidance_scale_2, + num_inference_steps=args.num_inference_steps, + num_outputs_per_prompt=args.num_images_per_prompt, + ) + + # Base layer: backward-compatible model-specific flags. New model params should + # instead be declared in vllm_omni/model_extras and passed via --extra-body, so + # this dict does not need to grow per-model. + user_extra = { + "cfg_scale": args.cfg_scale, + "cfg_text_scale": args.cfg_scale, + "negative_prompt": args.negative_prompt, + "timestep_shift": args.timesteps_shift, "timesteps_shift": args.timesteps_shift, "cfg_schedule": args.cfg_schedule, "use_norm": args.use_norm, "use_system_prompt": args.use_system_prompt, "system_prompt": args.system_prompt, } + # Override layer: generic JSON passthrough wins over the flags above. Keys are + # still filtered against the model's declared extra_body_params downstream. + if args.extra_body: + user_extra.update(args.extra_body) + if declared_extra_body_params: + apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra) + else: + diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None}) + if lora_request: - extra_args["lora_request"] = lora_request - extra_args["lora_scale"] = args.lora_scale - - outputs = omni.generate( - { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - }, - OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - generator=generator, - true_cfg_scale=args.cfg_scale, - guidance_scale=args.guidance_scale, - guidance_scale_2=args.guidance_scale_2, - num_inference_steps=args.num_inference_steps, - num_outputs_per_prompt=args.num_images_per_prompt, - extra_args=extra_args, - ), - ) + diffusion_params.lora_request = lora_request + diffusion_params.lora_scale = args.lora_scale + + # Build per-stage sampling params for multi-stage models (e.g. BAGEL), + # or wrap single diffusion params for single-stage models. + init_non_diffusion = should_init_extra_args_for_non_diffusion_stages( + model_class_name, + ) + defaults = list(omni.default_sampling_params_list or []) + sampling_params_list = [clone_sampling_params(p) for p in defaults] + if not sampling_params_list: + sampling_params_list = [diffusion_params] + + diffusion_replaced = False + for idx, params in enumerate(sampling_params_list): + if isinstance(params, OmniDiffusionSamplingParams): + sampling_params_list[idx] = diffusion_params + diffusion_replaced = True + elif init_non_diffusion and hasattr(params, "extra_args"): + if params.extra_args is None: + params.extra_args = {} + if args.seed is not None and hasattr(params, "seed"): + params.seed = args.seed + + if not diffusion_replaced and len(sampling_params_list) == 1: + sampling_params_list = [diffusion_params] + + outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list) generation_end = time.perf_counter() generation_time = generation_end - generation_start @@ -502,15 +591,22 @@ def main(): raise ValueError("No output generated from omni.generate()") logger.info(f"Outputs: {outputs}") - first_output = outputs[0] - if not hasattr(first_output, "request_output") or not first_output.request_output: - raise ValueError("No request_output found in OmniRequestOutput") - - req_out = first_output.request_output - if not hasattr(req_out, "images"): - raise ValueError("Invalid request_output structure or missing 'images'.") + images = None + for output in outputs: + images = getattr(output, "images", None) + if images: + break + req_out = getattr(output, "request_output", None) + images = getattr(req_out, "images", None) if req_out is not None else None + if images: + break + + # Fallback: generation-stage pipelines (e.g. MammothModa2's AR->DiT) return the + # generated image as a tensor under multimodal_output instead of populating the + # `images` field that diffusion-stage pipelines fill. + if not images: + images = _images_from_multimodal_output(outputs) - images = req_out.images if not images: raise ValueError("No images found in request_output") @@ -528,5 +624,36 @@ def main(): print(f"Saved generated image to {save_path}") +def _images_from_multimodal_output(outputs: list[Any]) -> list[Any]: + """Extract PIL images from multimodal_output tensors. + + Generation-stage pipelines (e.g. MammothModa2's AR->DiT) return the generated + image as a tensor (normalized to [-1, 1], CHW) under ``multimodal_output`` + rather than populating the ``images`` field. Convert any such tensors to PIL. + """ + from PIL import Image + + pil_images: list[Any] = [] + for output in outputs: + req_out = getattr(output, "request_output", output) + for completion in getattr(req_out, "outputs", None) or []: + # multimodal_output is a MultimodalPayload (a Mapping) keyed by modality, + # matching how omni examples (ming_flash_omni / magi_human / dynin) read it. + mm = getattr(completion, "multimodal_output", None) or {} + if "image" not in mm: + continue + payload = mm["image"] + for tensor in payload if isinstance(payload, list) else [payload]: + if not isinstance(tensor, torch.Tensor): + continue + img = tensor.detach().to("cpu", dtype=torch.float32) + if img.ndim == 4: + img = img[0] + img = (img / 2 + 0.5).clamp(0, 1).mul(255).to(torch.uint8) + img = img.permute(1, 2, 0).contiguous().numpy() + pil_images.append(Image.fromarray(img)) + return pil_images + + if __name__ == "__main__": main() diff --git a/examples/offline_inference/text_to_speech/README.md b/examples/offline_inference/text_to_speech/README.md index 15672b94f09..180ab995446 100644 --- a/examples/offline_inference/text_to_speech/README.md +++ b/examples/offline_inference/text_to_speech/README.md @@ -23,6 +23,7 @@ list of supported architectures across all modalities, see | OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz | | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz | | VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz | +| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz | ## Common Quick Start @@ -369,6 +370,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \ ``` Streaming requires `async_chunk: true` in the stage config. +### Word Timestamps +Generate a WAV offline and a JSON sidecar with word-level timestamps from +`Qwen/Qwen3-ForcedAligner-0.6B`: +```bash +python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ + --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \ + --text "Hello world." \ + --output-dir /tmp/qwen3_tts_timestamps +``` +The script writes `qwen3_tts_word_timestamps.wav` and +`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit, +set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT. + ### Batched decoding The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates): ```bash @@ -422,6 +437,50 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [ --- +## IndexTTS-2 + +2-stage TTS pipeline (GPT AR talker + S2Mel CFM DiT + BigVGAN vocoder) at 22.05 kHz. Every request requires reference audio for zero-shot voice cloning. Supports emotion conditioning via audio, text, or 8-dim vector. + +### Quick start +```bash +python examples/offline_inference/text_to_speech/indextts2/end2end.py \ + --model IndexTeam/IndexTTS-2 \ + --text "你好，这是一个语音合成测试。" \ + --ref-audio /path/to/reference.wav +``` + +### Emotion control +```bash +# Emotion from reference audio +python examples/offline_inference/text_to_speech/indextts2/end2end.py \ + --model IndexTeam/IndexTTS-2 \ + --text "今天天气真好！" \ + --ref-audio /path/to/ref.wav \ + --emo-audio /path/to/happy.wav + +# Emotion from 8-dim vector (happy angry sad afraid disgusted melancholy surprised calm) +python examples/offline_inference/text_to_speech/indextts2/end2end.py \ + --model IndexTeam/IndexTTS-2 \ + --text "今天天气真好！" \ + --ref-audio /path/to/ref.wav \ + --emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 + +# Emotion from text description +python examples/offline_inference/text_to_speech/indextts2/end2end.py \ + --model IndexTeam/IndexTTS-2 \ + --text "今天天气真好！" \ + --ref-audio /path/to/ref.wav \ + --emo-text "happy and excited" +``` + +### Notes +- `--ref-audio` is **required** — IndexTTS-2 does not support text-only synthesis. +- Stage 0 (AR Talker): GPT-2 generates mel codes from text + reference audio. +- Stage 1 (S2Mel + BigVGAN): CFM DiT converts mel codes to waveform at 22.05 kHz. +- Deploy config: `vllm_omni/deploy/indextts2.yaml`. Stage 1 runs with `enforce_eager: true` (DiT has dynamic shapes). + +--- + ## Voxtral TTS Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol. @@ -462,3 +521,100 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T - `--num-prompts N` replicates the prompt for performance measurement. - `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`. - Run `--help` for the full argument surface. + +--- + +## SoulX-Singer + +Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`. + +### Prerequisites + +Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it. + +```bash +# 1. DiT weights +export BASE=path/to/SoulX-Singer +export PREPROCESS=path/to/SoulX-Singer-Preprocess +export SVC_DIR=path/to/SoulX-Singer-svc + +huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE" + +# 2. Preprocess weights (required) +huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS" +export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS" + +# 3. SVS / SVC view directories +mkdir -p "$SVC_DIR" +cp $BASE/{config.yaml,README.md,assets} $SVC_DIR +mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt + +cat > "$BASE/config.json" <<'EOF' +{ + "model_type": "soulxsinger", + "architectures": ["SoulXSingerPipeline"], + "max_num_seqs": 1 +} +EOF + +cat > "$SVC_DIR/config.json" <<'EOF' +{ + "model_type": "soulxsinger", + "architectures": ["SoulXSingerSVCPipeline"], + "max_num_seqs": 1 +} +EOF +``` + +`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches. + +**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs: + +```bash +pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC +``` + +Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`: + +```bash +# install optional dependencies: +pip install -e ".[soulx-svs]" +``` + +English SVS adds NeMo ASR and NLTK data; pass `--language English`: + +```bash +pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2 +python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')" +``` + +**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni. + +### Quick start + +```bash +# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3 +python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \ + --model "$BASE" \ + --preprocess-weights-dir "$PREPROCESS" \ + --control score \ + --num-inference-steps 32 \ + -o output.wav + +python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \ + --model "$SVC_DIR" \ + --preprocess-weights-dir "$PREPROCESS" \ + --svc \ + --num-inference-steps 32 \ + -o output_svc.wav +``` + +`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options. + +### Notes + +- Output: 24 kHz mono WAV; batch only. +- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on. +- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only. + +--- diff --git a/examples/offline_inference/text_to_speech/fish_speech/end2end.py b/examples/offline_inference/text_to_speech/fish_speech/end2end.py index 4d6855246a5..5661cccd3b2 100644 --- a/examples/offline_inference/text_to_speech/fish_speech/end2end.py +++ b/examples/offline_inference/text_to_speech/fish_speech/end2end.py @@ -122,12 +122,48 @@ def _save_wav(output_dir: str, request_id: str, mm: dict) -> None: logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr) +def _collect_audio_chunk( + mm: dict, + all_audio_chunks: list[torch.Tensor], + sample_rate: int | None, +) -> tuple[int | None, int]: + """Extract audio and sample rate from a multimodal output chunk. + + Returns updated sample_rate and the number of audio segments in this chunk. + """ + audio = mm.get("audio") if mm else None + if audio is not None: + if isinstance(audio, list): + all_audio_chunks.extend(audio) + else: + all_audio_chunks.append(audio) + sr_raw = mm.get("sr") if mm else None + if sr_raw is not None and not (isinstance(sr_raw, (list, tuple)) and len(sr_raw) == 0): + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val) + n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1) + return sample_rate, n + + +def _build_omni_kwargs(args, model_name: str) -> dict: + """Build keyword arguments for Omni/AsyncOmni initialization.""" + omni_kwargs = dict( + model=model_name, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config + else: + omni_kwargs["stage_configs_path"] = args.stage_configs_path or DEFAULT_STAGE_CONFIG + return omni_kwargs + + def main(args): """Run offline inference with Omni.""" output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) - stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG model_name = args.model or DEFAULT_MODEL inputs = [ @@ -139,12 +175,7 @@ def main(args): ) ] - omni = Omni( - model=model_name, - stage_configs_path=stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) + omni = Omni(**_build_omni_kwargs(args, model_name)) t_start = time.perf_counter() for stage_outputs in omni.generate(inputs): @@ -165,7 +196,6 @@ async def main_streaming(args): output_dir = args.output_dir os.makedirs(output_dir, exist_ok=True) - stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG model_name = args.model or DEFAULT_MODEL prompt = build_prompt( @@ -175,24 +205,20 @@ async def main_streaming(args): model_name=model_name, ) - omni = AsyncOmni( - model=model_name, - stage_configs_path=stage_configs_path, - log_stats=args.log_stats, - stage_init_timeout=args.stage_init_timeout, - ) + omni = AsyncOmni(**_build_omni_kwargs(args, model_name)) request_id = "0" t_start = time.perf_counter() t_prev = t_start chunk_idx = 0 + all_audio_chunks: list[torch.Tensor] = [] + sample_rate = None async for stage_output in omni.generate(prompt, request_id=request_id): mm = stage_output.request_output.outputs[0].multimodal_output if not stage_output.finished: t_now = time.perf_counter() - audio = mm.get("audio") - n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1) + sample_rate, n = _collect_audio_chunk(mm, all_audio_chunks, sample_rate) dt_ms = (t_now - t_prev) * 1000 ttfa_ms = (t_now - t_start) * 1000 if chunk_idx == 0: @@ -202,10 +228,20 @@ async def main_streaming(args): t_prev = t_now chunk_idx += 1 else: + # Collect any remaining audio in the final message + sample_rate, _ = _collect_audio_chunk(mm, all_audio_chunks, sample_rate) + t_end = time.perf_counter() total_ms = (t_end - t_start) * 1000 logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx) - _save_wav(output_dir, request_id, mm) + + if all_audio_chunks and sample_rate: + audio_tensor = torch.cat(all_audio_chunks, dim=-1) + out_wav = os.path.join(output_dir, f"output_{request_id}.wav") + sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sample_rate, format="WAV") + logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sample_rate) + else: + logger.warning("Request %s: no audio chunks received", request_id) def parse_args(): @@ -236,11 +272,17 @@ def parse_args(): default=None, help=f"HuggingFace model path (default: {DEFAULT_MODEL}).", ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Deploy config YAML path (under vllm_omni/deploy/).", + ) parser.add_argument( "--stage-configs-path", type=str, default=None, - help="Path to stage configs YAML.", + help="Path to stage configs YAML (legacy, prefer --deploy-config).", ) parser.add_argument( "--output-dir", diff --git a/examples/offline_inference/text_to_speech/glm_tts/end2end.py b/examples/offline_inference/text_to_speech/glm_tts/end2end.py index d1f19f8a5c8..6624c2b35e6 100644 --- a/examples/offline_inference/text_to_speech/glm_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/glm_tts/end2end.py @@ -101,7 +101,7 @@ def _extract_sample_rate(audio_mm: dict) -> int: def _modify_deploy_config(base_path: str, async_chunk: bool) -> str: """Build deploy config with explicit sync/async mode and eager execution. - Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts.py`` + Mirrors the logic in ``tests/e2e/offline_inference/test_glm_tts_expansion.py`` (``_get_deploy_config``) so that example runs match CI behavior. """ with open(base_path) as f: diff --git a/examples/offline_inference/text_to_speech/indextts2/end2end.py b/examples/offline_inference/text_to_speech/indextts2/end2end.py new file mode 100644 index 00000000000..858487ae106 --- /dev/null +++ b/examples/offline_inference/text_to_speech/indextts2/end2end.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Offline inference example for IndexTTS2 via vLLM-Omni. + +Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1). +Output is 22050 Hz mono WAV. + +Usage: + python end2end.py \ + --model /path/to/IndexTeam/IndexTTS-2 \ + --text "你好，这是一个语音合成测试。" \ + --ref-audio /path/to/ref.wav + + # With emotion audio: + python end2end.py \ + --model /path/to/IndexTeam/IndexTTS-2 \ + --text "今天天气真好！" \ + --ref-audio /path/to/ref.wav \ + --emo-audio /path/to/happy.wav +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import soundfile as sf +import torch +from vllm import SamplingParams +from vllm.utils.argparse_utils import FlexibleArgumentParser + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + +from vllm_omni import Omni # noqa: E402 +from vllm_omni.model_executor.models.indextts2.prompt_utils import ( # noqa: E402 + build_indextts2_prefill_prompt_ids, +) + + +def build_request( + model: str, + text: str, + ref_audio_path: str | None = None, + emo_audio_path: str | None = None, + emo_text: str | None = None, + emo_vector: list[float] | None = None, + emo_alpha: float | None = None, + use_emo_text: bool = False, + use_random: bool = False, +) -> dict: + additional: dict = {"text": [text]} + if ref_audio_path: + additional["voice"] = [str(ref_audio_path)] + if emo_audio_path: + additional["emo_audio"] = [str(emo_audio_path)] + if emo_text: + additional["emo_text"] = [emo_text] + if emo_vector is not None: + additional["emo_vector"] = [emo_vector] + if emo_alpha is not None: + additional["emo_alpha"] = [emo_alpha] + if use_emo_text: + additional["use_emo_text"] = [True] + if use_random: + additional["use_random"] = [True] + return { + "prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text), + "additional_information": additional, + } + + +def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None: + audio_np = waveform.float().numpy() + sf.write(path, audio_np, sample_rate) + print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)") + + +def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]: + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if isinstance(audio, list): + chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0] + audio = torch.cat(chunks, dim=0) if chunks else None + + sr_val = mm.get("sr") + if isinstance(sr_val, list): + sr_val = sr_val[-1] if sr_val else None + if hasattr(sr_val, "item"): + sample_rate = int(sr_val.item()) + else: + sample_rate = int(sr_val) if sr_val is not None else 22050 + + return audio if isinstance(audio, torch.Tensor) else None, sample_rate + + +def main(args) -> None: + omni = Omni( + model=args.model, + deploy_config=args.deploy_config, + stage_init_timeout=args.stage_init_timeout, + ) + + # Stage 0: GPT AR (text → mel codes) + gpt_sampling = SamplingParams( + temperature=0.8, + top_p=0.8, + top_k=30, + max_tokens=1500, + repetition_penalty=10.0, + stop_token_ids=[8193], + seed=args.seed if args.seed is not None else 42, + detokenize=False, + ) + # Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored) + s2mel_sampling = SamplingParams( + temperature=0.0, + max_tokens=65536, + detokenize=True, + ) + sampling_params = [gpt_sampling, s2mel_sampling] + + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"Synthesizing: {args.text!r}") + if args.ref_audio: + print(f" ref_audio: {args.ref_audio}") + inputs = build_request( + model=args.model, + text=args.text, + ref_audio_path=args.ref_audio, + emo_audio_path=args.emo_audio, + emo_text=args.emo_text, + emo_vector=args.emo_vector, + emo_alpha=args.emo_alpha, + use_emo_text=args.use_emo_text, + use_random=args.use_random, + ) + + for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)): + mm = omni_out.multimodal_output + if not mm: + print(f" [req {i}] No multimodal output.") + continue + audio, sr = extract_audio(mm) + if audio is None: + print(f" [req {i}] No waveform in multimodal_output.") + continue + out_path = str(output_dir / f"output_{i}.wav") + save_audio(audio.cpu(), out_path, sr) + + print("Done.") + + +def parse_args(): + parser = FlexibleArgumentParser(description="IndexTTS2 offline inference") + parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.") + parser.add_argument("--text", default="你好，这是IndexTTS2语音合成测试。") + parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.") + parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.") + parser.add_argument("--emo-text", default=None, help="Emotion description text.") + parser.add_argument( + "--emo-vector", + type=float, + nargs=8, + default=None, + help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.", + ) + parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].") + parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.") + parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.") + parser.add_argument("--seed", type=int, default=None) + parser.add_argument( + "--output-dir", + default=os.path.join( + os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")), + "indextts2_output", + ), + ) + parser.add_argument("--deploy-config", default=None) + parser.add_argument("--stage-init-timeout", type=int, default=600) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/text_to_speech/moss_tts/end2end.py b/examples/offline_inference/text_to_speech/moss_tts/end2end.py index fb47ad8c9d3..51bf52168b3 100644 --- a/examples/offline_inference/text_to_speech/moss_tts/end2end.py +++ b/examples/offline_inference/text_to_speech/moss_tts/end2end.py @@ -1,11 +1,12 @@ """Offline inference example for MOSS-TTS variants. -Covers all five models: - - OpenMOSS-Team/MOSS-TTS (8B, general TTS) - - OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS) - - OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis) - - OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming) - - OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design) +Covers six models: + - OpenMOSS-Team/MOSS-TTS (8B, general TTS) + - OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS) + - OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis) + - OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming) + - OpenMOSS-Team/MOSS-VoiceGenerator (1.7B, voice design) + - OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 (n_vq=12, 48 kHz stereo) Usage ----- @@ -210,12 +211,18 @@ def run_tts(args: argparse.Namespace) -> None: "OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect", "OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime", "OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator", + "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5": "moss_tts_local", } deploy_config = deploy_map.get(args.model, "moss_tts") is_sound_effect = "SoundEffect" in args.model is_voice_gen = "VoiceGenerator" in args.model is_realtime = "Realtime" in args.model + # Local-v1.5 ships its own AutoProcessor (processor_config.json + + # processing_moss_tts.py), so it reuses the same _build_unified_codes + # path as the Delay variants below — just a different n_vq/codec. + is_local = "Local" in args.model + output_sr = 48000 if is_local else 24000 ref_audio = None builder_kwargs: dict = {} @@ -314,27 +321,33 @@ def run_tts(args: argparse.Namespace) -> None: wav = mm.get("model_outputs") if wav is None: continue + # Local-v1.5's codec emits stereo chunks shaped (C, T); keep the + # channel axis intact (don't flatten) so concatenation below stays + # along the time axis. Mono variants stay (T,) as before. if isinstance(wav, list): - non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0] + non_empty = [w for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0] if non_empty and t_first_audio is None: t_first_audio = _time.monotonic() chunks.extend(non_empty) elif isinstance(wav, torch.Tensor) and wav.numel() > 0: if t_first_audio is None: t_first_audio = _time.monotonic() - chunks.append(wav.reshape(-1)) + chunks.append(wav) t_gen_end = _time.monotonic() if not chunks: print("No audio in output — check model loading logs.") return - wav_t = torch.cat(chunks).float().cpu() - sf.write(args.output, wav_t.numpy(), 24000) - print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}") + cat_dim = -1 if chunks[0].dim() == 2 else 0 + wav_t = torch.cat(chunks, dim=cat_dim).float().cpu() + n_samples = wav_t.shape[-1] + # soundfile wants (T,) mono or (T, C) multi-channel. + sf.write(args.output, (wav_t.transpose(0, 1) if wav_t.dim() == 2 else wav_t).numpy(), output_sr) + print(f"Saved {n_samples / output_sr:.2f}s of audio to {args.output}") ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan") gen_total_ms = (t_gen_end - t_gen_start) * 1000.0 - print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}") + print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={n_samples / output_sr:.2f}") def main() -> None: diff --git a/examples/offline_inference/text_to_speech/omnivoice/end2end.py b/examples/offline_inference/text_to_speech/omnivoice/end2end.py index 619f8160b37..a6bdc7b4ef3 100644 --- a/examples/offline_inference/text_to_speech/omnivoice/end2end.py +++ b/examples/offline_inference/text_to_speech/omnivoice/end2end.py @@ -34,7 +34,7 @@ def run_e2e(): help="Model name or path (HuggingFace or local)", ) parser.add_argument( - "--stage-config", + "--deploy-config", type=str, default="vllm_omni/deploy/omnivoice.yaml", ) @@ -87,14 +87,14 @@ def run_e2e(): ) args = parser.parse_args() - if not os.path.exists(args.stage_config): - raise FileNotFoundError(f"Stage config not found: {args.stage_config}") + if not os.path.exists(args.deploy_config): + raise FileNotFoundError(f"Deploy config not found: {args.deploy_config}") print(f"Initializing OmniVoice with model={args.model}") omni = Omni( model=args.model, - stage_configs_path=args.stage_config, + deploy_config=args.deploy_config, log_stats=True, ) diff --git a/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py new file mode 100644 index 00000000000..581be2a7df6 --- /dev/null +++ b/examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py @@ -0,0 +1,172 @@ +"""Offline Qwen3-TTS example with word-level timestamps. + +This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with +the same shared forced aligner used by the streaming server path, then writes +the WAV and JSON sidecar: + + python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \ + --model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B + +On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0`` +to avoid FlashInfer sampler JIT during warmup. +""" + +from __future__ import annotations + +import asyncio +import json +import os +from pathlib import Path +from typing import Any + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") +# The in-process offline aligner consumes token_classify outputs. Keeping V1 +# multiprocessing off avoids msgspec serialization issues for token-wise logits. +os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + +import numpy as np +import soundfile as sf +import torch +from end2end import _estimate_prompt_len +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import Omni +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults +from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config + + +def _default_stage_config() -> str: + repo_root = Path(__file__).resolve().parents[4] + return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml") + + +def _build_custom_voice_input(args: Any) -> dict[str, Any]: + additional_information = { + "task_type": ["CustomVoice"], + "text": [args.text], + "language": [args.language], + "speaker": [args.speaker], + "instruct": [args.instructions], + "max_new_tokens": [args.max_new_tokens], + } + return { + "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model), + "additional_information": additional_information, + } + + +def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]: + audio_data = mm["audio"] + sr_raw = mm["sr"] + sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw + sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val) + audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data + return audio_tensor.float().detach().cpu().flatten(), sample_rate + + +def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes: + samples = audio.numpy() + pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2") + return pcm.tobytes() + + +async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None: + aligner_config = build_forced_aligner_config(args) + if aligner_config is None: + raise ValueError("--forced-aligner or --forced-aligner-config is required") + timestamps = await align( + audio=_float_audio_to_pcm16_bytes(audio), + text=args.text, + sample_rate=sample_rate, + config=aligner_config, + language=args.language, + ) + if timestamps is None: + return None + return [ + { + "word": item.word, + "start_ms": item.start_ms, + "end_ms": item.end_ms, + } + for item in timestamps + ] + + +def main(args: Any) -> None: + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + omni_kwargs = vars(args).copy() + for key in ( + "forced_aligner", + "text", + "language", + "speaker", + "instructions", + "max_new_tokens", + "output_dir", + ): + omni_kwargs.pop(key, None) + omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config() + omni_kwargs["log_stats"] = args.log_stats + omni = Omni(**omni_kwargs) + + prompt = _build_custom_voice_input(args) + final_output = None + for stage_outputs in omni.generate([prompt]): + final_output = stage_outputs.request_output + if final_output is None: + raise RuntimeError("Qwen3-TTS did not produce an output.") + + mm = final_output.outputs[0].multimodal_output + audio, sample_rate = _audio_tensor_and_sample_rate(mm) + + wav_path = output_dir / "qwen3_tts_word_timestamps.wav" + # Align the in-memory tensor, not the saved WAV. The WAV is only a demo + # artifact so users can listen to the same audio referenced by the sidecar. + timestamps = asyncio.run(_run_alignment(args, audio, sample_rate)) + sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV") + + sidecar = { + "text": args.text, + "sample_rate": sample_rate, + "audio_path": str(wav_path), + "timestamps": timestamps, + } + json_path = output_dir / "qwen3_tts_word_timestamps.json" + json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + + print(f"Saved audio: {wav_path}") + print(f"Saved timestamps: {json_path}") + print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2)) + + +def parse_args() -> Any: + parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example") + parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name") + parser.add_argument( + "--forced-aligner", + default=None, + help="Qwen3 forced aligner model path/name", + ) + parser.add_argument( + "--forced-aligner-config", + default=None, + help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)", + ) + parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML") + parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align") + parser.add_argument("--language", default="English", help="Qwen3-TTS language field") + parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name") + parser.add_argument("--instructions", default="", help="Optional speaking style instruction") + parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens") + parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar") + parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging") + nullify_stage_engine_defaults(parser) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py new file mode 100644 index 00000000000..97ed7eb9794 --- /dev/null +++ b/examples/offline_inference/text_to_speech/soulxsinger/benchmark.py @@ -0,0 +1,162 @@ +"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings). + +Usage: + python benchmark.py --model /path/to/SoulX-Singer --svs \\ + --prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\ + --target-audio ../SoulX-Singer/example/audio/music.mp3 \\ + --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ + -o benchmark.wav + + python benchmark.py ... --enable-diffusion-pipeline-profiler +""" + +import argparse +import statistics +import time +from pathlib import Path + +import soundfile as sf +from end2end import ( + SVC_DEPLOY_CONFIG, + SVS_DEPLOY_CONFIG, + add_inference_args, + build_sampling, + extract_audio, + resolve_preprocess_weights_dir, +) + +from vllm_omni.entrypoints.omni import Omni + + +def _audio_duration_sec(outputs) -> float: + audio_np, sr = extract_audio(outputs) + return audio_np.size / sr + + +def _stage_durations_from_output(outputs) -> dict[str, float]: + omni_out = outputs[0] + durations = getattr(omni_out, "stage_durations", None) or {} + if durations: + return dict(durations) + ro = getattr(omni_out, "request_output", None) + if ro is not None: + return dict(getattr(ro, "stage_durations", None) or {}) + return {} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark") + add_inference_args(parser) + parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)") + parser.add_argument("--runs", type=int, default=3, help="Measured runs") + parser.add_argument( + "--enable-diffusion-pipeline-profiler", + action="store_true", + help="Enable low-overhead stage timing in DiffusionOutput.stage_durations", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable torch.compile on diff_estimator (eager baseline for A/B)", + ) + parser.add_argument( + "--dtype", + type=str, + default=None, + choices=["bfloat16", "bf16", "float16", "fp16", "half"], + help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.", + ) + parser.add_argument( + "--output", + "-o", + type=str, + default=None, + help="Save WAV from the last measured run (after warmup)", + ) + parser.add_argument( + "--save-each-run", + action="store_true", + help="With --output, also write measured runs as <stem>_run1.wav, _run2.wav, ...", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir) + deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG) + mode = "SVS" if args.svs else "SVC" + + omni_kwargs: dict = { + "model": args.model, + "deploy_config": deploy_config, + "async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed) + } + if args.enable_diffusion_pipeline_profiler: + omni_kwargs["enable_diffusion_pipeline_profiler"] = True + if args.enforce_eager: + omni_kwargs["enforce_eager"] = True + if args.dtype is not None: + omni_kwargs["dtype"] = args.dtype + + compile_mode = "eager" if args.enforce_eager else "torch.compile" + dtype_label = args.dtype or "from deploy YAML" + print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]") + + omni = Omni(**omni_kwargs) + kind = "svs" if args.svs else "svc" + sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind) + prompt = {"prompt_token_ids": [0]} + + latencies_ms: list[float] = [] + rtfs: list[float] = [] + last_stages: dict[str, float] = {} + last_measured_outputs = None + + total_runs = args.warmup + args.runs + for run_idx in range(total_runs): + is_warmup = run_idx < args.warmup + label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}" + t0 = time.perf_counter() + outputs = list(omni.generate([prompt], sampling)) + elapsed_ms = (time.perf_counter() - t0) * 1000.0 + audio_sec = _audio_duration_sec(outputs) + rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf") + print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}") + if not is_warmup: + latencies_ms.append(elapsed_ms) + rtfs.append(rtf) + last_stages = _stage_durations_from_output(outputs) + last_measured_outputs = outputs + if args.output and args.save_each_run: + measured_idx = run_idx - args.warmup + 1 + out_path = Path(args.output) + run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}") + audio_np, sr = extract_audio(outputs) + sf.write(str(run_path), audio_np, sr) + print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)") + + omni.close() + + if args.output and last_measured_outputs is not None: + audio_np, sr = extract_audio(last_measured_outputs) + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + sf.write(str(out_path), audio_np, sr) + print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)") + + if latencies_ms: + print("\n=== Summary (measured runs) ===") + print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}") + print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}") + if last_stages: + print("\nStage durations (last measured run):") + for name, value in sorted(last_stages.items()): + if name.endswith("_ms"): + print(f" {name}: {value:.1f} ms") + else: + print(f" {name}: {value * 1000.0:.1f} ms") + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/text_to_speech/soulxsinger/end2end.py b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py new file mode 100644 index 00000000000..116b8da4b33 --- /dev/null +++ b/examples/offline_inference/text_to_speech/soulxsinger/end2end.py @@ -0,0 +1,285 @@ +"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline). + +Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference. + +Usage: + # SVS (default) + python end2end.py --model /path/to/SoulX-Singer \\ + --prompt-audio path/to/prompt/audio.mp3 \\ + --target-audio path/to/target/audio.mp3 \\ + --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ + -o output.wav + + # SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view) + python end2end.py --model /path/to/SoulX-Singer-svc \\ + --prompt-audio path/to/prompt/audio.mp3 \\ + --target-audio path/to/target/audio.mp3 \\ + --preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\ + -o output_svc.wav +""" + +import argparse +import os +from pathlib import Path + +import numpy as np +import soundfile as sf + +from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +REPO_ROOT = Path(__file__).resolve().parents[4] +SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml" +SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml" +DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3" +DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3" +_SAMPLE_RATE = 24000 + + +def _require_paths(paths: dict[str, str | None]) -> None: + missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)] + if missing: + raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing)) + + +def resolve_preprocess_weights_dir(explicit: str | None) -> str: + if explicit: + path = Path(explicit).expanduser().resolve() + if not (path / "rmvpe" / "rmvpe.pt").is_file(): + raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}") + return str(path) + env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR") + if env: + path = Path(env).expanduser().resolve() + if (path / "rmvpe" / "rmvpe.pt").is_file(): + return str(path) + raise FileNotFoundError( + "SoulX preprocess weights not found. Pass --preprocess-weights-dir or set " + "SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt " + "(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)." + ) + + +def build_sampling( + args: argparse.Namespace, + *, + preprocess_weights_dir: str, + kind: str, +) -> OmniDiffusionSamplingParams: + extra: dict = { + "prompt_audio": os.path.abspath(args.prompt_audio), + "target_audio": os.path.abspath(args.target_audio), + "preprocess_weights_dir": preprocess_weights_dir, + "vocal_sep": args.vocal_sep, + "auto_shift": args.auto_shift, + "pitch_shift": args.pitch_shift, + } + if kind == "svs": + extra["language"] = args.language + extra["control"] = args.control + if args.prompt_metadata_path: + extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path) + if args.target_metadata_path: + extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path) + if args.prompt_metadata_path or args.target_metadata_path: + extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio) + + extra = validate_soulx_extra_args(kind, extra) + + return OmniDiffusionSamplingParams( + num_inference_steps=args.num_inference_steps, + guidance_scale=args.guidance_scale, + seed=args.seed, + extra_args=extra, + ) + + +def extract_audio(outputs) -> tuple[np.ndarray, int]: + out = outputs[0] + mm = getattr(out, "multimodal_output", None) + if mm is None: + ro = getattr(out, "request_output", None) + if ro is not None: + mm = getattr(ro, "multimodal_output", None) + if mm is None and ro.outputs: + mm = getattr(ro.outputs[0], "multimodal_output", None) + if not mm or "audio" not in mm: + raise RuntimeError("No audio in multimodal_output") + + audio_val = mm["audio"] + if isinstance(audio_val, list): + chunks = [ + np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1) + for chunk in audio_val + if chunk is not None + ] + audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32) + elif hasattr(audio_val, "cpu"): + audio_np = audio_val.detach().cpu().numpy().squeeze() + else: + audio_np = np.asarray(audio_val, dtype=np.float32).squeeze() + + sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE)) + if isinstance(sr_val, list) and sr_val: + sr_val = sr_val[-1] + if hasattr(sr_val, "item"): + sr_val = sr_val.item() + return audio_np.astype(np.float32, copy=False), int(sr_val) + + +def add_inference_args(parser: argparse.ArgumentParser) -> None: + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--svc", + action="store_true", + help="Assert SVC mode (must match config.json architectures)", + ) + parser.add_argument( + "--model", + type=str, + required=True, + help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)", + ) + parser.add_argument( + "--deploy-config", + type=str, + default=None, + help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)", + ) + parser.add_argument( + "--prompt-audio", + type=str, + default=str(DEFAULT_PROMPT_AUDIO), + help="Prompt / reference audio (mp3 or wav)", + ) + parser.add_argument( + "--target-audio", + type=str, + default=str(DEFAULT_TARGET_AUDIO), + help="Target content audio (mp3 or wav)", + ) + parser.add_argument( + "--prompt-metadata-path", + type=str, + default=None, + help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)", + ) + parser.add_argument( + "--target-metadata-path", + type=str, + default=None, + help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)", + ) + parser.add_argument( + "--prompt-wav-path", + type=str, + default=None, + help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)", + ) + parser.add_argument( + "--preprocess-weights-dir", + type=str, + default=None, + help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)", + ) + parser.add_argument( + "--language", + type=str, + default="Mandarin", + help="SVS lyrics language (Mandarin, Cantonese, English)", + ) + parser.add_argument( + "--control", + type=str, + default="melody", + choices=["score", "melody"], + help="SVS control mode", + ) + parser.add_argument( + "--vocal-sep", + action=argparse.BooleanOptionalAction, + default=False, + help="Run vocal separation during preprocess", + ) + parser.add_argument( + "--auto-shift", + action=argparse.BooleanOptionalAction, + default=True, + help="Auto pitch shift to align target pitch with prompt", + ) + parser.add_argument("--pitch-shift", type=int, default=0) + parser.add_argument("--num-inference-steps", type=int, default=32) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Optional CFM RNG seed. Omit for non-deterministic sampling.", + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)") + add_inference_args(parser) + parser.add_argument("--output", "-o", type=str, default="output.wav") + args = parser.parse_args() + return args + + +def main() -> None: + args = parse_args() + kind = resolve_soulx_kind(args.model) + + if not (args.svc ^ (kind == "svs")): + raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}") + + required = { + "prompt_audio": args.prompt_audio, + "target_audio": args.target_audio, + } + if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path): + raise ValueError( + "SVS precomputed metadata requires both --prompt-metadata-path and " + "--target-metadata-path (or neither, for online preprocess)." + ) + if args.prompt_metadata_path: + required["prompt_metadata_path"] = args.prompt_metadata_path + required["target_metadata_path"] = args.target_metadata_path + if args.prompt_wav_path: + required["prompt_wav_path"] = args.prompt_wav_path + _require_paths(required) + preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir) + deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG) + mode = kind.upper() + + print(f"Loading SoulX-Singer {mode} from {args.model}") + print(f" deploy: {deploy_config}") + print(f" prompt: {args.prompt_audio}") + print(f" target: {args.target_audio}") + + omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only + sampling = build_sampling( + args, + preprocess_weights_dir=preprocess_dir, + kind=kind, + ) + prompt = {"prompt_token_ids": [0], "multi_modal_data": {}} + + print("Running generate (inline preprocess + DiT)...") + outputs = list(omni.generate([prompt], sampling)) + if not outputs: + raise RuntimeError("No output from omni.generate") + + err = getattr(outputs[0], "error", None) + if err is not None: + raise RuntimeError(f"Pipeline error: {err}") + + audio_np, sr = extract_audio(outputs) + sf.write(args.output, audio_np, sr) + print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)") + omni.close() + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index f852e980a78..876b8429b44 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -6,9 +6,13 @@ A unified script for text-to-video generation. Supports multiple models with mod | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | |---|---|---|---|---|---| +| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | +| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) | +| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — | ## Local CLI Usage @@ -29,6 +33,28 @@ python text_to_video.py \ --output t2v_out.mp4 ``` +### Wan2.1 VACE (T2V) + +VACE text-to-video uses this shared entrypoint. Conditional VACE tasks use +the shared [`image_to_video.py`](../image_to_video/README.md#wan21-vace-conditional-tasks) +entrypoint, which constructs the pipeline-native conditioning data from the +provided media inputs. No explicit mode parameter is required. + +```bash +python text_to_video.py \ + --model Wan-AI/Wan2.1-VACE-1.3B-diffusers \ + --prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \ + --seed 0 \ + --height 480 \ + --width 832 \ + --num-frames 81 \ + --num-inference-steps 30 \ + --guidance-scale 5.0 \ + --flow-shift 5.0 \ + --vae-use-tiling \ + --output vace_t2v_output.mp4 +``` + LTX2 example: ```bash @@ -45,6 +71,30 @@ python text_to_video.py \ --output ltx2_out.mp4 ``` +### LTX-2.3 + +```bash +python text_to_video.py \ + --model diffusers/LTX-2.3-Diffusers \ + --model-class-name LTX23Pipeline \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --num-inference-steps 20 \ + --guidance-scale 4.0 \ + --frame-rate 24 \ + --fps 24 \ + --audio-sample-rate 48000 \ + --output ltx23_t2v_output.mp4 +``` + +Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3 +text-to-video pipeline explicitly. + ### HunyuanVideo-1.5 (480p) ```bash @@ -100,11 +150,70 @@ python text_to_video.py \ --output quick_test.mp4 ``` +### Cosmos3 + +```bash +python text_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A robot arm is cleaning a plate in the kitchen." \ + --negative-prompt "blurry, distorted, low quality, jittery, deformed" \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, + "use_resolution_template": false, "use_duration_template": false}' \ + --output cosmos3_t2v.mp4 +``` + +### Helios (T2V) + +Helios ships three variants. Model-specific knobs (declared in +`vllm_omni/model_extras/helios.py`) are passed via the generic `--extra-body` +JSON flag rather than bespoke per-model flags. + +**Helios-Base** (Stage 1 only): + +```bash +python text_to_video.py \ + --model BestWishYsh/Helios-Base \ + --prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \ + --guidance-scale 5.0 \ + --output helios_t2v_base.mp4 +``` + +**Helios-Mid** (Stage 2 pyramid + CFG-Zero*): + +```bash +python text_to_video.py \ + --model BestWishYsh/Helios-Mid \ + --prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \ + --guidance-scale 5.0 \ + --extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}' \ + --output helios_t2v_mid.mp4 +``` + +**Helios-Distilled** (Stage 2 pyramid + DMD, few-step): + +```bash +python text_to_video.py \ + --model BestWishYsh/Helios-Distilled \ + --prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \ + --num-frames 240 \ + --guidance-scale 1.0 \ + --extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}' \ + --output helios_t2v_distilled.mp4 +``` + +> Helios image-to-video (I2V) and video-to-video (V2V) require image/video +> conditioning tensors that cannot be passed through the JSON `--extra-body` +> flag; they are out of scope for this text-to-video example. + ## Key Arguments ### Common - `--model`: Diffusers model ID or local path. +- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline` + for LTX-2.3 text-to-video. - `--prompt`: text description (string). - `--height/--width`: output resolution. Default depends on model. - `--num-frames`: number of frames. Default depends on model. @@ -112,6 +221,7 @@ python text_to_video.py \ - `--num-inference-steps`: sampling steps. Default depends on model. - `--fps`: frames per second for the saved MP4. - `--output`: path to save the generated video. +- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. See the Helios recipes above. - `--vae-use-slicing`: enable VAE slicing for memory optimization. - `--vae-use-tiling`: enable VAE tiling for memory optimization. - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel). @@ -119,9 +229,11 @@ python text_to_video.py \ - `--enable-cpu-offload`: enable CPU offloading for diffusion models. - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). -- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio). -- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). +- `--audio-sample-rate`: audio sample rate for embedded audio (when the + pipeline returns audio; LTX-2.3 outputs 48kHz audio). +- `--quantization`: quantization method (such as `fp8` for FP8). - `--flow-shift`: scheduler flow_shift parameter. +- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above). ### Wan2.2-specific diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 23b9afe8277..addf7cb7a29 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -11,12 +11,23 @@ import torch from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import get_extra_body_params, get_model_class_name from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform _MODEL_PRESETS = { + "vace": { + "height": 480, + "width": 832, + "num_frames": 81, + "num_inference_steps": 30, + "guidance_scale": 5.0, + "fps": 16, + "output": "vace_t2v_output.mp4", + }, "wan": { "height": 720, "width": 1280, @@ -35,13 +46,38 @@ "fps": 24, "output": "hunyuan_video_15_output.mp4", }, + "cosmos": { + "height": 720, + "width": 1280, + "num_frames": 189, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "fps": 24, + "flow_shift": 10.0, + "output": "cosmos3_t2v_output.mp4", + }, + "helios": { + "height": 384, + "width": 640, + "num_frames": 99, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "fps": 16, + "output": "helios_output.mp4", + }, } def _detect_preset(model: str) -> dict: model_lower = model.lower() + if "vace" in model_lower: + return _MODEL_PRESETS["vace"] + if "cosmos" in model_lower: + return _MODEL_PRESETS["cosmos"] if "hunyuan" in model_lower: return _MODEL_PRESETS["hunyuan"] + if "helios" in model_lower: + return _MODEL_PRESETS["helios"] return _MODEL_PRESETS["wan"] @@ -55,10 +91,26 @@ def parse_profiler_config(value: str) -> dict[str, Any]: return config +def parse_extra_body(value: str) -> dict[str, Any]: + """Parse a JSON object of model-specific extra_body params. + + Pipeline-declared knobs (see vllm_omni/model_extras/) are merged into + OmniDiffusionSamplingParams.extra_args, so a single generic example can + drive model-specific behaviour without bespoke per-model flags. + """ + try: + body = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e + if not isinstance(body, dict): + raise argparse.ArgumentTypeError("--extra-body must be a JSON object") + return body + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Generate a video from a text prompt. " - "Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models." + "Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models." ) parser.add_argument( "--model", @@ -72,8 +124,25 @@ def parse_args() -> argparse.Namespace: default=None, help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).", ) + parser.add_argument( + "--deploy-config", + default=None, + help="Optional deploy config YAML to use for pipeline-backed runs.", + ) parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.") parser.add_argument("--negative-prompt", default="", help="Negative prompt.") + parser.add_argument( + "--extra-body", + type=parse_extra_body, + default=None, + help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), " + "merged into sampling extra_args. Unknown keys for the chosen model are dropped. " + 'Cosmos3 example: \'{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, ' + '"use_resolution_template": false, "use_duration_template": false}\' ' + '(add "generate_sound": true and optional "sound_duration" for synchronized audio). ' + 'Helios-Distilled example: \'{"is_enable_stage2": true, ' + '"pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.', + ) parser.add_argument("--seed", type=int, default=42, help="Random seed.") parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.") parser.add_argument( @@ -163,7 +232,7 @@ def parse_args() -> argparse.Namespace: "--quantization", type=str, default=None, - choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8", "gguf"], + choices=["fp8", "mxfp8", "mxfp4", "mxfp4_dualscale", "int8"], help="Quantization method for the transformer. mxfp8: W8A8 MXFP8 (NPU). mxfp4: W4A4 MXFP4 (NPU). mxfp4_dualscale: W4A4 MXFP4 dual-scale + BF16 fallback mixed (NPU). fp8: online FP8 (GPU).", ) @@ -311,6 +380,8 @@ def main(): enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler, profiler_config=args.profiler_config, ) + if args.deploy_config: + omni_kwargs["deploy_config"] = args.deploy_config if args.boundary_ratio is not None: omni_kwargs["boundary_ratio"] = args.boundary_ratio if args.flow_shift is not None: @@ -322,7 +393,14 @@ def main(): omni_kwargs["cache_config"] = cache_config omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary + # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails + # gate is an engine-level config (offline analog of the server's --no-guardrails). + if args.extra_body and "guardrails" in args.extra_body: + omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])} + omni = Omni(**omni_kwargs) + model_class_name = get_model_class_name(omni) or model_class_name + declared_extra_body_params = get_extra_body_params(model_class_name) if profiler_enabled: print("[Profiler] Starting profiling...") @@ -358,10 +436,20 @@ def main(): if args.guidance_scale_high is not None: sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high + sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs) + # Route model-specific knobs through extra_body, filtered against the model's + # declared extra_body_params. Models without a declaration only forward explicit + # --extra-body JSON (preserving the generic flags' legacy behavior). + extra_body = dict(args.extra_body or {}) + if declared_extra_body_params: + apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body) + elif extra_body: + sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None}) + generation_start = time.perf_counter() frames = omni.generate( prompt_dict, - OmniDiffusionSamplingParams(**sampling_kwargs), + sampling_params, ) generation_end = time.perf_counter() generation_time = generation_end - generation_start @@ -374,6 +462,7 @@ def main(): print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") audio = None + audio_sample_rate = args.audio_sample_rate if isinstance(frames, list): frames = frames[0] if frames else None @@ -384,11 +473,13 @@ def main(): ) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] + audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) if frames.is_pipeline_output and frames.request_output is not None: inner_output = frames.request_output if isinstance(inner_output, OmniRequestOutput): if inner_output.multimodal_output and "audio" in inner_output.multimodal_output: audio = inner_output.multimodal_output["audio"] + audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate) frames = inner_output if isinstance(frames, OmniRequestOutput): if frames.images: @@ -396,6 +487,7 @@ def main(): frames, audio = frames.images[0] elif len(frames.images) == 1 and isinstance(frames.images[0], dict): audio = frames.images[0].get("audio") + audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate) frames = frames.images[0].get("frames") or frames.images[0].get("video") else: frames = frames.images @@ -408,6 +500,7 @@ def main(): frames, audio = first_item elif isinstance(first_item, dict): audio = first_item.get("audio") + audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate) frames = first_item.get("frames") or first_item.get("video") elif isinstance(first_item, list): frames = first_item @@ -416,6 +509,7 @@ def main(): frames, audio = frames elif isinstance(frames, dict): audio = frames.get("audio") + audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate) frames = frames.get("frames") or frames.get("video") if frames is None: @@ -530,7 +624,7 @@ def _ensure_frame_list(video_array): frames_u8, audio_np, fps=float(args.fps), - audio_sample_rate=args.audio_sample_rate, + audio_sample_rate=audio_sample_rate, ) with open(str(output_path), "wb") as f: f.write(video_bytes) diff --git a/examples/offline_inference/vace/vace_video_generation.md b/examples/offline_inference/vace/vace_video_generation.md deleted file mode 100644 index 8d8b00e48d9..00000000000 --- a/examples/offline_inference/vace/vace_video_generation.md +++ /dev/null @@ -1,148 +0,0 @@ -# VACE Video Generation - -Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline. - -- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options. - -## Table of Contents - -- [Overview](#overview) -- [Quick Start](#quick-start) -- [Key Arguments](#key-arguments) -- [More CLI Examples](#more-cli-examples) - -## Overview - -[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation. - -### Supported Models - -| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace | -|-------|-------------|-------------------|---------------------|-------------| -| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) | -| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) | - -!!! info -*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features. - -Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers` - -## Quick Start - -### Python API - -Text-to-video generation: - -```python -import torch - -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - -if __name__ == "__main__": - omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers") - prompt = "A sleek robot stands in a vast warehouse filled with boxes" - outputs = omni.generate( - prompt, - OmniDiffusionSamplingParams( - height=480, - width=832, - num_frames=81, - num_inference_steps=30, - guidance_scale=5.0, - ), - ) - video = outputs[0].images - - from diffusers.utils import export_to_video - export_to_video(list(video[0]), "t2v_output.mp4", fps=16) - omni.close() -``` - -### Local CLI Usage - -```bash -python vace_video_generation.py \ - --mode t2v \ - --prompt "A sleek robot stands in a vast warehouse filled with boxes" \ - --height 480 --width 832 --num-frames 81 \ - --output t2v_output.mp4 -``` - -## Key Arguments - -| Argument | Type | Default | Description | -| -------- | ---- | ------- | ----------- | -| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` | -| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path | -| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video | -| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance | -| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) | -| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) | -| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) | -| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) | -| `--num-frames` | int | `81` | Number of video frames to generate | -| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) | -| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale | -| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter | -| `--seed` | int | `42` | Random seed for deterministic sampling | -| `--fps` | int | `16` | Frames per second for the saved MP4 | -| `--output` | str | `"vace_output.mp4"` | Path to save the generated video | -| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization | -| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel | -| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference | -| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference | -| `--tensor-parallel-size` | int | — | Tensor parallel size | -| `--enforce-eager` | flag | off | Disable torch.compile | - -> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`). - -## More CLI Examples - -### Image-to-Video (I2V) - -First frame is kept, remaining frames are generated: - -```bash -python vace_video_generation.py \ - --mode i2v \ - --image astronaut.jpg \ - --prompt "An astronaut emerging from a cracked egg on the moon" \ - --height 480 --width 832 --num-frames 81 \ - --output i2v_output.mp4 -``` - -### First-Last-Frame Interpolation (FLF2V) - -```bash -python vace_video_generation.py \ - --mode flf2v \ - --image first_frame.jpg --last-image last_frame.jpg \ - --prompt "A bird takes off from a branch and lands on another" \ - --height 512 --width 512 --num-frames 81 \ - --output flf2v_output.mp4 -``` - -### Inpainting - -Center vertical stripe is masked and regenerated: - -```bash -python vace_video_generation.py \ - --mode inpaint \ - --image scene.jpg \ - --prompt "Shrek walks out of a building" \ - --height 480 --width 832 --num-frames 81 \ - --output inpaint_output.mp4 -``` - -### Reference Image-guided (R2V) - -```bash -python vace_video_generation.py \ - --mode r2v \ - --image reference.jpg \ - --prompt "Camera slowly zooms out from the character" \ - --height 480 --width 832 --num-frames 81 \ - --output r2v_output.mp4 -``` diff --git a/examples/offline_inference/vace/vace_video_generation.py b/examples/offline_inference/vace/vace_video_generation.py deleted file mode 100644 index 6ca0d74c52e..00000000000 --- a/examples/offline_inference/vace/vace_video_generation.py +++ /dev/null @@ -1,209 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""VACE video generation example. - -VACE (Video All-in-one Creation Engine) supports multiple video tasks: - - T2V: Text-to-Video - - I2V: Image-to-Video (first frame conditioning) - - V2LF: Video-to-Last-Frame - - FLF2V: First-Last-Frame interpolation - - Inpainting: Masked region generation - - R2V: Reference image-guided generation - -Usage examples: - # T2V (text-to-video) - python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse" - - # I2V (image-to-video, first frame kept) - python vace_video_generation.py --mode i2v --image input.jpg --prompt "..." - - # FLF2V (first-last frame interpolation) - python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg - - # R2V (reference image guided) - python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..." -""" - -import argparse -import time -from pathlib import Path - -import numpy as np -import PIL.Image -import torch - -from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams -from vllm_omni.platforms import current_omni_platform - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="VACE video generation.") - parser.add_argument( - "--model", - default="Wan-AI/Wan2.1-VACE-14B-diffusers", - help="VACE model ID or local path.", - ) - parser.add_argument( - "--mode", - default="t2v", - choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"], - help="Generation mode.", - ) - parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.") - parser.add_argument("--negative-prompt", default="", help="Negative prompt.") - parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).") - parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).") - parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).") - parser.add_argument("--seed", type=int, default=42, help="Random seed.") - parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.") - parser.add_argument("--height", type=int, default=480, help="Video height.") - parser.add_argument("--width", type=int, default=832, help="Video width.") - parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.") - parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.") - parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.") - parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.") - parser.add_argument("--fps", type=int, default=16, help="Output video FPS.") - parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.") - parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") - parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.") - parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.") - parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.") - return parser.parse_args() - - -def build_prompts(args): - """Build prompt dict with multi_modal_data based on mode.""" - h, w, nf = args.height, args.width, args.num_frames - - gray = PIL.Image.new("RGB", (w, h), (128, 128, 128)) - mask_black = PIL.Image.new("L", (w, h), 0) - mask_white = PIL.Image.new("L", (w, h), 255) - - prompt_data = { - "prompt": args.prompt, - "negative_prompt": args.negative_prompt, - } - - if args.mode == "t2v": - return prompt_data - - if args.mode == "r2v": - assert args.image, "--image required for R2V mode" - ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h)) - prompt_data["multi_modal_data"] = {"reference_images": [ref_img]} - return prompt_data - - if args.mode == "i2v": - assert args.image, "--image required for I2V mode" - img = PIL.Image.open(args.image).convert("RGB").resize((w, h)) - prompt_data["multi_modal_data"] = { - "video": [img] + [gray] * (nf - 1), - "mask": [mask_black] + [mask_white] * (nf - 1), - } - return prompt_data - - if args.mode == "v2lf": - assert args.image, "--image required for V2LF mode" - img = PIL.Image.open(args.image).convert("RGB").resize((w, h)) - prompt_data["multi_modal_data"] = { - "video": [gray] * (nf - 1) + [img], - "mask": [mask_white] * (nf - 1) + [mask_black], - } - return prompt_data - - if args.mode == "flf2v": - assert args.image and args.last_image, "--image and --last-image required for FLF2V" - first = PIL.Image.open(args.image).convert("RGB").resize((w, h)) - last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h)) - prompt_data["multi_modal_data"] = { - "video": [first] + [gray] * (nf - 2) + [last], - "mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black], - } - return prompt_data - - if args.mode == "inpaint": - assert args.image, "--image required for inpaint mode" - img = PIL.Image.open(args.image).convert("RGB").resize((w, h)) - d = 80 - frames, masks = [], [] - for _ in range(nf): - base = np.array(img).copy() - mask = PIL.Image.new("L", (w, h), 0) - stripe = PIL.Image.new("L", (2 * d, h), 255) - mask.paste(stripe, (w // 2 - d, 0)) - base[np.array(mask) > 128] = 128 - frames.append(PIL.Image.fromarray(base)) - masks.append(mask) - prompt_data["multi_modal_data"] = {"video": frames, "mask": masks} - return prompt_data - - raise ValueError(f"Unknown mode: {args.mode}") - - -def main(): - args = parse_args() - generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) - - parallel_config = DiffusionParallelConfig( - ulysses_degree=args.ulysses_degree, - ring_degree=args.ring_degree, - cfg_parallel_size=args.cfg_parallel_size, - ) - - omni = Omni( - model=args.model, - vae_use_tiling=args.vae_use_tiling, - flow_shift=args.flow_shift, - enforce_eager=args.enforce_eager, - parallel_config=parallel_config, - ) - - prompt_data = build_prompts(args) - - print(f"\n{'=' * 60}") - print(f"VACE {args.mode.upper()} Generation") - print(f" Model: {args.model}") - print(f" Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps") - print(f"{'=' * 60}\n") - - start = time.perf_counter() - outputs = omni.generate( - prompt_data, - OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_frames=args.num_frames, - num_inference_steps=args.num_inference_steps, - guidance_scale=args.guidance_scale, - generator=generator, - ), - ) - elapsed = time.perf_counter() - start - - video = outputs[0].images - if isinstance(video, list): - video = video[0] - if isinstance(video, torch.Tensor): - video = video.cpu().numpy() - if video.ndim == 5: - video = video[0] - print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s") - - output_path = Path(args.output) - output_path.parent.mkdir(parents=True, exist_ok=True) - - from diffusers.utils import export_to_video - - if np.issubdtype(video.dtype, np.integer): - video = video.astype(np.float32) / 255.0 - export_to_video(list(video), str(output_path), fps=args.fps) - print(f"Saved to {output_path}") - - omni.close() - - -if __name__ == "__main__": - main() diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index f5071ce6db4..8305c785f77 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md @@ -87,6 +87,63 @@ python x_to_video_audio.py \ ``` +## MagiHuman (`--model-type magi-human`) + +MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param +T5Gemma text encoder. A detailed text prompt is the only required input; an optional +image and/or audio file may be supplied for conditioning. Natively supports Tensor +Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard +the MoE weights and the text encoder. + +> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct +> attention-kernel behaviour (the pipeline otherwise falls back to stubs). + +Text-only generation: + +```bash +python x_to_video_audio.py \ + --model-type magi-human \ + --model /path/to/daVinci-MagiHuman \ + --prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \ + --tensor-parallel-size 4 \ + --height 256 --width 448 \ + --num-inference-steps 8 \ + --seed 52 \ + --extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \ + --output output_magihuman.mp4 +``` + +With optional image and audio conditioning: + +```bash +python x_to_video_audio.py \ + --model-type magi-human \ + --model /path/to/daVinci-MagiHuman \ + --prompt "A young woman..." \ + --tensor-parallel-size 4 \ + --height 256 --width 448 \ + --num-inference-steps 8 \ + --seed 52 \ + --extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \ + --output output_magihuman.mp4 +``` + +MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared +in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`): + +- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set, + because the audio length then determines the number of frames). +- `audio_path`: path to an audio file for audio-to-video conditioning. When provided, + the audio drives the frame count and is encoded as a latent condition into the DiT + (`is_a2v` mode). Omit for pure text-to-video+audio generation. +- `image_path`: path to an image file for visual conditioning. Applied at both the + base-resolution (BR) and super-resolution (SR) stages when SR is enabled. +- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped + when these are omitted. +- `sr_num_inference_steps`: denoising steps for the SR stage. + +## DreamID-Omni arguments + Key arguments: - `--prompt`: text description (string). - `--model`: path to the model local directory. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 30de8aa456f..0efbdc37d9a 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import argparse +import json import re import time from typing import Any @@ -20,14 +21,26 @@ def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).") - parser.add_argument("--model", required=True, help="DreamID ckpt root directory.") - parser.add_argument("--model-type", default="dreamid-omni", help="Model type.") + parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.") + parser.add_argument("--model", required=True, help="Model ckpt root directory.") + parser.add_argument( + "--model-type", + default="dreamid-omni", + choices=["dreamid-omni", "magi-human"], + help="Model type.", + ) parser.add_argument("--prompt", default=None, help="Text prompt.") parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path") parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path") parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.") + parser.add_argument( + "--extra-body", + type=str, + default=None, + help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), " + 'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.', + ) parser.add_argument("--height", type=int, default=704, help="Video height.") parser.add_argument("--width", type=int, default=1280, help="Video width.") @@ -142,8 +155,6 @@ def main() -> None: text_prompt = args.prompt if args.prompt_file: - import json - with open(args.prompt_file) as f: text_prompt = json.load(f) text_prompt = re.sub( @@ -155,25 +166,36 @@ def main() -> None: text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt) text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip() - image, audio = load_image_and_audio(args.image_path, args.audio_path) + if args.model_type == "magi-human": + prompt = text_prompt + extra_args = json.loads(args.extra_body) if args.extra_body else {} + sampling_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_inference_steps=args.num_inference_steps, + seed=args.seed, + extra_args=extra_args, + ) + else: + image, audio = load_image_and_audio(args.image_path, args.audio_path) - prompt = { - "prompt": text_prompt, - "video_negative_prompt": args.video_negative_prompt, - "audio_negative_prompt": args.audio_negative_prompt, - "multi_modal_data": {"image": image, "audio": audio}, - } + prompt = { + "prompt": text_prompt, + "video_negative_prompt": args.video_negative_prompt, + "audio_negative_prompt": args.audio_negative_prompt, + "multi_modal_data": {"image": image, "audio": audio}, + } - sampling_params = OmniDiffusionSamplingParams( - height=args.height, - width=args.width, - num_inference_steps=args.num_inference_steps, - seed=args.seed, - extra_args={ - "solver_name": args.solver_name, - "shift": args.shift, - }, - ) + sampling_params = OmniDiffusionSamplingParams( + height=args.height, + width=args.width, + num_inference_steps=args.num_inference_steps, + seed=args.seed, + extra_args={ + "solver_name": args.solver_name, + "shift": args.shift, + }, + ) parallel_config = DiffusionParallelConfig( cfg_parallel_size=args.cfg_parallel_size, @@ -215,22 +237,28 @@ def main() -> None: elapsed = time.perf_counter() - start if not outputs: - raise RuntimeError("No output returned from DreamID-Omni.") + raise RuntimeError("No output returned from the model.") result = outputs[0] if not result.images: - raise RuntimeError("No video frames found in DreamID-Omni output.") + raise RuntimeError("No video frames found in model output.") generated_video = result.images[0] mm = result.multimodal_output or {} generated_audio = mm.get("audio") - fps = int(mm.get("fps", 24)) - sample_rate = int(mm.get("audio_sample_rate", 16000)) - # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1]. - # mux_video_audio_bytes expects (F, H, W, C) uint8. - if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4: - raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}") - frames = generated_video.transpose(1, 2, 3, 0) - frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8) + if args.model_type == "magi-human": + # MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux. + fps = float(mm.get("fps", 25)) + sample_rate = int(mm.get("audio_sample_rate", 24000)) + frames = generated_video + else: + fps = float(mm.get("fps", 24)) + sample_rate = int(mm.get("audio_sample_rate", 16000)) + # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1]. + # mux_video_audio_bytes expects (F, H, W, C) uint8. + if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4: + raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}") + frames = generated_video.transpose(1, 2, 3, 0) + frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8) audio_np = None if generated_audio is not None: diff --git a/examples/online_serving/audiox/README.md b/examples/online_serving/audiox/README.md index aad65eb301e..bf068b3eb4f 100644 --- a/examples/online_serving/audiox/README.md +++ b/examples/online_serving/audiox/README.md @@ -37,11 +37,10 @@ python openai_chat_client.py --task tv2a \ The client sends: -- `num_inference_steps`, `guidance_scale`, `seed` as first-class OpenAI chat-completion fields -- `audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max` nested under - `extra_args` (a reserved dict on the request body that the server forwards verbatim into - the pipeline's `sampling_params.extra_args` — the same escape hatch `serving_video.py` exposes - as `extra_params` on /v1/videos) +- `num_inference_steps`, `guidance_scale`, `seed`, `audiox_task`, `seconds_start`, + `seconds_total`, `sigma_min`, `sigma_max` as flattened keys under `extra_body`. The server + routes the declared knobs (see `vllm_omni/model_extras/audiox.py`) into the pipeline's + `sampling_params.extra_args`. - For `v2*` / `tv2*` tasks, the video as a `video_url` content item (data URI for local files) ## curl @@ -52,10 +51,10 @@ curl -sS -X POST http://localhost:8099/v1/chat/completions \ -d '{ "model": "zhangj1an/AudioX", "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}], - "num_inference_steps": 250, - "guidance_scale": 7.0, - "seed": 42, - "extra_args": { + "extra_body": { + "num_inference_steps": 250, + "guidance_scale": 7.0, + "seed": 42, "audiox_task": "t2m", "seconds_total": 10.0, "sigma_min": 0.3, diff --git a/examples/online_serving/audiox/openai_chat_client.py b/examples/online_serving/audiox/openai_chat_client.py index 8c8b572e50d..fd9aa8cc550 100755 --- a/examples/online_serving/audiox/openai_chat_client.py +++ b/examples/online_serving/audiox/openai_chat_client.py @@ -4,8 +4,8 @@ AudioX supports 6 tasks (t2a, t2m, v2a, v2m, tv2a, tv2m). Text-only tasks send the prompt as the chat message; video-conditioned tasks additionally attach the video via a ``video_url`` content item (data URI for local files). Task + generation knobs (steps, cfg, sigma range, seconds, seed) -are sent via the OpenAI SDK's ``extra_body`` as ``extra_args`` — the same pipeline-agnostic escape -hatch used by the /v1/videos endpoint's ``extra_params`` field. +are sent as flattened keys under ``extra_body``; the server routes the declared knobs into +``sampling_params.extra_args`` (see ``vllm_omni/model_extras/audiox.py``). Usage: python openai_chat_client.py --task t2a --prompt "Fireworks burst twice..." --output t2a.wav @@ -88,10 +88,10 @@ def main() -> int: payload = { "model": args.model, "messages": [{"role": "user", "content": content}], - "num_inference_steps": args.steps, - "guidance_scale": args.guidance_scale, - "seed": args.seed, - "extra_args": { + "extra_body": { + "num_inference_steps": args.steps, + "guidance_scale": args.guidance_scale, + "seed": args.seed, "audiox_task": args.task, "seconds_start": args.seconds_start, "seconds_total": args.seconds_total, diff --git a/examples/online_serving/aura_omni/README.md b/examples/online_serving/aura_omni/README.md new file mode 100644 index 00000000000..5f246c33677 --- /dev/null +++ b/examples/online_serving/aura_omni/README.md @@ -0,0 +1,146 @@ +# AURA Omni: Online serving + +`aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline: + +```text +Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav +``` + +The pipeline has three semantic modules, but four engine stages because the +existing Qwen3-TTS implementation is natively split into Talker and Code2Wav. + +Start the server with the deploy profile: + +```bash +vllm serve aurateam/AURA \ + --omni \ + --port 8091 \ + --deploy-config vllm_omni/deploy/aura_omni.yaml \ + --served-model-name aurateam/AURA \ + --trust-remote-code +``` + +The deploy file sets per-stage model repos: + +- Stage 0 ASR: `Qwen/Qwen3-ASR-1.7B` +- Stage 1 AURA: `aurateam/AURA` +- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base` + +For local weights, edit the `model` value on each stage in +`vllm_omni/deploy/aura_omni.yaml`. The deploy profile includes +`pipeline: aura_omni`, so the server uses this four-stage topology even when +the command-line model path points at one component checkpoint. + +Expected request shape: + +- Send microphone audio as the Stage 0 multimodal audio input. +- Include video frames in the original request `multi_modal_data`; the + `asr2aura` processor carries them forward to AURA. +- Optional `additional_information` keys: + - `aura_system_prompt` + - `tts_task_type` + - `tts_language` + - `tts_speaker` + - `tts_instruct` + - `tts_ref_audio` + - `tts_ref_text` + - `tts_x_vector_only_mode` + - `tts_pass_token_ids` + +If AURA emits `<|silent|>`, the `aura2tts` processor returns no TTS request, so +the TTS stages are skipped for that turn. + +## GPU Utilization Recommendation + +Tune `gpu_memory_utilization` per stage in `vllm_omni/deploy/aura_omni.yaml`. +Recommended baseline on one GPU for H200 + +- Stage 0 (ASR): `0.10` +- Stage 1 (AURA): `0.4` +- Stage 2 (Qwen3-TTS Talker): `0.20` +- Stage 3 (Qwen3-TTS Code2Wav): `0.20` + +## Python Client + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --host localhost \ + --port 8091 \ + --model aurateam/AURA \ + --modalities text,audio +``` + +Use local media: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --audio-path /path/to/input.wav \ + --video-path /path/to/video.mp4 \ + --output-dir output_aura_omni_online +``` + +Base voice clone mode (default, recommended as x-vector while debugging ICL): + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-task-type Base \ + --tts-ref-audio vllm-omni/tests/assets/qwen3_tts/clone_2.wav \ + --tts-ref-text "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +``` + +Enable AURA token-id passthrough explicitly: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-pass-token-ids +``` + +CustomVoice mode requires stages 2 and 3 in `aura_omni.yaml` to point at a +Qwen3-TTS CustomVoice checkpoint: + +```bash +python examples/online_serving/aura_omni/openai_chat_completion_client.py \ + --tts-task-type CustomVoice \ + --tts-speaker Vivian +``` + +By default, AURA responses are passed to Qwen3-TTS as text. Set +`tts_pass_token_ids=true` to pass AURA-generated assistant token ids directly +to Qwen3-TTS instead. The processor still uses AURA token ids, when available, +to estimate the Talker prompt length in the default text path. + +## Curl + +```bash +cd examples/online_serving/aura_omni +bash run_curl_multimodal_generation.sh +``` + +Set `PORT`, `MODEL`, or `OUTPUT_DIR` to override defaults: + +```bash +PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh +TTS_PASS_TOKEN_IDS=true PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh +``` + +## Gradio + +Launch the server and Gradio UI together: + +```bash +cd examples/online_serving/aura_omni +bash run_gradio_demo.sh +``` + +If the server is already running: + +```bash +python examples/online_serving/aura_omni/gradio_demo.py \ + --model aurateam/AURA \ + --api-base http://localhost:8091/v1 +``` + +## Offline + +For offline inference, see +[`examples/offline_inference/aura_omni`](../../offline_inference/aura_omni/). diff --git a/examples/online_serving/aura_omni/gradio_demo.py b/examples/online_serving/aura_omni/gradio_demo.py new file mode 100644 index 00000000000..c6badee7826 --- /dev/null +++ b/examples/online_serving/aura_omni/gradio_demo.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gradio demo for AURA Omni online serving.""" + +from __future__ import annotations + +import argparse +import base64 +import io +from pathlib import Path +from typing import Any + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None + +import numpy as np +import soundfile as sf +from openai import OpenAI + +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + DEFAULT_QWEN3_TTS_REF_TEXT, + default_qwen3_tts_ref_audio_path, +) + +SEED = 42 +DEFAULT_TTS_REF_AUDIO = default_qwen3_tts_ref_audio_path() + + +def parse_args(): + parser = argparse.ArgumentParser(description="Gradio demo for AURA Omni online inference.") + parser.add_argument("--model", default="aurateam/AURA", help="Model name served by vLLM.") + parser.add_argument("--api-base", default="http://localhost:8091/v1", help="OpenAI-compatible API base.") + parser.add_argument("--ip", default="127.0.0.1", help="Gradio host.") + parser.add_argument("--port", type=int, default=7862, help="Gradio port.") + parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.") + return parser.parse_args() + + +def _audio_to_data_url(audio_file: Any | None) -> str | None: + if audio_file is None: + return None + + def _path_to_data_url(path_str: str) -> str | None: + if not path_str: + return None + path = Path(path_str) + if not path.exists(): + return None + with open(path, "rb") as f: + return f"data:audio/wav;base64,{base64.b64encode(f.read()).decode('utf-8')}" + + if isinstance(audio_file, str): + return _path_to_data_url(audio_file) + + # Gradio >=5 may return dict payloads with local temp path. + if isinstance(audio_file, dict): + path_val = audio_file.get("path") or audio_file.get("name") + if isinstance(path_val, str): + return _path_to_data_url(path_val) + return None + + if isinstance(audio_file, tuple): + # Common shape: (sample_rate, np.ndarray) + if len(audio_file) == 2 and isinstance(audio_file[0], (int, float)): + sample_rate, audio_np = audio_file + audio_np = np.asarray(audio_np) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + if audio_np.dtype != np.int16: + audio_np = np.clip(audio_np.astype(np.float32), -1.0, 1.0) + buf = io.BytesIO() + sf.write(buf, audio_np, int(sample_rate), format="WAV") + return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}" + # Compatibility shape from some gradio builds: (filepath, (sr, np.ndarray)) + if len(audio_file) >= 1 and isinstance(audio_file[0], str): + path_url = _path_to_data_url(audio_file[0]) + if path_url is not None: + return path_url + return None + + +def _video_to_data_url(video_file: str | None) -> str | None: + if not video_file: + return None + path = Path(video_file) + mime = { + ".mp4": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".avi": "video/x-msvideo", + ".mkv": "video/x-matroska", + }.get(path.suffix.lower(), "video/mp4") + with open(path, "rb") as f: + return f"data:{mime};base64,{base64.b64encode(f.read()).decode('utf-8')}" + + +def _sampling_params_list() -> list[dict[str, Any]]: + return [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED}, + { + "temperature": 0.5, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 256, + "seed": SEED, + "repetition_penalty": 1.0, + }, + { + "temperature": 0.9, + "top_k": 50, + "max_tokens": 4096, + "seed": SEED, + "detokenize": False, + "repetition_penalty": 1.05, + "stop_token_ids": [2150], + }, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 65536, + "seed": SEED, + "repetition_penalty": 1.0, + }, + ] + + +def _decode_audio_choice(audio_data: str | None) -> tuple[int, np.ndarray] | None: + if not audio_data: + return None + wav_bytes = base64.b64decode(audio_data) + audio_np, sample_rate = sf.read(io.BytesIO(wav_bytes)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return int(sample_rate), audio_np.astype(np.float32) + + +def build_interface(client: OpenAI, model: str): + def run( + audio_file, + video_file, + prompt: str, + aura_system_prompt: str, + tts_task_type: str, + tts_language: str, + tts_speaker: str, + tts_instruct: str, + tts_ref_audio: str, + tts_ref_text: str, + tts_x_vector_only_mode: bool, + tts_pass_token_ids: bool, + ): + audio_url = _audio_to_data_url(audio_file) + video_url = _video_to_data_url(video_file) + if not audio_url: + return "Please provide audio input for the ASR stage.", None + if tts_task_type == "Base" and (not (tts_ref_audio or "").strip() or not (tts_ref_text or "").strip()): + return "Base TTS requires both reference audio and reference transcript.", None + content: list[dict[str, Any]] = [{"type": "audio_url", "audio_url": {"url": audio_url}}] + if video_url: + content.append({"type": "video_url", "video_url": {"url": video_url}}) + content.append( + {"type": "text", "text": prompt or "Use the audio and video together to decide whether a reply is needed."} + ) + + additional_information: dict[str, Any] = { + "aura_system_prompt": aura_system_prompt, + "tts_task_type": tts_task_type, + "tts_instruct": tts_instruct, + "tts_pass_token_ids": bool(tts_pass_token_ids), + } + if tts_task_type == "CustomVoice": + additional_information.update( + { + "tts_language": tts_language or "English", + "tts_speaker": tts_speaker or "Vivian", + } + ) + else: + additional_information.update( + { + "tts_ref_audio": tts_ref_audio.strip() if tts_ref_audio else None, + "tts_ref_text": tts_ref_text.strip() if tts_ref_text else None, + "tts_x_vector_only_mode": bool(tts_x_vector_only_mode), + } + ) + + response = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": content}], + modalities=["text", "audio"], + extra_body={ + "sampling_params_list": _sampling_params_list(), + "additional_information": additional_information, + }, + timeout=600.0, + ) + + text_parts: list[str] = [] + audio_output = None + for choice in response.choices: + message = choice.message + if message.content: + text_parts.append(str(message.content)) + if getattr(message, "audio", None): + audio_output = _decode_audio_choice(message.audio.data) + return "\n\n".join(text_parts) or "No text output.", audio_output + + default_system = ( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English." + ) + + def _toggle_tts_advanced(task_type: str): + is_base = task_type == "Base" + return ( + gr.update(visible=not is_base), + gr.update(visible=is_base), + ) + + with gr.Blocks() as demo: + gr.Markdown("# AURA Omni") + gr.Markdown("ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav") + with gr.Row(): + audio_input = gr.Audio(label="Audio input", sources=["upload", "microphone"], type="numpy") + video_input = gr.Video(label="Video input (optional)", sources=["upload"]) + prompt = gr.Textbox( + label="Instruction", + value="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.", + lines=2, + ) + with gr.Accordion("Advanced", open=False): + aura_system_prompt = gr.Textbox(label="AURA system prompt", value=default_system, lines=4) + tts_task_type = gr.Radio(choices=["Base", "CustomVoice"], value="Base", label="TTS task type") + tts_instruct = gr.Textbox(label="TTS instruction", value="") + tts_pass_token_ids = gr.Checkbox( + label="Pass AURA token ids directly to TTS", + value=False, + ) + + with gr.Group(visible=False) as customvoice_config_group: + tts_language = gr.Dropdown( + choices=["Chinese", "English", "Japanese", "Korean", "Cantonese"], + value="Chinese", + label="CustomVoice language", + ) + tts_speaker = gr.Textbox(label="CustomVoice speaker", value="Vivian") + + with gr.Group(visible=True) as base_config_group: + tts_ref_audio = gr.Textbox( + label="Base reference audio path/URL", + value=DEFAULT_TTS_REF_AUDIO, + ) + tts_ref_text = gr.Textbox( + label="Base reference transcript", + value=DEFAULT_QWEN3_TTS_REF_TEXT, + lines=3, + ) + tts_x_vector_only_mode = gr.Checkbox( + label="Base x-vector only mode (disable ICL)", + value=False, + ) + + tts_task_type.change( + _toggle_tts_advanced, + inputs=[tts_task_type], + outputs=[customvoice_config_group, base_config_group], + ) + button = gr.Button("Generate", variant="primary") + with gr.Row(): + text_output = gr.Textbox(label="Text output", lines=8) + audio_output = gr.Audio(label="Audio output", interactive=False) + button.click( + run, + inputs=[ + audio_input, + video_input, + prompt, + aura_system_prompt, + tts_task_type, + tts_language, + tts_speaker, + tts_instruct, + tts_ref_audio, + tts_ref_text, + tts_x_vector_only_mode, + tts_pass_token_ids, + ], + outputs=[text_output, audio_output], + ) + demo.queue() + return demo + + +def main() -> None: + args = parse_args() + client = OpenAI(api_key="EMPTY", base_url=args.api_base) + demo = build_interface(client, args.model) + demo.launch(server_name=args.ip, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/aura_omni/openai_chat_completion_client.py b/examples/online_serving/aura_omni/openai_chat_completion_client.py new file mode 100644 index 00000000000..c8225bf0d1d --- /dev/null +++ b/examples/online_serving/aura_omni/openai_chat_completion_client.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OpenAI-compatible client for the AURA Omni pipeline.""" + +from __future__ import annotations + +import base64 +import io +import os + +import soundfile as sf +from openai import OpenAI +from vllm.assets.audio import AudioAsset +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + DEFAULT_QWEN3_TTS_REF_TEXT, + default_qwen3_tts_ref_audio_path, +) + +SEED = 42 +DEFAULT_MODEL = "aurateam/AURA" +DEFAULT_VIDEO_URL = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + + +def _encode_file(path: str) -> str: + with open(path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + + +def _data_url(path: str, default_mime: str) -> str: + suffix = os.path.splitext(path)[1].lower() + mime_by_suffix = { + ".wav": "audio/wav", + ".mp3": "audio/mpeg", + ".ogg": "audio/ogg", + ".flac": "audio/flac", + ".m4a": "audio/mp4", + ".mp4": "video/mp4", + ".webm": "video/webm", + ".mov": "video/quicktime", + ".avi": "video/x-msvideo", + ".mkv": "video/x-matroska", + } + return f"data:{mime_by_suffix.get(suffix, default_mime)};base64,{_encode_file(path)}" + + +def media_url(path_or_url: str | None, *, kind: str) -> str: + if path_or_url: + if path_or_url.startswith(("http://", "https://", "data:")): + return path_or_url + if not os.path.exists(path_or_url): + raise FileNotFoundError(f"{kind} file not found: {path_or_url}") + return _data_url(path_or_url, "audio/wav" if kind == "audio" else "video/mp4") + if kind == "audio": + return AudioAsset("mary_had_lamb").url + return DEFAULT_VIDEO_URL + + +def sampling_params_list() -> list[dict]: + return [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED}, + { + "temperature": 0.5, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 256, + "seed": SEED, + "repetition_penalty": 1.0, + }, + { + "temperature": 0.9, + "top_k": 50, + "max_tokens": 4096, + "seed": SEED, + "detokenize": False, + "repetition_penalty": 1.05, + "stop_token_ids": [2150], + }, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 65536, + "seed": SEED, + "repetition_penalty": 1.0, + }, + ] + + +def parse_modalities(value: str | None) -> list[str] | None: + if not value: + return None + return [item.strip() for item in value.split(",") if item.strip()] + + +def save_response(response, output_dir: str) -> None: + os.makedirs(output_dir, exist_ok=True) + for idx, choice in enumerate(response.choices): + message = choice.message + if message.content: + out_txt = os.path.join(output_dir, f"choice_{idx}.txt") + with open(out_txt, "w", encoding="utf-8") as f: + f.write(str(message.content).strip() + "\n") + print(f"Text saved to {out_txt}") + print(message.content) + if getattr(message, "audio", None): + audio_bytes = base64.b64decode(message.audio.data) + audio_np, sample_rate = sf.read(io.BytesIO(audio_bytes)) + out_wav = os.path.join(output_dir, f"choice_{idx}.wav") + sf.write(out_wav, audio_np, int(sample_rate), format="WAV") + print(f"Audio saved to {out_wav}") + + +def main(args) -> None: + client = OpenAI(base_url=f"http://{args.host}:{args.port}/v1", api_key="EMPTY") + messages = [ + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": media_url(args.audio_path, kind="audio")}}, + {"type": "video_url", "video_url": {"url": media_url(args.video_path, kind="video")}}, + {"type": "text", "text": args.prompt}, + ], + } + ] + response = client.chat.completions.create( + model=args.model, + messages=messages, + modalities=parse_modalities(args.modalities), + extra_body={ + "sampling_params_list": sampling_params_list(), + "additional_information": { + "aura_system_prompt": args.aura_system_prompt, + "tts_task_type": args.tts_task_type, + "tts_language": args.tts_language, + "tts_speaker": args.tts_speaker, + "tts_instruct": args.tts_instruct, + "tts_ref_audio": args.tts_ref_audio, + "tts_ref_text": args.tts_ref_text, + "tts_x_vector_only_mode": args.tts_x_vector_only_mode, + "tts_pass_token_ids": args.tts_pass_token_ids, + }, + }, + timeout=args.timeout, + ) + save_response(response, args.output_dir) + + +def parse_args(): + parser = FlexibleArgumentParser(description="AURA Omni online serving client") + parser.add_argument("--host", default="localhost") + parser.add_argument("--port", type=int, default=8091) + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--audio-path", default=None, help="Audio file, URL, or data URL.") + parser.add_argument("--video-path", default=None, help="Video file, URL, or data URL.") + parser.add_argument( + "--prompt", + default="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.", + ) + parser.add_argument("--modalities", default="text,audio") + parser.add_argument("--output-dir", default="output_aura_omni_online") + parser.add_argument( + "--aura-system-prompt", + default=( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English." + ), + ) + parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"]) + parser.add_argument("--tts-language", default="English") + parser.add_argument("--tts-speaker", default="Vivian") + parser.add_argument("--tts-instruct", default="") + parser.add_argument( + "--tts-ref-audio", + default=default_qwen3_tts_ref_audio_path(), + help="Base-mode reference audio path/URL visible to server.", + ) + parser.add_argument( + "--tts-ref-text", + default=DEFAULT_QWEN3_TTS_REF_TEXT, + help="Base-mode reference audio transcript.", + ) + parser.add_argument( + "--tts-x-vector-only-mode", + action="store_true", + help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).", + ) + parser.add_argument( + "--tts-pass-token-ids", + action="store_true", + help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.", + ) + parser.add_argument("--timeout", type=float, default=600.0) + return parser.parse_args() + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/online_serving/aura_omni/run_curl_multimodal_generation.sh b/examples/online_serving/aura_omni/run_curl_multimodal_generation.sh new file mode 100755 index 00000000000..7510781093f --- /dev/null +++ b/examples/online_serving/aura_omni/run_curl_multimodal_generation.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +set -euo pipefail + +PORT="${PORT:-8091}" +MODEL="${MODEL:-aurateam/AURA}" +OUTPUT_DIR="${OUTPUT_DIR:-output_aura_omni_online}" +TTS_PASS_TOKEN_IDS="${TTS_PASS_TOKEN_IDS:-false}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +VLLM_OMNI_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" +CLONE_REF_AUDIO="${VLLM_OMNI_ROOT}/tests/assets/qwen3_tts/clone_2.wav" +CLONE_REF_TEXT="Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +mkdir -p "$OUTPUT_DIR" + +MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" +SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + +request_body=$(cat <<EOF +{ + "model": "$MODEL", + "modalities": ["text", "audio"], + "sampling_params_list": [ + {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42}, + {"temperature": 0.5, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42, "repetition_penalty": 1.0}, + {"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "seed": 42, "detokenize": false, "repetition_penalty": 1.05, "stop_token_ids": [2150]}, + {"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 65536, "seed": 42, "repetition_penalty": 1.0} + ], + "additional_information": { + "aura_system_prompt": "You are receiving a live video stream where the final frame is the present moment. Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English.", + "tts_task_type": "Base", + "tts_ref_audio": "file://${CLONE_REF_AUDIO}", + "tts_ref_text": "${CLONE_REF_TEXT}", + "tts_language": "English", + "tts_speaker": "Vivian", + "tts_instruct": "", + "tts_pass_token_ids": ${TTS_PASS_TOKEN_IDS} + }, + "messages": [{ + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "$MARY_HAD_LAMB_AUDIO_URL"}}, + {"type": "video_url", "video_url": {"url": "$SAMPLE_VIDEO_URL"}}, + {"type": "text", "text": "Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English."} + ] + }] +} +EOF +) + +response=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \ + -X POST "http://localhost:${PORT}/v1/chat/completions" \ + -H "Content-Type: application/json" \ + -d "$request_body") + +echo "$response" | jq '.choices[].message.content' + +audio_b64=$(echo "$response" | jq -r '.choices[]?.message.audio.data // empty' | head -n 1) +if [[ -n "$audio_b64" ]]; then + echo "$audio_b64" | base64 -d > "${OUTPUT_DIR}/aura_omni_output.wav" + echo "Audio saved to ${OUTPUT_DIR}/aura_omni_output.wav" +fi diff --git a/examples/online_serving/aura_omni/run_gradio_demo.sh b/examples/online_serving/aura_omni/run_gradio_demo.sh new file mode 100755 index 00000000000..bf591c84321 --- /dev/null +++ b/examples/online_serving/aura_omni/run_gradio_demo.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL="aurateam/AURA" +SERVER_MODEL="aurateam/AURA" +DEPLOY_CONFIG="/data/yrr/vllm-omni/vllm_omni/deploy/aura_omni.yaml" +SERVER_PORT=8091 +GRADIO_PORT=7862 +SERVER_HOST="0.0.0.0" +GRADIO_IP="127.0.0.1" +GRADIO_SHARE=false + + +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL="$2"; shift 2 ;; + --server-model) SERVER_MODEL="$2"; shift 2 ;; + --deploy-config) DEPLOY_CONFIG="$2"; shift 2 ;; + --server-port) SERVER_PORT="$2"; shift 2 ;; + --gradio-port) GRADIO_PORT="$2"; shift 2 ;; + --server-host) SERVER_HOST="$2"; shift 2 ;; + --gradio-ip) GRADIO_IP="$2"; shift 2 ;; + --share) GRADIO_SHARE=true; shift ;; + --help) + echo "Usage: $0 [--model SERVED_MODEL_NAME] [--server-model MODEL_PATH] [--deploy-config YAML] [--server-port PORT] [--gradio-port PORT] [--share]" + exit 0 + ;; + *) echo "Unknown option: $1"; exit 1 ;; + esac +done + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +API_BASE="http://localhost:${SERVER_PORT}/v1" +LOG_FILE="/tmp/aura_omni_vllm_${SERVER_PORT}.log" + +cleanup() { + echo "Shutting down..." + [[ -n "${SERVER_PID:-}" ]] && kill "$SERVER_PID" 2>/dev/null || true + [[ -n "${GRADIO_PID:-}" ]] && kill "$GRADIO_PID" 2>/dev/null || true +} +trap cleanup SIGINT SIGTERM EXIT + +vllm serve "$SERVER_MODEL" \ + --omni \ + --host "$SERVER_HOST" \ + --port "$SERVER_PORT" \ + --deploy-config "$DEPLOY_CONFIG" \ + --served-model-name "$MODEL" \ + --trust-remote-code 2>&1 | tee "$LOG_FILE" & +SERVER_PID=$! + +echo "Waiting for server startup..." +for _ in $(seq 1 600); do + if grep -q "Application startup complete" "$LOG_FILE" 2>/dev/null; then + break + fi + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "vLLM server exited before startup completed" + wait "$SERVER_PID" || true + exit 1 + fi + sleep 1 +done + +# cd "$SCRIPT_DIR" +GRADIO_CMD=(python gradio_demo.py --model "$MODEL" --api-base "$API_BASE" --ip "$GRADIO_IP" --port "$GRADIO_PORT") +if [[ "$GRADIO_SHARE" == "true" ]]; then + GRADIO_CMD+=(--share) +fi +"${GRADIO_CMD[@]}" & +GRADIO_PID=$! + +echo "vLLM server: http://${SERVER_HOST}:${SERVER_PORT}" +echo "Gradio demo: http://${GRADIO_IP}:${GRADIO_PORT}" +if [[ -n "${SERVER_PID:-}" ]]; then + wait "$SERVER_PID" "$GRADIO_PID" +else + wait "$GRADIO_PID" +fi diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md deleted file mode 100644 index b7a54da356c..00000000000 --- a/examples/online_serving/bagel/README.md +++ /dev/null @@ -1,404 +0,0 @@ -# BAGEL-7B-MoT - -## Installation - -Please refer to [README.md](../../../README.md) - -## Architecture - -BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies: - -| Topology | Stages | Description | -| :------- | :----- | :---------- | -| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. | -| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. | - -Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`. - -> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices. - -## Launch the Server - -### Two-Stage (Default) - -The default pipeline is auto-detected from the model. No extra flags needed: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 -``` - -Or use the convenience script: - -```bash -cd examples/online_serving/bagel -bash run_server.sh - -# Initialize each stage in a discrete isolated process terminal -bash run_server_stage_cli.sh --stage 0 -bash run_server_stage_cli.sh --stage 1 -``` - -To use a custom deploy YAML, pass it via `--deploy-config`: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config /path/to/deploy_config.yaml -``` - -See [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration. - -### Single-Stage - -The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml -``` - -See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry. - -### Tensor Parallelism (TP) - -For larger models or multi-GPU environments, enable TP via CLI: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2 -``` - -Or set `tensor_parallel_size` per stage in a custom deploy YAML. - -### VAE Patch Parallelism - -[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck. - -**Scope for Bagel:** - -| Topology | VAE patch parallel | -| :------- | :----------------- | -| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) | -| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated | - -**Requirements:** - -- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline). -- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage). -- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup. - -VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`). - -**Online serving (single-stage, 2 GPUs):** - -```bash -CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config vllm_omni/deploy/bagel_single_stage.yaml \ - --tensor-parallel-size 2 \ - --vae-patch-parallel-size 2 \ - --vae-use-tiling -``` - -**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example: - -```yaml -stages: - - stage_id: 0 - devices: "0" - # Thinker (AR) — no VAE patch parallel here - - - stage_id: 1 - devices: "0,1" - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 2 - vae_patch_parallel_size: 2 -``` - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \ - --deploy-config /path/to/bagel_vae_pp.yaml -``` - -**Verify it is active** (check server logs at startup): - -```text -INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it. -``` - -| CLI flag | Default | Description | -| :------- | :------ | :---------- | -| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). | -| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). | - -#### Hybrid Sharded Data Parallel (HSDP) - -For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage. - -1. **Enable HSDP**: Set `use_hsdp: true`. -2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`). -3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups. -4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`). - -Example configuration for HSDP across 4 GPUs: -```yaml - - stage_id: 1 - devices: "0,1,2,3" - parallel_config: - use_hsdp: true - hsdp_shard_size: 4 - hsdp_replicate_size: 1 -``` - -### Multi-Node Deployment - -Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node. - -**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node: - -```bash -# API server port for client requests: 8000 -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \ - --port 8000 \ - --stage-id 0 \ - --omni-master-address <ORCHESTRATOR_IP> \ - --omni-master-port 8091 -``` - -**2. Launch Stage 1 (DiT)** on the remote node in headless mode: - -```bash -vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \ - --stage-id 1 \ - --headless \ - --omni-master-address <ORCHESTRATOR_IP> \ - --omni-master-port 8091 -``` - -Or use the convenience script: - -```bash -# Terminal 1: Stage 0 -bash run_server_stage_cli.sh --stage 0 - -# Terminal 2: Stage 1 -bash run_server_stage_cli.sh --stage 1 - -# With extra args -bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2 -bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9 -``` - -**vllm serve arguments:** - -| Argument | Description | -| :------- | :---------- | -| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) | -| `--headless` | Run without the API server (worker-only mode) | -| `-oma` / `--omni-master-address` | Orchestrator master address | -| `-omp` / `--omni-master-port` | Orchestrator master port | - -> [!IMPORTANT] -> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless). -> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior. - -### Inter-Stage Connectors - -When deploying stages across nodes, configure the connector type in the deploy YAML: - -- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed. -- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`. - -To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section. - -## Send Requests - -```bash -cd examples/online_serving/bagel -``` - -### Text to Image (text2img) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "A beautiful sunset over mountains" \ - --modality text2img \ - --output sunset.png \ - --steps 50 -``` - -**curl:** - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}], - "modalities": ["image"], - "height": 512, - "width": 512, - "num_inference_steps": 50, - "seed": 42 - }' -``` - -### Image to Image (img2img) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "Make the cat stand up" \ - --modality img2img \ - --image-url /path/to/input.jpg \ - --output transformed.png -``` - -**curl:** - -```bash -IMAGE_BASE64=$(base64 -w 0 cat.jpg) - -cat <<EOF > payload.json -{ - "messages": [{ - "role": "user", - "content": [ - {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"}, - {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}} - ] - }], - "modalities": ["image"], - "height": 512, - "width": 512, - "num_inference_steps": 50, - "seed": 42 -} -EOF - -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d @payload.json -``` - -### Image to Text (img2text) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "Describe this image in detail" \ - --modality img2text \ - --image-url /path/to/image.jpg -``` - -**curl:** - -```bash -IMAGE_BASE64=$(base64 -w 0 cat.jpg) - -cat <<EOF > payload.json -{ - "messages": [{ - "role": "user", - "content": [ - {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"}, - {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}} - ] - }], - "modalities": ["text"] -} -EOF - -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d @payload.json -``` - -### Text to Text (text2text) - -**Python client:** - -```bash -python openai_chat_client.py \ - --prompt "What is the capital of France?" \ - --modality text2text -``` - -**curl:** - -```bash -curl http://localhost:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}], - "modalities": ["text"] - }' -``` - -### Python Client Arguments - -| Argument | Default | Description | -| :------- | :------ | :---------- | -| `--prompt` / `-p` | `A cute cat` | Text prompt | -| `--output` / `-o` | `bagel_output.png` | Output file path | -| `--server` / `-s` | `http://localhost:8091` | Server URL | -| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) | -| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` | -| `--height` | `512` | Image height in pixels | -| `--width` | `512` | Image width in pixels | -| `--steps` | `25` | Number of inference steps | -| `--seed` | `42` | Random seed | -| `--negative` | `None` | Negative prompt for CFG | - -Example with custom parameters: - -```bash -python openai_chat_client.py \ - --prompt "A futuristic city" \ - --modality text2img \ - --height 768 \ - --width 768 \ - --steps 50 \ - --seed 42 \ - --negative "blurry, low quality" -``` - -## Configuration Reference - -### Deploy YAML Files - -| File | Description | -| :--- | :---------- | -| [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) | -| [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) | - -### Key Deploy YAML Fields - -| Field | Scope | Description | -| :---- | :---- | :---------- | -| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) | -| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) | -| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) | -| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences | -| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use | -| `stages[].enforce_eager` | per-stage | Disable CUDA graphs | -| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage | -| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) | -| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) | -| `connectors` | top-level | Define available connector instances (SHM, Mooncake) | -| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) | - -## FAQ - -- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML. - -**Two-stage VRAM usage:** - -| Stage | VRAM | -| :---- | :--- | -| Stage 0 (Thinker) | **15.04 GiB + KV Cache** | -| Stage 1 (DiT) | **26.50 GiB** | -| Total | **~42 GiB + KV Cache** | - -**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process. diff --git a/examples/online_serving/bagel/openai_chat_client.py b/examples/online_serving/bagel/openai_chat_client.py deleted file mode 100755 index cc9ec32db91..00000000000 --- a/examples/online_serving/bagel/openai_chat_client.py +++ /dev/null @@ -1,184 +0,0 @@ -#!/usr/bin/env python3 -""" -Bagel OpenAI-compatible chat client for image generation and multimodal tasks. - -Usage: - python openai_chat_client.py --prompt "A cute cat" --output output.png - python openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png -""" - -import argparse -import base64 -from pathlib import Path - -import requests - - -def generate_image( - prompt: str, - server_url: str = "http://localhost:8091", - image_url: str | None = None, - height: int | None = None, - width: int | None = None, - steps: int | None = None, - seed: int | None = None, - negative_prompt: str | None = None, - modality: str = "text2img", # "text2img" (default), "img2img", "img2text", "text2text" -) -> bytes | str | None: - """Generate an image or text using the chat completions API. - - Args: - prompt: Text description or prompt - server_url: Server URL - image_url: URL or path to input image (for img2img/img2text) - height: Image height in pixels - width: Image width in pixels - steps: Number of inference steps - seed: Random seed - negative_prompt: Negative prompt - modality: Task modality hint - - Returns: - Image bytes (for image outputs) or Text string (for text outputs) or None if failed - """ - - # Construct Message Content - content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}] - - if image_url: - # Check if local file - if Path(image_url).exists(): - with open(image_url, "rb") as f: - b64_data = base64.b64encode(f.read()).decode("utf-8") - final_image_url = f"data:image/jpeg;base64,{b64_data}" - else: - final_image_url = image_url - - content.append({"type": "image_url", "image_url": {"url": final_image_url}}) - - messages = [{"role": "user", "content": content}] - - # Build request payload with all parameters at top level - # Note: vLLM ignores "extra_body", so we put parameters directly in the payload - payload = {"messages": messages} - - # Set output modalities at top level - if modality == "text2img" or modality == "img2img": - payload["modalities"] = ["image"] - elif modality == "img2text" or modality == "text2text": - payload["modalities"] = ["text"] - - # Add generation parameters directly to payload - if height is not None: - payload["height"] = height - if width is not None: - payload["width"] = width - if steps is not None: - payload["num_inference_steps"] = steps - if seed is not None: - payload["seed"] = seed - if negative_prompt: - payload["negative_prompt"] = negative_prompt - - # Send request - try: - print(f"Sending request to {server_url} with modality {modality}...") - response = requests.post( - f"{server_url}/v1/chat/completions", - headers={"Content-Type": "application/json"}, - json=payload, - timeout=300, - ) - response.raise_for_status() - data = response.json() - - # Extract content - check ALL choices since server may return multiple - # (e.g., text in choices[0], image in choices[1]) - choices = data.get("choices", []) - - # First pass: look for image output in any choice - for choice in choices: - choice_content = choice.get("message", {}).get("content") - - # Handle Image Output - if isinstance(choice_content, list) and len(choice_content) > 0: - first_item = choice_content[0] - if isinstance(first_item, dict) and "image_url" in first_item: - img_url_str = first_item["image_url"].get("url", "") - if img_url_str.startswith("data:image"): - _, b64_data = img_url_str.split(",", 1) - return base64.b64decode(b64_data) - - # Second pass: look for text output if no image found - for choice in choices: - choice_content = choice.get("message", {}).get("content") - if isinstance(choice_content, str) and choice_content: - return choice_content - - print(f"Unexpected response format: {choices}") - return None - - except Exception as e: - print(f"Error: {e}") - return None - - -def main(): - parser = argparse.ArgumentParser(description="Bagel multimodal chat client") - parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt") - parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)") - parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL") - - # Modality Control - parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path") - parser.add_argument( - "--modality", - "-m", - default="text2img", - choices=["text2img", "img2img", "img2text", "text2text"], - help="Task modality", - ) - - # Generation Params - parser.add_argument("--height", type=int, default=512, help="Image height") - parser.add_argument("--width", type=int, default=512, help="Image width") - parser.add_argument("--steps", type=int, default=25, help="Inference steps") - parser.add_argument("--seed", type=int, default=42, help="Random seed") - parser.add_argument("--negative", help="Negative prompt") - - args = parser.parse_args() - - print(f"Mode: {args.modality}") - if args.image_url: - print(f"Input Image: {args.image_url}") - - result = generate_image( - prompt=args.prompt, - server_url=args.server, - image_url=args.image_url, - height=args.height, - width=args.width, - steps=args.steps, - seed=args.seed, - negative_prompt=args.negative, - modality=args.modality, - ) - - if result: - if isinstance(result, bytes): - # It's an image - output_path = Path(args.output) - output_path.write_bytes(result) - print(f"Image saved to: {output_path}") - print(f"Size: {len(result) / 1024:.1f} KB") - elif isinstance(result, str): - # It's text - print("Response:") - print(result) - else: - print("Failed to generate response") - exit(1) - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/bagel/run_server.sh b/examples/online_serving/bagel/run_server.sh deleted file mode 100755 index a64057ef033..00000000000 --- a/examples/online_serving/bagel/run_server.sh +++ /dev/null @@ -1,12 +0,0 @@ -#!/bin/bash -# Bagel online serving startup script - -MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}" -PORT="${PORT:-8091}" - -echo "Starting Bagel server..." -echo "Model: $MODEL" -echo "Port: $PORT" - -vllm serve "$MODEL" --omni \ - --port "$PORT" diff --git a/examples/online_serving/bagel/run_server_stage_cli.sh b/examples/online_serving/bagel/run_server_stage_cli.sh deleted file mode 100644 index 912e212f97e..00000000000 --- a/examples/online_serving/bagel/run_server_stage_cli.sh +++ /dev/null @@ -1,164 +0,0 @@ -#!/bin/bash -# Bagel multi-stage online serving startup script. -# -# Usage: -# ./run_server_stage_cli.sh --stage 0 -# ./run_server_stage_cli.sh --stage 1 -# ./run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2 -# ./run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9 -# -# By default, `--stage all` keeps the old behavior and launches both stages in -# one session. Use `--stage 0` / `--stage 1` to launch each stage separately in -# different terminal sessions, with stage-specific extra CLI arguments passed -# after `--`. - -set -euo pipefail - -MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}" -PORT="${PORT:-8091}" -MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}" -MASTER_PORT="${MASTER_PORT:-8092}" -STAGE="all" -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -DEPLOY_CONFIG="${DEPLOY_CONFIG:-$SCRIPT_DIR/../../../vllm_omni/deploy/bagel.yaml}" -EXTRA_ARGS=() - -usage() { - cat <<EOF -Usage: $0 [OPTIONS] [-- EXTRA_VLLM_ARGS...] - -Options: - --stage {0|1|all} Stage to launch (default: all) - --model MODEL Model name/path (default: $MODEL) - --port PORT API port for stage 0 (default: $PORT) - --master-address ADDRESS Master/orchestrator address (default: $MASTER_ADDRESS) - --master-port PORT Master/orchestrator port (default: $MASTER_PORT) - --deploy-config PATH Deploy config YAML path (default: $DEPLOY_CONFIG) - --help Show this help message - -Examples: - $0 --stage 0 - $0 --stage 1 - $0 --stage 0 -- --tensor-parallel-size 2 - $0 --stage 1 -- --gpu-memory-utilization 0.9 - -Notes: - - Use different terminal sessions to launch stage 0 and stage 1 separately. - - Extra args after '--' are forwarded only to the selected stage. - - When using '--stage all', the extra args are forwarded to both stages. -EOF -} - -while [[ $# -gt 0 ]]; do - case "$1" in - --stage) - STAGE="$2" - shift 2 - ;; - --model) - MODEL="$2" - shift 2 - ;; - --port) - PORT="$2" - shift 2 - ;; - --master-address) - MASTER_ADDRESS="$2" - shift 2 - ;; - --master-port) - MASTER_PORT="$2" - shift 2 - ;; - --deploy-config) - DEPLOY_CONFIG="$2" - shift 2 - ;; - --help|-h) - usage - exit 0 - ;; - --) - shift - EXTRA_ARGS=("$@") - break - ;; - *) - echo "Unknown option: $1" >&2 - usage - exit 1 - ;; - esac -done - -if [[ "$STAGE" != "0" && "$STAGE" != "1" && "$STAGE" != "all" ]]; then - echo "Invalid --stage value: $STAGE" >&2 - usage - exit 1 -fi - -print_config() { - echo "Model: $MODEL" - echo "API Port: $PORT" - echo "Master Address: $MASTER_ADDRESS" - echo "Master Port: $MASTER_PORT" - echo "Deploy Config: $DEPLOY_CONFIG" - echo "Selected Stage: $STAGE" - if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then - echo "Extra Args: ${EXTRA_ARGS[*]}" - fi -} - -run_stage_0() { - echo "Starting Stage 0 (Thinker) as master..." - vllm serve "$MODEL" --omni \ - --port "$PORT" \ - --deploy-config "$DEPLOY_CONFIG" \ - --stage-id 0 \ - --omni-master-address "$MASTER_ADDRESS" \ - --omni-master-port "$MASTER_PORT" \ - "${EXTRA_ARGS[@]}" -} - -run_stage_1() { - echo "Starting Stage 1 (DiT) in headless mode..." - vllm serve "$MODEL" --omni \ - --deploy-config "$DEPLOY_CONFIG" \ - --stage-id 1 \ - --headless \ - --omni-master-address "$MASTER_ADDRESS" \ - --omni-master-port "$MASTER_PORT" \ - "${EXTRA_ARGS[@]}" -} - -echo "Starting Bagel multi-stage server..." -print_config - -case "$STAGE" in - 0) - run_stage_0 - ;; - 1) - run_stage_1 - ;; - all) - echo "Launching both stages in one session (legacy mode)..." - echo "Starting Stage 0 (Thinker) in background first..." - run_stage_0 & - STAGE_0_PID=$! - - cleanup() { - if [[ -n "${STAGE_0_PID:-}" ]]; then - kill "$STAGE_0_PID" 2>/dev/null || true - wait "$STAGE_0_PID" 2>/dev/null || true - fi - } - - trap cleanup EXIT INT TERM - - echo "Waiting briefly for Stage 0 to initialize..." - sleep 2 - run_stage_1 - ;; -esac diff --git a/examples/online_serving/covo_audio/README.md b/examples/online_serving/covo_audio/README.md index 6894507efc4..b6a4ad74a8b 100644 --- a/examples/online_serving/covo_audio/README.md +++ b/examples/online_serving/covo_audio/README.md @@ -1,4 +1,4 @@ -# Covo-Audio-Chat +# Covo-Audio-Chat: Online serving ## Installation diff --git a/examples/online_serving/diffusers_pipeline_adapter/README.md b/examples/online_serving/diffusers_pipeline_adapter/README.md index d6f41f6328f..37bf48aa007 100644 --- a/examples/online_serving/diffusers_pipeline_adapter/README.md +++ b/examples/online_serving/diffusers_pipeline_adapter/README.md @@ -6,7 +6,8 @@ vLLM-Omni supports running diffusion models with the diffusers backend, directly The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online. Currently, the following features are NOT yet supported. -It is not guaranteed whether they will be supported in the future. +Community contributions are welcome to add these features when they can be +implemented with clear Diffusers or vLLM-Omni behavior. - CFG parallel execution - Sequence parallel execution @@ -68,6 +69,35 @@ But if that parameter is simultaneously set in both the vLLM-Omni interface and Consider referring to [`OmniDiffusionSamplingParams`](../../../vllm_omni/inputs/data.py) for the default sampling parameters in the vLLM-Omni interface, and the corresponding [diffusers pipeline](https://huggingface.co/docs/diffusers/main/en/api/pipelines/overview)'s `__call__` function documentation. +### Quantization + +Use `diffusers_load_kwargs` for Diffusers-native quantization options. If +`diffusers_load_kwargs["quantization_config"]` is provided as a dictionary, the +diffusers backend builds a Diffusers quantization config and lets Diffusers +validate it before calling `DiffusionPipeline.from_pretrained()`. + +The diffusers backend also provides a small compatibility shortcut for +vLLM-Omni quantization configs when `diffusers_load_kwargs` does not already +contain `quantization_config`. Currently this maps online/dynamic `fp8` and +online/dynamic `int8` to Diffusers/TorchAO dynamic quantization for transformer +components such as `transformer` or `transformer_2`. This path requires +`torchao` to be installed. + +For example, the CLI can request dynamic FP8 through the vLLM-Omni interface: + +```bash +vllm serve "Qwen/Qwen-Image" \ + --omni \ + --diffusion-load-format diffusers \ + --quantization-config '{"method": "fp8"}' +``` + +Other vLLM-Omni quantization methods, such as `gguf`, `modelopt`, `mxfp4`, +`mxfp8`, serialized checkpoints, static FP8 configs, and layer-name skip lists +such as `ignored_layers`, are intentionally not translated by this compatibility +shortcut. Use Diffusers-native configuration through `diffusers_load_kwargs` or +a native vLLM-Omni pipeline for those cases. + ### Attention Backends The diffusers backend converts diff --git a/examples/online_serving/dynin_omni/README.md b/examples/online_serving/dynin_omni/README.md index 50e2f3f8039..6e78a5c1fc4 100644 --- a/examples/online_serving/dynin_omni/README.md +++ b/examples/online_serving/dynin_omni/README.md @@ -1,4 +1,4 @@ -# Dynin-Omni Online Serving Example +# Dynin-Omni: Online serving ## Installation diff --git a/examples/online_serving/helios/README.md b/examples/online_serving/helios/README.md deleted file mode 100644 index 5e13df19d63..00000000000 --- a/examples/online_serving/helios/README.md +++ /dev/null @@ -1,175 +0,0 @@ -# Helios Video Generation via Online API - -This example demonstrates how to use the `/v1/videos` API with Helios models using the generic `extra_params` field. - -## Overview - -The `/v1/videos` API now supports model-specific parameters through the `extra_params` field, which accepts a JSON object containing any model-specific configuration. This allows supporting new models like Helios without modifying the API. - -## Helios Model Variants - -- **Helios-Base**: Basic T2V/I2V/V2V generation (Stage 1 only) -- **Helios-Mid**: Advanced generation with CFG-Zero* (Stage 2) -- **Helios-Distilled**: Fast generation with DMD (Stage 2) - -## API Usage Examples - -### 1. T2V (Text-to-Video) - Helios-Base - -Basic text-to-video generation: - -```bash -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=A serene lakeside sunrise with mist over the water." \ - -F "model=BestWishYsh/Helios-Base" \ - -F "width=640" \ - -F "height=384" \ - -F "num_frames=99" \ - -F "num_inference_steps=50" \ - -F "guidance_scale=5.0" \ - -F "seed=42" -``` - -### 2. I2V (Image-to-Video) - Helios-Base - -Generate video from an input image: - -```bash -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=The lake water gently ripples as morning mist rises." \ - -F "model=BestWishYsh/Helios-Base" \ - -F "input_reference=@/path/to/image.jpg" \ - -F "width=640" \ - -F "height=384" \ - -F "num_frames=99" \ - -F "guidance_scale=5.0" -``` - -### 3. Helios-Mid with Stage 2 + CFG-Zero* - -Advanced generation with pyramid multi-stage denoising: - -```bash -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=A serene lakeside sunrise with mist over the water." \ - -F "model=BestWishYsh/Helios-Mid" \ - -F "width=640" \ - -F "height=384" \ - -F "guidance_scale=5.0" \ - -F 'extra_params={ - "is_enable_stage2": true, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [20, 20, 20], - "use_cfg_zero_star": true, - "use_zero_init": true, - "zero_steps": 1 - }' -``` - -### 4. Helios-Distilled with DMD - -Fast generation with Distribution Matching Distillation: - -```bash -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=A serene lakeside sunrise with mist over the water." \ - -F "model=BestWishYsh/Helios-Distilled" \ - -F "width=640" \ - -F "height=384" \ - -F "guidance_scale=1.0" \ - -F 'extra_params={ - "is_enable_stage2": true, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [2, 2, 2], - "is_amplify_first_chunk": true - }' -``` - -## Model-Specific Parameters - -The `extra_params` field accepts a JSON object with model-specific parameters. For Helios models, supported parameters include: - -### Stage 2 Parameters -- `is_enable_stage2` (bool): Enable pyramid multi-stage denoising -- `pyramid_num_stages` (int): Number of pyramid stages (default: 3) -- `pyramid_num_inference_steps_list` (array): Steps per stage, e.g., `[20, 20, 20]` - -### CFG Zero Star Parameters (Helios-Mid) -- `use_cfg_zero_star` (bool): Enable CFG Zero Star guidance -- `use_zero_init` (bool): Use zero initialization for first steps -- `zero_steps` (int): Number of initial zero prediction steps (default: 1) - -### DMD Parameters (Helios-Distilled) -- `is_amplify_first_chunk` (bool): Enable DMD amplification for first chunk - -### Video Input (V2V mode) -For video-to-video generation, upload a video file via `input_reference`: - -```bash -curl -X POST http://localhost:8000/v1/videos \ - -F "prompt=Transform the video into a watercolor painting style." \ - -F "model=BestWishYsh/Helios-Base" \ - -F "input_reference=@/path/to/video.mp4" \ - -F "guidance_scale=5.0" -``` - -## Python Example - -```python -import requests -import json - -url = "http://localhost:8000/v1/videos" - -# Helios-Mid with Stage 2 -data = { - "prompt": "A serene lakeside sunrise with mist over the water.", - "model": "BestWishYsh/Helios-Mid", - "width": 640, - "height": 384, - "guidance_scale": 5.0, - "extra_params": json.dumps({ - "is_enable_stage2": True, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [20, 20, 20], - "use_cfg_zero_star": True, - "use_zero_init": True, - "zero_steps": 1 - }) -} - -response = requests.post(url, data=data) -video_job = response.json() -print(f"Video job created: {video_job['id']}") - -# Poll for completion -import time -while True: - status_response = requests.get(f"{url}/{video_job['id']}") - status = status_response.json() - - if status['status'] == 'completed': - print(f"Video generated: {status['file_name']}") - break - elif status['status'] == 'failed': - print(f"Generation failed: {status.get('error')}") - break - - print(f"Progress: {status['progress']}%") - time.sleep(2) -``` - -## Benefits of This Approach - -1. **No API Changes**: New models can be supported without modifying the API -2. **Backward Compatible**: Existing clients continue to work -3. **Flexible**: Any model-specific parameter can be passed -4. **Type-Safe**: Parameters are validated at the model level -5. **Future-Proof**: Supports models that don't exist yet - -## Notes - -- The `extra_params` field must be a valid JSON object -- Parameters are passed directly to the model's `extra_args` -- Invalid parameters will be caught by the model implementation -- Tensor inputs (images/videos) should use `input_reference`, not `extra_params` diff --git a/examples/online_serving/helios/helios_client.py b/examples/online_serving/helios/helios_client.py deleted file mode 100644 index d5024996620..00000000000 --- a/examples/online_serving/helios/helios_client.py +++ /dev/null @@ -1,212 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Helios video generation client example using the /v1/videos API. - -This example demonstrates how to use the generic extra_params field -to pass Helios-specific parameters without modifying the API. -""" - -import argparse -import json -import time -from pathlib import Path - -import requests - - -def create_video_job( - api_url: str, - prompt: str, - model: str, - width: int = 640, - height: int = 384, - num_frames: int = 99, - guidance_scale: float = 5.0, - seed: int = 42, - extra_params: dict | None = None, - input_image: Path | None = None, -) -> dict: - """Create a video generation job.""" - data = { - "prompt": prompt, - "model": model, - "width": width, - "height": height, - "num_frames": num_frames, - "guidance_scale": guidance_scale, - "seed": seed, - } - - files = {} - if input_image: - files["input_reference"] = open(input_image, "rb") - - if extra_params: - data["extra_params"] = json.dumps(extra_params) - - response = requests.post(f"{api_url}/v1/videos", data=data, files=files) - response.raise_for_status() - - if files: - files["input_reference"].close() - - return response.json() - - -def poll_video_status(api_url: str, video_id: str, poll_interval: int = 2) -> dict: - """Poll video generation status until completion.""" - print(f"Polling video job: {video_id}") - - while True: - response = requests.get(f"{api_url}/v1/videos/{video_id}") - response.raise_for_status() - status_data = response.json() - - status = status_data["status"] - progress = status_data.get("progress", 0) - - print(f"Status: {status}, Progress: {progress}%") - - if status == "completed": - print("Video generation completed!") - return status_data - elif status == "failed": - error = status_data.get("error", {}) - raise RuntimeError(f"Video generation failed: {error}") - - time.sleep(poll_interval) - - -def main(): - parser = argparse.ArgumentParser(description="Helios video generation client") - parser.add_argument( - "--api-url", - default="http://localhost:8000", - help="API server URL", - ) - parser.add_argument( - "--model", - default="BestWishYsh/Helios-Base", - help="Model name (Helios-Base, Helios-Mid, Helios-Distilled)", - ) - parser.add_argument( - "--prompt", - default="A serene lakeside sunrise with mist over the water.", - help="Text prompt", - ) - parser.add_argument( - "--width", - type=int, - default=640, - help="Video width", - ) - parser.add_argument( - "--height", - type=int, - default=384, - help="Video height", - ) - parser.add_argument( - "--num-frames", - type=int, - default=99, - help="Number of frames", - ) - parser.add_argument( - "--guidance-scale", - type=float, - default=5.0, - help="Guidance scale", - ) - parser.add_argument( - "--seed", - type=int, - default=42, - help="Random seed", - ) - parser.add_argument( - "--input-image", - type=Path, - help="Input image for I2V mode", - ) - parser.add_argument( - "--preset", - choices=["base", "mid-stage2", "distilled"], - default="base", - help="Helios preset configuration", - ) - - args = parser.parse_args() - - # Define preset configurations - presets = { - "base": None, # No extra params for base model - "mid-stage2": { - "is_enable_stage2": True, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [20, 20, 20], - "use_cfg_zero_star": True, - "use_zero_init": True, - "zero_steps": 1, - }, - "distilled": { - "is_enable_stage2": True, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [2, 2, 2], - "is_amplify_first_chunk": True, - }, - } - - extra_params = presets[args.preset] - - print("=" * 50) - print("Helios Video Generation") - print("=" * 50) - print(f"API URL: {args.api_url}") - print(f"Model: {args.model}") - print(f"Preset: {args.preset}") - print(f"Prompt: {args.prompt}") - print(f"Size: {args.width}x{args.height}") - print(f"Frames: {args.num_frames}") - print(f"Guidance Scale: {args.guidance_scale}") - if extra_params: - print(f"Extra Params: {json.dumps(extra_params, indent=2)}") - print() - - # Create video job - print("Creating video generation job...") - job = create_video_job( - api_url=args.api_url, - prompt=args.prompt, - model=args.model, - width=args.width, - height=args.height, - num_frames=args.num_frames, - guidance_scale=args.guidance_scale, - seed=args.seed, - extra_params=extra_params, - input_image=args.input_image, - ) - - video_id = job["id"] - print(f"Video job created: {video_id}") - print() - - # Poll for completion - result = poll_video_status(args.api_url, video_id) - - print() - print("=" * 50) - print("Generation Complete") - print("=" * 50) - print(f"Video ID: {result['id']}") - print(f"File: {result.get('file_name', 'N/A')}") - print(f"Inference Time: {result.get('inference_time_s', 'N/A')}s") - print() - - -if __name__ == "__main__": - main() diff --git a/examples/online_serving/helios/run_helios_distilled.sh b/examples/online_serving/helios/run_helios_distilled.sh deleted file mode 100644 index 15b6d8c3032..00000000000 --- a/examples/online_serving/helios/run_helios_distilled.sh +++ /dev/null @@ -1,79 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Helios-Distilled with DMD Example -# Fast generation with Distribution Matching Distillation - -API_URL="${API_URL:-http://localhost:8000}" -MODEL="${MODEL:-BestWishYsh/Helios-Distilled}" -PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}" - -echo "===================================" -echo "Helios-Distilled with DMD" -echo "===================================" -echo "API URL: $API_URL" -echo "Model: $MODEL" -echo "Prompt: $PROMPT" -echo "" - -# Model-specific parameters for Helios-Distilled -extra_params='{ - "is_enable_stage2": true, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [2, 2, 2], - "is_amplify_first_chunk": true -}' - -echo "Model extra params: $extra_params" -echo "" - -# Create video generation job -echo "Creating video generation job..." -RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \ - -F "prompt=$PROMPT" \ - -F "model=$MODEL" \ - -F "width=640" \ - -F "height=384" \ - -F "guidance_scale=1.0" \ - -F "seed=42" \ - -F "extra_params=$extra_params") - -echo "Response: $RESPONSE" -echo "" - -# Extract video ID -VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4) - -if [ -z "$VIDEO_ID" ]; then - echo "Error: Failed to create video job" - exit 1 -fi - -echo "Video job created: $VIDEO_ID" -echo "" - -# Poll for completion -echo "Polling for completion..." -while true; do - STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID") - STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4) - PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2) - - echo "Status: $STATUS, Progress: $PROGRESS%" - - if [ "$STATUS" = "completed" ]; then - echo "" - echo "Video generation completed!" - echo "Full response:" - echo "$STATUS_RESPONSE" | jq '.' - break - elif [ "$STATUS" = "failed" ]; then - echo "" - echo "Video generation failed!" - echo "$STATUS_RESPONSE" | jq '.' - exit 1 - fi - - sleep 2 -done diff --git a/examples/online_serving/helios/run_helios_mid_stage2.sh b/examples/online_serving/helios/run_helios_mid_stage2.sh deleted file mode 100644 index d5da9b22c0b..00000000000 --- a/examples/online_serving/helios/run_helios_mid_stage2.sh +++ /dev/null @@ -1,81 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Helios-Mid with Stage 2 + CFG-Zero* Example -# This demonstrates advanced generation with pyramid multi-stage denoising - -API_URL="${API_URL:-http://localhost:8000}" -MODEL="${MODEL:-BestWishYsh/Helios-Mid}" -PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}" - -echo "===================================" -echo "Helios-Mid Stage 2 + CFG-Zero*" -echo "===================================" -echo "API URL: $API_URL" -echo "Model: $MODEL" -echo "Prompt: $PROMPT" -echo "" - -# Model-specific parameters for Helios-Mid -extra_params='{ - "is_enable_stage2": true, - "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [20, 20, 20], - "use_cfg_zero_star": true, - "use_zero_init": true, - "zero_steps": 1 -}' - -echo "Model extra params: $extra_params" -echo "" - -# Create video generation job -echo "Creating video generation job..." -RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \ - -F "prompt=$PROMPT" \ - -F "model=$MODEL" \ - -F "width=640" \ - -F "height=384" \ - -F "guidance_scale=5.0" \ - -F "seed=42" \ - -F "extra_params=$extra_params") - -echo "Response: $RESPONSE" -echo "" - -# Extract video ID -VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4) - -if [ -z "$VIDEO_ID" ]; then - echo "Error: Failed to create video job" - exit 1 -fi - -echo "Video job created: $VIDEO_ID" -echo "" - -# Poll for completion -echo "Polling for completion..." -while true; do - STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID") - STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4) - PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2) - - echo "Status: $STATUS, Progress: $PROGRESS%" - - if [ "$STATUS" = "completed" ]; then - echo "" - echo "Video generation completed!" - echo "Full response:" - echo "$STATUS_RESPONSE" | jq '.' - break - elif [ "$STATUS" = "failed" ]; then - echo "" - echo "Video generation failed!" - echo "$STATUS_RESPONSE" | jq '.' - exit 1 - fi - - sleep 2 -done diff --git a/examples/online_serving/helios/run_helios_t2v.sh b/examples/online_serving/helios/run_helios_t2v.sh deleted file mode 100644 index e920fc872b6..00000000000 --- a/examples/online_serving/helios/run_helios_t2v.sh +++ /dev/null @@ -1,69 +0,0 @@ -#!/bin/bash -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -# Helios Text-to-Video Generation Example -# This script demonstrates how to use the /v1/videos API with Helios models - -API_URL="${API_URL:-http://localhost:8000}" -MODEL="${MODEL:-BestWishYsh/Helios-Base}" -PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}" - -echo "===================================" -echo "Helios T2V Generation" -echo "===================================" -echo "API URL: $API_URL" -echo "Model: $MODEL" -echo "Prompt: $PROMPT" -echo "" - -# Create video generation job -echo "Creating video generation job..." -RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \ - -F "prompt=$PROMPT" \ - -F "model=$MODEL" \ - -F "width=640" \ - -F "height=384" \ - -F "num_frames=99" \ - -F "num_inference_steps=50" \ - -F "guidance_scale=5.0" \ - -F "seed=42") - -echo "Response: $RESPONSE" -echo "" - -# Extract video ID -VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4) - -if [ -z "$VIDEO_ID" ]; then - echo "Error: Failed to create video job" - exit 1 -fi - -echo "Video job created: $VIDEO_ID" -echo "" - -# Poll for completion -echo "Polling for completion..." -while true; do - STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID") - STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4) - PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2) - - echo "Status: $STATUS, Progress: $PROGRESS%" - - if [ "$STATUS" = "completed" ]; then - echo "" - echo "Video generation completed!" - echo "Full response:" - echo "$STATUS_RESPONSE" | jq '.' - break - elif [ "$STATUS" = "failed" ]; then - echo "" - echo "Video generation failed!" - echo "$STATUS_RESPONSE" | jq '.' - exit 1 - fi - - sleep 2 -done diff --git a/examples/online_serving/image_to_image/README.md b/examples/online_serving/image_to_image/README.md index 59b1f0e2c15..8a6c6a0e6c2 100644 --- a/examples/online_serving/image_to_image/README.md +++ b/examples/online_serving/image_to_image/README.md @@ -1,6 +1,8 @@ # Image-To-Image -This example demonstrates how to deploy Qwen-Image-Edit model for online image editing service using vLLM-Omni. +This example demonstrates how to deploy image-to-image models for online image editing service using vLLM-Omni. + +Supported models include Qwen-Image-Edit, BAGEL, and other image-to-image pipelines. For **multi-image** input editing, use **Qwen-Image-Edit-2509** (QwenImageEditPlusPipeline) and send multiple images in the user message content. @@ -21,6 +23,12 @@ vllm serve Qwen/Qwen-Image-Edit --omni --port 8092 vllm serve Qwen/Qwen-Image-Edit-2509 --omni --port 8092 ``` +### BAGEL + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 +``` + ### Start with Parameters @@ -119,6 +127,16 @@ python openai_chat_client.py --input input.png --prompt "Convert to oil painting python openai_chat_client.py --input input1.png input2.png --prompt "Combine these images into a single scene" --output output.png ``` +Pass model-specific parameters through `--extra-body` (e.g. for BAGEL): + +```bash +python openai_chat_client.py \ + --input input.png \ + --prompt "Make the scene look like a watercolor painting" \ + --server http://localhost:8091 \ + --extra-body '{"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}' +``` + ### Method 4: Using Gradio Demo ```bash @@ -318,6 +336,9 @@ count, use `size` and `n` rather than `height`, `width`, or | `layers` | int | 4 | Number of layers (Qwen-Image-Layered) | | `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) | +Models like BAGEL accept additional parameters via `extra_body` (e.g. +`cfg_text_scale`, `cfg_img_scale`). See the [BAGEL recipe](../../../recipes/Bagel/BAGEL-7B-MoT.md) for the full list. + ## Response Format ```json diff --git a/examples/online_serving/image_to_image/openai_chat_client.py b/examples/online_serving/image_to_image/openai_chat_client.py index 7aac2b214e2..46f2ad5a7ef 100644 --- a/examples/online_serving/image_to_image/openai_chat_client.py +++ b/examples/online_serving/image_to_image/openai_chat_client.py @@ -38,6 +38,7 @@ def edit_image( guidance_scale: float | None = None, seed: int | None = None, negative_prompt: str | None = None, + extra_body: dict | None = None, ) -> bytes | None: """Edit an image using the chat completions API. @@ -51,6 +52,7 @@ def edit_image( guidance_scale: CFG guidance scale seed: Random seed negative_prompt: Negative prompt + extra_body: Additional model-specific params (e.g. cfg_text_scale for BAGEL) Returns: Edited image bytes or None if failed @@ -75,24 +77,26 @@ def edit_image( ] # Build extra_body with generation parameters - extra_body = {} + merged_extra_body: dict[str, object] = {} if height is not None: - extra_body["height"] = height + merged_extra_body["height"] = height if width is not None: - extra_body["width"] = width + merged_extra_body["width"] = width if steps is not None: - extra_body["num_inference_steps"] = steps + merged_extra_body["num_inference_steps"] = steps if guidance_scale is not None: - extra_body["guidance_scale"] = guidance_scale + merged_extra_body["guidance_scale"] = guidance_scale if seed is not None: - extra_body["seed"] = seed + merged_extra_body["seed"] = seed if negative_prompt: - extra_body["negative_prompt"] = negative_prompt + merged_extra_body["negative_prompt"] = negative_prompt + if extra_body: + merged_extra_body.update(extra_body) # Build request payload - payload = {"messages": messages} - if extra_body: - payload["extra_body"] = extra_body + payload: dict[str, object] = {"messages": messages} + if merged_extra_body: + payload["extra_body"] = merged_extra_body # Send request try: @@ -121,6 +125,19 @@ def edit_image( return None +def parse_extra_body(value: str) -> dict: + """Parse a JSON string into a dict for --extra-body.""" + import json + + try: + obj = json.loads(value) + except json.JSONDecodeError as e: + raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e + if not isinstance(obj, dict): + raise argparse.ArgumentTypeError("--extra-body must be a JSON object") + return obj + + def main(): parser = argparse.ArgumentParser(description="Qwen-Image-Edit chat client") parser.add_argument("--input", "-i", required=True, nargs="+", help="Input image path(s)") @@ -133,6 +150,12 @@ def main(): parser.add_argument("--guidance", type=float, default=7.5, help="Guidance scale") parser.add_argument("--seed", type=int, default=0, help="Random seed") parser.add_argument("--negative", help="Negative prompt") + parser.add_argument( + "--extra-body", + type=parse_extra_body, + default=None, + help="JSON object merged into request extra_body, e.g. '{\"cfg_text_scale\": 4.0}'.", + ) args = parser.parse_args() @@ -152,6 +175,7 @@ def main(): guidance_scale=args.guidance, seed=args.seed, negative_prompt=args.negative, + extra_body=args.extra_body, ) if image_bytes: diff --git a/examples/online_serving/image_to_video/README.md b/examples/online_serving/image_to_video/README.md index 285eeb27983..f18cea1c8ce 100644 --- a/examples/online_serving/image_to_video/README.md +++ b/examples/online_serving/image_to_video/README.md @@ -120,14 +120,17 @@ Use `/v1/videos/sync` if you want to write the MP4 directly to a file. `POST /v1 Generated video files are stored on local disk by the async video API. Local file storage behavior can be controlled via the following environment variables: -- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`) -- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`) +- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files (default: `/tmp/storage`) +- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete operations (default: `4`) + +`VLLM_OMNI_STORAGE_PATH` and `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` are deprecated and will be +removed in a future release; use the names above instead. Example: ```bash -export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos -export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8 +export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos +export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8 ``` ## API Calls diff --git a/examples/online_serving/joyvl_interaction/bridges/asr_bridge.py b/examples/online_serving/joyvl_interaction/bridges/asr_bridge.py new file mode 100644 index 00000000000..042e5808fc8 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/bridges/asr_bridge.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import struct +import wave + +import aiohttp +from aiohttp import web + +_HEADER = struct.Struct(">iii") +_SAMPLE_RATE = 16000 +_DEFAULT_PROMPT = "Transcribe the audio." +_ASR_TEXT_TAG = "<asr_text>" + + +def _clean_transcript(raw: str) -> str: + """Qwen3-ASR emits ``language <Lang><asr_text><transcript>``; keep the transcript.""" + text = raw or "" + marker = text.rfind(_ASR_TEXT_TAG) + if marker != -1: + text = text[marker + len(_ASR_TEXT_TAG) :] + return text.replace("</asr_text>", "").strip() + + +def _pcm_to_wav_b64(pcm: bytes, sample_rate: int) -> str: + buf = io.BytesIO() + with wave.open(buf, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(sample_rate) + wav.writeframes(pcm) + return base64.b64encode(buf.getvalue()).decode() + + +def _result(text: str, *, final: bool) -> dict: + return { + "asr_response": { + "event_type": "IS_FINAL" if final else "IS_PARTIAL", + "recognition_result": {"hypothesis": [{"text": text, "confidence": None}]}, + }, + "mid": "", + "code": 0, + "msg": "ok", + } + + +async def _transcribe(cfg: dict, pcm: bytes) -> str: + if not pcm: + return "" + wav_b64 = _pcm_to_wav_b64(pcm, cfg["sample_rate"]) + body = { + "model": cfg["model"], + "messages": [ + { + "role": "user", + "content": [ + {"type": "input_audio", "input_audio": {"data": wav_b64, "format": "wav"}}, + {"type": "text", "text": cfg["prompt"]}, + ], + } + ], + "temperature": 0.0, + "max_tokens": cfg["max_tokens"], + } + url = cfg["backend_url"].rstrip("/") + "/v1/chat/completions" + async with aiohttp.ClientSession() as session: + async with session.post(url, json=body) as resp: + payload = await resp.json() + choices = payload.get("choices") or [] + if not choices: + return "" + return _clean_transcript(choices[0].get("message", {}).get("content") or "") + + +async def _handle(request: web.Request) -> web.WebSocketResponse: + ws = web.WebSocketResponse(max_msg_size=0) + await ws.prepare(request) + cfg = request.app["cfg"] + pcm = bytearray() + try: + async for msg in ws: + if msg.type == aiohttp.WSMsgType.BINARY and len(msg.data) >= _HEADER.size: + seqid, _, _ = _HEADER.unpack(msg.data[: _HEADER.size]) + pcm.extend(msg.data[_HEADER.size :]) + if seqid < 0: + text = await _transcribe(cfg, bytes(pcm)) + await ws.send_str(json.dumps(_result(text, final=True), ensure_ascii=False)) + pcm.clear() + elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR): + break + finally: + if not ws.closed: + await ws.close() + return ws + + +def create_app(backend_url: str, model: str, sample_rate: int, prompt: str, max_tokens: int) -> web.Application: + app = web.Application() + app["cfg"] = { + "backend_url": backend_url, + "model": model, + "sample_rate": sample_rate, + "prompt": prompt, + "max_tokens": max_tokens, + } + app.router.add_get("/v1/asr", _handle) + return app + + +def main() -> None: + parser = argparse.ArgumentParser(description="JoyVL webui <-> Qwen3-ASR (chat input_audio) bridge") + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8093) + parser.add_argument("--backend-url", default="http://127.0.0.1:8094") + parser.add_argument("--model", default="Qwen/Qwen3-ASR-1.7B") + parser.add_argument("--sample-rate", type=int, default=_SAMPLE_RATE) + parser.add_argument("--prompt", default=_DEFAULT_PROMPT) + parser.add_argument("--max-tokens", type=int, default=256) + args = parser.parse_args() + app = create_app(args.backend_url, args.model, args.sample_rate, args.prompt, args.max_tokens) + web.run_app(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/joyvl_interaction/bridges/tts_bridge.py b/examples/online_serving/joyvl_interaction/bridges/tts_bridge.py new file mode 100644 index 00000000000..d6567081a77 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/bridges/tts_bridge.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging + +import aiohttp +from aiohttp import web + +try: + import cn2an # optional: Arabic-digit -> Chinese number normalization +except ImportError: + cn2an = None + +logger = logging.getLogger("tts_bridge") + + +def _normalize(text: str) -> str: + """Spell out Arabic digits in Chinese so short counting outputs synthesize + intelligibly: bare "1次" garbles on Qwen3-TTS (and others), while "一次" + is clean. ``cn2an.transform`` rewrites digits in place ("第3次"->"第三次") + and is pure-Python (~microseconds). No-op if cn2an is unavailable.""" + if not cn2an or not text: + return text + try: + return cn2an.transform(text, "an2cn") + except Exception as err: # never let TN break synthesis + logger.debug("cn2an normalize skipped: %s", err) + return text + + +async def _pump_backend_to_front(back: aiohttp.ClientWebSocketResponse, front: web.WebSocketResponse) -> None: + async for msg in back: + if msg.type == aiohttp.WSMsgType.BINARY: + await front.send_bytes(msg.data) + elif msg.type == aiohttp.WSMsgType.TEXT: + event = json.loads(msg.data).get("type") + if event == "session.done": + await front.send_str(json.dumps({"type": "response.done"})) + return + if event == "error": + await front.send_str(json.dumps({"type": "error", "error": "tts backend error"})) + return + elif msg.type in (aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.ERROR): + return + + +async def _handle(request: web.Request) -> web.WebSocketResponse: + front = web.WebSocketResponse(heartbeat=20, max_msg_size=0) + await front.prepare(request) + cfg = request.app["cfg"] + + session = aiohttp.ClientSession() + back: aiohttp.ClientWebSocketResponse | None = None + pump: asyncio.Task | None = None + parts: list[str] = [] + try: + async for msg in front: + if msg.type != web.WSMsgType.TEXT: + continue + data = json.loads(msg.data) + + if "config" in data and back is None: + voice = (data["config"] or {}).get("voice") or cfg["voice"] + back = await session.ws_connect(cfg["backend_url"], max_msg_size=0) + await back.send_str( + json.dumps({"type": "session.config", "response_format": cfg["response_format"], "voice": voice}) + ) + pump = asyncio.create_task(_pump_backend_to_front(back, front)) + elif data.get("type") == "input_text.append" and back is not None: + parts.append(data.get("text", "")) + elif data.get("type") == "input_text.commit" and back is not None: + # Normalize the full utterance once (digits may span chunks), then send. + await back.send_str(json.dumps({"type": "input.text", "text": _normalize("".join(parts))})) + parts = [] + await back.send_str(json.dumps({"type": "input.done"})) + if pump is not None: + await pump + break + except Exception as err: + logger.warning("tts bridge error: %s", err) + finally: + if pump is not None and not pump.done(): + pump.cancel() + if back is not None and not back.closed: + await back.close() + await session.close() + if not front.closed: + await front.close() + return front + + +def create_app(backend_url: str, voice: str, response_format: str) -> web.Application: + app = web.Application() + app["cfg"] = {"backend_url": backend_url, "voice": voice, "response_format": response_format} + app.router.add_get("/v1/tts", _handle) + return app + + +def main() -> None: + parser = argparse.ArgumentParser(description="JoyVL webui <-> Qwen3-TTS websocket bridge") + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8092) + parser.add_argument("--backend-url", default="ws://127.0.0.1:8091/v1/audio/speech/stream") + parser.add_argument("--voice", default="vivian") + parser.add_argument("--response-format", default="pcm") + args = parser.parse_args() + logging.basicConfig(level=logging.INFO) + web.run_app(create_app(args.backend_url, args.voice, args.response_format), host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/joyvl_interaction/cli/run_cli_demo.py b/examples/online_serving/joyvl_interaction/cli/run_cli_demo.py new file mode 100644 index 00000000000..8d86e033bec --- /dev/null +++ b/examples/online_serving/joyvl_interaction/cli/run_cli_demo.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import argparse + +from stream_client import stream + +ICON = {"response": "SPEAK ", "silence": "silence ", "delegate": "DELEGATE"} + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("video") + parser.add_argument("--server", default="http://127.0.0.1:8070") + parser.add_argument("--query", default=None) + parser.add_argument("--fps", type=float, default=1.0) + parser.add_argument("--session-id", default="cli") + args = parser.parse_args() + + for tick in stream(args.video, args.server, args.session_id, args.query, args.fps): + if tick.action == "silence": + continue + line = f"[{tick.t:6.1f}s] {ICON.get(tick.action)} {tick.text}" + if tick.delegation: + line += f" (-> {tick.delegation.get('question', '')})" + print(line, flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/joyvl_interaction/cli/stream_client.py b/examples/online_serving/joyvl_interaction/cli/stream_client.py new file mode 100644 index 00000000000..82053021c10 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/cli/stream_client.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import base64 +import time +from collections.abc import Iterator +from dataclasses import dataclass + +import cv2 +import requests + + +@dataclass +class Tick: + index: int + t: float + action: str + text: str + skipped: bool + latency_ms: float + delegation: dict | None + + +def iter_frames(video_path: str, fps: float) -> Iterator[tuple[float, bytes]]: + cap = cv2.VideoCapture(video_path) + src_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + step = max(1, round(src_fps / fps)) + idx = 0 + while True: + ok, frame = cap.read() + if not ok: + break + if idx % step == 0: + ok, buf = cv2.imencode(".jpg", frame) + if ok: + yield idx / src_fps, buf.tobytes() + idx += 1 + cap.release() + + +def stream( + video_path: str, + server: str, + session_id: str, + query: str | None = None, + fps: float = 1.0, + realtime: bool = False, +) -> Iterator[Tick]: + """Send each sampled frame as one chat completion; yield the parsed action. + + ``query`` is attached to the first tick to arm the session; inject later + queries by posting another frame turn with a new text part. + """ + url = server.rstrip("/") + "/v1/chat/completions" + requests.post(server.rstrip("/") + "/reset", json={"session_id": session_id}, timeout=10) + for i, (t, jpeg) in enumerate(iter_frames(video_path, fps)): + data_url = "data:image/jpeg;base64," + base64.b64encode(jpeg).decode() + content = [{"type": "image_url", "image_url": {"url": data_url}}] + if i == 0 and query: + content.insert(0, {"type": "text", "text": query}) + resp = requests.post( + url, + json={"session_id": session_id, "messages": [{"role": "user", "content": content}]}, + timeout=120, + ).json() + info = resp.get("interaction", {}) + yield Tick( + index=info.get("frame_index", i + 1), + t=t, + action=info.get("action", "?"), + text=resp["choices"][0]["message"]["content"], + skipped=info.get("inference_skipped", False), + latency_ms=info.get("latency_ms", 0.0), + delegation=info.get("delegation"), + ) + if realtime: + time.sleep(1.0 / fps) diff --git a/examples/online_serving/joyvl_interaction/rtsp/README.md b/examples/online_serving/joyvl_interaction/rtsp/README.md new file mode 100644 index 00000000000..6f6fcc73ddb --- /dev/null +++ b/examples/online_serving/joyvl_interaction/rtsp/README.md @@ -0,0 +1,46 @@ +# RTSP local streaming + +Simulate an RTSP camera from a local video file so the WebUI (or any RTSP client) can +test the RTSP input path without a physical IP camera. RTSP is a **webui-side input** — +the browser pulls the stream and feeds frames to the orchestrator over the normal API; +no serving-layer code is involved. + +Needs [`ffmpeg`](https://ffmpeg.org/) and [MediaMTX](https://github.com/bluenviron/mediamtx/releases) +(download the binary + `mediamtx.yml` and keep them next to these scripts, or point +`MEDIAMTX_BIN` / `MEDIAMTX_CONFIG` at them). + +## 1. Start the RTSP server + +```bash +bash ./mediamtx.sh # MediaMTX, listens on :8554 +``` + +## 2. Push a local video as an RTSP stream + +```bash +bash ./rtsp.sh ./videos/example.mp4 rtsp://127.0.0.1:8554/fire1 +``` + +`rtsp.sh` takes `[video-path] [rtsp-url]` (or set `VIDEO_PATH` / `RTSP_URL`). To stream a +whole directory of clips, one stream per file: + +```bash +bash ./rtsp_all.sh ./videos rtsp://127.0.0.1:8554 +``` + +## 3. Use it in the WebUI + +Enter the URL in the WebUI RTSP box: + +``` +rtsp://127.0.0.1:8554/fire1 +``` + +If the WebUI runs on another machine, replace `127.0.0.1` with the MediaMTX host's IP. + +## Notes + +- `rtsp.sh` transcodes audio to AAC; replace `-c:a aac -b:a 128k -ar 44100` with `-an` + if the source clip has no audio track and your `ffmpeg` errors on it. +- If a client cannot connect, check that port `8554` is reachable and not blocked by a + firewall, container network, or security group. diff --git a/examples/online_serving/joyvl_interaction/rtsp/mediamtx.sh b/examples/online_serving/joyvl_interaction/rtsp/mediamtx.sh new file mode 100755 index 00000000000..1af9c0b6360 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/rtsp/mediamtx.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Start a local MediaMTX RTSP server (listens on :8554 by default). +# Download the binary + mediamtx.yml from https://github.com/bluenviron/mediamtx/releases +# and either keep them next to this script or point the vars below at them. + +MEDIAMTX_BIN="${MEDIAMTX_BIN:-./mediamtx}" +MEDIAMTX_CONFIG="${MEDIAMTX_CONFIG:-./mediamtx.yml}" + +exec "$MEDIAMTX_BIN" "$MEDIAMTX_CONFIG" diff --git a/examples/online_serving/joyvl_interaction/rtsp/rtsp.sh b/examples/online_serving/joyvl_interaction/rtsp/rtsp.sh new file mode 100755 index 00000000000..03f8052dd41 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/rtsp/rtsp.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Loop a local video file and push it to an RTSP URL, simulating an RTSP camera. +# bash ./rtsp.sh [video-path] [rtsp-output-url] +# Defaults stream ./videos/example.mp4 to rtsp://127.0.0.1:8554/fire1 + +VIDEO_PATH="${1:-${VIDEO_PATH:-./videos/example.mp4}}" +RTSP_URL="${2:-${RTSP_URL:-rtsp://127.0.0.1:8554/fire1}}" + +if [[ ! -f "$VIDEO_PATH" ]]; then + echo "Video file not found: $VIDEO_PATH" >&2 + exit 1 +fi + +echo "Input video: $VIDEO_PATH" +echo "RTSP output: $RTSP_URL" + +# Drop "-c:a aac -b:a 128k -ar 44100" for "-an" if the source has no audio track. +exec ffmpeg \ + -hide_banner \ + -loglevel info \ + -re \ + -stream_loop -1 \ + -i "$VIDEO_PATH" \ + -vf "scale='min(1280,iw)':-2" \ + -c:v libx264 \ + -preset veryfast \ + -tune zerolatency \ + -b:v 2500k \ + -maxrate 2500k \ + -bufsize 5000k \ + -c:a aac \ + -b:a 128k \ + -ar 44100 \ + -f rtsp \ + -rtsp_transport udp \ + "$RTSP_URL" diff --git a/examples/online_serving/joyvl_interaction/rtsp/rtsp_all.sh b/examples/online_serving/joyvl_interaction/rtsp/rtsp_all.sh new file mode 100755 index 00000000000..b73f1d87168 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/rtsp/rtsp_all.sh @@ -0,0 +1,64 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Push every .mp4 in a directory as its own RTSP stream (one ffmpeg per file), +# named rtsp://<base>/<filename-without-extension>. +# bash ./rtsp_all.sh [video-dir] [rtsp-base-url] + +VIDEO_DIR="${1:-./videos}" +RTSP_BASE="${2:-rtsp://127.0.0.1:8554}" + +if [[ ! -d "$VIDEO_DIR" ]]; then + echo "Video directory not found: $VIDEO_DIR" >&2 + exit 1 +fi + +shopt -s nullglob +VIDEO_FILES=("$VIDEO_DIR"/*.mp4) +shopt -u nullglob + +if [[ ${#VIDEO_FILES[@]} -eq 0 ]]; then + echo "No video files found in: $VIDEO_DIR" >&2 + exit 1 +fi + +echo "Found ${#VIDEO_FILES[@]} video(s) in $VIDEO_DIR" + +PIDS=() +cleanup() { + echo "Stopping all streams..." + for pid in "${PIDS[@]}"; do + kill "$pid" 2>/dev/null || true + done + wait 2>/dev/null +} +trap cleanup EXIT INT TERM + +for video in "${VIDEO_FILES[@]}"; do + stream_name="$(basename "${video%.*}")" + rtsp_url="${RTSP_BASE}/${stream_name}" + echo "Starting stream: $rtsp_url <- $(basename "$video")" + ffmpeg \ + -hide_banner \ + -loglevel error \ + -re \ + -stream_loop -1 \ + -i "$video" \ + -vf "format=yuv420p,scale='min(1280,iw)':'min(720,ih)':force_original_aspect_ratio=decrease:force_divisible_by=2" \ + -c:v libx264 \ + -preset veryfast \ + -tune zerolatency \ + -b:v 4000k \ + -maxrate 4500k \ + -bufsize 9000k \ + -c:a aac \ + -b:a 128k \ + -ar 44100 \ + -f rtsp \ + -rtsp_transport tcp \ + "$rtsp_url" & + PIDS+=($!) +done + +echo "All ${#VIDEO_FILES[@]} streams started. Press Ctrl+C to stop." +wait diff --git a/examples/online_serving/joyvl_interaction/scripts/start_all.sh b/examples/online_serving/joyvl_interaction/scripts/start_all.sh new file mode 100755 index 00000000000..d2f512b26f0 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/scripts/start_all.sh @@ -0,0 +1,108 @@ +#!/bin/bash +# One-shot JoyVL demo: model -> interaction orchestrator -> JD webui. +# +# The model is the only fixed component; every external module is pluggable and +# configured here via env (unset = disabled), matching JoyVL's design: +# +# TTS_URL voice out (point at our tts_bridge, or any OpenAI-Realtime-style TTS WS) +# ASR_URL voice in (point at our asr_bridge, or any compatible ASR WS) +# DELEGATION_BACKEND_URL orchestrator background brain for </delegation> (OpenAI-compatible; DELEGATION_MODEL/DELEGATION_API_KEY pick model+key) +# DELEGATION_KIND=chat (default) text/VL brain answers; =image a text-to-image model generates a picture +# chat+Claude: DELEGATION_BACKEND_URL=https://api.anthropic.com/v1/ DELEGATION_MODEL=claude-... DELEGATION_API_KEY=sk-ant-... +# image+Qwen-Image: DELEGATION_KIND=image DELEGATION_BACKEND_URL=http://127.0.0.1:8091/v1 +# BACKGROUND_CODEX_API_URL webui-side background agent +# ENABLE_MEMORY (default 1) 3-tier memory; SUMMARIZER_BACKEND_URL/MODEL pick the summarizer (default: reuse main) +# +# Usage: bash start_all.sh +# TTS_URL=ws://host/v1/tts ASR_URL=ws://host/v1/asr GPU=0 bash start_all.sh +set -o pipefail + +CODE="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)" +VENV="${VENV:-$CODE/.venv}" +WEBUI_DIR="${WEBUI_DIR:-$CODE/joyvl-interaction/joyvl-interaction-webui}" +MODEL="${MODEL:-jdopensource/JoyAI-VL-Interaction-Preview}" +SERVED_NAME="${SERVED_NAME:-JoyAI-VL-Interaction-Preview}" +GPU="${GPU:-0}" +MODEL_PORT="${MODEL_PORT:-8061}" +ORCH_PORT="${ORCH_PORT:-8070}" +WEBUI_PORT="${WEBUI_PORT:-8999}" +WEBUI_USERNAME="${WEBUI_USERNAME:-admin}" +WEBUI_PASSWORD="${WEBUI_PASSWORD:-changeme}" # override via env for any exposed deployment +ENABLE_MEMORY="${ENABLE_MEMORY:-1}" +RESPONSE_DEDUP="${RESPONSE_DEDUP:-0.85}" # demo drops near-duplicate narration; set 1.0 for reference (exact-only) +TTS_URL="${TTS_URL:-}" +ASR_URL="${ASR_URL:-}" +BACKGROUND_CODEX_API_URL="${BACKGROUND_CODEX_API_URL:-}" + +LOG=/tmp/joyvl_demo +mkdir -p "$LOG" + +[ -d "$WEBUI_DIR" ] || { echo "! WEBUI_DIR not found: $WEBUI_DIR (set WEBUI_DIR=...)"; exit 1; } + +kill_port() { + local p + p=$(ss -tlnp 2>/dev/null | grep ":$1 " | grep -oE 'pid=[0-9]+' | head -1 | cut -d= -f2) + [ -n "${p:-}" ] && kill "$p" 2>/dev/null + return 0 +} +wait_http() { + for _ in $(seq 1 150); do + curl -sk -m2 "$1" -o /dev/null 2>/dev/null && return 0 + sleep 2 + done + echo " ! timeout waiting for $1" + return 1 +} + +echo "Resetting ports ${MODEL_PORT} ${ORCH_PORT} ${WEBUI_PORT}…" +for p in "$MODEL_PORT" "$ORCH_PORT" "$WEBUI_PORT"; do kill_port "$p"; done +sleep 3 + +[ -f "$VENV/bin/activate" ] && source "$VENV/bin/activate" + +# detach into a new session so servers outlive this launcher +spawn() { local log="$1"; shift; setsid "$@" > "$log" 2>&1 < /dev/null & } + +echo "[1/3] model on GPU ${GPU}…" +spawn "$LOG/model.log" env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 CUDA_VISIBLE_DEVICES="$GPU" \ + vllm serve "$MODEL" --served-model-name "$SERVED_NAME" --port "$MODEL_PORT" \ + --gpu-memory-utilization 0.85 --max-model-len 131072 --enable-prefix-caching \ + --limit-mm-per-prompt '{"image":256,"video":1}' +wait_http "http://127.0.0.1:${MODEL_PORT}/health" && echo " up :${MODEL_PORT}" + +echo "[2/3] orchestrator (memory=${ENABLE_MEMORY})…" +ORCH_ARGS=(--port "$ORCH_PORT" --main-backend-url "http://127.0.0.1:${MODEL_PORT}/v1" + --main-model "$SERVED_NAME" + --response-dedup-threshold "$RESPONSE_DEDUP") +# delegation: point at a stronger background brain to enable it; otherwise keep it off +if [ -n "${DELEGATION_BACKEND_URL:-}" ]; then + ORCH_ARGS+=(--delegation-backend-url "$DELEGATION_BACKEND_URL") + ORCH_ARGS+=(--delegation-kind "${DELEGATION_KIND:-chat}") + [ -n "${DELEGATION_MODEL:-}" ] && ORCH_ARGS+=(--delegation-model "$DELEGATION_MODEL") + [ -n "${DELEGATION_API_KEY:-}" ] && ORCH_ARGS+=(--delegation-api-key "$DELEGATION_API_KEY") +else + ORCH_ARGS+=(--no-delegation) +fi +[ "$ENABLE_MEMORY" = "0" ] && ORCH_ARGS+=(--no-memory) +[ -n "${SUMMARIZER_BACKEND_URL:-}" ] && ORCH_ARGS+=(--summarizer-backend-url "$SUMMARIZER_BACKEND_URL") +[ -n "${SUMMARIZER_MODEL:-}" ] && ORCH_ARGS+=(--summarizer-model "$SUMMARIZER_MODEL") +spawn "$LOG/orch.log" env PYTHONPATH="$CODE" python -m vllm_omni.experimental.fullduplex.joyvl.serving.server "${ORCH_ARGS[@]}" +wait_http "http://127.0.0.1:${ORCH_PORT}/health" && echo " up :${ORCH_PORT}" + +echo "[3/3] JD webui…" +spawn "$LOG/webui.log" env WEBUI_USERNAME="$WEBUI_USERNAME" WEBUI_PASSWORD="$WEBUI_PASSWORD" \ + TTS_URL="$TTS_URL" ASR_URL="$ASR_URL" BACKGROUND_CODEX_API_URL="$BACKGROUND_CODEX_API_URL" \ + bash -c "cd '$WEBUI_DIR' && exec bash scripts/start_server.sh --api-base 'http://127.0.0.1:${ORCH_PORT}/v1'" +webui_up=0 +for _ in $(seq 1 60); do grep -q "Server startup complete" "$LOG/webui.log" 2>/dev/null && { webui_up=1; break; }; sleep 1; done +[ "$webui_up" = 1 ] || { echo "! webui did not start — see $LOG/webui.log"; exit 1; } + +IP=$(hostname -I 2>/dev/null | awk '{print $1}') +echo +echo "=== JoyVL demo ready ===" +echo " webui : https://${IP}:${WEBUI_PORT} (user ${WEBUI_USERNAME} / pass ${WEBUI_PASSWORD})" +echo " memory : $([ "$ENABLE_MEMORY" = 0 ] && echo off || echo on)" +echo " TTS : ${TTS_URL:-<unset>}" +echo " ASR : ${ASR_URL:-<unset>}" +echo " agent : ${BACKGROUND_CODEX_API_URL:-<unset>}" +echo " logs : ${LOG}/{model,orch,webui}.log" diff --git a/examples/online_serving/joyvl_interaction/scripts/start_asr.sh b/examples/online_serving/joyvl_interaction/scripts/start_asr.sh new file mode 100755 index 00000000000..cbbb9547e2e --- /dev/null +++ b/examples/online_serving/joyvl_interaction/scripts/start_asr.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# Serve Qwen3-ASR (the same audio-LLM AURA uses as its ASR stage) and run the +# webui<->ASR bridge. Then start the webui with: +# ASR_URL=ws://127.0.0.1:8093/v1/asr bash scripts/start_server.sh +set -euo pipefail + +ASR_MODEL="${ASR_MODEL:-Qwen/Qwen3-ASR-1.7B}" +ASR_PORT="${ASR_PORT:-8094}" +BRIDGE_PORT="${BRIDGE_PORT:-8093}" +GPU="${GPU:-2}" +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +# Qwen3-ASR is audio-in / text-out; served as a normal multimodal chat model. +CUDA_VISIBLE_DEVICES="${GPU}" vllm serve "${ASR_MODEL}" \ + --host 0.0.0.0 --port "${ASR_PORT}" \ + --gpu-memory-utilization 0.3 --trust-remote-code & +ASR_PID=$! +trap 'kill ${ASR_PID} 2>/dev/null || true' EXIT + +until curl -s "http://127.0.0.1:${ASR_PORT}/health" >/dev/null 2>&1; do sleep 2; done +python "${SCRIPT_DIR}/../bridges/asr_bridge.py" \ + --port "${BRIDGE_PORT}" \ + --backend-url "http://127.0.0.1:${ASR_PORT}" \ + --model "${ASR_MODEL}" diff --git a/examples/online_serving/joyvl_interaction/scripts/start_interaction.sh b/examples/online_serving/joyvl_interaction/scripts/start_interaction.sh new file mode 100644 index 00000000000..b97e5e9c6e7 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/scripts/start_interaction.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Launch the interaction orchestrator in front of a running model server. +set -euo pipefail + +PORT="${PORT:-8070}" +MAIN_BACKEND_URL="${MAIN_BACKEND_URL:-http://127.0.0.1:8061/v1}" +MAIN_MODEL="${MAIN_MODEL:-JoyAI-VL-Interaction-Preview}" +PERSONA="${PERSONA:-default}" + +# By default memory reuses the main model as its own summarizer. Set +# SUMMARIZER_BACKEND_URL/SUMMARIZER_MODEL to a dedicated Qwen3-VL-4B server, +# or pass --no-memory for the lightest setup. +EXTRA=() +[[ -n "${SUMMARIZER_BACKEND_URL:-}" ]] && EXTRA+=(--summarizer-backend-url "${SUMMARIZER_BACKEND_URL}") +[[ -n "${SUMMARIZER_MODEL:-}" ]] && EXTRA+=(--summarizer-model "${SUMMARIZER_MODEL}") +[[ "${NO_MEMORY:-0}" == "1" ]] && EXTRA+=(--no-memory) + +python -m vllm_omni.experimental.fullduplex.joyvl.serving.server \ + --port "${PORT}" \ + --main-backend-url "${MAIN_BACKEND_URL}" \ + --main-model "${MAIN_MODEL}" \ + --persona "${PERSONA}" \ + "${EXTRA[@]}" diff --git a/examples/online_serving/joyvl_interaction/scripts/start_model.sh b/examples/online_serving/joyvl_interaction/scripts/start_model.sh new file mode 100644 index 00000000000..4cfa95d539e --- /dev/null +++ b/examples/online_serving/joyvl_interaction/scripts/start_model.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# Serve the JoyVL interaction model as a plain VLM (NOT --omni: it is a +# standard Qwen3-VL autoregressive model, not an omni/diffusion model). +set -euo pipefail + +MODEL="${MODEL:-jdopensource/JoyAI-VL-Interaction-Preview}" +SERVED_NAME="${SERVED_NAME:-JoyAI-VL-Interaction-Preview}" +PORT="${PORT:-8061}" +GPU="${GPU:-0}" +# Covers the default short-term window (chunk_frames=100) + key_frames=0 (all frames +# summarized). Lower MAX_MODEL_LEN / IMAGE_LIMIT (and --chunk-frames) together for ~24GB. +MAX_MODEL_LEN="${MAX_MODEL_LEN:-131072}" +IMAGE_LIMIT="${IMAGE_LIMIT:-256}" + +CUDA_VISIBLE_DEVICES="${GPU}" vllm serve "${MODEL}" \ + --served-model-name "${SERVED_NAME}" \ + --port "${PORT}" \ + --max-model-len "${MAX_MODEL_LEN}" \ + --gpu-memory-utilization 0.85 \ + --enable-prefix-caching \ + --enable-chunked-prefill \ + --limit-mm-per-prompt "{\"image\":${IMAGE_LIMIT},\"video\":1}" diff --git a/examples/online_serving/joyvl_interaction/scripts/start_tts.sh b/examples/online_serving/joyvl_interaction/scripts/start_tts.sh new file mode 100755 index 00000000000..2c6c2ab3536 --- /dev/null +++ b/examples/online_serving/joyvl_interaction/scripts/start_tts.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Serve Qwen3-TTS (vLLM-Omni) and run the webui<->TTS bridge. +# Then start the webui with: TTS_URL=ws://127.0.0.1:8092/v1/tts bash scripts/start_server.sh +set -euo pipefail + +TTS_MODEL="${TTS_MODEL:-Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice}" +TTS_PORT="${TTS_PORT:-8091}" +BRIDGE_PORT="${BRIDGE_PORT:-8092}" +GPU="${GPU:-1}" +VOICE="${VOICE:-vivian}" +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" + +CUDA_VISIBLE_DEVICES="${GPU}" vllm-omni serve "${TTS_MODEL}" \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ + --host 0.0.0.0 --port "${TTS_PORT}" \ + --gpu-memory-utilization 0.9 --trust-remote-code --omni & +TTS_PID=$! +trap 'kill ${TTS_PID} 2>/dev/null || true' EXIT + +# wait for the TTS server, then run the bridge in the foreground +until curl -s "http://127.0.0.1:${TTS_PORT}/health" >/dev/null 2>&1; do sleep 2; done +python "${SCRIPT_DIR}/../bridges/tts_bridge.py" \ + --port "${BRIDGE_PORT}" \ + --backend-url "ws://127.0.0.1:${TTS_PORT}/v1/audio/speech/stream" \ + --voice "${VOICE}" diff --git a/examples/online_serving/lance/README.md b/examples/online_serving/lance/README.md index 5f2febce924..614f4990c0e 100644 --- a/examples/online_serving/lance/README.md +++ b/examples/online_serving/lance/README.md @@ -1,4 +1,4 @@ -# Lance — online serving +# Lance: Online serving OpenAI-compatible chat completions API for Lance, served by `vllm-omni`. diff --git a/examples/online_serving/mimo_audio/README.md b/examples/online_serving/mimo_audio/README.md index 9c1be7f21c8..7c5e2db0f05 100644 --- a/examples/online_serving/mimo_audio/README.md +++ b/examples/online_serving/mimo_audio/README.md @@ -1,4 +1,4 @@ -# Online serving Example of vLLM-Omni for MiMo-Audio +# MiMo-Audio: Online serving ## 🛠️ Installation diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md index 76b0bef9df7..377acc566e7 100644 --- a/examples/online_serving/ming_flash_omni/README.md +++ b/examples/online_serving/ming_flash_omni/README.md @@ -1,4 +1,4 @@ -# Ming-flash-omni 2.0 +# Ming-flash-omni 2.0: Online serving ## Installation @@ -49,22 +49,14 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene --modalities text ``` -Parameterized curl wrapper in this directory: -```bash -bash run_curl_multimodal_generation.sh text -bash run_curl_multimodal_generation.sh use_image -bash run_curl_multimodal_generation.sh use_audio -bash run_curl_multimodal_generation.sh use_video -bash run_curl_multimodal_generation.sh use_mixed_modalities -bash run_curl_multimodal_generation.sh use_image_gen -``` +## Image generation (text-to-image / img2img) -## Image generation (text-to-image) +Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker. -Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. -Launch with the image deploy YAML, which adds an image-gen stage behind -the thinker: +The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below). + +### Launch ```bash vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \ @@ -74,73 +66,143 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \ --port 8091 ``` -Then request image output by passing `"modalities": ["image"]`: + +### Text-to-image + +Request image output with `"modalities": ["image"]`: ```bash curl -s http://127.0.0.1:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "messages": [{"role": "user", "content": "Please draw a cute cat."}], - "modalities": ["image"] - }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen.png + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + { + "role": "user", + "content": "Please draw a cute cat." + } + ], + "modalities": ["image"] + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_imagegen.png ``` -### Optional knobs +Pass generation knobs under a literal `extra_body` object (the OpenAI Python client's `extra_body=` kwarg produces the same request). Keys are filtered against the declared set and routed into every stage's `extra_args`: -Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`: - -| Key | Default | Description | -| --- | --- | --- | -| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). | -| `steps` | 30 | Number of FlowMatchEuler denoise steps. | -| `cfg` | 2.0 | Classifier-free guidance scale. | -| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). | -| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. | -| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). | +```bash +curl http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "modalities": ["image"], + "messages": [ + { + "role": "user", + "content": "Draw a poster." + } + ], + "extra_body": { + "steps": 6, + "cfg": 1.5, + "height": 512, + "width": 512, + "seed": 123, + "byte5_text": ["理解与生成统一"], + "negative_prompt": "ugly, blurry, distorted" + } + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_imagegen_extra_body.png +``` -Example with all the knobs: +NOTE: `extra_body` does **not** set the thinker's own sampling params. To tune the thinker (stage-0) sampling (`temperature` / `top_p` / `top_k` / `max_tokens`) or place knobs explicitly per stage, use `sampling_params_list` (`[thinker, imagegen]`). +`negative_prompt` must sit on the **stage-0** entry to trigger the real-CFG companion; the imagegen knobs go on **stage-1**: ```bash curl http://127.0.0.1:8091/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "modalities": ["image"], - "sampling_params_list": [ - {"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42, - "extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}}, - {"seed":42,"extra_args":{"image_gen":{ - "steps":6,"cfg":1.5,"height":512,"width":512,"seed":123, - "byte5_text":["理解与生成统一"]}}} - ], - "messages": [{"role": "user", "content": "Draw a poster."}] - }' | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' | base64 -d > ming_imagegen_knobs.png + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "modalities": ["image"], + "sampling_params_list": [ + { + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "max_tokens": 1, + "seed": 42, + "extra_args": { + "negative_prompt": "ugly, blurry, distorted" + } + }, + { + "seed": 42, + "extra_args": { + "steps": 6, + "cfg": 1.5, + "height": 512, + "width": 512, + "seed": 123, + "byte5_text": ["理解与生成统一"] + } + } + ], + "messages": [ + { + "role": "user", + "content": "Draw a poster." + } + ] + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_imagegen_knobs.png ``` ### img2img (reference image) -Add an `image_url` content part to the user message; the parser routes it -into the diffusion stage as `extra[reference_image]`: - -```jsonc -"messages": [{ - "role": "user", - "content": [ - {"type": "text", "text": "Change the background to a sandy beach at sunset."}, - {"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}} - ] -}] +Add an `image_url` content part (a base64 data URL) to the user message; it is routed into the DiT stage as `extra[reference_image]`. base64-encode a local file and stream it through `jq` via stdin - piping `base64 -> jq -> curl` so that avoids the shell `ARG_MAX` limit that inlining a large base64 string in `-d '…'` hits. + +```bash +# Reference image: figures/cases/person_gen_05.png from the upstream Ming repo +# Check https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/examples/vllm_demo.py +wget https://raw.githubusercontent.com/inclusionAI/Ming/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/person_gen_05.png + +base64 -w0 person_gen_05.png \ +| jq -Rs --arg prompt "Put a pair of sunglasses on the person." '{ + model: "Jonathan1909/Ming-flash-omni-2.0", + modalities: ["image"], + messages: [ + { + role: "user", + content: [ + { type: "text", text: $prompt }, + { type: "image_url", image_url: { url: ("data:image/png;base64," + (. | rtrimstr("\n"))) } } + ] + } + ] + }' \ +| curl http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d @- \ +| jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ +| base64 -d > ming_img2img.png ``` -### GPU layout +The reference image can also be a public URL (`"url": "https://…/photo.jpg"`) or the simplified `{"image": "<base64>"}` content-part form — see the [image-to-image request formats](../image_to_image/README.md#request-format). + +### Knobs (declared `extra_body` params) + +| Key | Default | Description | +| --- | --- | --- | +| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). | +| `steps` | 30 | Number of FlowMatchEuler denoise steps. | +| `cfg` | 2.0 | Classifier-free guidance scale. | +| `seed` | 42 | Per-request RNG seed. | +| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. | +| `negative_prompt` | (empty) | Real CFG negative conditioning (text-to-image only). | -The shipped `ming_flash_omni_image.yaml` allocates the thinker on GPUs 0–3 -(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit -`devices` per stage to relocate; with fewer GPUs available, drop -the thinker TP to 2 and run the diffusion stage on a free card. Image-gen -warmup takes roughly an extra 30–60 s on top of the thinker — set -`--stage-init-timeout 1800` if the default 300 s is too tight. +For the offline `text_to_image.py` / `image_edit.py` scripts and the full knob reference, see the [image-generation section in the recipe](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md#image-generation-text-to-image--img2img). ## Modality control diff --git a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh deleted file mode 100755 index 92bb32c471b..00000000000 --- a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh +++ /dev/null @@ -1,169 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -# Server port -PORT="${PORT:-8091}" -# Default query type -QUERY_TYPE="${1:-text}" - -# Validate query type -if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then - echo "Error: Invalid query type '$QUERY_TYPE'" - echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]" - echo " text: Text-only query" - echo " use_audio: Audio + Text query" - echo " use_image: Image + Text query" - echo " use_video: Video + Text query" - echo " use_mixed_modalities: Audio + Image + Video + Text query" - echo " use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally" - exit 1 -fi - -# Define URLs for assets -MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" -CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg" -SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" - -# Build user content based on query type -case "$QUERY_TYPE" in - text) - user_content='[ - { - "type": "text", - "text": "请详细介绍鹦鹉的生活习性。" - } - ]' - ;; - use_image) - user_content='[ - { - "type": "image_url", - "image_url": { - "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" - } - }, - { - "type": "text", - "text": "Describe this image in detail." - } - ]' - ;; - use_audio) - user_content='[ - { - "type": "audio_url", - "audio_url": { - "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" - } - }, - { - "type": "text", - "text": "Please recognize the language of this speech and transcribe it. Format: oral." - } - ]' - ;; - use_video) - user_content='[ - { - "type": "video_url", - "video_url": { - "url": "'"$SAMPLE_VIDEO_URL"'" - } - }, - { - "type": "text", - "text": "Describe what is happening in this video." - } - ]' - ;; - use_mixed_modalities) - user_content='[ - { - "type": "image_url", - "image_url": { - "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" - } - }, - { - "type": "audio_url", - "audio_url": { - "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" - } - }, - { - "type": "text", - "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral" - } - ]' - ;; -esac - -echo "Running query type: $QUERY_TYPE" -echo "" - -if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then - # Image-gen branch: dual-stage diffusion endpoint, modalities=["image"]. - # The shipped ming_flash_omni_image.yaml runs the thinker on cards 0-3 - # (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen - # section in README.md for knobs and img2img. - out_path="${OUT_PATH:-/tmp/ming_imagegen.png}" - request_body=$(cat <<'EOF' -{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "modalities": ["image"], - "messages": [ - {"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."} - ] -} -EOF -) - output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \ - -X POST "http://localhost:${PORT}/v1/chat/completions" \ - -H "Content-Type: application/json" \ - -d "$request_body") - # Server returns base64-encoded PNG either as a data URL in - # ``choices[0].message.content[0].image_url.url`` or as a raw string. - b64=$(echo "$output" | jq -r '.choices[0].message.content - | if type == "array" - then (map(select(.image_url? != null)) | .[0].image_url.url // "") - elif type == "string" then . else "" end') - if [[ -z "$b64" ]]; then - echo "Error: no image returned. Raw response:" - echo "$output" | jq '.' - exit 1 - fi - if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi - printf '%s' "$b64" | base64 --decode > "$out_path" - echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)" - exit 0 -fi - -request_body=$(cat <<EOF -{ - "model": "Jonathan1909/Ming-flash-omni-2.0", - "modalities": ["text"], - "messages": [ - { - "role": "system", - "content": [ - { - "type": "text", - "text": "你是一个友好的AI助手。\n\ndetailed thinking off" - } - ] - }, - { - "role": "user", - "content": $user_content - } - ] -} -EOF -) - -output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \ - -X POST http://localhost:${PORT}/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d "$request_body") - -echo "Output of request: $(echo "$output" | jq '.choices[0].message.content')" diff --git a/examples/online_serving/minicpmo/README.md b/examples/online_serving/minicpmo/README.md index a3476465255..334b5ed1154 100644 --- a/examples/online_serving/minicpmo/README.md +++ b/examples/online_serving/minicpmo/README.md @@ -11,17 +11,18 @@ The UI supports: ## 1. Start the backend server The deploy config auto-loads via `--omni`; the default -`vllm_omni/deploy/minicpmo_4_5.yaml` targets a 2-GPU layout (thinker on -GPU 0, talker + t2w sharing GPU 1). For other hardware layouts pick +`vllm_omni/deploy/minicpmo_4_5.yaml` targets a single-GPU layout (thinker +and talker + t2w co-located on GPU 0). For other hardware layouts pick one of the deploy variants below. | deploy config | GPUs | Notes | |---|---|---| -| `minicpmo_4_5.yaml` (default) | 2 | Thinker on GPU0, talker+t2w on GPU1. | +| `minicpmo_4_5.yaml` (default) | 1 | Thinker and talker+t2w co-located on GPU0. | +| `minicpmo_4_5_2gpu.yaml` | 2 | Thinker on GPU0, talker+t2w on GPU1. | | `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. | | `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. | -Default (2-GPU): +Default (single-GPU): ```bash vllm serve openbmb/MiniCPM-o-4_5 --omni \ diff --git a/examples/online_serving/qwen2_5_omni/README.md b/examples/online_serving/qwen2_5_omni/README.md index c528732064a..48406003d98 100644 --- a/examples/online_serving/qwen2_5_omni/README.md +++ b/examples/online_serving/qwen2_5_omni/README.md @@ -1,4 +1,4 @@ -# Qwen2.5-Omni +# Qwen2.5-Omni: Online serving ## 🛠️ Installation diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index a13ab52dfd6..74a18af322c 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -1,4 +1,4 @@ -# Qwen3-Omni +# Qwen3-Omni: Online serving ## 🛠️ Installation @@ -24,12 +24,15 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ --deploy-config /path/to/your_deploy_config.yaml ``` -For the bundled 3x-GPU multi-replica layout (talker/code2wav scale-out), -use: +For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2), +use `--stage-overrides` on top of the default config: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ - --deploy-config vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml + --stage-overrides '{ + "1": {"num_replicas": 2, "devices": "1,2"}, + "2": {"num_replicas": 2, "devices": "1,2"} + }' ``` ### Launch individual stages (stage-based CLI) @@ -350,13 +353,16 @@ curl http://localhost:8091/v1/chat/completions \ #### Text + Audio ```bash -curl -s http://localhost:8091/v1/chat/completions \ +response=$(curl -s http://localhost:8091/v1/chat/completions \ -H "Content-Type: application/json" \ -d '{ "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", "messages": [{"role": "user", "content": "Describe vLLM in brief."}], - "modalities": ["audio"] - }' | jq -r '.choices[0].message.audio.data' | base64 -d > output.wav + "modalities": ["text", "audio"] + }') + +echo "$response" | jq -r '.choices[0].message.content' +echo "$response" | jq -r '.choices[1].message.audio.data' | base64 -d > output.wav ``` ### Using Python client @@ -453,8 +459,8 @@ response = client.chat.completions.create( modalities=["audio"], extra_body={"speaker": "chelsie"} ) -# Audio uses the specified speaker -print(response.choices[1].message.audio) +# Audio-only responses contain a single audio choice +print(response.choices[0].message.audio) ``` Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default. diff --git a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml index a32ca7d2dec..7ce88f80aac 100644 --- a/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml +++ b/examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml @@ -20,7 +20,7 @@ stage_args: trust_remote_code: true engine_output_type: text distributed_executor_backend: "mp" - enable_prefix_caching: false + enable_prefix_caching: true hf_config_name: thinker_config tensor_parallel_size: 2 final_output: true diff --git a/examples/online_serving/replica_data_parallel/README.md b/examples/online_serving/replica_data_parallel/README.md new file mode 100644 index 00000000000..4bc98d6f718 --- /dev/null +++ b/examples/online_serving/replica_data_parallel/README.md @@ -0,0 +1,91 @@ +# Replica Data Parallelism for Video Generation + +A runnable recipe + benchmark for **replica data parallelism** (replica-DP) on a +video DiT. Replica-DP runs `N` independent diffusion engine replicas, one per GPU, +and routes each request to a single replica. It scales **request throughput** +near-linearly; it does **not** change single-request latency (for that, use an +intra-request axis such as tensor/sequence/CFG parallelism). + +This example uses `Wan-AI/Wan2.2-TI2V-5B-Diffusers` with one GPU per replica. + +## Files + +| File | Purpose | +|------|---------| +| `wan2_2_ti2v_dp.yaml` | Diffusion stage config; `runtime.num_replicas` + `runtime.devices` drive the fan-out. | +| `run_server.sh` | Substitutes `NUM_REPLICAS` / `DEVICES` into the config and serves. | +| `bench_replica_dp.py` | Replica-agnostic load driver; reports throughput, latency, and a per-input isolation check. | + +## Run + +Start the server with the replica count you want (one GPU per replica): + +```bash +# baseline +NUM_REPLICAS=1 DEVICES=0 ./run_server.sh +# 2 replicas +NUM_REPLICAS=2 DEVICES=0,1 ./run_server.sh +# 4 replicas +NUM_REPLICAS=4 DEVICES=0,1,2,3 ./run_server.sh +``` + +In another shell, drive a fixed workload and read the throughput: + +```bash +python bench_replica_dp.py --url http://127.0.0.1:8098 \ + --num-requests 28 --concurrency 8 --label "replicas=4" +``` + +Re-run the server at `NUM_REPLICAS=1,2,4` (same client workload each time) and +compare `videos/min`. Pass `--baseline-tpm <N=1 value>` to print a scaling factor. + +### Isolation check + +Identical outputs under a *shared* prompt+seed prove nothing -- every request +would be identical regardless of isolation. Instead each request gets a distinct +input (a prompt from the pool + a unique seed `base_seed+i`), and we check that a +request's output is unchanged whether it runs alone or concurrently amid other +distinct requests. Capture a per-input baseline on `N=1`, then verify it under +`N>=2` concurrency: + +```bash +# 1) against NUM_REPLICAS=1 (serial), record each input's output hash +python bench_replica_dp.py --num-requests 8 --concurrency 1 \ + --write-baseline baseline.json + +# 2) against NUM_REPLICAS>=2, replay the SAME inputs concurrently and verify +python bench_replica_dp.py --num-requests 8 --concurrency 8 \ + --check-baseline baseline.json +``` + +Step 2 exits non-zero if any request's output differs from its own baseline (a +mismatch means cross-request state bleed or misrouting). Keep `--num-requests`, +`--base-seed`, and the prompt source identical between the two runs so the input +sets match. + +## Measured scaling + +Wan2.2-TI2V-5B (832×480, 33 frames, 30 steps), 4× A800-80GB (NVLink), one GPU per replica: + +| Replicas | Throughput (videos/min) | Scaling | Efficiency | +|----------|-------------------------|---------|------------| +| 1 | 4.71 | 1.00× | — | +| 2 | 9.20 | 1.95× | 98% | +| 4 | 18.02 | 3.83× | 96% | + +Per-request latency stayed flat (~13–14 s) across all replica counts. Isolation +is checked per-input via the baseline/verify protocol above: each request's +output under `N>=2` concurrency must match its own single-replica baseline. + +## Notes on the config surface + +- `runtime.num_replicas: N` fans out `N` replicas; `runtime.devices` assigns their + GPUs. With `tensor_parallel_size = 1`, list one GPU per replica + (`num_replicas * tensor_parallel_size` entries, pool mode). +- Replica fan-out here comes from the stage config, not a CLI flag. (The headless / + multi-runtime launch path uses the process-local `--omni-dp-size-local`, which + requires `--stage-id`.) +- This recipe serves via `--stage-configs-path`. Expressing replica-DP under the + newer `--deploy-config` schema is part of the ongoing serving-config alignment + (see the diffusion parallelism guide and the `#4707` discussion); `num_replicas` + is accepted by both paths. diff --git a/examples/online_serving/replica_data_parallel/bench_replica_dp.py b/examples/online_serving/replica_data_parallel/bench_replica_dp.py new file mode 100755 index 00000000000..977f475ac9e --- /dev/null +++ b/examples/online_serving/replica_data_parallel/bench_replica_dp.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""Throughput-scaling + isolation benchmark for diffusion replica data parallelism. + +Sends `--num-requests` text-to-video requests (up to `--concurrency` in flight) +to a running vLLM-Omni video server via `POST /v1/videos/sync` (blocks until the +clip is produced) and reports: + + completed / failed, makespan (wall clock), throughput (videos/min), + per-request latency p50/p90/max (nearest-rank), and -- when a baseline + manifest is supplied -- a per-request isolation check. + +Each request gets a *distinct* input (a prompt from the pool + a unique seed), +so an output is a deterministic function of its own input and nothing else. + +Measuring replica scaling +------------------------- +This client is replica-agnostic. Vary the *server's* replica count +(`runtime.num_replicas` in the stage config; see run_server.sh) and keep the +client workload fixed, then compare throughput: + + N=1 -> baseline + N=2 -> expect ~2x + N=4 -> expect ~4x (near-linear is the goal) + +Pass the single-replica videos/min via `--baseline-tpm` to print a scaling +factor for the current run. + +Isolation / correctness +----------------------- +Byte-identical outputs under a *shared* prompt+seed prove nothing -- every +request would be identical regardless of isolation. Instead, give each request a +distinct input and check that its output is unchanged whether it runs alone or +concurrently amid other distinct requests: + + 1. Against an N=1 server, capture a per-input baseline: + bench_replica_dp.py --num-requests 8 --concurrency 1 \ + --write-baseline baseline.json + 2. Against an N>=2 server, replay the *same* inputs concurrently and verify + each output matches its own baseline: + bench_replica_dp.py --num-requests 8 --concurrency 8 \ + --check-baseline baseline.json + +A mismatch means a request's output changed under concurrency -- cross-request +state bleed or misrouting -- and the run exits non-zero. + +Requires: requests (`pip install requests`). +""" + +import argparse +import concurrent.futures as cf +import hashlib +import json +import math +import sys +import time +from pathlib import Path + +import requests + +# Distinct prompts, cycled across requests; combined with a per-request unique +# seed this makes every request a distinct input (see input_key()). +DEFAULT_PROMPTS = [ + "A cat playing piano, cinematic, high detail", + "A drone shot over a snowy mountain range at sunrise", + "A red sports car driving through a neon-lit city at night", + "Ocean waves crashing on a rocky shore, slow motion", + "A hot air balloon floating over green rolling hills", + "Timelapse of clouds moving over a desert canyon at dusk", + "A steam train crossing a stone viaduct in autumn", + "A jellyfish drifting through deep blue water, backlit", +] + + +def input_key(prompt, seed, args): + """Stable identity of a request's input; matches baseline<->verify by content, not order.""" + return "\x1f".join([prompt, str(seed), args.size, str(args.num_frames), str(args.steps)]) + + +def nearest_rank(xs, q): + """Nearest-rank percentile; q in [0, 1]. Correct for tiny samples (unlike int(q*n)-1).""" + if not xs: + return float("nan") + s = sorted(xs) + rank = max(1, math.ceil(q * len(s))) + return s[rank - 1] + + +def one_request(args, spec, idx, save): + """Fire one /v1/videos/sync request; return (ok, latency, nbytes, md5, key, err).""" + prompt, seed = spec["prompt"], spec["seed"] + key = input_key(prompt, seed, args) + # multipart form fields match POST /v1/videos/sync in the omni API server + form = { + "prompt": (None, prompt), + "size": (None, args.size), # e.g. "832x480" + "num_frames": (None, str(args.num_frames)), + "num_inference_steps": (None, str(args.steps)), + "seed": (None, str(seed)), # per-request unique -> distinct, reproducible input + } + if args.model: + form["model"] = (None, args.model) + if args.negative_prompt: + form["negative_prompt"] = (None, args.negative_prompt) + + t0 = time.perf_counter() + try: + r = requests.post(f"{args.url}/v1/videos/sync", files=form, timeout=args.timeout) + dt = time.perf_counter() - t0 + if r.status_code != 200: + return (False, dt, 0, "", key, f"HTTP {r.status_code}: {r.text[:160]}") + body = r.content + md5 = hashlib.md5(body).hexdigest() + if save and args.save_dir: + Path(args.save_dir).mkdir(parents=True, exist_ok=True) + (Path(args.save_dir) / f"req{idx:03d}_seed{seed}.mp4").write_bytes(body) + return (True, dt, len(body), md5, key, "") + except Exception as e: # noqa: BLE001 - report any client-side failure + return (False, time.perf_counter() - t0, 0, "", key, repr(e)) + + +def build_specs(args): + """Deterministic list of distinct (prompt, seed) inputs; identical across baseline/verify runs.""" + prompts = DEFAULT_PROMPTS + if args.prompts_file: + prompts = [ln.strip() for ln in Path(args.prompts_file).read_text().splitlines() if ln.strip()] + return [{"prompt": prompts[i % len(prompts)], "seed": args.base_seed + i} for i in range(args.num_requests)] + + +def run_batch(args, specs, save): + """Fire specs concurrently; return (results, makespan). results = list of one_request tuples.""" + results = [] + wall0 = time.perf_counter() + with cf.ThreadPoolExecutor(max_workers=args.concurrency) as ex: + futs = {ex.submit(one_request, args, spec, i, save): i for i, spec in enumerate(specs)} + for done, f in enumerate(cf.as_completed(futs), 1): + res = f.result() + results.append(res) + good, dt, nbytes, _md5, _key, err = res + print(f" [{done}/{len(specs)}] {'ok' if good else 'FAIL'} {dt:6.1f}s {(nbytes / 1e6):5.1f}MB {err}") + return results, time.perf_counter() - wall0 + + +def main(): + ap = argparse.ArgumentParser(description="Diffusion replica-DP throughput + isolation benchmark") + ap.add_argument("--url", default="http://127.0.0.1:8098", help="server base URL (no path)") + ap.add_argument("--num-requests", type=int, default=16, help="total requests to send") + ap.add_argument("--concurrency", type=int, default=4, help="max in-flight (set >= replica count)") + ap.add_argument("--warmup", type=int, default=-1, help="warmup requests before timing (default: =concurrency)") + ap.add_argument("--base-seed", type=int, default=42, help="request i uses seed base_seed+i (distinct inputs)") + ap.add_argument("--prompts-file", default="", help="optional file, one prompt per line (else a built-in pool)") + ap.add_argument("--negative-prompt", default="") + ap.add_argument("--size", default="832x480", help="WIDTHxHEIGHT") + ap.add_argument("--num-frames", type=int, default=33) + ap.add_argument("--steps", type=int, default=30, help="num_inference_steps") + ap.add_argument("--model", default="", help="optional; defaults to the served model") + ap.add_argument("--timeout", type=float, default=1800, help="per-request timeout (s)") + ap.add_argument("--save-dir", default="", help="save returned clips") + ap.add_argument("--label", default="", help="run label, e.g. 'replicas=2'") + ap.add_argument( + "--baseline-tpm", type=float, default=0.0, help="single-replica videos/min, to print a scaling factor" + ) + ap.add_argument("--write-baseline", default="", help="run against N=1: write per-input md5 manifest to this path") + ap.add_argument("--check-baseline", default="", help="run against N>=2: verify each output matches this manifest") + args = ap.parse_args() + + if args.write_baseline and args.check_baseline: + ap.error("--write-baseline and --check-baseline are mutually exclusive") + + specs = build_specs(args) + warmup = args.concurrency if args.warmup < 0 else args.warmup + + print( + f"-> {args.url}/v1/videos/sync | {args.num_requests} reqs " + f"| concurrency {args.concurrency} | warmup {warmup} | {args.size} {args.num_frames}f {args.steps}steps " + f"| {args.label}" + ) + + # Warmup: absorb each replica's first-call compile/init cost so it does not + # skew the measured makespan. Results are discarded and never saved. + if warmup > 0: + print(f"warming up ({warmup} reqs, not timed)...") + run_batch(args, [specs[i % len(specs)] for i in range(warmup)], save=False) + + results, makespan = run_batch(args, specs, save=True) + + ok = sum(1 for r in results if r[0]) + fail = len(results) - ok + lat = [r[1] for r in results if r[0]] + errs = [r[5] for r in results if not r[0]] + got = {r[4]: r[3] for r in results if r[0]} # input_key -> md5 + + print(f"\n===== result {args.label} =====") + print(f"completed {ok} / failed {fail} | makespan {makespan:.1f}s") + if ok: + tpm = ok / makespan * 60.0 + print(f"throughput: {tpm:.2f} videos/min ({ok / makespan:.4f} videos/s)") + print( + f"latency (s): p50={nearest_rank(lat, 0.50):.1f} " + f"p90={nearest_rank(lat, 0.90):.1f} " + f"min={min(lat):.1f} max={max(lat):.1f}" + ) + if args.baseline_tpm > 0: + print(f"scaling vs baseline ({args.baseline_tpm:.2f}/min): {tpm / args.baseline_tpm:.2f}x") + if errs: + print("sample failures:", errs[:3]) + + # Isolation check: compare each request's output against its own single-replica baseline. + isolation_failed = False + if args.write_baseline: + Path(args.write_baseline).write_text(json.dumps(got, indent=2)) + print(f"baseline: wrote {len(got)} per-input hashes -> {args.write_baseline}") + elif args.check_baseline: + baseline = json.loads(Path(args.check_baseline).read_text()) + matched = mismatched = unknown = 0 + for key, md5 in got.items(): + if key not in baseline: + unknown += 1 + print(f" isolation: NO BASELINE for input {key!r} (not in the baseline manifest)") + elif baseline[key] == md5: + matched += 1 + else: + mismatched += 1 + print(f" isolation: MISMATCH for input {key!r} (output changed under concurrency)") + # PASS requires the run to cover the *whole* baseline: a subset (fewer requests, or + # some failed) must not pass just because the inputs it did run happened to match. + unverified = sorted(set(baseline) - set(got)) + for key in unverified: + print(f" isolation: NOT VERIFIED for baseline input {key!r} (absent from this run)") + isolation_failed = mismatched > 0 or unknown > 0 or len(unverified) > 0 + verdict = "PASS (isolated)" if not isolation_failed else "FAIL" + print( + f"isolation: {matched} matched / {mismatched} mismatched / {unknown} unknown / " + f"{len(unverified)} unverified vs baseline({len(baseline)}) -> {verdict}" + ) + + if fail > 0 or ok == 0 or isolation_failed: + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/replica_data_parallel/run_server.sh b/examples/online_serving/replica_data_parallel/run_server.sh new file mode 100755 index 00000000000..915732eee4b --- /dev/null +++ b/examples/online_serving/replica_data_parallel/run_server.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Serve N independent Wan2.2 TI2V-5B replicas (replica data parallelism). +# +# Usage: +# NUM_REPLICAS=1 DEVICES=0 ./run_server.sh # baseline +# NUM_REPLICAS=2 DEVICES=0,1 ./run_server.sh +# NUM_REPLICAS=4 DEVICES=0,1,2,3 ./run_server.sh +# +# Then drive load with bench_replica_dp.py (see README.md) and compare +# throughput across replica counts. +set -euo pipefail + +MODEL="${MODEL:-Wan-AI/Wan2.2-TI2V-5B-Diffusers}" +HOST="${HOST:-127.0.0.1}" +PORT="${PORT:-8098}" +NUM_REPLICAS="${NUM_REPLICAS:-1}" +DEVICES="${DEVICES:-0}" # one GPU per replica for tensor_parallel_size=1 + +HERE="$(cd "$(dirname "$0")" && pwd)" +TEMPLATE="$HERE/wan2_2_ti2v_dp.yaml" +CFG="$(mktemp --suffix=_wan2_2_dp.yaml)" +trap 'rm -f "$CFG"' EXIT + +sed -e "s/__NUM_REPLICAS__/${NUM_REPLICAS}/" \ + -e "s/__DEVICES__/${DEVICES}/" \ + "$TEMPLATE" > "$CFG" + +echo "Serving ${MODEL}: num_replicas=${NUM_REPLICAS} devices=${DEVICES} (config ${CFG})" + +vllm serve "$MODEL" --omni \ + --host "$HOST" \ + --port "$PORT" \ + --stage-configs-path "$CFG" diff --git a/examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml b/examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml new file mode 100644 index 00000000000..b472d7e2d99 --- /dev/null +++ b/examples/online_serving/replica_data_parallel/wan2_2_ti2v_dp.yaml @@ -0,0 +1,43 @@ +# Replica data-parallel recipe for Wan2.2 TI2V-5B (video DiT). +# +# Fans out `num_replicas` independent DiffusionEngine replicas, one per GPU, +# and lets the serving layer route each request to a single replica. This +# scales request throughput near-linearly; it does not change single-request +# latency. See ./README.md for measured numbers and the `run_server.sh` driver. +# +# run_server.sh substitutes __NUM_REPLICAS__ / __DEVICES__ from env before serving. +# For tensor_parallel_size = 1, `devices` lists one GPU per replica (pool mode): +# num_replicas=4 -> devices "0,1,2,3". + +stage_args: + - stage_id: 0 + stage_type: diffusion + runtime: + num_replicas: __NUM_REPLICAS__ + devices: "__DEVICES__" + max_batch_size: 1 + engine_args: + model_stage: dit + # `model_class_name` takes the diffusion registry KEY (WanPipeline), which + # resolves to the Wan22Pipeline class. Wan2.2 TI2V-5B is served by the base + # WanPipeline, which auto-detects TI2V-5B mode from model_index.json. + model_class_name: WanPipeline + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + distributed_executor_backend: "mp" + vae_use_tiling: true + parallel_config: + tensor_parallel_size: 1 + + final_output: true + final_output_type: video + is_comprehension: false + default_sampling_params: + seed: 42 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 diff --git a/examples/online_serving/sensenova_u1/README.md b/examples/online_serving/sensenova_u1/README.md index 86eadb1d654..d529eedad3e 100644 --- a/examples/online_serving/sensenova_u1/README.md +++ b/examples/online_serving/sensenova_u1/README.md @@ -1,4 +1,4 @@ -# SenseNova-U1 Online Serving +# SenseNova-U1: Online serving ## Launch the Server @@ -6,11 +6,11 @@ vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 ``` -With Cache-DiT acceleration: +With cache acceleration (`cache_dit` or `tea_cache`): ```bash vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \ - --cache-backend cache_dit + --cache-backend tea_cache ``` Or use the convenience script: @@ -19,8 +19,8 @@ Or use the convenience script: cd examples/online_serving/sensenova_u1 bash run_server.sh -# Cache-DiT -CACHE_BACKEND=cache_dit bash run_server.sh +# Cache acceleration (cache_dit or tea_cache) +CACHE_BACKEND=tea_cache bash run_server.sh ``` ### Tensor Parallelism (TP) diff --git a/examples/online_serving/sensenova_u1/openai_chat_client.py b/examples/online_serving/sensenova_u1/openai_chat_client.py index df9afa83d7f..50ff202d4d7 100644 --- a/examples/online_serving/sensenova_u1/openai_chat_client.py +++ b/examples/online_serving/sensenova_u1/openai_chat_client.py @@ -48,7 +48,7 @@ def generate( forwarded as top-level fields in the request payload. The serving layer maps standard keys (height, width, seed, num_inference_steps) to sampling params and forwards model-specific keys (think, cfg_norm, img_cfg_scale, - …) to ``extra_args`` based on the pipeline's ``EXTRA_BODY_PARAMS``. + …) to ``extra_args`` based on the model extra registry. Returns: Image bytes (for image outputs) or text string (for text outputs). @@ -131,7 +131,7 @@ def main(): parser.add_argument("--width", type=int, default=2048, help="Image width") parser.add_argument("--num-steps", type=int, default=50, help="Inference steps") parser.add_argument("--seed", type=int, default=42, help="Random seed") - # Model-specific parameters (forwarded via EXTRA_BODY_PARAMS) + # Model-specific parameters forwarded via the model extra registry. parser.add_argument("--cfg-scale", type=float, default=4.0, help="CFG scale") parser.add_argument("--img-cfg-scale", type=float, help="Image CFG scale (img2img)") parser.add_argument("--cfg-norm", type=str, help="CFG normalization mode") diff --git a/examples/online_serving/step_audio2/README.md b/examples/online_serving/step_audio2/README.md index 3a04f4565cb..25e2a678506 100644 --- a/examples/online_serving/step_audio2/README.md +++ b/examples/online_serving/step_audio2/README.md @@ -1,4 +1,4 @@ -# Step-Audio2 Online Serving +# Step-Audio2: Online serving This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API. @@ -10,15 +10,15 @@ Please refer to [README.md](../../../README.md) ```bash # Async chunk mode (recommended — lower first-packet latency for TTS) -vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \ - --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml \ +vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 \ + --deploy-config vllm_omni/deploy/step_audio_2_async_chunk.yaml \ --trust-remote-code --enforce-eager ``` Sequential mode: ```bash -vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \ - --stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2.yaml \ +vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 \ + --deploy-config vllm_omni/deploy/step_audio_2.yaml \ --trust-remote-code --enforce-eager ``` @@ -51,7 +51,7 @@ Or via curl: ```bash curl -X POST http://localhost:8092/v1/audio/speech \ -H "Content-Type: application/json" \ - -d '{"model":"stepfun-ai/Step-Audio2-mini","input":"你好世界","voice":"default"}' \ + -d '{"model":"stepfun-ai/Step-Audio-2-mini","input":"你好世界","voice":"default"}' \ --output output.wav ``` @@ -139,7 +139,7 @@ Step-Audio2 uses the OpenAI-compatible chat completions API: ```json { - "model": "stepfun-ai/Step-Audio2-mini", + "model": "stepfun-ai/Step-Audio-2-mini", "messages": [ { "role": "system", diff --git a/examples/online_serving/step_audio2/openai_chat_completion_client.py b/examples/online_serving/step_audio2/openai_chat_completion_client.py index 2b91dacb21d..94acb553359 100644 --- a/examples/online_serving/step_audio2/openai_chat_completion_client.py +++ b/examples/online_serving/step_audio2/openai_chat_completion_client.py @@ -9,7 +9,7 @@ Usage examples: # Start server first: - vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 + vllm serve stepfun-ai/Step-Audio-2-mini --omni --port 8092 # Audio to text (ASR) python openai_chat_completion_client.py --query-type audio_to_text @@ -311,7 +311,7 @@ def parse_args(): "--model", "-m", type=str, - default="stepfun-ai/Step-Audio2-mini", + default="stepfun-ai/Step-Audio-2-mini", help="Model name", ) parser.add_argument( diff --git a/examples/online_serving/step_audio2/openai_speech_client.py b/examples/online_serving/step_audio2/openai_speech_client.py index 740b368055e..762ca72a6fb 100644 --- a/examples/online_serving/step_audio2/openai_speech_client.py +++ b/examples/online_serving/step_audio2/openai_speech_client.py @@ -87,7 +87,7 @@ def parse_args(): "--model", "-m", type=str, - default="stepfun-ai/Step-Audio2-mini", + default="stepfun-ai/Step-Audio-2-mini", help="Model name/path", ) parser.add_argument( diff --git a/examples/online_serving/step_audio2/run_curl.sh b/examples/online_serving/step_audio2/run_curl.sh index 3811308e735..3a3b1c9f4bf 100755 --- a/examples/online_serving/step_audio2/run_curl.sh +++ b/examples/online_serving/step_audio2/run_curl.sh @@ -110,7 +110,7 @@ output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \ -H "Content-Type: application/json" \ -d @- <<EOF { - "model": "stepfun-ai/Step-Audio2-mini", + "model": "stepfun-ai/Step-Audio-2-mini", "sampling_params_list": $sampling_params_list, "messages": [ { diff --git a/examples/online_serving/streaming_video_generation/README.md b/examples/online_serving/streaming_video_generation/README.md new file mode 100644 index 00000000000..1cfa29dfd59 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/README.md @@ -0,0 +1,94 @@ +# Streaming Video Generation + +This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced. +It covers text-only video generation. Image/reference input is intentionally not included for now. + +## Start The Server + +Start a diffusion video model with streaming output enabled: + +```bash +vllm serve BestWishYsh/Helios-Distilled \ + --omni \ + --diffusion-streaming-output \ + --port 8000 +``` + +The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=True` in the default diffusion stage `engine_args`, then loaded by `OmniDiffusionConfig.from_kwargs()`. + +## WebSocket Protocol + +| Direction | Message | Format | Description | +| --- | --- | --- | --- | +| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. | +| Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. | +| Server to client | Video chunk | Binary WebSocket frame | Fragmented MP4 (`m4s`) video bytes. | +| Client to server | `session.stop` | JSON text: `{"type":"session.stop"}` | Requests cancellation of the active session. | +| Client to server | `session.ping` | JSON text: `{"type":"session.ping"}` | Optional keepalive; refreshes the server stall clock. | +| Server to client | `session.done` | JSON text: `{"type":"session.done","request_id":"...","chunks":3,"stopped":false}` | Ends a completed or stopped session. | +| Server to client | `session.pong` | JSON text: `{"type":"session.pong"}` | Reply to `session.ping`. | +| Server to client | `error` | JSON text: `{"type":"error","message":"..."}` | Reports invalid input, unsupported formats, generation failures, control-message errors, or stall timeout. | + +During generation the client normally sends only `session.start` and then receives binary chunks; silence on the client socket is expected. The server closes the session with a stall error only when there is no engine progress and no `session.ping` for about 60 seconds. + +## Install Client Dependency + +```bash +pip install websockets +# For the Gradio demo: +pip install 'vllm-omni[demo]' websockets +``` + +## Run The Client + +```bash +python streaming_video_client.py \ + --host 127.0.0.1 \ + --port 8000 \ + --model BestWishYsh/Helios-Distilled \ + --prompt "A serene lakeside sunrise with mist over the water." \ + --width 640 \ + --height 384 \ + --fps 16 \ + --num-frames 99 \ + --guidance-scale 1.0 \ + --seed 42 \ + --output helios_stream.mp4 +``` + +The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`. +The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration. + +## Run The Gradio Demo + +```bash +python gradio_demo.py \ + --host 127.0.0.1 \ + --port 7860 +``` + +The Gradio demo requests fMP4 (`m4s`) chunks and appends them directly in the browser with a Media Source Extensions player. + +## Model Choice + +### Helios + +The example uses `BestWishYsh/Helios-Distilled` model by default. + +To ensure streaming-level generation speed, `pyramid_num_inference_steps_list` is suggested to be as low as `[1, 1, 1]`. Both example clients uses the following Helios-Distilled preset by default: + +```json +{ + "is_enable_stage2": true, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": true +} +``` + +Disable it in the CLI example with `--no-helios-distilled-preset`, or override/extend it with `--extra-params`: + +```bash +python streaming_video_client.py \ + --extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}' +``` diff --git a/examples/online_serving/streaming_video_generation/gradio_demo.py b/examples/online_serving/streaming_video_generation/gradio_demo.py new file mode 100644 index 00000000000..2f7a301f14b --- /dev/null +++ b/examples/online_serving/streaming_video_generation/gradio_demo.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Gradio demo for `/v1/realtime/video`. + +Usage: + python gradio_demo.py --host 127.0.0.1 --port 7860 + +Start a vLLM-Omni server with streaming output enabled first, for example: + vllm serve BestWishYsh/Helios-Distilled --omni --diffusion-streaming-output --port 8000 +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None + +DEFAULT_PROMPT = "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." +VIDEO_STREAM_VIEW_HTML = Path(__file__).with_name("video-stream-view.html") +VIDEO_STREAM_VIEW_JS = Path(__file__).with_name("video-stream-view.js") + + +def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None: + if value is None: + return + if isinstance(value, str) and not value.strip(): + return + payload[key] = value + + +def _optional_int(value: int | float | None) -> int | None: + return None if value is None else int(value) + + +def _optional_float(value: int | float | None) -> float | None: + return None if value is None else float(value) + + +def _build_session_start( + *, + model: str, + prompt: str, + negative_prompt: str, + width: int | float | None, + height: int | float | None, + fps: int | float | None, + num_frames: int | float | None, + num_inference_steps: int | float | None, + guidance_scale: int | float | None, + guidance_scale_2: int | float | None, + boundary_ratio: int | float | None, + flow_shift: int | float | None, + true_cfg_scale: int | float | None, + seed: int | float | None, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "type": "session.start", + "model": model, + "prompt": prompt, + "format": "m4s", + } + + _maybe_set(payload, "negative_prompt", negative_prompt) + _maybe_set(payload, "width", _optional_int(width)) + _maybe_set(payload, "height", _optional_int(height)) + _maybe_set(payload, "fps", _optional_int(fps)) + _maybe_set(payload, "num_frames", _optional_int(num_frames)) + _maybe_set(payload, "num_inference_steps", _optional_int(num_inference_steps)) + _maybe_set(payload, "guidance_scale", _optional_float(guidance_scale)) + _maybe_set(payload, "guidance_scale_2", _optional_float(guidance_scale_2)) + _maybe_set(payload, "boundary_ratio", _optional_float(boundary_ratio)) + _maybe_set(payload, "flow_shift", _optional_float(flow_shift)) + _maybe_set(payload, "true_cfg_scale", _optional_float(true_cfg_scale)) + _maybe_set(payload, "seed", _optional_int(seed)) + + if "Helios" in model: + payload["extra_params"] = { + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": True, + } + + return payload + + +def _set_optional_field_enabled(enabled: bool) -> Any: + return gr.update(interactive=enabled) + + +def build_browser_stream_config( + host: str, + port: int | float, + model: str, + prompt: str, + negative_prompt: str, + width: int | float | None, + height: int | float | None, + fps: int | float, + num_frames: int | float | None, + num_inference_steps: int | float | None, + include_guidance_scale: bool, + guidance_scale: int | float | None, + include_guidance_scale_2: bool, + guidance_scale_2: int | float | None, + include_seed: bool, + seed: int | float | None, + include_boundary_ratio: bool, + boundary_ratio: int | float | None, + include_flow_shift: bool, + flow_shift: int | float | None, + include_true_cfg_scale: bool, + true_cfg_scale: int | float | None, +) -> tuple[str, Any]: + """Build the browser-side WebSocket request config for the HTML/MSE player.""" + if not model.strip(): + raise gr.Error("Model is required.") + if not prompt.strip(): + raise gr.Error("Prompt is required.") + + payload = _build_session_start( + model=model.strip(), + prompt=prompt.strip(), + negative_prompt=negative_prompt, + width=width, + height=height, + fps=fps, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale if include_guidance_scale else None, + guidance_scale_2=guidance_scale_2 if include_guidance_scale_2 else None, + boundary_ratio=boundary_ratio if include_boundary_ratio else None, + flow_shift=flow_shift if include_flow_shift else None, + true_cfg_scale=true_cfg_scale if include_true_cfg_scale else None, + seed=seed if include_seed else None, + ) + config = { + "url": f"ws://{host}:{int(port)}/v1/realtime/video", + "payload": payload, + } + return json.dumps(config, ensure_ascii=False), gr.update(value="Streaming...", interactive=False) + + +def create_demo() -> gr.Blocks: + with gr.Blocks(title="Streaming Video Generation") as demo: + gr.Markdown("# Streaming Video Generation") + gr.Markdown( + "Connects to `WS /v1/realtime/video` in the browser and appends fMP4 chunks " + "directly to a Media Source Extensions video player." + ) + + with gr.Row(): + with gr.Column(scale=1): + host = gr.Textbox(label="Server Host", value="127.0.0.1") + port = gr.Number(label="Server Port", value=8000, precision=0) + model = gr.Textbox(label="Model", value="BestWishYsh/Helios-Distilled") + prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=3) + negative_prompt = gr.Textbox(label="Negative Prompt", value="", lines=2) + + with gr.Row(): + width = gr.Number(label="Width", value=640, precision=0) + height = gr.Number(label="Height", value=384, precision=0) + + with gr.Row(): + fps = gr.Number(label="FPS", value=16, precision=0) + num_frames = gr.Number(label="Num Frames", value=132, precision=0) + num_inference_steps = gr.Number(label="Num Inference Steps", value=50, precision=0) + + with gr.Row(): + with gr.Column(): + guidance_scale = gr.Number( + label="Guidance Scale", + value=1.0, + interactive=True, + scale=4, + ) + include_guidance_scale = gr.Checkbox(label="Send", value=True, scale=1, min_width=70) + with gr.Column(): + guidance_scale_2 = gr.Number( + label="Guidance Scale 2", + value=None, + interactive=False, + scale=4, + ) + include_guidance_scale_2 = gr.Checkbox(label="Send", value=False, scale=1, min_width=70) + with gr.Column(): + true_cfg_scale = gr.Number( + label="True CFG Scale", + value=None, + interactive=False, + scale=4, + ) + include_true_cfg_scale = gr.Checkbox(label="Send", value=False, scale=1, min_width=70) + # I want it to wrap after 3, but it auto wrap after 2. I'll rather fully delegate to auto wrapping. + with gr.Column(): + boundary_ratio = gr.Number( + label="Boundary Ratio", + value=None, + interactive=False, + scale=4, + ) + include_boundary_ratio = gr.Checkbox(label="Send", value=False, scale=1, min_width=70) + with gr.Column(): + flow_shift = gr.Number(label="Flow Shift", value=None, interactive=False, scale=4) + include_flow_shift = gr.Checkbox(label="Send", value=False, scale=1, min_width=70) + with gr.Column(): + seed = gr.Number(label="Seed", value=42, precision=0, interactive=True, scale=4) + include_seed = gr.Checkbox(label="Send", value=True, scale=1, min_width=70) + + with gr.Row(): + start_button = gr.Button("Start", variant="primary", elem_id="streaming-video-start") + hidden_stream_config = gr.Textbox(value="", visible=False) + + with gr.Column(scale=1): + gr.HTML( + VIDEO_STREAM_VIEW_HTML.read_text(encoding="utf-8"), + label="Streaming Preview", + js_on_load=VIDEO_STREAM_VIEW_JS.read_text(encoding="utf-8"), + ) + + inputs = [ + host, + port, + model, + prompt, + negative_prompt, + width, + height, + fps, + num_frames, + num_inference_steps, + include_guidance_scale, + guidance_scale, + include_guidance_scale_2, + guidance_scale_2, + include_seed, + seed, + include_boundary_ratio, + boundary_ratio, + include_flow_shift, + flow_shift, + include_true_cfg_scale, + true_cfg_scale, + ] + for toggle, field in [ + (include_guidance_scale, guidance_scale), + (include_guidance_scale_2, guidance_scale_2), + (include_seed, seed), + (include_boundary_ratio, boundary_ratio), + (include_flow_shift, flow_shift), + (include_true_cfg_scale, true_cfg_scale), + ]: + toggle.change( + fn=_set_optional_field_enabled, + inputs=toggle, + outputs=field, + ) + start_button.click( + fn=build_browser_stream_config, + inputs=inputs, + outputs=[hidden_stream_config, start_button], + ).then( + fn=lambda config: config, + inputs=hidden_stream_config, + outputs=hidden_stream_config, + js=""" + (config) => { + if (window.vllmStreamingVideoStart) { + window.vllmStreamingVideoStart(config); + } + return config; + } + """, + ) + + return demo + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Streaming video generation Gradio demo") + parser.add_argument("--host", default="127.0.0.1", help="Host/IP for gradio launch") + parser.add_argument("--port", type=int, default=7860, help="Port for gradio launch") + parser.add_argument("--share", action="store_true", help="Create a public Gradio link") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + demo = create_demo() + demo.queue().launch(server_name=args.host, server_port=args.port, share=args.share) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/streaming_video_generation/streaming_video_client.py b/examples/online_serving/streaming_video_generation/streaming_video_client.py new file mode 100644 index 00000000000..c77113420d6 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/streaming_video_client.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Text-only WebSocket client for `/v1/realtime/video`. + +The client prints a line for every binary video chunk it receives and writes +the received bytes to disk when the session finishes or is interrupted. +Image/reference input is intentionally not supported in this example yet. + +Requirements: + pip install websockets +""" + +import argparse +import asyncio +import io +import json +import os +import time +from pathlib import Path +from typing import Any + +import av + +try: + from websockets.asyncio.client import connect # pyright: ignore[reportMissingImports] +except ImportError: + try: + from websockets import connect # pyright: ignore[reportMissingImports] + except ImportError: + print("Please install websockets: pip install websockets") + raise SystemExit(1) + + +HELIOS_DISTILLED_EXTRA_PARAMS = { + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": True, +} + + +def _json_object(value: str, *, argument_name: str) -> dict[str, Any]: + try: + parsed = json.loads(value) + except json.JSONDecodeError as exc: + raise argparse.ArgumentTypeError(f"{argument_name} must be valid JSON: {exc}") from exc + if not isinstance(parsed, dict): + raise argparse.ArgumentTypeError(f"{argument_name} must be a JSON object") + return parsed + + +def _maybe_set(payload: dict[str, Any], key: str, value: Any) -> None: + if value is not None: + payload[key] = value + + +def build_session_start(args: argparse.Namespace) -> dict[str, Any]: + payload: dict[str, Any] = { + "type": "session.start", + "model": args.model, + "prompt": args.prompt, + } + + _maybe_set(payload, "negative_prompt", args.negative_prompt) + _maybe_set(payload, "width", args.width) + _maybe_set(payload, "height", args.height) + _maybe_set(payload, "size", args.size) + _maybe_set(payload, "seconds", args.seconds) + _maybe_set(payload, "fps", args.fps) + _maybe_set(payload, "num_frames", args.num_frames) + _maybe_set(payload, "num_inference_steps", args.num_inference_steps) + _maybe_set(payload, "guidance_scale", args.guidance_scale) + _maybe_set(payload, "guidance_scale_2", args.guidance_scale_2) + _maybe_set(payload, "boundary_ratio", args.boundary_ratio) + _maybe_set(payload, "flow_shift", args.flow_shift) + _maybe_set(payload, "true_cfg_scale", args.true_cfg_scale) + _maybe_set(payload, "seed", args.seed) + + extra_params: dict[str, Any] = {} + if "Helios-Distilled" in args.model and args.helios_distilled_preset: + extra_params.update(HELIOS_DISTILLED_EXTRA_PARAMS) + if args.extra_params: + extra_params.update(args.extra_params) + if extra_params: + payload["extra_params"] = extra_params + + return payload + + +def _count_decoded_frames(video_bytes: bytes, *, stream_format: str) -> int | None: + """Return the number of decodable video frames in a streamed byte sequence, if known.""" + if not video_bytes: + return 0 + + av_format = "mp4" if stream_format == "m4s" else None + try: + with av.open(io.BytesIO(video_bytes), format=av_format, mode="r") as container: + stream = container.streams.video[0] + return sum(1 for _ in container.decode(stream)) + except Exception: + return None + + +def save_video( + chunks: list[bytes], + output_path: Path, + *, + fps: float, + stream_format: str, +) -> None: + if not chunks: + print("No video chunks received; nothing to save.") + return + + from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes + + streamed_bytes = b"".join(chunks) + playback_bytes = finalize_streaming_video_bytes( + streamed_bytes, + input_format="m4s", + fps=fps, + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = output_path.with_suffix(output_path.suffix + ".tmp") + with tmp_path.open("wb") as f: + f.write(playback_bytes) + os.replace(tmp_path, output_path) + print( + f"Saved {len(chunks)} streamed chunk(s), {len(streamed_bytes)} streamed bytes " + f"-> {len(playback_bytes)} playback bytes at {output_path}" + ) + + +async def stream_video(args: argparse.Namespace) -> None: + payload = build_session_start(args) + output_path = Path(args.output) + url = f"ws://{args.host}:{args.port}/v1/realtime/video" + chunks: list[bytes] = [] + total_bytes = 0 + total_frames = 0 + accumulated_stream = bytearray() + stream_format = "m4s" + done = False + ws = None + + try: + async with connect(url, max_size=None) as websocket: + ws = websocket + await websocket.send(json.dumps(payload)) + print(f"Connected: {url}") + print(f"Sent session.start: {json.dumps(payload, ensure_ascii=False)}") + started_at = time.perf_counter() + last_chunk_at = started_at + + while True: + message = await websocket.recv() + now = time.perf_counter() + + if isinstance(message, bytes): + chunks.append(message) + chunk_bytes = len(message) + chunk_elapsed = now - last_chunk_at + last_chunk_at = now + + total_bytes += chunk_bytes + accumulated_stream.extend(message) + decoded_total = _count_decoded_frames(bytes(accumulated_stream), stream_format=stream_format) + if decoded_total is not None: + chunk_frames = max(decoded_total - total_frames, 0) + total_frames = decoded_total + else: + chunk_frames = 0 + + total_elapsed = now - started_at + print( + f"[chunk {len(chunks):03d}] " + f"bytes={chunk_bytes} " + f"frames={chunk_frames} " + f"elapsed={chunk_elapsed:.2f}s " + f"total_bytes={total_bytes} " + f"total_frames={total_frames} " + f"total_elapsed={total_elapsed:.2f}s" + ) + continue + + msg = json.loads(message) + msg_type = msg.get("type") + if msg_type == "video.start": + stream_format = msg.get("format") or stream_format + print(f"Video session started: request_id={msg.get('request_id')} format={msg.get('format')}") + elif msg_type == "session.done": + print(f"Session complete: {json.dumps(msg, ensure_ascii=False)}") + done = True + break + elif msg_type == "error": + print(f"ERROR: {msg.get('message', msg)}") + else: + print(f"Received control message: {msg}") + except (KeyboardInterrupt, asyncio.CancelledError): + print("\nInterrupted. Asking server to stop the active session...") + if ws is not None: + try: + await ws.send(json.dumps({"type": "session.stop"})) + await ws.close(code=1000, reason="client interrupted") + except Exception: + pass + raise + finally: + print( + "Saving video... (May take a while. Remuxing concatenated chunks into one progressive MP4 file with total duration metadata in file header)" + ) + save_video( + chunks, + output_path, + fps=float(payload.get("fps") or args.fps), + stream_format=stream_format, + ) + if not done: + print("Saved partial output because the stream did not finish normally.") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Streaming video generation WebSocket client") + parser.add_argument("--host", default="127.0.0.1", help="API server host") + parser.add_argument("--port", type=int, default=8000, help="API server port") + parser.add_argument( + "--model", + default="BestWishYsh/Helios-Distilled", + help="Model name served by the API server", + ) + parser.add_argument( + "--prompt", + default="A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement.", + help="Text prompt for video generation", + ) + parser.add_argument("--output", default="streaming_video_output.mp4", help="Path for the received video bytes") + parser.add_argument("--width", type=int, default=640, help="Video width") + parser.add_argument("--height", type=int, default=384, help="Video height") + parser.add_argument("--size", default=None, help="Video size as WIDTHxHEIGHT; alternative to width/height") + parser.add_argument("--seconds", default=None, help="Clip duration in seconds") + parser.add_argument("--fps", type=int, default=16, help="Frames per second") + parser.add_argument("--num-frames", type=int, default=99, help="Number of generated frames") + + parser.add_argument("--negative-prompt", default=None, help="Negative prompt") + parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps") + parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale") + parser.add_argument("--guidance-scale-2", type=float, default=None, help="High-noise CFG scale for video models") + parser.add_argument("--boundary-ratio", type=float, default=None, help="Boundary split ratio for video models") + parser.add_argument("--flow-shift", type=float, default=None, help="Scheduler flow_shift for video models") + parser.add_argument("--true-cfg-scale", type=float, default=None, help="True CFG scale for models that support it") + parser.add_argument("--seed", type=int, default=42, help="Random seed") + + parser.add_argument( + "--helios-distilled-preset", + action=argparse.BooleanOptionalAction, + default=True, + help="Send the Helios-Distilled extra_params used by the bundled example.", + ) + parser.add_argument( + "--extra-params", + type=lambda value: _json_object(value, argument_name="--extra-params"), + default=None, + help="JSON object merged into extra_params; overrides preset keys on conflict.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + try: + asyncio.run(stream_video(args)) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/streaming_video_generation/video-stream-view.html b/examples/online_serving/streaming_video_generation/video-stream-view.html new file mode 100644 index 00000000000..bbef95c691b --- /dev/null +++ b/examples/online_serving/streaming_video_generation/video-stream-view.html @@ -0,0 +1,11 @@ +<!-- This file and the corresponding JS file are needed to display a streamable video view in the Gradio demo. + Gradio's built-in video player does not support modern M4S streaming format. --> +<div id="vllm-streaming-video-view" style="display:flex; flex-direction:column; gap:10px;"> + <video id="vllm-streaming-video-player" autoplay muted controls playsinline + style="width:100%; min-height:320px; background:#111; border-radius:8px;"></video> + <div style="display:flex; align-items:center; gap:10px;"> + <span id="vllm-streaming-video-status" style="font-weight:600;">Ready</span> + </div> + <pre id="vllm-streaming-video-log" style="height:220px; overflow:auto; margin:0; padding:10px; border-radius:8px; + background:#0f172a; color:#dbeafe; font-size:12px; white-space:pre-wrap;">Ready.</pre> +</div> diff --git a/examples/online_serving/streaming_video_generation/video-stream-view.js b/examples/online_serving/streaming_video_generation/video-stream-view.js new file mode 100644 index 00000000000..1b0a5024b37 --- /dev/null +++ b/examples/online_serving/streaming_video_generation/video-stream-view.js @@ -0,0 +1,198 @@ +/** + * This file and the corresponding HTML file are needed to display a streamable video view in the Gradio demo. + * Gradio's built-in video player does not support modern M4S streaming format. + */ +(function () { + const playerId = "vllm-streaming-video-player"; + const statusId = "vllm-streaming-video-status"; + const logId = "vllm-streaming-video-log"; + let ws = null; + let mediaSource = null; + let sourceBuffer = null; + let queue = []; + let chunkCount = 0; + let totalBytes = 0; + let done = false; + let objectUrl = null; + + function el(id) { + return document.getElementById(id); + } + + function setStatus(text) { + const node = el(statusId); + if (node) node.textContent = text; + } + + function log(line) { + const node = el(logId); + if (!node) return; + node.textContent += "\n" + line; + node.scrollTop = node.scrollHeight; + } + + function setStartButton(disabled, label) { + const btn = document.querySelector("#streaming-video-start button") || document.getElementById("streaming-video-start"); + if (!btn) return; + btn.disabled = disabled; + btn.textContent = label; + } + + function chooseMime() { + const candidates = [ + 'video/mp4; codecs="avc1.42E01E"', + 'video/mp4; codecs="avc1.4D401F"', + 'video/mp4; codecs="avc1.64001F"', + "video/mp4" + ]; + for (const mime of candidates) { + if (window.MediaSource && MediaSource.isTypeSupported(mime)) return mime; + } + return null; + } + + function pump() { + if (!sourceBuffer || sourceBuffer.updating || queue.length === 0) return; + try { + sourceBuffer.appendBuffer(queue.shift()); + } catch (err) { + done = true; + setStatus("SourceBuffer error"); + setStartButton(false, "Restart"); + log("SourceBuffer append error: " + err.message); + finishStream(); + } + } + + function finishStream() { + if (!mediaSource || mediaSource.readyState !== "open") return; + if (sourceBuffer && sourceBuffer.updating) return; + if (queue.length > 0) return; + try { + mediaSource.endOfStream(); + } catch (_) { } + } + + function stopCurrent() { + if (ws && ws.readyState === WebSocket.OPEN) { + try { ws.send(JSON.stringify({ type: "session.stop" })); } catch (_) { } + try { ws.close(); } catch (_) { } + } + ws = null; + sourceBuffer = null; + mediaSource = null; + queue = []; + if (objectUrl) URL.revokeObjectURL(objectUrl); + objectUrl = null; + } + + window.vllmStreamingVideoStart = function (configJson) { + stopCurrent(); + chunkCount = 0; + totalBytes = 0; + done = false; + const logNode = el(logId); + if (logNode) logNode.textContent = "Starting..."; + setStatus("Starting..."); + + let config; + try { + config = JSON.parse(configJson); + } catch (err) { + setStatus("Invalid request"); + setStartButton(false, "Restart"); + log("Invalid request config: " + err.message); + return configJson; + } + + const mime = chooseMime(); + if (!mime) { + setStatus("MSE unsupported"); + setStartButton(false, "Restart"); + log("This browser does not support MP4 Media Source Extensions playback."); + return configJson; + } + + const video = el(playerId); + if (!video) return configJson; + mediaSource = new MediaSource(); + objectUrl = URL.createObjectURL(mediaSource); + video.src = objectUrl; + + mediaSource.addEventListener("sourceopen", () => { + try { + sourceBuffer = mediaSource.addSourceBuffer(mime); + sourceBuffer.mode = "segments"; + sourceBuffer.addEventListener("updateend", () => { + pump(); + if (done) finishStream(); + }); + } catch (err) { + setStatus("SourceBuffer error"); + setStartButton(false, "Restart"); + log("SourceBuffer error: " + err.message); + return; + } + + ws = new WebSocket(config.url); + ws.binaryType = "arraybuffer"; + ws.onopen = () => { + setStatus("Streaming..."); + setStartButton(true, "Streaming..."); + log("Connected: " + config.url); + ws.send(JSON.stringify(config.payload)); + log("Sent session.start: " + JSON.stringify(config.payload)); + }; + ws.onmessage = (event) => { + if (typeof event.data === "string") { + const msg = JSON.parse(event.data); + if (msg.type === "video.start") { + log("Video session started: request_id=" + (msg.request_id || "") + " format=" + (msg.format || "")); + } else if (msg.type === "session.done") { + done = true; + setStatus("Done"); + setStartButton(false, "Restart"); + log("Session complete: " + JSON.stringify(msg)); + finishStream(); + } else if (msg.type === "error") { + done = true; + setStatus("Error"); + setStartButton(false, "Restart"); + log("ERROR: " + (msg.message || JSON.stringify(msg))); + finishStream(); + } else { + log("Control message: " + JSON.stringify(msg)); + } + return; + } + + const data = event.data; + queue.push(data); + chunkCount += 1; + totalBytes += data.byteLength; + log(`[chunk ${String(chunkCount).padStart(3, "0")}] bytes=${data.byteLength} total_bytes=${totalBytes}`); + pump(); + }; + ws.onerror = () => { + done = true; + setStatus("Error"); + setStartButton(false, "Restart"); + log("ERROR: WebSocket error"); + finishStream(); + }; + ws.onclose = () => { + if (!done) { + done = true; + setStatus("Closed"); + setStartButton(false, "Restart"); + log("WebSocket closed before session.done"); + finishStream(); + } + }; + }, { once: true }); + + return configJson; + }; + + console.log("video-stream-view.html loaded"); +})(); diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index c9498596eef..201cb5d781b 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) | — | ✓ | +| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | compat only, non-chunk | uploaded audio voice only; no presets | — | | 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) | — | ✓ | @@ -23,6 +24,7 @@ For the full list of supported architectures across all modalities, see | Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) | | VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ | | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ | +| SoulX-Singer | `Soul-AILab/SoulX-Singer` | ✓ (prompt audio) | — (batch only) | — (prompt + target audio) | — (chat client) | CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](../../offline_inference/text_to_speech/README.md#cosyvoice3) instead. @@ -34,7 +36,7 @@ Launch the server (defaults shown — adjust `--port`, `--gpu-memory-utilization vllm serve <hf-repo-or-local-path> --omni --port 8091 ``` -Send a TTS request via curl: +Send a TTS request via curl. These generic snippets assume a model with a preset/default voice; voice-cloning-only models such as IndexTTS-2 require `ref_audio` or an uploaded audio `voice` (see model-specific sections below). ```bash curl -X POST http://localhost:8091/v1/audio/speech \ @@ -77,7 +79,7 @@ response = client.audio.speech.create( response.stream_to_file("output.wav") ``` -Streaming PCM output (where supported) — set `stream=true` with `response_format="pcm"`: +Streaming PCM output (where supported) — set `stream=true`, `stream_format="audio"`, and `response_format="pcm"`: ```bash curl -X POST http://localhost:8091/v1/audio/speech \ @@ -86,11 +88,12 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "input": "Hello, how are you?", "voice": "default", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others). +Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others). For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md). @@ -135,6 +138,40 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh --- +## IndexTTS-2 + +2-stage TTS (GPT AR + S2Mel CFM DiT + BigVGAN) at 22.05 kHz. Requests use `ref_audio` for voice cloning, or an uploaded audio `voice` from `/v1/audio/voices`. Supports emotion conditioning via `emo_audio`, `emo_text`, or `emo_vector` passed in `extra_params`. + +### Launch +```bash +vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 +# or, to pass the bundled deploy config explicitly: +bash examples/online_serving/text_to_speech/indextts2/run_server.sh +``` + +### Sending requests +```bash +# Voice cloning (ref_audio required) +python examples/online_serving/text_to_speech/indextts2/speech_client.py \ + --text "你好，世界！" \ + --ref-audio /path/to/reference.wav + +# With emotion audio +python examples/online_serving/text_to_speech/indextts2/speech_client.py \ + --text "今天心情很好！" \ + --ref-audio /path/to/ref.wav \ + --emo-audio /path/to/happy.wav +``` + +### Notes +- Output: 22.05 kHz mono WAV. +- Provide `ref_audio` on the documented raw request path, or pass `voice` only when it names an uploaded audio voice; IndexTTS-2 does not provide a built-in text-only preset voice. +- Emotion params (`emo_audio`, `emo_text`, `emo_vector`, `emo_alpha`, `use_emo_text`, `use_random`) are passed via the `extra_params` field. Official precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as the speaker reference. +- `stream=true` is accepted as an OpenAI-compatible response path, but IndexTTS-2 is not async-chunk streaming; audio is produced after S2Mel receives the full mel-code sequence. +- Deploy config: `vllm_omni/deploy/indextts2.yaml` (auto-loaded). + +--- + ## Fish Speech S2 Pro 4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec. @@ -327,6 +364,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ \"input\": \"Hello, streaming output from MOSS-TTS-Nano.\", \"ref_audio\": \"data:audio/wav;base64,${REF_AUDIO}\", \"stream\": true, + \"stream_format\": \"audio\", \"response_format\": \"pcm\" }" --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 - ``` @@ -492,10 +530,11 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "voice": "vivian", "language": "English", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` -Streaming requires `response_format="pcm"` and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. +Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. ### Streaming WebSocket The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence: @@ -504,6 +543,35 @@ python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 ``` +To receive word-level timestamps, launch the server with a forced aligner: +```bash +vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ + --omni \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ + --trust-remote-code \ + --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B +``` +Then request PCM JSON sidecar chunks: +```bash +python qwen3_tts/streaming_speech_client.py \ + --text "Hello world. How are you?" \ + --stream-audio \ + --response-format pcm \ + --word-timestamps +``` +The client writes one PCM file per sentence and a matching +`sentence_XXX_timestamps.json` sidecar. + +To *see* the alignment instead of reading a JSON sidecar, run the +word-timestamp Gradio demo (server must be launched with `--forced-aligner`): +```bash +python qwen3_tts/word_timestamps_demo.py --api-base http://localhost:8091 +``` +Each sentence's audio plays in an `<audio>` element while its text is rendered +as inline word spans; the current word highlights as `audio.currentTime` +crosses each `start_ms`. The **Stop (barge-in)** button cuts playback and +reports the last-spoken word, useful for the voice-agent barge-in case. + ### Gradio demos ```bash ./qwen3_tts/run_gradio_demo.sh # CustomVoice (default) @@ -567,7 +635,7 @@ After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use ```bash python voxcpm2/gradio_demo.py ``` -Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Audio is streamed from the OpenAI Speech endpoint with `stream=true`. +Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Raw PCM audio is streamed from the OpenAI Speech endpoint with `stream=true` and `stream_format="audio"`. --- @@ -597,3 +665,65 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts - Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`). - Voice cloning is gated upstream and may require a recent `mistral_common`. - A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now. + +--- + +## SoulX-Singer + +Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Single-stage DiT with inline preprocess. Uses the `/v1/chat/completions` endpoint with multimodal input (`prompt_audio` + `target_audio`). + +### Prerequisites + +Download DiT and preprocess weights, then set up separate SVS / SVC view directories and install dependencies as described in the [offline README](../../offline_inference/text_to_speech/README.md#soulx-singer). `config.json` `architectures` field is the single source of truth for SVS vs SVC — point `MODEL` at the matching directory. + +### Launch + +```bash +# SVS (default) +export MODEL=/path/to/SoulX-Singer +export PREPROCESS=/path/to/SoulX-Singer-Preprocess +bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh + +# SVC +export MODE=svc +export MODEL=/path/to/SoulX-Singer-svc +bash examples/online_serving/text_to_speech/soulxsinger/run_server.sh +``` + +Or equivalently, set `SOULX_PREPROCESS_WEIGHTS_DIR` and launch directly: +```bash +export SOULX_PREPROCESS_WEIGHTS_DIR=$PREPROCESS +vllm serve $MODEL --omni \ + --deploy-config vllm_omni/deploy/soulxsinger_${MODE}.yaml \ + --port 8192 --trust-remote-code --enforce-eager +``` + +### Sending requests + +Audio paths must be reachable from the server host (local filesystem or data URL). The client sends prompt vocal via `input_audio` and target accompaniment via `extra_args['target_audio']`. + +```bash +# Default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3 +python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \ + --prompt-audio /path/on/server/zh_prompt.mp3 \ + --target-audio /path/on/server/music.mp3 \ + --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \ + -o output.wav +``` + +Use precomputed metadata to skip online preprocess with following command: +```bash +python examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py \ + --prompt-metadata-path /path/on/server/zh_prompt.json \ + --target-metadata-path /path/on/server/music.json \ + --audio-path /path/on/server/zh_prompt.mp3 \ + -o output.wav +``` + +`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. See `openai_chat_client.py --help` for `--vocal-sep`, `--language`, `--num-inference-steps`, `--guidance-scale`, and `--seed`. + +### Notes + +- Output: 24 kHz mono WAV; batch only. +- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on. +- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only. diff --git a/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py b/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py index a837d924345..5ac11cb97e9 100644 --- a/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py +++ b/examples/online_serving/text_to_speech/cosyvoice3/speech_client.py @@ -60,6 +60,7 @@ def run_tts(args) -> None: if args.stream: payload["stream"] = True + payload["stream_format"] = "audio" payload["response_format"] = "pcm" print(f"Model: {args.model}") 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 cd227947626..ed22aa007e5 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 @@ -66,6 +66,8 @@ def build_payload( "response_format": "pcm" if stream else response_format, "stream": stream, } + if stream: + payload["stream_format"] = "audio" # Voice cloning: ref_audio takes priority over URL. ref_url_stripped = ref_audio_url.strip() if ref_audio_url else "" 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 43acffdc73d..726ccb0a6e6 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 @@ -53,6 +53,7 @@ def run_tts(args) -> None: if args.stream: payload["stream"] = True + payload["stream_format"] = "audio" payload["response_format"] = "pcm" print(f"Model: {args.model}") diff --git a/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py b/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py index 12e5d4b1181..3c5d3e206c9 100644 --- a/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py +++ b/examples/online_serving/text_to_speech/glm_tts/gradio_demo.py @@ -64,6 +64,8 @@ def build_payload( "stream": stream, "max_new_tokens": max_new_tokens, } + if stream: + payload["stream_format"] = "audio" # Voice cloning: ref_audio takes priority over URL. ref_url_stripped = ref_audio_url.strip() if ref_audio_url else "" diff --git a/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py b/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py index a0ff181c3f6..904d0190933 100644 --- a/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py +++ b/examples/online_serving/text_to_speech/glm_tts/openai_speech_client.py @@ -40,6 +40,9 @@ def run_tts_generation(args) -> None: "ref_audio": args.ref_audio, "ref_text": args.ref_text, } + if args.stream: + payload["stream_format"] = "audio" + payload["response_format"] = "pcm" if args.max_new_tokens: payload["max_new_tokens"] = args.max_new_tokens diff --git a/examples/online_serving/text_to_speech/indextts2/run_server.sh b/examples/online_serving/text_to_speech/indextts2/run_server.sh new file mode 100644 index 00000000000..4a17b3b04ae --- /dev/null +++ b/examples/online_serving/text_to_speech/indextts2/run_server.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Launch vLLM-Omni server for IndexTTS2 +# +# Usage from repository root: +# examples/online_serving/text_to_speech/indextts2/run_server.sh +# CUDA_VISIBLE_DEVICES=0 PORT=8092 MODEL=/path/to/IndexTeam/IndexTTS-2 examples/online_serving/text_to_speech/indextts2/run_server.sh + +set -e + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd -- "$SCRIPT_DIR/../../../.." && pwd)" + +MODEL="${MODEL:-IndexTeam/IndexTTS-2}" +PORT="${PORT:-8092}" +DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT_DIR/vllm_omni/deploy/indextts2.yaml}" + +echo "Starting IndexTTS2 server with model: $MODEL" + +FLASHINFER_DISABLE_VERSION_CHECK=1 \ +vllm serve "$MODEL" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --omni \ + --trust-remote-code \ + --deploy-config "$DEPLOY_CONFIG" diff --git a/examples/online_serving/text_to_speech/indextts2/speech_client.py b/examples/online_serving/text_to_speech/indextts2/speech_client.py new file mode 100644 index 00000000000..b938f0bd3cc --- /dev/null +++ b/examples/online_serving/text_to_speech/indextts2/speech_client.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OpenAI-compatible client for IndexTTS2 TTS via /v1/audio/speech endpoint. + +Examples: + # With reference audio for voice cloning + python speech_client.py --text "你好，世界！" \ + --ref-audio /path/to/reference.wav + + # With emotion audio + python speech_client.py --text "今天心情很好！" \ + --ref-audio /path/to/ref.wav \ + --emo-audio /path/to/happy.wav + +Server setup: + vllm serve IndexTeam/IndexTTS-2 --omni --host 0.0.0.0 --port 8092 +""" + +from __future__ import annotations + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8092" +DEFAULT_API_KEY = "sk-empty" + + +def encode_audio_to_base64(audio_path: str) -> str: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + ext = audio_path.lower().rsplit(".", 1)[-1] + mime = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac"}.get(ext, "audio/wav") + with open(audio_path, "rb") as f: + b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime};base64,{b64}" + + +def main() -> None: + parser = argparse.ArgumentParser(description="IndexTTS2 OpenAI speech client") + parser.add_argument("--text", type=str, required=True) + parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio for voice cloning") + parser.add_argument("--emo-audio", type=str, default=None, help="Emotion reference audio") + parser.add_argument("--emo-text", type=str, default=None, help="Emotion description text") + parser.add_argument( + "--emo-vector", + type=float, + nargs=8, + default=None, + help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm", + ) + parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1]") + parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text") + parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes") + parser.add_argument("--model", type=str, default="IndexTeam/IndexTTS-2") + parser.add_argument("--voice", type=str, default=None, help="Uploaded voice name to use instead of --ref-audio") + parser.add_argument("--output", type=str, default="output.wav") + parser.add_argument("--api-base", type=str, default=DEFAULT_API_BASE) + parser.add_argument("--api-key", type=str, default=DEFAULT_API_KEY) + parser.add_argument("--response-format", type=str, default="wav") + args = parser.parse_args() + + if not args.ref_audio and not args.voice: + parser.error("IndexTTS2 requires --ref-audio or --voice for voice cloning") + + payload: dict = { + "model": args.model, + "input": args.text, + "response_format": args.response_format, + } + if args.voice: + payload["voice"] = args.voice + + if args.ref_audio: + ref = args.ref_audio + if ref.startswith(("http://", "https://", "data:")): + payload["ref_audio"] = ref + else: + payload["ref_audio"] = encode_audio_to_base64(ref) + + extra_params = {} + if args.emo_audio: + emo = args.emo_audio + if emo.startswith(("http://", "https://", "data:")): + extra_params["emo_audio"] = emo + else: + extra_params["emo_audio"] = encode_audio_to_base64(emo) + if args.emo_text: + extra_params["emo_text"] = args.emo_text + if args.emo_vector is not None: + extra_params["emo_vector"] = args.emo_vector + if args.emo_alpha is not None: + extra_params["emo_alpha"] = args.emo_alpha + if args.use_emo_text: + extra_params["use_emo_text"] = True + if args.use_random: + extra_params["use_random"] = True + if extra_params: + payload["extra_params"] = extra_params + + url = f"{args.api_base}/v1/audio/speech" + print(f"POST {url}") + print(f" text: {args.text}") + if args.ref_audio: + print(f" ref_audio: {args.ref_audio[:80]}...") + + with httpx.Client(timeout=300) as client: + resp = client.post( + url, + json=payload, + headers={"Authorization": f"Bearer {args.api_key}"}, + ) + + if resp.status_code != 200: + print(f"Error {resp.status_code}: {resp.text[:500]}") + return + + with open(args.output, "wb") as f: + f.write(resp.content) + print(f"Saved: {args.output} ({len(resp.content):,} bytes)") + + +if __name__ == "__main__": + main() 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 index 536bfb04877..cd7bfafbd24 100644 --- 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 @@ -123,6 +123,7 @@ def run_tts(args) -> None: payload["max_new_tokens"] = args.max_new_tokens if args.stream: payload["stream"] = True + payload["stream_format"] = "audio" payload["response_format"] = "pcm" api_url = f"{args.api_base}/v1/audio/speech" 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 index 10fcaf59ae0..35f7956f2a8 100755 --- a/examples/online_serving/text_to_speech/ming_tts/run_curl.sh +++ b/examples/online_serving/text_to_speech/ming_tts/run_curl.sh @@ -65,6 +65,7 @@ PY \"model\": \"${MODEL}\", \"input\": \"${TEXT}\", \"stream\": true, + \"stream_format\": \"audio\", \"response_format\": \"pcm\" }" "$STREAM_OUTPUT" ;; diff --git a/examples/online_serving/text_to_speech/moss_tts_nano/gradio_demo.py b/examples/online_serving/text_to_speech/moss_tts_nano/gradio_demo.py index 4a8e4896d20..0278db09f03 100644 --- a/examples/online_serving/text_to_speech/moss_tts_nano/gradio_demo.py +++ b/examples/online_serving/text_to_speech/moss_tts_nano/gradio_demo.py @@ -335,6 +335,7 @@ def build_payload( "ref_audio": encode_audio_to_base64(ref_audio), "response_format": "pcm" if stream else response_format, "stream": stream, + **({"stream_format": "audio"} if stream else {}), } diff --git a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py index 76c348bbbae..92a40739ef0 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py @@ -18,6 +18,21 @@ --text "Hello world. How are you? I am fine." \ --simulate-stt --stt-delay 0.1 + # Receive JSON sidecar chunks with word-level timestamps + python streaming_speech_client.py \ + --text "Hello world. How are you?" \ + --stream-audio \ + --response-format pcm \ + --word-timestamps + + # Also save every received audio.chunk frame as individual PCM/WAV files + python streaming_speech_client.py \ + --text "Hello world. How are you?" \ + --stream-audio \ + --response-format pcm \ + --word-timestamps \ + --save-chunks + # VoiceDesign task python streaming_speech_client.py \ --text "Today is a great day. The weather is nice." \ @@ -37,8 +52,10 @@ import argparse import asyncio +import base64 import json import os +import wave try: import websockets @@ -47,6 +64,37 @@ raise SystemExit(1) +def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str: + audio_bytes = b"".join(chunks) + if response_format == "pcm": + raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm") + with open(raw_path, "wb") as f: + f.write(audio_bytes) + + wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav") + save_pcm_wav(wav_path, audio_bytes) + return wav_path + + filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}") + with open(filename, "wb") as f: + f.write(audio_bytes) + return filename + + +def save_pcm_wav(path: str, audio_bytes: bytes, sample_rate: int = 24000) -> None: + with wave.open(path, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(audio_bytes) + + +def timestamp_words_text(timestamps: list[dict] | None) -> str: + if not timestamps: + return "" + return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word")) + + async def stream_tts( url: str, text: str, @@ -54,6 +102,7 @@ async def stream_tts( output_dir: str, simulate_stt: bool = False, stt_delay: float = 0.1, + save_chunks: bool = False, ) -> None: """Connect to the streaming TTS endpoint and process audio responses.""" os.makedirs(output_dir, exist_ok=True) @@ -99,8 +148,14 @@ async def send_text(): sender_task = asyncio.create_task(send_text()) response_format = config.get("response_format", "wav") + word_timestamps = bool(config.get("word_timestamps", False)) current_sentence_index = 0 + current_sentence_text = "" current_chunks: list[bytes] = [] + current_timestamps: list[dict] = [] + # Distinguish `null` (aligner failed) from `[]` (silence / no tokens): + # the sentence is a failure only if no frame ever carried timestamps. + alignment_frame_seen = False try: while True: @@ -116,22 +171,94 @@ async def send_text(): if msg_type == "audio.start": current_sentence_index = msg["sentence_index"] + current_sentence_text = msg.get("sentence_text", "") current_chunks = [] - print(f" [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}") + current_timestamps = [] + alignment_frame_seen = False + print(f" [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}") + elif msg_type == "audio.chunk": + audio = base64.b64decode(msg["audio_b64"]) + current_chunks.append(audio) + chunk_timestamps = msg.get("timestamps") + # Sentence-level alignment sends a trailing timestamp-only + # frame with empty audio; don't write an empty chunk file. + if save_chunks and audio: + chunk_base = os.path.join( + output_dir, + f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}", + ) + chunk_pcm = f"{chunk_base}.pcm" + chunk_wav = f"{chunk_base}.wav" + chunk_json = f"{chunk_base}_timestamps.json" + with open(chunk_pcm, "wb") as f: + f.write(audio) + save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000))) + with open(chunk_json, "w", encoding="utf-8") as f: + json.dump( + { + "sentence_index": msg["sentence_index"], + "chunk_id": msg["chunk_id"], + "chunk_start_ms": msg.get("chunk_start_ms"), + "chunk_end_ms": msg.get("chunk_end_ms"), + "sentence_text": current_sentence_text, + "covered_text": timestamp_words_text(chunk_timestamps), + "timestamps": chunk_timestamps, + }, + f, + ensure_ascii=False, + indent=2, + ) + if chunk_timestamps is None: + print( + f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: " + f"{len(audio)} bytes, timestamps unavailable" + ) + else: + alignment_frame_seen = True + current_timestamps.extend(chunk_timestamps) + print( + f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: " + f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)" + ) + if save_chunks and audio: + print(f" saved chunk -> {chunk_wav}") elif msg_type == "audio.done": - filename = os.path.join( + filename = save_audio_file( output_dir, - f"sentence_{msg['sentence_index']:03d}.{response_format}", + msg["sentence_index"], + response_format, + current_chunks, ) - with open(filename, "wb") as f: - f.write(b"".join(current_chunks)) + if word_timestamps: + # `null` on failure, else the accumulated list (`[]` for silence). + sentence_timestamps = current_timestamps if alignment_frame_seen else None + ts_filename = os.path.join( + output_dir, + f"sentence_{msg['sentence_index']:03d}_timestamps.json", + ) + with open(ts_filename, "w", encoding="utf-8") as f: + json.dump( + { + "sentence_index": msg["sentence_index"], + "sentence_text": current_sentence_text, + "covered_text": timestamp_words_text(sentence_timestamps), + "timestamps": sentence_timestamps, + }, + f, + ensure_ascii=False, + indent=2, + ) print( f" [sentence {msg['sentence_index']}] Done" f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}" f" error={msg.get('error', False)}" f" -> {filename}" ) + if word_timestamps: + print(f" [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}") current_chunks = [] + current_timestamps = [] + alignment_frame_seen = False elif msg_type == "session.done": print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated") break @@ -189,6 +316,16 @@ def main(): action="store_true", help="Receive one or more PCM chunks per sentence (requires --response-format pcm)", ) + parser.add_argument( + "--word-timestamps", + action="store_true", + help="Receive JSON sidecar chunks with word-level timestamps (requires --stream-audio --response-format pcm)", + ) + parser.add_argument( + "--save-chunks", + action="store_true", + help="Save each audio.chunk frame as sentence_XXX_chunk_YYY.{pcm,wav,json}", + ) parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)") parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens") @@ -238,6 +375,10 @@ def main(): config[key] = val if args.stream_audio: config["stream_audio"] = True + if args.word_timestamps: + config["word_timestamps"] = True + config["stream_audio"] = True + config["response_format"] = "pcm" if args.x_vector_only_mode: config["x_vector_only_mode"] = True @@ -249,6 +390,7 @@ def main(): output_dir=args.output_dir, simulate_stt=args.simulate_stt, stt_delay=args.stt_delay, + save_chunks=args.save_chunks, ) ) diff --git a/examples/online_serving/text_to_speech/qwen3_tts/tts_common.py b/examples/online_serving/text_to_speech/qwen3_tts/tts_common.py index 44ba7c3ab7b..df66d452b76 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/tts_common.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/tts_common.py @@ -93,6 +93,8 @@ def build_payload( "response_format": "pcm" if stream else response_format, "stream": stream, } + if stream: + payload["stream_format"] = "audio" if not stream: payload["speed"] = speed diff --git a/examples/online_serving/text_to_speech/qwen3_tts/word_timestamps_demo.py b/examples/online_serving/text_to_speech/qwen3_tts/word_timestamps_demo.py new file mode 100644 index 00000000000..704a64a063d --- /dev/null +++ b/examples/online_serving/text_to_speech/qwen3_tts/word_timestamps_demo.py @@ -0,0 +1,457 @@ +"""Gradio demo for streaming TTS word-level timestamps. + +Connects to the WebSocket endpoint ``/v1/audio/speech/stream`` with +``word_timestamps: true`` and visualizes the alignment: each sentence's +audio plays in a native ``<audio>`` element while its text is rendered as +inline word spans, the current word highlighting as ``audio.currentTime`` +crosses each ``start_ms``. + +A "Stop (barge-in)" button cuts playback and reports the last-spoken word, +demonstrating the voice-agent barge-in use case. + +Timestamps are sentence-level: audio streams in real time, then the forced +aligner runs once per sentence and the complete word timestamps arrive in a +trailing frame. The demo therefore plays one sentence block at a time. + +Usage: + # Launch the server with a forced aligner first, e.g.: + # vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \ + # --deploy-config vllm_omni/deploy/qwen3_tts.yaml --trust-remote-code \ + # --forced-aligner Qwen/Qwen3-ForcedAligner-0.6B + # Then: + python word_timestamps_demo.py --api-base http://localhost:8000 +""" + +import argparse +import asyncio +import json +import logging + +try: + import gradio as gr +except ImportError: + raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None +try: + import websockets +except ImportError: + raise ImportError("websockets is required to run this demo. Install it with: pip install websockets") from None +from fastapi import FastAPI, WebSocket, WebSocketDisconnect +from tts_common import ( + PCM_SAMPLE_RATE, + SUPPORTED_LANGUAGES, + TASK_TYPES, + add_common_args, + fetch_voices, +) + +logger = logging.getLogger(__name__) + + +def _upstream_ws_url(api_base: str) -> str: + """Map an http(s) API base to its WebSocket streaming endpoint.""" + base = api_base.rstrip("/") + if base.startswith("https://"): + base = "wss://" + base[len("https://") :] + elif base.startswith("http://"): + base = "ws://" + base[len("http://") :] + else: + base = "ws://" + base + return base + "/v1/audio/speech/stream" + + +# ── Front-end ──────────────────────────────────────────────────────── +PLAYER_HTML = """ +<div id="wt-root"> + <div style="display:flex; align-items:center; gap:10px;"> + <div id="wt-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div> + <span id="wt-status" style="font-weight:600;font-size:1.05em;">Ready</span> + <button id="wt-stop" onclick="window.wtStop()" + style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552; + background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em;">Stop (barge-in)</button> + </div> + <div id="wt-barge" style="display:none; margin-top:10px; padding:8px 12px; border-radius:6px; + background:#fff5f5; border:1px solid #f3c2c2; color:#b02a2a; font-size:0.9em;"></div> + <div id="wt-sentences" style="margin-top:12px; display:flex; flex-direction:column; gap:14px;"></div> +</div> +""" + + +def _build_head_js(sample_rate: int) -> str: + """Browser-side player: WebSocket -> per-sentence <audio> + word highlight.""" + return f""" +<script> +const WT_SR = {sample_rate}; +let wtWS = null, wtQueue = [], wtPlaying = false, wtCur = null, wtLastWord = ""; + +function wtSet(text, color) {{ + const s = document.getElementById('wt-status'); + const d = document.getElementById('wt-dot'); + if (s) s.textContent = text; + if (d) d.style.background = color || '#ccc'; +}} + +function wtStopBtn(show) {{ + const b = document.getElementById('wt-stop'); + if (b) b.style.display = show ? 'inline-block' : 'none'; +}} + +// Build a WAV Blob (16-bit mono PCM) from concatenated Int16 samples. +function wtPcmToWav(int16, sr) {{ + const n = int16.length; + const buf = new ArrayBuffer(44 + n * 2); + const v = new DataView(buf); + const w = (off, s) => {{ for (let i = 0; i < s.length; i++) v.setUint8(off + i, s.charCodeAt(i)); }}; + w(0, 'RIFF'); v.setUint32(4, 36 + n * 2, true); w(8, 'WAVE'); + w(12, 'fmt '); v.setUint32(16, 16, true); v.setUint16(20, 1, true); v.setUint16(22, 1, true); + v.setUint32(24, sr, true); v.setUint32(28, sr * 2, true); v.setUint16(32, 2, true); v.setUint16(34, 16, true); + w(36, 'data'); v.setUint32(40, n * 2, true); + let off = 44; + for (let i = 0; i < n; i++) {{ v.setInt16(off, int16[i], true); off += 2; }} + return new Blob([buf], {{ type: 'audio/wav' }}); +}} + +function wtConcat(chunks) {{ + let total = 0; + for (const c of chunks) total += c.length; + const out = new Int16Array(total); + let off = 0; + for (const c of chunks) {{ out.set(c, off); off += c.length; }} + return out; +}} + +// Render one finished sentence: <audio> + inline word spans. +function wtRenderSentence(sentence) {{ + const wrap = document.getElementById('wt-sentences'); + if (!wrap) return; + const block = document.createElement('div'); + block.style.cssText = 'border:1px solid #e3e8ef;border-radius:8px;padding:10px 12px;'; + + const head = document.createElement('div'); + head.style.cssText = 'font-size:0.75em;color:#888;margin-bottom:6px;'; + head.textContent = 'Sentence ' + sentence.index + + (sentence.timestamps ? ' (' + sentence.timestamps.length + ' words)' : ' (no timestamps)'); + block.appendChild(head); + + const textEl = document.createElement('div'); + textEl.style.cssText = 'font-size:1.15em;line-height:1.8;margin-bottom:8px;'; + if (sentence.timestamps && sentence.timestamps.length) {{ + sentence.spans = []; + for (const ts of sentence.timestamps) {{ + const span = document.createElement('span'); + span.textContent = ts.word; + span.dataset.start = ts.start_ms; + span.dataset.end = ts.end_ms; + span.style.cssText = 'padding:1px 2px;border-radius:3px;transition:background 0.1s;'; + textEl.appendChild(span); + // Keep spacing for spaced languages; CJK chars sit flush. + textEl.appendChild(document.createTextNode(' ')); + sentence.spans.push(span); + }} + }} else {{ + textEl.textContent = sentence.text || '(timestamps unavailable)'; + }} + block.appendChild(textEl); + + const audio = document.createElement('audio'); + audio.controls = true; + audio.style.cssText = 'width:100%;'; + audio.src = URL.createObjectURL(wtPcmToWav(wtConcat(sentence.chunks), WT_SR)); + block.appendChild(audio); + sentence.audio = audio; + + // Highlight the current word. Driven by requestAnimationFrame (~60fps) + // rather than the <audio> 'timeupdate' event, which only fires ~4x/sec + // and is too coarse for short words. We light the latest word whose + // start_ms has passed, so there is always exactly one active word (no + // inter-word flicker) and wtLastWord stays well-defined for barge-in. + let raf = null; + function highlight() {{ + if (!sentence.spans) return; + const t = audio.currentTime * 1000; + let idx = -1; + for (let i = 0; i < sentence.spans.length; i++) {{ + if (t >= +sentence.spans[i].dataset.start) idx = i; + else break; + }} + sentence.spans.forEach((sp, i) => {{ + const on = i === idx; + sp.style.background = on ? '#ffe680' : 'transparent'; + sp.style.fontWeight = on ? '700' : '400'; + }}); + if (idx >= 0) wtLastWord = sentence.spans[idx].textContent; + }} + function loop() {{ + highlight(); + if (!audio.paused && !audio.ended) raf = requestAnimationFrame(loop); + }} + audio.addEventListener('play', () => {{ + wtCur = audio; + wtSet('Playing sentence ' + sentence.index + '...', '#64dd17'); + cancelAnimationFrame(raf); + raf = requestAnimationFrame(loop); + }}); + audio.addEventListener('pause', () => cancelAnimationFrame(raf)); + audio.addEventListener('seeked', highlight); + audio.addEventListener('ended', () => {{ + cancelAnimationFrame(raf); + highlight(); + wtPlaying = false; + wtPlayNext(); + }}); + + wrap.appendChild(block); + wtQueue.push(audio); + wtPlayNext(); +}} + +function wtPlayNext() {{ + if (wtPlaying) return; + const next = wtQueue.shift(); + if (!next) {{ + if (wtWS === null) {{ wtSet('Done', '#64dd17'); wtStopBtn(false); }} + return; + }} + wtPlaying = true; + next.play().catch(() => {{ wtPlaying = false; }}); +}} + +window.wtStop = function() {{ + // Barge-in: cut playback, stop generation, report the last-spoken word. + if (wtCur) {{ try {{ wtCur.pause(); }} catch (e) {{}} }} + wtQueue = []; + wtPlaying = false; + if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} wtWS = null; }} + wtSet('Stopped (barge-in)', '#999'); + wtStopBtn(false); + const b = document.getElementById('wt-barge'); + if (b) {{ + b.style.display = 'block'; + b.textContent = wtLastWord + ? 'Barge-in at word: "' + wtLastWord + '" (a voice agent would resume the conversation from here)' + : 'Barge-in before any word was spoken.'; + }} +}}; + +window.wtGenerate = function(payload) {{ + if (wtWS) {{ try {{ wtWS.close(); }} catch (e) {{}} }} + wtQueue = []; wtPlaying = false; wtCur = null; wtLastWord = ""; + document.getElementById('wt-sentences').innerHTML = ''; + const barge = document.getElementById('wt-barge'); + if (barge) {{ barge.style.display = 'none'; barge.textContent = ''; }} + wtStopBtn(true); + wtSet('Connecting...', '#4A90D9'); + + const proto = location.protocol === 'https:' ? 'wss:' : 'ws:'; + const ws = new WebSocket(proto + '//' + location.host + '/proxy/stream'); + wtWS = ws; + const sentences = {{}}; + + ws.onopen = () => {{ + wtSet('Streaming...', '#4A90D9'); + ws.send(JSON.stringify(Object.assign({{ type: 'session.config' }}, payload.config))); + ws.send(JSON.stringify({{ type: 'input.text', text: payload.text }})); + ws.send(JSON.stringify({{ type: 'input.done' }})); + }}; + + ws.onmessage = (ev) => {{ + let msg; + try {{ msg = JSON.parse(ev.data); }} catch (e) {{ return; }} + if (msg.type === 'audio.start') {{ + sentences[msg.sentence_index] = {{ + index: msg.sentence_index, text: msg.sentence_text || '', + chunks: [], timestamps: null + }}; + }} else if (msg.type === 'audio.chunk') {{ + const s = sentences[msg.sentence_index]; + if (!s) return; + if (msg.audio_b64) {{ + const bin = atob(msg.audio_b64); + const bytes = new Uint8Array(bin.length); + for (let i = 0; i < bin.length; i++) bytes[i] = bin.charCodeAt(i); + const usable = bytes.length - (bytes.length % 2); + if (usable > 0) s.chunks.push(new Int16Array(bytes.buffer, 0, usable / 2)); + }} + if (msg.timestamps != null) s.timestamps = msg.timestamps; + }} else if (msg.type === 'audio.done') {{ + const s = sentences[msg.sentence_index]; + if (s) wtRenderSentence(s); + }} else if (msg.type === 'session.done') {{ + wtWS = null; + try {{ ws.close(); }} catch (e) {{}} + wtPlayNext(); + }} else if (msg.type === 'error') {{ + wtSet('Error: ' + (msg.message || 'unknown'), '#EF5552'); + }} + }}; + + ws.onerror = () => wtSet('WebSocket error', '#EF5552'); + ws.onclose = () => {{ if (wtWS === ws) wtWS = null; }}; +}}; +</script> +""" + + +def create_app(api_base: str) -> FastAPI: + """FastAPI app with a same-origin WebSocket proxy + Gradio UI.""" + fastapi_app = FastAPI() + upstream_url = _upstream_ws_url(api_base) + + @fastapi_app.websocket("/proxy/stream") + async def proxy_stream(client_ws: WebSocket): + await client_ws.accept() + try: + async with websockets.connect(upstream_url, max_size=None) as upstream: + + async def client_to_upstream(): + try: + while True: + await upstream.send(await client_ws.receive_text()) + except (WebSocketDisconnect, Exception): + await upstream.close() + + async def upstream_to_client(): + try: + async for msg in upstream: + if isinstance(msg, bytes): + await client_ws.send_bytes(msg) + else: + await client_ws.send_text(msg) + except Exception: + pass + finally: + try: + await client_ws.close() + except Exception: + pass + + await asyncio.gather(client_to_upstream(), upstream_to_client()) + except Exception as exc: + logger.exception("WS proxy error") + try: + await client_ws.send_text(json.dumps({"type": "error", "message": f"proxy: {exc}"})) + await client_ws.close() + except Exception: + pass + + voices = fetch_voices(api_base) + + with gr.Blocks(title="Qwen3-TTS Word Timestamps") as demo: + gr.HTML(f""" + <div style="margin-bottom:8px;"> + <h1 style="margin:0; font-size:1.5em;">Qwen3-TTS Streaming Word Timestamps</h1> + <span style="font-size:0.85em; color:#666;"> + Served by <a href="https://github.com/vllm-project/vllm-omni" target="_blank">vLLM-Omni</a> +  ·  <code>{api_base}</code>  ·  + launch the server with <code>--forced-aligner</code>. + </span> + </div> + """) + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Text to Synthesize", + value="Hello world. How are you today? I am doing great.", + lines=3, + ) + with gr.Row(): + task_type = gr.Radio(choices=TASK_TYPES, value="CustomVoice", label="Task Type", scale=2) + language = gr.Dropdown(choices=SUPPORTED_LANGUAGES, value="Auto", label="Language", scale=1) + voice = gr.Dropdown( + choices=voices, + value=voices[0] if voices else None, + label="Speaker", + allow_custom_value=True, + ) + instructions = gr.Textbox( + label="Instructions", + placeholder="Optional style/emotion (required for VoiceDesign)", + lines=1, + ) + with gr.Column(visible=False) as ref_group: + ref_audio_url = gr.Textbox( + label="Reference Audio URL (Base task)", + placeholder="https://example.com/reference.wav", + lines=1, + ) + ref_text = gr.Textbox(label="Reference Transcript", lines=1) + generate_btn = gr.Button("Generate + Align", variant="primary", size="lg") + + with gr.Column(scale=2): + gr.HTML(value=PLAYER_HTML) + + hidden_payload = gr.Textbox(visible=False, elem_id="wt-payload") + + def on_task_change(tt): + return gr.update(visible=(tt == "Base")) + + task_type.change(fn=on_task_change, inputs=[task_type], outputs=[ref_group]) + + def on_generate(text, tt, voice_v, lang_v, instr_v, ref_url, ref_t): + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + config: dict = { + "task_type": tt, + "word_timestamps": True, + "stream_audio": True, + "response_format": "pcm", + } + if lang_v: + config["language"] = lang_v + if tt == "CustomVoice": + if voice_v: + config["voice"] = voice_v + if instr_v and instr_v.strip(): + config["instructions"] = instr_v.strip() + elif tt == "VoiceDesign": + if not instr_v or not instr_v.strip(): + raise gr.Error("VoiceDesign requires voice style instructions.") + config["instructions"] = instr_v.strip() + elif tt == "Base": + if not ref_url or not ref_url.strip(): + raise gr.Error("Base (voice clone) needs a reference audio URL.") + config["ref_audio"] = ref_url.strip() + if ref_t and ref_t.strip(): + config["ref_text"] = ref_t.strip() + return json.dumps({"config": config, "text": text.strip()}) + + generate_btn.click( + fn=on_generate, + inputs=[text_input, task_type, voice, language, instructions, ref_audio_url, ref_text], + outputs=[hidden_payload], + ).then( + fn=lambda p: p, + inputs=[hidden_payload], + outputs=[hidden_payload], + js="(p) => { if (p && p.trim()) window.wtGenerate(JSON.parse(p)); return p; }", + ) + + demo.queue() + + return gr.mount_gradio_app( + fastapi_app, + demo, + path="/", + head=_build_head_js(PCM_SAMPLE_RATE), + ) + + +def parse_args(): + parser = argparse.ArgumentParser(description="Gradio demo for streaming TTS word-level timestamps.") + add_common_args(parser) + return parser.parse_args() + + +def main(): + args = parse_args() + logging.basicConfig(level=logging.INFO) + print(f"Connecting to vLLM server at: {args.api_base}") + app = create_app(args.api_base) + + import uvicorn + + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py b/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py new file mode 100644 index 00000000000..91035e62510 --- /dev/null +++ b/examples/online_serving/text_to_speech/soulxsinger/openai_chat_client.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +"""SoulX-Singer OpenAI-compatible chat client (SVS / SVC). + +Sends prompt audio via ``input_audio`` and target accompaniment via +``extra_args['target_audio']`` (server-local path). For integrated preprocess, +also pass ``preprocess_weights_dir`` in ``extra_args``. + +Usage: + python openai_chat_client.py \\ + --prompt-audio /path/on/server/zh_prompt.mp3 \\ + --target-audio /path/on/server/music.mp3 \\ + --preprocess-weights-dir /path/on/server/SoulX-Singer-Preprocess \\ + -o output.wav +""" + +from __future__ import annotations + +import argparse +import base64 +import io +import sys +from pathlib import Path + +import requests +import soundfile +import torch + + +def _audio_to_data_url(path: Path) -> str: + with path.open("rb") as handle: + data = base64.b64encode(handle.read()).decode("ascii") + return f"data:audio/mpeg;base64,{data}" + + +def _save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None: + audio = audio.to(torch.float32) + peak = audio.abs().max().clamp(min=1e-8) + audio = audio / peak + path.parent.mkdir(parents=True, exist_ok=True) + soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16") + + +def _decode_audio_from_response(body: dict) -> tuple[torch.Tensor, int]: + for choice in body.get("choices", []): + audio_obj = choice.get("message", {}).get("audio") + if isinstance(audio_obj, dict) and audio_obj.get("data"): + data, sr = soundfile.read( + io.BytesIO(base64.b64decode(audio_obj["data"])), + dtype="float32", + always_2d=True, + ) + return torch.from_numpy(data).transpose(0, 1), sr + brief = {k: v for k, v in body.items() if k != "choices"} + raise RuntimeError(f"no audio in response message.audio: {brief}") + + +def main() -> int: + repo_root = Path(__file__).resolve().parents[4] + default_assets = repo_root / "tests" / "assets" / "soulxsinger" + + parser = argparse.ArgumentParser(description="SoulX-Singer online chat client") + parser.add_argument("--port", type=int, default=8192) + parser.add_argument("--model", default="Soul-AILab/SoulX-Singer") + parser.add_argument( + "--prompt-audio", + default=str(default_assets / "zh_prompt.mp3"), + help="Prompt vocal audio (path on server if using extra_args, or local for input_audio)", + ) + parser.add_argument( + "--target-audio", + default=str(default_assets / "music.mp3"), + help="Target accompaniment path on the server (extra_args['target_audio'])", + ) + parser.add_argument( + "--prompt-metadata-path", + default=None, + help="SVS precomputed prompt metadata.json", + ) + parser.add_argument( + "--target-metadata-path", + default=None, + help="SVS precomputed target metadata.json", + ) + parser.add_argument( + "--audio-path", + default=None, + help="SVS prompt vocal wav for precomputed metadata", + ) + parser.add_argument("--preprocess-weights-dir", default=None) + parser.add_argument("--output", "-o", default="soulxsinger_out.wav") + parser.add_argument("--svc", action="store_true", help="Use SVC mode knobs") + parser.add_argument("--language", default="Mandarin") + parser.add_argument("--num-inference-steps", type=int, default=32) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument( + "--seed", + type=int, + default=42, + help="Optional CFM seed. Omit for non-deterministic sampling.", + ) + parser.add_argument( + "--auto-shift", + action=argparse.BooleanOptionalAction, + default=True, + help="Auto pitch shift (default: on, original upstream infer.sh)", + ) + parser.add_argument( + "--control", + default="melody", + choices=["melody", "score"], + help="SVS control mode", + ) + parser.add_argument("--vocal-sep", action="store_true") + args = parser.parse_args() + + meta_paths = (args.prompt_metadata_path, args.target_metadata_path, args.audio_path) + if any(meta_paths) and not all(meta_paths): + print( + "ERROR: precomputed metadata requires --prompt-metadata-path, " + "--target-metadata-path, and --audio-path together.", + file=sys.stderr, + ) + return 2 + + extra_args: dict = { + "vocal_sep": args.vocal_sep, + "auto_shift": args.auto_shift, + "pitch_shift": 0, + } + if all(meta_paths): + extra_args.update( + { + "prompt_metadata_path": str(Path(args.prompt_metadata_path).expanduser().resolve()), + "target_metadata_path": str(Path(args.target_metadata_path).expanduser().resolve()), + "audio_path": str(Path(args.audio_path).expanduser().resolve()), + } + ) + content = [{"type": "text", "text": "soulx-singer"}] + else: + prompt_path = Path(args.prompt_audio).expanduser().resolve() + if not prompt_path.is_file(): + print(f"ERROR: prompt audio not found: {prompt_path}", file=sys.stderr) + return 2 + extra_args["prompt_audio"] = str(prompt_path) + extra_args["target_audio"] = str(Path(args.target_audio).expanduser().resolve()) + if args.preprocess_weights_dir: + extra_args["preprocess_weights_dir"] = str(Path(args.preprocess_weights_dir).expanduser().resolve()) + content = [ + {"type": "text", "text": "soulx-singer"}, + { + "type": "input_audio", + "input_audio": {"data": _audio_to_data_url(prompt_path), "format": "mp3"}, + }, + ] + if not args.svc: + extra_args["language"] = args.language + extra_args["control"] = args.control + + payload = { + "model": args.model, + "modalities": ["audio"], + "messages": [{"role": "user", "content": content}], + "num_inference_steps": args.num_inference_steps, + "guidance_scale": args.guidance_scale, + "extra_args": extra_args, + } + if args.seed is not None: + payload["seed"] = args.seed + + print(f"POST http://localhost:{args.port}/v1/chat/completions") + response = requests.post( + f"http://localhost:{args.port}/v1/chat/completions", + headers={"Content-Type": "application/json"}, + json=payload, + timeout=1800, + ) + response.raise_for_status() + audio, sample_rate = _decode_audio_from_response(response.json()) + _save_wav(audio, Path(args.output), sample_rate) + duration = audio.shape[-1] / sample_rate + print(f"saved {args.output} sr={sample_rate}Hz duration={duration:.2f}s") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/online_serving/text_to_speech/soulxsinger/run_server.sh b/examples/online_serving/text_to_speech/soulxsinger/run_server.sh new file mode 100755 index 00000000000..3e20de39bdf --- /dev/null +++ b/examples/online_serving/text_to_speech/soulxsinger/run_server.sh @@ -0,0 +1,39 @@ +#!/bin/bash +# Launch vLLM-Omni server for SoulX-Singer (single-stage DiT, preprocess inline). +# +# Usage: +# MODEL=/path/to/SoulX-Singer PREPROCESS=/path/to/Preprocess \ +# +# Audio paths in client extra_args must be readable on the server host. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)" + +MODEL="${MODEL:-Soul-AILab/SoulX-Singer}" +MODE="${MODE:-svs}" +PORT="${PORT:-8192}" +GPUS="${GPUS:-0}" + +if [[ "$MODE" == "svc" ]]; then + DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svc.yaml}" +else + DEPLOY_CONFIG="${DEPLOY_CONFIG:-$REPO_ROOT/vllm_omni/deploy/soulxsinger_svs.yaml}" +fi + +echo "Starting SoulX-Singer server" +echo " MODEL=$MODEL" +echo " MODE=$MODE" +echo " PORT=$PORT" +echo " DEPLOY_CONFIG=$DEPLOY_CONFIG" +echo " CUDA_VISIBLE_DEVICES=$GPUS" + +CUDA_VISIBLE_DEVICES="$GPUS" \ +vllm serve "$MODEL" \ + --omni \ + --deploy-config "$DEPLOY_CONFIG" \ + --host 0.0.0.0 \ + --port "$PORT" \ + --trust-remote-code \ + --enforce-eager diff --git a/examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py b/examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py index c6706198ae4..f9df10d528d 100644 --- a/examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py +++ b/examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py @@ -515,6 +515,8 @@ def on_generate(stream_enabled, text, ctrl_instr, ref_a, ref_url, ult_clone, p_t "response_format": "pcm" if stream_enabled else "wav", "stream": stream_enabled, } + if stream_enabled: + payload["stream_format"] = "audio" # Reference audio for cloning ref_url_s = ref_url.strip() if ref_url else "" diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index c01e0602ff9..93de44a88e5 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/examples/online_serving/text_to_video/README.md @@ -10,6 +10,7 @@ This example demonstrates how to deploy text-to-video models for online video ge | Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | | Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | | LTX-2 | `Lightricks/LTX-2` | +| Helios (Base / Mid / Distilled) | `BestWishYsh/Helios-Base`, `Helios-Mid`, `Helios-Distilled` | ## Wan2.2 T2V @@ -87,14 +88,17 @@ curl -X POST http://localhost:8091/v1/videos/sync \ Generated video files are stored on local disk by the async video API. Local file storage behavior can be controlled via the following environment variables: -- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`) -- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`) +- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files (default: `/tmp/storage`) +- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete operations (default: `4`) + +`VLLM_OMNI_STORAGE_PATH` and `VLLM_OMNI_STORAGE_MAX_CONCURRENCY` are deprecated and will be +removed in a future release; use the names above instead. Example: ```bash -export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos -export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8 +export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos +export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8 ``` ## API Calls @@ -320,3 +324,45 @@ curl -sS -X POST http://localhost:8098/v1/videos \ -F "guidance_scale=3.0" \ -F "seed=42" ``` + +## Helios + +Helios ships three variants (`Helios-Base`, `Helios-Mid`, `Helios-Distilled`) that +share the same server launch. Variant-specific knobs (declared in +`vllm_omni/model_extras/helios.py`) are sent per request through the generic +`extra_params` JSON form field — no per-model server flags required. + +### Start Server + +```bash +vllm serve BestWishYsh/Helios-Base --omni --port 8098 +# or: MODEL=BestWishYsh/Helios-Mid bash run_server_helios.sh +``` + +### Send Requests (curl) + +```bash +# Helios-Base (Stage 1 only) +bash run_curl_helios.sh + +# Helios-Mid (Stage 2 pyramid + CFG-Zero*) +PRESET=mid-stage2 MODEL=BestWishYsh/Helios-Mid bash run_curl_helios.sh + +# Helios-Distilled (Stage 2 pyramid + DMD, few-step) +PRESET=distilled MODEL=BestWishYsh/Helios-Distilled bash run_curl_helios.sh +``` + +The `mid-stage2` and `distilled` presets attach an `extra_params` field, e.g. for Helios-Distilled: + +```bash +curl -sS -X POST http://localhost:8098/v1/videos \ + -H "Accept: application/json" \ + -F "prompt=A dynamic time-lapse of scenery rushing past the window of a speeding train." \ + -F "model=BestWishYsh/Helios-Distilled" \ + -F "size=640x384" \ + -F "num_frames=99" \ + -F "fps=16" \ + -F "guidance_scale=1.0" \ + -F "seed=42" \ + -F 'extra_params={"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}' +``` diff --git a/examples/online_serving/text_to_video/run_curl_helios.sh b/examples/online_serving/text_to_video/run_curl_helios.sh new file mode 100755 index 00000000000..6a1c6f6ba03 --- /dev/null +++ b/examples/online_serving/text_to_video/run_curl_helios.sh @@ -0,0 +1,93 @@ +#!/bin/bash +# Helios text-to-video curl example using the async video job API. +# +# Helios-specific knobs (declared in vllm_omni/model_extras/helios.py) are passed +# through the generic `extra_params` JSON form field. Select a variant via PRESET: +# PRESET=base -> Helios-Base, Stage 1 only (default) +# PRESET=mid-stage2 -> Helios-Mid, Stage 2 pyramid + CFG-Zero* +# PRESET=distilled -> Helios-Distilled, Stage 2 pyramid + DMD (few-step) + +set -euo pipefail + +BASE_URL="${BASE_URL:-http://localhost:8098}" +MODEL="${MODEL:-BestWishYsh/Helios-Base}" +PROMPT="${PROMPT:-A dynamic time-lapse of scenery rushing past the window of a speeding train.}" +PRESET="${PRESET:-base}" +OUTPUT_PATH="${OUTPUT_PATH:-helios_t2v_${PRESET}.mp4}" +POLL_INTERVAL="${POLL_INTERVAL:-2}" + +case "${PRESET}" in + base) + EXTRA_PARAMS="" + GUIDANCE_SCALE="5.0" + ;; + mid-stage2) + EXTRA_PARAMS='{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}' + GUIDANCE_SCALE="5.0" + ;; + distilled) + EXTRA_PARAMS='{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}' + GUIDANCE_SCALE="1.0" + ;; + *) + echo "Unknown PRESET '${PRESET}' (expected base|mid-stage2|distilled)" + exit 1 + ;; +esac + +create_args=( + -sS -X POST "${BASE_URL}/v1/videos" + -H "Accept: application/json" + -F "prompt=${PROMPT}" + -F "model=${MODEL}" + -F "size=640x384" + -F "num_frames=99" + -F "fps=16" + -F "num_inference_steps=50" + -F "guidance_scale=${GUIDANCE_SCALE}" + -F "seed=42" +) +if [ -n "${EXTRA_PARAMS}" ]; then + create_args+=(-F "extra_params=${EXTRA_PARAMS}") +fi + +create_response=$(curl "${create_args[@]}") + +video_id="$(echo "${create_response}" | jq -r '.id')" +if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then + echo "Failed to create video job:" + echo "${create_response}" | jq . + exit 1 +fi + +echo "Created video job ${video_id} (preset=${PRESET})" +echo "${create_response}" | jq . + +while true; do + status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")" + status="$(echo "${status_response}" | jq -r '.status')" + + case "${status}" in + queued|in_progress) + echo "Video job ${video_id} status: ${status}" + sleep "${POLL_INTERVAL}" + ;; + completed) + echo "${status_response}" | jq . + break + ;; + failed) + echo "Video generation failed:" + echo "${status_response}" | jq . + exit 1 + ;; + *) + echo "Unexpected status response:" + echo "${status_response}" | jq . + exit 1 + ;; + esac +done + +curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}" +echo "Saved video to ${OUTPUT_PATH}" diff --git a/examples/online_serving/text_to_video/run_server_helios.sh b/examples/online_serving/text_to_video/run_server_helios.sh new file mode 100755 index 00000000000..3f18a12a3fd --- /dev/null +++ b/examples/online_serving/text_to_video/run_server_helios.sh @@ -0,0 +1,15 @@ +#!/bin/bash +# Helios online serving startup script. +# All three variants (Helios-Base / Helios-Mid / Helios-Distilled) share the same +# server launch; variant-specific knobs are sent per-request via `extra_params` +# (see run_curl_helios.sh). + +MODEL="${MODEL:-BestWishYsh/Helios-Base}" +PORT="${PORT:-8098}" + +echo "Starting Helios server..." +echo "Model: $MODEL" +echo "Port: $PORT" + +vllm serve "$MODEL" --omni \ + --port "$PORT" diff --git a/midstream/vllm-version b/midstream/vllm-version index 0c2a959e86e..70449b9e7ab 100644 --- a/midstream/vllm-version +++ b/midstream/vllm-version @@ -1 +1 @@ -v0.23.0 +v0.25.1 diff --git a/midstream/vllm-wheels.yml b/midstream/vllm-wheels.yml index ffd6818dbaa..534f5812346 100644 --- a/midstream/vllm-wheels.yml +++ b/midstream/vllm-wheels.yml @@ -1,3 +1,8 @@ +v0.25.1: + run_id: "29766140891" + branch: "sync-v0.25.1" + note: "v0.25.1 wheel, built 2026-07-20 for upstream rebase" + v0.23.0: run_id: "27562765781" branch: "doug/v0.23.0" diff --git a/mkdocs.yml b/mkdocs.yml index 74f2145fe2f..8d5ffaa1bf1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -102,6 +102,7 @@ plugins: - "vllm_omni.entrypoints.async_diffusion" # avoid importing vllm in mkdocs building - "vllm_omni.entrypoints.openai" # avoid importing vllm in mkdocs building - "vllm_omni.entrypoints.openai.protocol" # avoid importing vllm in mkdocs building + - "vllm_omni.experimental" # avoid importing optional serving deps (openai, fastapi, uvicorn, httpx) in mkdocs building - mkdocstrings: handlers: python: diff --git a/pyproject.toml b/pyproject.toml index 21b0ff42e57..4a39c45c2f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,8 @@ dev = [ "pytest-asyncio>=0.21.0", "pytest-cov>=4.0.0", "pytest-mock>=3.10.0", + "websockets>=15.0.1", + "pytest-subtests>=0.11.0", "datasets>=2.14.0", "mypy==1.11.1", "pre-commit==4.0.1", @@ -67,11 +69,18 @@ dev = [ "torchdiffeq>=0.2.5", # for covo_audio Test "FlagEmbedding", # for hunyuan_image3 tests "kernels<0.15.0", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0) - "s3tokenizer>=0.3.0" # for tts test + "s3tokenizer>=0.3.0", # for tts test + "step-audio2>=1.0.0", # Step-Audio2 Token2Wav (flashcosyvoice / hyperpyyaml) ] demo = [ - "gradio>=6.7.0", + "opencv-python>=4.12.0.88", # cli/stream_client frame reading + "requests>=2.28.0", # cli/stream_client HTTP + "cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次") +] + +forced-aligner = [ + "qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner ] # Local-test-only extras @@ -112,6 +121,12 @@ docs = [ "pydantic", ] +soulx-svs = [ + "funasr>=1.0.0", + "g2pM>=0.1.2.5", + "g2p-en>=2.1.0", + "ToJyutping>=3.2.0", +] [project.urls] Homepage = "https://github.com/vllm-project/vllm-omni" @@ -136,7 +151,6 @@ include = ["vllm_omni*"] [tool.setuptools.package-data] "vllm_omni" = ["_version.py", "py.typed"] -"vllm_omni.model_executor.stage_configs" = ["*.yaml"] [tool.setuptools_scm] # no extra settings needed, presence enables setuptools-scm @@ -181,6 +195,7 @@ ignore = [ # These files keep upstream conventions (star imports, long lines, math-symbol # argument names) so they stay easy to diff against upstream releases. "vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"] +"vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py" = ["E501"] # System/summarizer prompts are long prose lines [tool.ruff.lint.flake8-tidy-imports.banned-api] "librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead" @@ -228,6 +243,7 @@ markers = [ "tts: TTS model tests", "cache: Cache backend tests", "parallel: Parallelism/distributed tests", + "sp: Sequence parallelism tests", "example: Doc example code tests", # platform markers "cpu: Tests that run on CPU", @@ -256,6 +272,7 @@ markers = [ "skipif_npu: Skip if the num of NPU cards is less than the required", "skipif_musa: Skip if the num of MUSA cards is less than the required", # more detailed markers + "sp: Sequence parallelism tests (multi-GPU)", "slow: Slow tests (may skip in quick CI)", "benchmark: Benchmark tests", ] @@ -300,3 +317,6 @@ CANN = "CANN" thi = "thi" ND = "ND" byt5 = "byt5" +# "seeked" is the standard HTML media element DOM event name (fires when a +# seek completes), used by the word-timestamps Gradio demo. Not a typo. +seeked = "seeked" diff --git a/recipes/Bagel/BAGEL-7B-MoT.md b/recipes/Bagel/BAGEL-7B-MoT.md new file mode 100644 index 00000000000..a7b0524e47f --- /dev/null +++ b/recipes/Bagel/BAGEL-7B-MoT.md @@ -0,0 +1,227 @@ +# BAGEL-7B-MoT + +> BAGEL image generation through the shared online and offline image examples + +## Summary + +- Vendor: ByteDance Seed +- Model: `ByteDance-Seed/BAGEL-7B-MoT` +- Task: Text-to-image and image-to-image generation +- Mode: Offline inference and OpenAI-compatible online serving +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to run BAGEL through the shared image examples +instead of model-specific example clients. The generic examples can format +BAGEL text-to-image and image-to-image prompts, select the image output +modality, attach reference images, and forward BAGEL-specific generation +parameters through the pipeline-declared `extra_args` contract. + +## References + +- Upstream model: + [`ByteDance-Seed/BAGEL-7B-MoT`](https://huggingface.co/ByteDance-Seed/BAGEL-7B-MoT) +- Related offline example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Related offline image-to-image example: + [`examples/offline_inference/image_to_image/image_edit.py`](../../examples/offline_inference/image_to_image/image_edit.py) +- Related online example: + [`examples/online_serving/text_to_image/openai_chat_client.py`](../../examples/online_serving/text_to_image/openai_chat_client.py) +- Related online image-to-image example: + [`examples/online_serving/image_to_image/openai_chat_client.py`](../../examples/online_serving/image_to_image/openai_chat_client.py) +- Default deploy configs: + [`vllm_omni/deploy/bagel.yaml`](../../vllm_omni/deploy/bagel.yaml), + [`vllm_omni/deploy/bagel_single_stage.yaml`](../../vllm_omni/deploy/bagel_single_stage.yaml) + +## Hardware Support + +This recipe documents the CUDA layouts used by the in-repo BAGEL deploy +configs. The default two-stage config shares one 80 GB GPU; for more headroom, +move the diffusion stage to a second GPU in a custom deploy config. + +## GPU + +### 1x A100 80GB + +#### Environment + +- OS: Linux +- Python: Match the repository requirements for your checkout +- Driver / runtime: NVIDIA CUDA environment with one A100 80 GB GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Offline Commands + +Run text-to-image with the shared offline example from the repository root. +Pick the topology with `--stage-configs-path` and forward BAGEL-specific +generation parameters as a JSON object through `--extra-body`: + +```bash +# Two-stage (Thinker + DiT), shares one GPU by default +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --stage-configs-path vllm_omni/deploy/bagel.yaml \ + --prompt "A beautiful sunset over mountains" \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 50 \ + --negative-prompt "blurry, low quality" \ + --seed 42 \ + --extra-body '{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_img_scale": 1.5, "cfg_interval": [0.4, 1.0], "cfg_renorm_type": "global", "cfg_renorm_min": 0.0, "think": false}' \ + --output /tmp/bagel_text2img.png +``` + +```bash +# Single-stage (DiT only, with internal LLM/ViT/VAE) +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --stage-configs-path vllm_omni/deploy/bagel_single_stage.yaml \ + --prompt "A beautiful sunset over mountains" \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 50 \ + --seed 42 \ + --extra-body '{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0], "cfg_renorm_type": "global", "cfg_renorm_min": 0.0}' \ + --output /tmp/bagel_text2img_single.png +``` + +The `--extra-body` JSON forwards BAGEL-specific parameters into +`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's +declared `extra_body_params` (see +[`vllm_omni/model_extras/bagel.py`](../../vllm_omni/model_extras/bagel.py)), so +unknown keys for BAGEL are silently dropped. This is the preferred way to pass +model-specific params — `timestep_shift` (use `3.0` for 1024×1024; the +deprecated `--timesteps-shift` flag defaults to `1.0` and is NextStep-oriented), +`cfg_text_scale`, `cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`, +`cfg_renorm_min`, and `think`. `--seed` makes generation reproducible. + +> **Note**: the two-stage `bagel.yaml` runs both stages on GPU 0. If you hit +> CUDA OOM during warmup, lower stage-0 `gpu_memory_utilization` in the YAML or +> split the stages across two GPUs (see the `2x CUDA GPUs` section). + +Run image-to-image with the shared offline image-to-image example: + +```bash +python examples/offline_inference/image_to_image/image_edit.py \ + --model ByteDance-Seed/BAGEL-7B-MoT \ + --prompt "Make the scene look like a watercolor painting" \ + --image /tmp/bagel_text2img.png \ + --height 512 \ + --width 512 \ + --extra-args '{"cfg_text_scale": 4.0, "cfg_img_scale": 1.5}' \ + --output /tmp/bagel_img2img.png +``` + +The `--extra-args` JSON forwards BAGEL-specific parameters (e.g. `cfg_text_scale`, +`cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`) into +`OmniDiffusionSamplingParams.extra_args` via the model-extras registry. + +#### Online Commands + +Start the OpenAI-compatible server: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 +``` + +To use the single-stage topology online: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT \ + --omni \ + --port 8091 \ + --deploy-config vllm_omni/deploy/bagel_single_stage.yaml +``` + +Send a text-to-image request with BAGEL-specific generation parameters: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"} + ] + } + ], + "modalities": ["image"], + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 50, + "cfg_text_scale": 4.0, + "cfg_img_scale": 1.5, + "negative_prompt": "blurry, low quality", + "seed": 42 + } + }' +``` + +The important part is that BAGEL-specific keys such as `cfg_text_scale`, +`cfg_img_scale`, `cfg_interval`, `cfg_renorm_type`, and `cfg_renorm_min` belong +in `extra_body`. The serving layer routes the declared keys into +`OmniDiffusionSamplingParams.extra_args` for the BAGEL pipeline. + +#### Verification + +Decode the returned data URL into an image: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "<|im_start|>A ceramic teapot on a wooden table<|im_end|>"} + ] + } + ], + "modalities": ["image"], + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 25, + "cfg_text_scale": 4.0, + "seed": 42 + } + }' | jq -r '.choices[]?.message.content[]? | select(.image_url.url) | .image_url.url' \ + | head -n 1 \ + | cut -d',' -f2- \ + | base64 -d > /tmp/bagel_online.png + +ls -lh /tmp/bagel_online.png +``` + +### 2x CUDA GPUs + +Create a custom deploy config from `vllm_omni/deploy/bagel.yaml` and move the +diffusion stage to GPU 1: + +```yaml +stages: + - stage_id: 0 + devices: "0" + # keep the remaining stage-0 settings from bagel.yaml + - stage_id: 1 + devices: "1" + # keep the remaining stage-1 settings from bagel.yaml +``` + +Then start serving with that config: + +```bash +vllm serve ByteDance-Seed/BAGEL-7B-MoT \ + --omni \ + --port 8091 \ + --deploy-config /path/to/custom_bagel_2gpu.yaml +``` + +Use the online curl request from the `1x A100 80GB` section to verify that the +server returns an image. diff --git a/recipes/Helios/Helios.md b/recipes/Helios/Helios.md index ebc446b745e..fdc0660bb68 100644 --- a/recipes/Helios/Helios.md +++ b/recipes/Helios/Helios.md @@ -4,27 +4,27 @@ - Vendor: Helios - Model: `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` -- Task: Text-to-video, image-to-video, and video-to-video generation -- Mode: Offline inference with the bundled Helios example +- Task: Text-to-video generation +- Mode: Offline inference with the shared text_to_video example - Maintainer: Community ## When to use this recipe Use this recipe when you want a known-good starting point for running Helios video generation with vLLM-Omni. The concrete command below focuses on -`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU. The -same example script also supports Helios-Mid, Helios-Distilled, image-to-video, -and video-to-video; see the linked example docs for those model-specific flags. +`BestWishYsh/Helios-Base` text-to-video generation on one NVIDIA H20 GPU via the +shared `text_to_video.py` example. The same example also supports Helios-Mid and +Helios-Distilled through the generic `--extra-body` flag (see below). +Image-to-video and video-to-video require image/video conditioning inputs and +are not covered by the text-to-video example. ## References - Upstream repository: <https://github.com/PKU-YuanGroup/Helios> - Model weights: <https://huggingface.co/BestWishYsh/Helios-Base> -- Upstream or canonical docs: - [`docs/user_guide/examples/offline_inference/helios.md`](../../docs/user_guide/examples/offline_inference/helios.md) - Related example under `examples/`: - [`examples/offline_inference/helios/README.md`](../../examples/offline_inference/helios/README.md) + [`examples/offline_inference/text_to_video/text_to_video.md`](../../examples/offline_inference/text_to_video/text_to_video.md) ## Hardware Support @@ -46,11 +46,10 @@ and video-to-video; see the linked example docs for those model-specific flags. Run the baseline Helios-Base text-to-video example from the repository root: ```bash -cd examples/offline_inference/helios +cd examples/offline_inference/text_to_video -python end2end.py \ +python text_to_video.py \ --model BestWishYsh/Helios-Base \ - --sample-type t2v \ --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ --guidance-scale 5.0 \ --output helios_t2v_base.mp4 @@ -59,13 +58,12 @@ python end2end.py \ To use cache-dit acceleration, run on a vLLM-Omni checkout that includes Helios cache-dit support and enable the cache backend: ```bash -cd examples/offline_inference/helios +cd examples/offline_inference/text_to_video -python end2end.py \ +python text_to_video.py \ --cache-backend cache_dit \ --enable-cache-dit-summary \ --model BestWishYsh/Helios-Base \ - --sample-type t2v \ --prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \ --guidance-scale 5.0 \ --output helios_t2v_base.mp4 @@ -85,22 +83,17 @@ Saved generated video to helios_t2v_base.mp4 - `--model BestWishYsh/Helios-Base` selects the base Helios checkpoint used by this recipe. -- `--sample-type t2v` selects text-to-video generation. Use `i2v` with - `--image-path` for image-to-video, or `v2v` with `--video-path` for - video-to-video. - `--guidance-scale 5.0` matches the Helios-Base recommendation. Use `--guidance-scale 1.0` for Helios-Distilled. - `--cache-backend cache_dit` enables the cache-dit acceleration path. - `--enable-cache-dit-summary` prints cache-dit summary information after diffusion forward passes. -- No separate stage config is required for this offline recipe; the bundled - Helios example configures the pipeline through `end2end.py` arguments. -- For Helios-Mid, add `--is-enable-stage2`, - `--pyramid-num-inference-steps-list 20 20 20`, `--use-cfg-zero-star`, - `--use-zero-init`, and `--zero-steps 1`. -- For Helios-Distilled, add `--is-enable-stage2`, - `--pyramid-num-inference-steps-list 2 2 2`, and - `--is-amplify-first-chunk`. +- No separate stage config is required for this offline recipe; the shared + `text_to_video.py` example configures the pipeline through its arguments. +- Helios-specific knobs (declared in `vllm_omni/model_extras/helios.py`) are + passed via the generic `--extra-body` JSON flag: + - Helios-Mid: `--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}'` + - Helios-Distilled: `--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}'` #### Known limitations diff --git a/recipes/IndexTeam/IndexTTS-2.md b/recipes/IndexTeam/IndexTTS-2.md new file mode 100644 index 00000000000..54c1119dd26 --- /dev/null +++ b/recipes/IndexTeam/IndexTTS-2.md @@ -0,0 +1,139 @@ +# IndexTTS-2 for voice-cloned TTS on 1x GPU + +## Summary + +- Vendor: IndexTeam +- Model: `IndexTeam/IndexTTS-2` +- Task: Text-to-speech with voice cloning and optional emotion control +- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API +- Maintainer: Community + +## When to use this recipe + +Use this recipe to serve IndexTTS-2 as a two-stage TTS system for +voice-cloned speech synthesis. Each request provides synthesis text plus either +a reference audio clip or an uploaded audio voice; the model returns 22.05 kHz +mono speech. IndexTTS-2 also supports optional emotion conditioning through +model-specific `extra_params`. + +IndexTTS-2 is **not an async-chunk streaming model** in the bundled deployment +configuration: the S2Mel flow-matching stage consumes the full mel-code sequence +from the AR talker before it produces audio. The OpenAI-compatible endpoint can +still be called with `stream=true` as an HTTP compatibility path, but that should +not be interpreted as incremental model streaming. + +## References + +- Upstream model: + [IndexTeam/IndexTTS-2 on Hugging Face](https://huggingface.co/IndexTeam/IndexTTS-2) +- Online serving example: + [`examples/online_serving/text_to_speech/README.md#indextts-2`](../../examples/online_serving/text_to_speech/README.md#indextts-2) +- Offline inference example: + [`examples/offline_inference/text_to_speech/README.md#indextts-2`](../../examples/offline_inference/text_to_speech/README.md#indextts-2) +- OpenAI-compatible client: + [`examples/online_serving/text_to_speech/indextts2/speech_client.py`](../../examples/online_serving/text_to_speech/indextts2/speech_client.py) +- Deploy config: + [`vllm_omni/deploy/indextts2.yaml`](../../vllm_omni/deploy/indextts2.yaml) + +## Hardware Support + +### GPU + +### 1x L4 24GB or larger CUDA GPU + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA CUDA environment +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +Start the server from the repository root: + +```bash +vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 +``` + +The bundled deployment config is selected automatically for the `indextts2` +model type. To make the config choice explicit, pass it directly: + +```bash +vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code --port 8092 \ + --deploy-config vllm_omni/deploy/indextts2.yaml +``` + +#### Verification + +Run the bundled client. Local audio paths passed to this client are converted to +base64 data URLs before they are sent to the HTTP API: + +```bash +python examples/online_serving/text_to_speech/indextts2/speech_client.py \ + --api-base http://localhost:8092 \ + --text "你好，这是 IndexTTS-2 语音合成测试。" \ + --ref-audio /path/to/reference.wav \ + --output indextts2.wav +``` + +For a raw HTTP request, pass `ref_audio` as an HTTP(S) URL, a base64 data URL, +or a `file://` URI allowed by the server's `--allowed-local-media-path`. Do not +pass an arbitrary local filesystem path directly in JSON unless your client +converts it first: + +```bash +curl http://localhost:8092/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "IndexTeam/IndexTTS-2", + "input": "你好，这是 IndexTTS-2 语音合成测试。", + "response_format": "wav", + "ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_AUDIO>" + }' \ + --output indextts2.wav +``` + +Optional emotion conditioning is passed through `extra_params`. Use one emotion +source mode at a time: + +```bash +curl http://localhost:8092/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "IndexTeam/IndexTTS-2", + "input": "今天心情很好！", + "response_format": "wav", + "ref_audio": "data:audio/wav;base64,<BASE64_ENCODED_REFERENCE_AUDIO>", + "extra_params": { + "emo_vector": [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "emo_alpha": 0.8 + } + }' \ + --output indextts2_happy.wav +``` + +#### Notes + +- Hardware scope: the default config is a single-GPU CUDA deployment. Stage 0 + and Stage 1 share GPU 0 by default. +- Audio output: 22.05 kHz mono WAV. +- Voice cloning: `ref_audio` is required for the documented raw request path. + Alternatively, pass `voice` only when it names an uploaded audio voice from + `/v1/audio/voices`; IndexTTS-2 does not provide a built-in text-only preset + voice in this serving recipe. +- Reference audio format: use an HTTP(S) URL, a base64 data URL, or a `file://` + URI allowed by the server's `--allowed-local-media-path` for raw API calls. + The bundled Python client accepts local paths and converts them to base64 data + URLs. +- Emotion control: `extra_params.emo_audio`, `extra_params.emo_vector`, and + `extra_params.use_emo_text` are alternative conditioning modes. Official + precedence is `use_emo_text` > `emo_vector` > `emo_audio` > same emotion as + the speaker reference. The `emo_vector` order is `[happy, angry, sad, afraid, + disgusted, melancholic, surprised, calm]`. +- Streaming limitation: `async_chunk` is disabled in + `vllm_omni/deploy/indextts2.yaml`; S2Mel requires the full mel-code sequence + before audio generation. +- First request can be slower because reference encoders and auxiliary models + are lazily loaded and cached. diff --git a/recipes/JD/JoyAI-VL-Interaction.md b/recipes/JD/JoyAI-VL-Interaction.md new file mode 100644 index 00000000000..ed767830fc5 --- /dev/null +++ b/recipes/JD/JoyAI-VL-Interaction.md @@ -0,0 +1,221 @@ +# JoyAI-VL-Interaction + +> Real-time streaming video-language interaction (proactive speak / silence / delegate) + +## Summary + +- Vendor: JD (Joy Future Academy) +- Model: [`jdopensource/JoyAI-VL-Interaction-Preview`](https://huggingface.co/jdopensource/JoyAI-VL-Interaction-Preview) (8B, Qwen3-VL architecture with weights retrained for streaming interaction) +- Task: Per-tick proactive interaction over a live video stream — the model decides + on its own each second to speak, stay silent, or delegate a hard question +- Mode: Online serving — an OpenAI-compatible interaction orchestrator in front of + a plain `vllm serve` backend +- Maintainer: Community + +## When to use this recipe + +Use this to stand up the streaming-interaction serving layer (`vllm_omni/experimental/fullduplex/`). +The model is served unchanged by `vllm serve`; this layer adds session state, 3-tier +summary memory, the per-tick decision, and pluggable ASR / TTS / delegation. For the +framework internals and how to add another full-duplex model, see +[`vllm_omni/experimental/fullduplex/README.md`](../../vllm_omni/experimental/fullduplex/README.md). + +## Environment + +- OS: Linux +- Python: 3.10+ +- Hardware: 1x GPU. The default `T_s=100` frame window + `--max-model-len 131072` wants + ≈48GB+; for ~24GB, lower `--chunk-frames`, `--max-model-len`, and the image limit together, + and/or load the weights in fp8 (see "Smaller GPU: fp8 weights" below — ~9.9 GiB vs 16.8 GiB) +- vLLM / vLLM-Omni: versions from your current checkout + +## Start server + +From repository root: + +```bash +# 1. Serve the model (plain `vllm serve`, NOT --omni; it uses the Qwen3-VL architecture). +# The image limit must cover the short-term frame window (chunk_frames, default 100 = T_s); +# prefix caching keeps the accumulating window cheap. Lower both for smaller GPUs. +vllm serve jdopensource/JoyAI-VL-Interaction-Preview \ + --served-model-name JoyAI-VL-Interaction-Preview --port 8061 \ + --max-model-len 131072 --enable-prefix-caching \ + --limit-mm-per-prompt '{"image":256,"video":1}' + +# 2. Interaction orchestrator (OpenAI-compatible, :8070) +python -m vllm_omni.experimental.fullduplex.joyvl.serving.server --port 8070 \ + --main-backend-url http://127.0.0.1:8061/v1 --main-model JoyAI-VL-Interaction-Preview +``` + +### Smaller GPU: fp8 weights + +Add `--quantization fp8` to the step-1 `vllm serve` to load the weights in fp8 online +from the bf16 checkpoint (no separate quantized checkpoint needed; uses vLLM's generic +fp8 path, since the model is a standard Qwen3-VL VLM): + +```bash +vllm serve jdopensource/JoyAI-VL-Interaction-Preview \ + --served-model-name JoyAI-VL-Interaction-Preview --port 8061 \ + --quantization fp8 --max-model-len 131072 --enable-prefix-caching \ + --limit-mm-per-prompt '{"image":256,"video":1}' +``` + +Measured here: model weights **9.9 GiB** vs **16.8 GiB** bf16 (**−41%**), which lets the 8B +fit a much smaller card. This is a **memory** option, not a speed one: at one request per +tick (the streaming regime) fp8 was **not** faster (≈19% slower in our test) because +single-stream decode is memory-latency-bound and pays fp8's dynamic-scale overhead; its +throughput gain needs batching. Decisions stay coherent, but fp8 can shift the +speak/silence boundary, so for timing-critical alerting validate output against bf16 first. + +Optional one-shot launch (model + orchestrator + JD webui + ASR/TTS/background, +all env-configurable). The JD webui frontend is external — set `WEBUI_DIR` to point at it: + +```bash +bash examples/online_serving/joyvl_interaction/scripts/start_all.sh +``` + +## Using the model + +The orchestrator is OpenAI-compatible: send **one user turn per video frame** (~1 fps) to +`/v1/chat/completions` with an `x-session-id` header, and attach an optional **standing +instruction** as a text part. Each reply carries an `interaction` block — `action` is +`silence` / `response` / `delegate` and `text` is what to say: + +```bash +curl -s http://127.0.0.1:8070/v1/chat/completions \ + -H 'x-session-id: s1' -H 'content-type: application/json' -d '{ + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "Alert me if a fire breaks out"}, + {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}} + ]}]}' | jq .interaction +``` + +Send the standing instruction once (it persists for the session); subsequent turns send +just the frame. Ready-made headless client: `cli/run_cli_demo.py`. Reset a session with +`POST /reset {"session_id": "..."}`. + +**What to ask it** — give a standing task and let it act on its own each second: + +- **Proactive alerting** — "tell me when someone enters", "alert me if a fire breaks out" +- **Streaming Q&A** — ask about what's on screen; it answers once it has the evidence +- **Time & memory** — "how many people have walked past?", "what did you see earlier?" + (the 3-tier summary memory lets it recall beyond the live frame window) +- **Read & translate** — "translate the on-screen text as it appears" +- **Delegate hard questions** — handed to a background brain (see below) +- **Live commentary** — `--persona talkative` for continuous, danmaku-style narration + +**Personas** (`--persona`): `default` (speak on meaningful events or to answer), `silent` +(answer only when asked, never delegate), `talkative` (proactively narrate). + +**Tuning:** `--chunk-frames` (short-term window `T_s`), `--response-dedup-threshold` +(lower drops more near-duplicate narration), and `--no-memory` to disable the 3-tier +summaries. The model stays silent until the first instruction arrives +(`force_silence_before_query`, on by default). + +## Delegation (background brain) + +When the model judges a question too hard, it emits `</delegation> <question>` and the +orchestrator hands it to a **background brain** — any OpenAI-compatible endpoint you +self-host. Enable it by pointing the orchestrator at one: + +```bash +python -m vllm_omni.experimental.fullduplex.joyvl.serving.server --port 8070 \ + --main-backend-url http://127.0.0.1:8061/v1 --main-model JoyAI-VL-Interaction-Preview \ + --delegation-backend-url <brain-endpoint>/v1 \ + --delegation-model <brain-model> --delegation-kind chat +``` + +`--delegation-kind` picks the bridge: + +- `chat` — a stronger text/VL model answers (`/chat/completions`) +- `image` — a text-to-image model generates a picture (`/images/generations`, e.g. Qwen-Image) +- `edit` — an image-edit model restyles the current frame (e.g. Qwen-Image-Edit) +- `router` — dispatch each request to chat / image / edit by inspecting it (set + `--delegation-image-url` / `--delegation-edit-url` for the latter two) +- `stub` — canned answers for tests/demos only (no backend needed) + +`chat`/`image`/`edit`/`router` each need a backend URL — **without one, delegation stays +off** (the model's delegate note is still spoken, but nothing is folded back). The brain is +**bring-your-own**: a larger vLLM you serve, or any OpenAI-compatible API (e.g. +`--delegation-backend-url https://api.anthropic.com/v1/ --delegation-model claude-... +--delegation-api-key …`). The reference deployment instead drives the `codex` CLI as the +brain via a separate background-agent service; that agent runs with its own credentials +and bypasses its sandbox, so it is **not bundled here** — self-host a plain +OpenAI-compatible endpoint instead. + +## Host the WebUI demo + +For the full browser experience — live webcam / RTSP input, voice (ASR/TTS), and the +per-tick decision stream — run JD's official WebUI in front of the orchestrator. Clone the +model repo and start its WebUI pointed at the orchestrator (`:8070`): + +```bash +git clone https://github.com/jd-opensource/JoyAI-VL-Interaction.git +cd JoyAI-VL-Interaction/services/webui +uv venv && uv pip install -e . +bash scripts/start_server.sh --api-base http://127.0.0.1:8070/v1 +``` + +Open the printed HTTPS URL, allow the camera (or enter an RTSP URL), and give a standing +instruction. `examples/online_serving/joyvl_interaction/scripts/start_all.sh` can launch +the model + orchestrator + WebUI together — point `WEBUI_DIR` at your clone. + +## Verification + +```bash +# headless: stream a clip and print the per-tick decision timeline +# (the CLI reads video frames via OpenCV: `uv pip install opencv-python`) +python examples/online_serving/joyvl_interaction/cli/run_cli_demo.py \ + path/to/video.mp4 --query "Alert me if a fire breaks out" + +pytest tests/fullduplex # framework + JoyVL unit tests +``` + +## Testing with an RTSP stream (optional) + +RTSP is a **webui-side input** — the browser pulls the stream and feeds frames to the +orchestrator over the normal API; no serving-layer code is involved. To simulate an RTSP +camera from a local video file (no physical IP camera needed), use the helper scripts in +[`examples/online_serving/joyvl_interaction/rtsp/`](../../examples/online_serving/joyvl_interaction/rtsp/), +which wrap [MediaMTX](https://github.com/bluenviron/mediamtx/releases) + `ffmpeg`: + +```bash +cd examples/online_serving/joyvl_interaction/rtsp + +# 1. Local RTSP server (MediaMTX, listens on :8554) +bash ./mediamtx.sh + +# 2. Push a local video file as an RTSP stream (another terminal) +bash ./rtsp.sh ./videos/example.mp4 rtsp://127.0.0.1:8554/fire1 + +# 3. In the WebUI RTSP box, enter: rtsp://127.0.0.1:8554/fire1 +# (replace 127.0.0.1 with the MediaMTX host IP if the webui runs on another machine) +``` + +See the directory's `README.md` for streaming a whole video folder (`rtsp_all.sh`) and +the audio-track caveat. + +## Notes + +- `--omni` is **not** used: the model keeps the Qwen3-VL architecture (only the + weights are retrained), so stock `vllm serve` runs the forward pass; this recipe + only adds the interaction/serving layer. +- On a host without `nvcc` / `ninja`, `vllm serve` of the 8B can crash engine-core in the + FlashInfer sampler JIT (`FileNotFoundError: 'ninja'`) during `profile_run`. Set + `VLLM_USE_FLASHINFER_SAMPLER=0` (or install `ninja`) to work around it. +- `force_silence_before_query` is on by default — the model stays silent until an + instruction arrives; give a standing task (e.g. "translate the on-screen text") + to arm proactive output. +- **Frame resolution is the main latency knob.** The system/memory prefix is prefix-cached, + so per-tick *new* compute is dominated by the new frame's vision tokens. Measured on the + 8B: ~256×192 frames hit Qwen3-VL's min-pixel floor (~72 vision tokens, ~17 ms/tick) vs + ~302 tokens / ~38 ms at 640×480 — about 2× cheaper. Downsample frames to ~256×192 for the + tightest latency / highest concurrency; one GPU then sustains ~150–180 concurrent 1 fps + streams with p95 < 200 ms. Per-tick latency is already far inside the 1 fps budget, so + serving is rarely the bottleneck — resolution is the lever if you need more headroom. +- Speech is external and pluggable: point `ASR_URL` / `TTS_URL` at the bridges in + `examples/online_serving/joyvl_interaction/bridges/` or any compatible service. +- The decision prompts, sampling, and 3-tier summary memory (`T_s=100`, mid→long every + 5 chunks, `key_frames=0` = summarize all chunk frames) are aligned to the JoyVL + reference adapter so per-tick behavior matches the released model; the framework only + supplies the serving structure. diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md index 269ea642594..471ee7ae040 100644 --- a/recipes/LTX/LTX-2.3.md +++ b/recipes/LTX/LTX-2.3.md @@ -5,22 +5,24 @@ ## Summary - Vendor: Lightricks -- Model: `dg845/LTX-2.3-Diffusers` -- Task: Text-to-video with synchronized audio generation +- Model: `diffusers/LTX-2.3-Diffusers` +- Task: Text-to-video and image-to-video with synchronized audio generation - Mode: Online serving (pure diffusion) - Maintainer: @oglok ## When to use this recipe -Use this recipe when you want to serve LTX-2.3 for text-to-video generation -with audio. The model generates videos up to 20+ seconds at 768x512 resolution -with 48kHz audio, all from a single text prompt. Requires a GPU with at least -96GB VRAM due to the 22B parameter transformer (~44GB weights) plus text +Use this recipe when you want to serve LTX-2.3 for text-to-video or +image-to-video generation with audio. The model generates videos up to 20+ +seconds at 768x512 resolution with 48kHz audio from text prompts, and can +condition video generation on an initial image. Start with a 96GB-class GPU for +validation because the model combines a 22B parameter transformer with text encoder, VAE, and vocoder components. ## References -- Model: <https://huggingface.co/dg845/LTX-2.3-Diffusers> +- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3> +- Diffusers-format checkpoint: <https://huggingface.co/diffusers/LTX-2.3-Diffusers> - Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`) ## Serving @@ -28,13 +30,22 @@ encoder, VAE, and vocoder components. ### Command ```bash -vllm serve dg845/LTX-2.3-Diffusers \ +vllm serve diffusers/LTX-2.3-Diffusers \ --omni \ --model-class-name LTX23Pipeline \ --stage-init-timeout 600 ``` -### Verification +For image-to-video: + +```bash +vllm serve diffusers/LTX-2.3-Diffusers \ + --omni \ + --model-class-name LTX23ImageToVideoPipeline \ + --stage-init-timeout 600 +``` + +### T2V Verification ```bash # Health check @@ -44,7 +55,7 @@ curl http://localhost:8000/health curl -X POST http://localhost:8000/v1/videos \ -F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \ -F "negative_prompt=blurry, low quality, distorted, watermark" \ - -F "model=dg845/LTX-2.3-Diffusers" \ + -F "model=diffusers/LTX-2.3-Diffusers" \ -F "num_frames=81" \ -F "fps=24" \ -F "size=768x512" \ @@ -55,7 +66,7 @@ curl -X POST http://localhost:8000/v1/videos \ # Generate a 10-second video (241 frames) curl -X POST http://localhost:8000/v1/videos \ -F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \ - -F "model=dg845/LTX-2.3-Diffusers" \ + -F "model=diffusers/LTX-2.3-Diffusers" \ -F "num_frames=241" \ -F "fps=24" \ -F "size=768x512" \ @@ -65,7 +76,7 @@ curl -X POST http://localhost:8000/v1/videos \ # Generate a 20-second video (481 frames) curl -X POST http://localhost:8000/v1/videos \ -F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \ - -F "model=dg845/LTX-2.3-Diffusers" \ + -F "model=diffusers/LTX-2.3-Diffusers" \ -F "num_frames=481" \ -F "fps=24" \ -F "size=768x512" \ @@ -73,19 +84,48 @@ curl -X POST http://localhost:8000/v1/videos \ -F "guidance_scale=4.0" ``` +### I2V Verification + +Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass +an image reference. The I2V pipeline requires `input_reference` or +`image_reference` because the first frame is the conditioning input. + +```bash +curl http://localhost:8000/health + +curl -X POST http://localhost:8000/v1/videos/sync \ + -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \ + -F "input_reference=@/absolute/path/to/reference.png" \ + -F "model=diffusers/LTX-2.3-Diffusers" \ + -F "num_frames=81" \ + -F "fps=24" \ + -F "size=768x512" \ + -F "num_inference_steps=30" \ + -F "guidance_scale=4.0" \ + -F "seed=42" \ + -o ltx23_i2v.mp4 +``` + ### Notes -- Memory usage: Model loads at ~36 GiB, peaks at ~62 GiB during inference +- Memory usage: run a latest-head serving smoke or DFX benchmark on the target + GPU before publishing load/peak VRAM numbers. +- Checkpoint format: use a Diffusers-format checkpoint such as + `diffusers/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships + raw safetensors and does not contain the subfolder configs required by this + pipeline loader. - Key flags: - - `--stage-init-timeout 600`: Required for the initial `torch.compile` warmup (~90-140 seconds on first request) - - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 pipeline (not LTX-2) + - `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints + - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2) + - `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline + - `input_reference=@...`: Uploads the I2V reference image directly - Audio: 48kHz AAC via BWE vocoder, automatically synced with video - CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed - Supported resolutions: 768x512, 512x384 (must be divisible by 32) - Frame rate: 24 fps - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1) - Known limitations: - - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder) + - LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract. - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git) ## Hardware Support @@ -99,13 +139,16 @@ curl -X POST http://localhost:8000/v1/videos \ - OS: Ubuntu 22.04 - Python: 3.10+ - Driver / runtime: CUDA 13.0, Driver 580.126.09 -- vLLM version: 0.19.x -- vLLM-Omni version: 0.19.x +- vLLM version: Match the PR or release checkout being validated +- vLLM-Omni version or commit: Use the commit you are deploying from + +### Benchmarking -### Validated configurations +Do not copy latency or VRAM numbers between commits or machines. Record them +only after running a latest-head sweep on the target hardware. -| Duration | Frames | Resolution | Steps | Guidance | Inference Time | Peak VRAM | -|----------|--------|------------|-------|----------|----------------|-----------| -| 3s | 81 | 768x512 | 30 | 4.0 | ~110s | ~62 GB | -| 10s | 241 | 768x512 | 30 | 4.0 | ~130s | ~62 GB | -| 20s | 481 | 768x512 | 30 | 4.0 | ~420s | ~62 GB | +For formal PR benchmarking, reuse +`tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with +`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the +single-device eager baseline and CFG-parallel=2 case for a small 384x512, +25-frame, 20-step workload. diff --git a/recipes/MammothModa2/MammothModa2-Preview.md b/recipes/MammothModa2/MammothModa2-Preview.md new file mode 100644 index 00000000000..0733b063068 --- /dev/null +++ b/recipes/MammothModa2/MammothModa2-Preview.md @@ -0,0 +1,110 @@ +# MammothModa2-Preview + +> MammothModa2-Preview text-to-image generation through the shared offline image example + +## Summary + +- Vendor: ByteDance Research +- Model: `bytedance-research/MammothModa2-Preview` +- Task: Text-to-image generation (AR → DiT two-stage pipeline) +- Mode: Offline inference +- Maintainer: Community + +## When to use this recipe + +Use this recipe to run MammothModa2-Preview text-to-image through the shared +offline image example (`text_to_image.py`) instead of a model-specific script. +The generic example formats the AR prompt, drives the AR → DiT stage pipeline, +and forwards MammothModa2-specific generation parameters through the +pipeline-declared `extra_body` contract. + +MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs +interface (not `OmniDiffusionRequest`), so its generation knobs +(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via +`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale` +flags. Image size uses the standard `--height` / `--width` flags. + +## References + +- Upstream model: + [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview) +- Related offline example: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) +- Declared parameters: + [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py) +- Deploy config: + [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml) + +## Hardware Support + +The default deploy config runs both the AR and DiT stages on a single GPU +(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`, +stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU +after rebalancing the split so the AR weights (~23 GB) leave room for the KV +cache — see the note under *1x L40S 48GB*. + +## GPU + +### 1x L40S 48GB + +> **48 GB config adjustment:** the committed +> `vllm_omni/deploy/mammoth_moda2.yaml` uses +> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S, +> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16` +> before running. (On an ~80 GB GPU, leave the defaults unchanged.) + +#### Environment + +- OS: Linux +- Python: Match the repository requirements for your checkout +- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an + ~80 GB GPU for the default config +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Offline Commands + +Download the model: + +```bash +hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview +``` + +Run text-to-image with the shared offline example from the repository root. The +deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the +MammothModa2 generation parameters as a JSON object through `--extra-body`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model ./MammothModa2-Preview \ + --stage-configs-path vllm_omni/deploy/mammoth_moda2.yaml \ + --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \ + --height 1024 \ + --width 1024 \ + --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \ + --output mammoth_t2i.png +``` + +The `--extra-body` JSON forwards MammothModa2-specific parameters into +`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's +declared `extra_body_params` (see +[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)), +so unknown keys for MammothModa2 are silently dropped: + +- `text_guidance_scale` — classifier-free guidance scale for the DiT stage + (default `9.0`; CFG is active only when `> 1.0`). +- `cfg_range` — relative step range `[start, end]` over which CFG is applied + (default `[0.0, 1.0]`). +- `num_inference_steps` — number of DiT denoising steps (default `50`). + +`--height` and `--width` must be multiples of 16. + +#### Verification + +The example writes the generated image to the `--output` path. Confirm the file +exists and is a valid image: + +```bash +ls -lh mammoth_t2i.png +python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)" +``` diff --git a/recipes/NVIDIA/GR00T-N1.7.md b/recipes/NVIDIA/GR00T-N1.7.md new file mode 100644 index 00000000000..91f676a524a --- /dev/null +++ b/recipes/NVIDIA/GR00T-N1.7.md @@ -0,0 +1,101 @@ +# GR00T-N1.7 + +> NVIDIA Isaac GR00T-N1.7-3B robot VLA policy served over the OpenPI WebSocket protocol + +## Summary + +- Vendor: NVIDIA +- Model: `nvidia/GR00T-N1.7-3B` +- Task: Vision-Language-Action (VLA) inference for robot manipulation +- Mode: Online serving via OpenPI WebSocket endpoint +- Maintainer: timzsu + +## When to use this recipe + +Use this recipe when you need to serve GR00T-N1.7 as a real-time robot policy +over the OpenPI WebSocket API. It configures the DROID embodiment +(`OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT`) and exposes the standard DROID action +keys (`eef_9d`, `gripper_position`, `joint_position`) with action horizon 40. + +## References + +- Upstream model: <https://huggingface.co/nvidia/GR00T-N1.7-3B> +- Upstream codebase: <https://github.com/NVIDIA/Isaac-GR00T> +- OpenPI client library: <https://github.com/Physical-Intelligence/openpi> +- Pipeline: `vllm_omni.diffusion.models.gr00t.pipeline_gr00t.Gr00tN1d7Pipeline` +- Deploy config: [`vllm_omni/deploy/Gr00tN1d7.yaml`](../../vllm_omni/deploy/Gr00tN1d7.yaml) +- E2E test: [`tests/e2e/online_serving/test_gr00t_openpi.py`](../../tests/e2e/online_serving/test_gr00t_openpi.py) + +## Environment + +- OS: Linux +- Python: 3.11+ +- Driver / runtime: NVIDIA CUDA +- Hardware: 1 NVIDIA GPU. The upstream model card lists 16 GB+ VRAM for inference (e.g. RTX 4090, L40, H100); in practice this serving path uses ~6 GiB peak (bf16, TP=1, `max_num_seqs: 1`). +- vLLM-Omni version or commit: use versions from your current checkout + +## Start server + +From repository root: + +```bash +vllm serve nvidia/GR00T-N1.7-3B \ + --omni \ + --host 127.0.0.1 \ + --port 8000 \ + --served-model-name gr00t-n1d7 \ + --stage-configs-path vllm_omni/deploy/Gr00tN1d7.yaml +``` + +Notes: + +- Only `max_num_seqs: 1` is supported (configured in the deploy YAML). The + policy is markovian (`state_history_length=1`, `reset()` returns `{}`) — the + reason for the cap is that `Gr00tPolicy` does its own per-sample + (un)batching inside `get_action` and is not integrated with vLLM's + continuous batching path. +- This pipeline is a thin wrapper around the upstream HF policy + (`AutoModel.from_pretrained` + `enforce_eager`, `tensor_parallel_size=1`, + pipeline `load_weights` is a no-op). The standard diffusion accelerators + (SP, CFG, TeaCache, VAE tiling) do not transfer to a flow-matching action + policy at batch size 1, so a native-kernel port is intentionally out of + scope for this recipe — the value is the OpenPI serving integration, not + kernel-level acceleration. +- The WebSocket endpoint is `ws://127.0.0.1:8000/v1/realtime/robot/openpi`. + The server handshake message (first frame after connect) is a msgpack-encoded + dict with `action_horizon`, `action_keys`, `embodiment_tag`, and + `needs_session_id`. + +## Verification + +```python +from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result +validate_session_result(run_policy_session(host="127.0.0.1", port=8000)) +``` + +Or run the e2e test suite: + +```bash +python -m pytest tests/e2e/online_serving/test_gr00t_openpi.py -v +``` + +The test sends a synthetic two-frame DROID observation and checks: + +- GR00T metadata contract: `image_resolution`, `action_horizon`, `action_keys`, `embodiment_tag` +- Action shapes: `eef_9d (1,40,9)`, `gripper_position (1,40,1)`, `joint_position (1,40,7)` +- All action values are finite float32 +- Reset response is `"reset successful"` + +## Notes + +- **Do not change the model-specific `policy_server_config` values.** `action_horizon`, + `action_keys`, and `supported_embodiments` are fixed by the GR00T-N1.7 checkpoint and + are validated against the loaded policy at startup (the server refuses to start on a + mismatch). Only `image_resolution` and `needs_session_id` are deployment knobs. +- To switch embodiment, edit `embodiment_tag` under both `model_config` and + `policy_server_config` in `vllm_omni/deploy/Gr00tN1d7.yaml`. Supported values: + `OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT` (default), `XDOF`, `XDOF_SUBTASK`, + `REAL_G1`, `REAL_R1_PRO_SHARPA`, `LIBERO_PANDA`, `SIMPLER_ENV_GOOGLE`, + `SIMPLER_ENV_WIDOWX`. +- GR00T weights are loaded directly by `Gr00tPolicy` via `AutoModel.from_pretrained`; + the pipeline's `load_weights` is intentionally a no-op. diff --git a/recipes/OpenBMB/MiniCPM-o-4_5.md b/recipes/OpenBMB/MiniCPM-o-4_5.md index efc1699b373..4bbf416ba86 100644 --- a/recipes/OpenBMB/MiniCPM-o-4_5.md +++ b/recipes/OpenBMB/MiniCPM-o-4_5.md @@ -19,15 +19,18 @@ Use this recipe as a known-good starting point for serving of the MiniCPM-o family — it pairs a multimodal-understanding thinker LLM with a streaming `MiniCPMTTS + Token2Wav` talker so a single `/v1/chat/completions` call can return text and 24 kHz speech in one -shot. The recipe covers three shipped GPU layouts (2 / 3 / 8 GPUs) -selected via `--deploy-config`. +shot. The recipe covers the shipped GPU layouts (single / 2 / 3 / 8 GPUs): + the default co-locates both stages on one GPU, and the larger scale-out + layouts (2 / 3 / 8 GPUs) are selected via `--deploy-config`. ## References - Default deploy configs (auto-loaded by HF `model_type=minicpmo` + `hf_config.version="4.5"`): - - 2-GPU layout (default): + - Single-GPU layout (default): [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) + - 2-GPU layout (thinker on GPU 0, talker on GPU 1): + [`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml) - 3-GPU layout (thinker TP=2): [`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml) - 8x RTX 4090 layout: @@ -45,28 +48,30 @@ selected via `--deploy-config`. ## Hardware Support -Three GPU layouts ship with default deploy configs. Pick the layout that -matches your hardware and pass it via `--deploy-config`; the talker -(`MiniCPMTTS + Token2Wav`) always lives on its own GPU because of the -in-process vocoder, and the thinker is the part that scales out via TP. +Four GPU layouts ship with default deploy configs. Pick the layout that +matches your hardware and pass it via `--deploy-config`. The default +co-locates the thinker and the talker (`MiniCPMTTS + Token2Wav`, with its +in-process vocoder) on a single GPU; the multi-GPU layouts split the +talker onto its own GPU and scale the thinker out via TP. | Layout | Thinker | Talker + Token2Wav | Typical hardware | | --- | --- | --- | --- | -| 2-GPU (default) | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB | +| Single-GPU (default) | GPU 0 | GPU 0 | 1x A100/H100/H200 80GB | +| 2-GPU | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB | | 3-GPU (thinker TP=2) | GPU 0,1 (TP=2) | GPU 2 | 3x mid-tier GPUs | | 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | 8x RTX 4090 consumer | ## GPU -### 2 x GPU (default — single command) +### 1 x GPU (default — single command) The default [`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml) -puts the thinker on GPU 0 (`~70 %` memory, `enforce_eager: true`, -`max_num_seqs: 1`) and the talker + Token2Wav vocoder on GPU 1 -(`~75 %` memory). This is the recommended starting layout — works on -any pair of 80GB-class GPUs (A100, H100, H200) and on most 40GB+ -pairs as long as the thinker model weights fit. +co-locates both stages on GPU 0: the thinker (`~80 %` memory) and the +talker + Token2Wav vocoder (`~10 %` memory, `max_num_seqs: 1`). This is +the recommended starting layout — works on a single 80GB-class GPU +(A100, H100, H200) as long as the thinker model weights and the talker's +in-process vocoder fit together. #### Environment @@ -85,7 +90,7 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ ``` The deploy config is auto-loaded by the model registry — no -`--deploy-config` flag needed for this default 2-GPU layout. +`--deploy-config` flag needed for this default single-GPU layout. #### Verification @@ -102,8 +107,9 @@ curl http://localhost:8099/v1/chat/completions \ ``` **Text + speech in one response** (the headline 4.5 feature). The TTS -path is gated by a Jinja flag on the chat template, passed through -`extra_body.chat_template_kwargs.use_tts_template=true`: +path is gated by a Jinja flag on the chat template. Pass +`use_tts_template=true` via the **top-level** `chat_template_kwargs` +field (curl does not flatten nested `extra_body`): ```bash curl http://localhost:8099/v1/chat/completions \ @@ -112,12 +118,18 @@ curl http://localhost:8099/v1/chat/completions \ "model": "openbmb/MiniCPM-o-4_5", "messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}], "modalities": ["text", "audio"], - "extra_body": {"chat_template_kwargs": {"use_tts_template": true}} + "chat_template_kwargs": {"use_tts_template": true} }' ``` -Response carries text in `choices[0].message.content` and base64 WAV -in `choices[0].message.audio.data` (24 kHz mono, see Notes). +When using the OpenAI Python SDK, the same flag can also be sent as +`extra_body={"chat_template_kwargs": {"use_tts_template": True}}` +because the client merges `extra_body` into the request root. + +Response carries text in one choice's `message.content` and base64 WAV +in another choice's `message.audio.data` (24 kHz mono, see Notes). With +`modalities: ["text", "audio"]` you typically get two `choices` entries +(one text, one audio). **Gradio demo (text + image + audio + video UI)**: @@ -135,17 +147,40 @@ speech output (TTS)"** checkbox on / off. #### Notes -- Memory budget: thinker weights occupy GPU 0 at `gpu_memory_utilization: - 0.7`; talker + Token2Wav vocoder share GPU 1 at `0.75`. +- Memory budget: both stages share GPU 0 — the thinker at + `gpu_memory_utilization: 0.8`, the talker + Token2Wav vocoder at `0.1`. - `--trust-remote-code` is required — the HF repo ships a custom `MiniCPMO` config / model class. -- Pin: `enforce_eager: true` on both stages (CUDA graph capture is off - by design for the talker's Token2Wav path). - Stage 1 (talker) is hard-capped to `max_num_seqs: 1`: the talker only consumes `runtime_additional_information[0]`, so any value > 1 makes concurrent requests share request-0's audio. This is the same cap baked into the deploy config. +### 2 x GPU (talker on its own GPU) + +Use +[`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml) +when you have two GPUs and want to give the talker + Token2Wav vocoder a +dedicated card instead of sharing GPU 0 with the thinker. The thinker +runs on GPU 0 (`~90 %` mem, TP=1) and the talker on GPU 1 (`~75 %` mem, +`max_num_seqs: 1`). This relieves the memory pressure of the default +single-GPU co-located layout and is the recommended step up when a +second 80GB-class card is available but full 3-way TP scale-out is not +needed. + +#### Command + +```bash +vllm serve openbmb/MiniCPM-o-4_5 --omni \ + --deploy-config vllm_omni/deploy/minicpmo_4_5_2gpu.yaml \ + --trust-remote-code \ + --host 0.0.0.0 --port 8099 +``` + +Verification and Notes mirror the single-GPU section; the only +difference is that the talker no longer competes with the thinker for +GPU 0 memory. + ### 3 x GPU (thinker TP=2) Use @@ -164,7 +199,7 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ --host 0.0.0.0 --port 8099 ``` -Verification and Notes mirror the 2-GPU section; thinker latency +Verification and Notes mirror the single-GPU section; thinker latency roughly halves under load thanks to TP=2. ### 8 x RTX 4090 24GB (consumer-GPU layout) @@ -190,8 +225,8 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ 4090s. Raise it if your cards have more headroom (e.g. 4090 D / custom 32 GB SKUs), but verify with a long-prompt run before promoting. -- All other knobs match the 2-GPU section; the only difference is the - per-card memory pressure on the thinker shards. +- All other knobs match the single-GPU section; the only difference is + the per-card memory pressure on the thinker shards. ## Notes (applies to all layouts) @@ -209,9 +244,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \ missing dep raises `ImportError` at first request with the same install hint instead of silently emitting empty audio. -- **TTS trigger**: speech output is only emitted when the client passes - `extra_body.chat_template_kwargs.use_tts_template=true`. Without it, - the response is text-only (which is also faster). +- **TTS trigger**: speech output requires + `chat_template_kwargs.use_tts_template=true` so the chat template + appends `<|tts_bos|>` before generation. Without it, Stage-1 talker + receives no TTS token span and returns silent audio (not text-only). + For **curl**, put `chat_template_kwargs` at the request root; nested + `extra_body.chat_template_kwargs` is ignored. The OpenAI Python SDK + may use `extra_body` because it flattens those fields into the root. - **Output audio**: 24 kHz mono WAV inside the OpenAI-style `message.audio.data` (base64). The Gradio demo's WAV player decodes diff --git a/recipes/OpenBMB/VoxCPM2.md b/recipes/OpenBMB/VoxCPM2.md index 9414d8a4804..0fd8e210ecd 100644 --- a/recipes/OpenBMB/VoxCPM2.md +++ b/recipes/OpenBMB/VoxCPM2.md @@ -149,6 +149,7 @@ curl -X POST http://127.0.0.1:8000/v1/audio/speech \ "input": "Streaming PCM test on a single RTX 4090 with VoxCPM2.", "voice": "default", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 - ``` diff --git a/recipes/OpenMOSS/MOSS-TTS.md b/recipes/OpenMOSS/MOSS-TTS.md index c88bf7281bf..e9cb798fe27 100644 --- a/recipes/OpenMOSS/MOSS-TTS.md +++ b/recipes/OpenMOSS/MOSS-TTS.md @@ -104,7 +104,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "voice": "default", "ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav", "response_format": "wav", - "stream": true + "stream": true, + "stream_format": "audio" }' --output output.wav ``` diff --git a/recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md b/recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md new file mode 100644 index 00000000000..f98d78b6296 --- /dev/null +++ b/recipes/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers.md @@ -0,0 +1,109 @@ +# Wan2.2 VACE-Fun Diffusers + +> All-in-one video creation and editing with Wan2.2 VACE-Fun A14B in diffusers format. + +## Summary + +- Vendor: Pyros13 +- Model: `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` +- Base model: `alibaba-pai/Wan2.2-VACE-Fun-A14B` +- Task: VACE video generation and editing (T2V / I2V / R2V / V2V / inpainting) +- Mode: Offline inference and OpenAI-compatible online serving +- Maintainer: Community + +## When To Use This Recipe + +Use this recipe to run Wan2.2-VACE-Fun-A14B through the standard diffusers-style +`Wan22VACEPipeline` path in vLLM-Omni. The checkpoint is a diffusers conversion of +`alibaba-pai/Wan2.2-VACE-Fun-A14B`, with `model_index.json`, `transformer/`, +`transformer_2/`, `text_encoder/`, `tokenizer/`, `vae/`, and `scheduler/` +subfolders. Because the model already uses the expected diffusers layout, vLLM-Omni +does not need any original-format key remapping or `.pth` component loaders. + +The model is a two-expert architecture. The high-noise `transformer` and low-noise +`transformer_2` are switched by the Wan2.2 boundary during denoising. + +## References + +- Model card: <https://huggingface.co/Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers> +- Base model card: <https://huggingface.co/alibaba-pai/Wan2.2-VACE-Fun-A14B> +- Related example: [`examples/offline_inference/vace/vace_video_generation.md`](../../examples/offline_inference/vace/vace_video_generation.md) +- Related issue: [vllm-project/vllm-omni#4206](https://github.com/vllm-project/vllm-omni/issues/4206) +- PR: [vllm-project/vllm-omni#4667](https://github.com/vllm-project/vllm-omni/pull/4667) + +## Offline Inference + +```bash +python examples/offline_inference/vace/vace_video_generation.py \ + --model Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers \ + --mode i2v --image ./i2v_input.JPG \ + --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard" \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \ + --seed 42 --guidance-scale 5.0 \ + --vae-use-tiling --enforce-eager \ + --output vace_fun_i2v.mp4 +``` + +See [`examples/offline_inference/vace/vace_video_generation.md`](../../examples/offline_inference/vace/vace_video_generation.md) +for other VACE modes and their input flags. + +## Online Serving + +### Server + +```bash +vllm serve Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers --omni \ + --model-class-name Wan22VACEPipeline \ + --vae-use-tiling \ + --enforce-eager \ + --port 8091 +``` + +For multi-GPU sequence-parallel serving, add the same parallelism flags you would +use offline, for example `--ulysses-degree 4`. The explicit `--model-class-name` +is optional for this diffusers checkpoint, but keeping it in examples makes the +selected pipeline obvious. + +### Client + +```bash +no_proxy=127.0.0.1 \ +curl -X POST http://127.0.0.1:8091/v1/videos/sync \ + -F "prompt=Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard" \ + -F "input_reference=@/absolute/path/to/i2v_input.JPG" \ + -F "width=832" \ + -F "height=480" \ + -F "num_frames=81" \ + -F "fps=16" \ + -F "num_inference_steps=30" \ + -F "guidance_scale=5.0" \ + --output vace_fun_i2v_serve.mp4 +``` + +For video-to-video or inpainting-style requests, upload the reference video with +`input_reference=@/absolute/path/to/input.mp4` or use `video_reference` with a URL +or JSON-safe data URI. + +## Cache-DiT / TaylorSeer Notes + +`Wan22VACEPipeline` uses the shared Wan2.2 Cache-DiT enabler for both the +high-noise and low-noise transformers. The adapter intentionally wraps only the +main denoising `blocks`. VACE `vace_blocks` form a conditioning branch: each step +they combine the current latent with the VACE context to produce hints that are +then injected into selected main blocks. Recomputing that branch preserves the +control signal; caching the main backbone still gives the acceleration target. + +## Validation + +Validated on NVIDIA B300 with `mkt_part_000.mp4`, 736x1280, 61 frames, 30 +inference steps, `seed=1`, and `boundary_ratio=0.875`. + +The Pyros13 diffusers checkpoint produced a byte-identical MP4 to the +`alibaba-pai/Wan2.2-VACE-Fun-A14B` original-format checkpoint under the same +vLLM-Omni pipeline and inputs: + +- SHA256: `2e2db74a612a009b0d6ad6d739e039883f0de84ce5cdf2676324daf78fb05ef8` +- Frame shape: 61 x 736 x 1280 +- MAE / RMSE / max pixel diff: 0 / 0 / 0 +- PSNR: infinite +- Mean SSIM: 1.0 diff --git a/recipes/Qwen/Qwen-Image-2512.md b/recipes/Qwen/Qwen-Image-2512.md new file mode 100644 index 00000000000..3b7d23302cc --- /dev/null +++ b/recipes/Qwen/Qwen-Image-2512.md @@ -0,0 +1,289 @@ +# Qwen-Image-2512 + +> Text-to-image serving for the 2512 (Dec-2025) refresh of Qwen-Image + +## Summary + +- Vendor: Qwen +- Model: `Qwen/Qwen-Image-2512` +- Task: Text-to-image generation +- Mode: Online serving with optional step-wise continuous batching +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to serve the newer `Qwen/Qwen-Image-2512` +checkpoint instead of the original `Qwen/Qwen-Image`. `Qwen-Image-2512` is a +drop-in updated checkpoint that runs on the **same** `QwenImagePipeline`, so the +serving path, flags, and clients are identical to +[`Qwen-Image.md`](./Qwen-Image.md) — only the model id changes. Pick this recipe +to validate the normal online-serving path on a single 80 GB A800, or to run the +validated 2-GPU ModelOpt mixed FP8/NVFP4 quantized variant on B200. + +## Difference from Qwen-Image + +`Qwen-Image-2512` requires **no code, registry, or pipeline changes** relative to +`Qwen-Image`: + +- Same architecture / pipeline: `QwenImagePipeline` + (`vllm_omni/diffusion/registry.py`, see + [`docs/models/supported_models.md`](../../docs/models/supported_models.md)). +- Same serving entrypoint, CLI flags, OpenAI image-generation API, example + clients, and benchmark scripts. + +The only change is the **model id**: swap `Qwen/Qwen-Image` for +`Qwen/Qwen-Image-2512` everywhere (serve command, API `model` field, client +`--model`, benchmark `--model`). Everything else in +[`Qwen-Image.md`](./Qwen-Image.md) applies unchanged. + +## References + +- Sibling recipe (canonical, identical serving path): + [`recipes/Qwen/Qwen-Image.md`](./Qwen-Image.md) +- Upstream or canonical docs: + [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md) +- Related example under `examples/`: + [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md) +- Related benchmark: + [`benchmarks/diffusion/diffusion_benchmark_serving.py`](../../benchmarks/diffusion/diffusion_benchmark_serving.py) +- Supported models table: + [`docs/models/supported_models.md`](../../docs/models/supported_models.md) + +## Hardware Support + +This recipe documents a single-GPU BF16 serving configuration on A800 and a +validated 2-GPU ModelOpt mixed FP8/NVFP4 configuration on B200. Extend it with +more hardware sections as community validation lands. + +## GPU + +### 1x A800-SXM4-80GB + +#### Environment + +- OS: Linux +- Python: 3.12.13 +- PyTorch: 2.11.0+cu130 +- CUDA runtime: 13.0 (driver 580.126.09) +- GPU: NVIDIA A800-SXM4-80GB (81920 MiB) +- vLLM: 0.22.0 +- Transformers: 5.8.1 +- vLLM-Omni: v0.22.0rc1-70-ge91efc24 + +#### Command + +Start the baseline server: + +```bash +vllm serve Qwen/Qwen-Image-2512 --omni --port 8091 +``` + +To enable the step-wise runtime without batching: + +```bash +vllm serve Qwen/Qwen-Image-2512 --omni \ + --port 8091 \ + --step-execution \ + --max-num-seqs 1 +``` + +To enable experimental compatible-request batching: + +```bash +vllm serve Qwen/Qwen-Image-2512 --omni \ + --port 8091 \ + --step-execution \ + --max-num-seqs 8 +``` + +You can also use the example launcher and pass the model id and extra flags +through: + +```bash +bash examples/online_serving/text_to_image/run_server.sh \ + --model Qwen/Qwen-Image-2512 --step-execution --max-num-seqs 8 +``` + +#### Verification + +Run the existing client example after the server is ready: + +```bash +python examples/online_serving/text_to_image/openai_chat_client.py \ + --server http://localhost:8091 \ + --prompt "A ceramic teapot on a wooden table" \ + --output /tmp/qwen_image_2512_recipe.png +``` + +For a direct API smoke test: + +```bash +curl http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Qwen/Qwen-Image-2512", + "prompt": "A ceramic teapot on a wooden table", + "size": "1024x1024", + "num_inference_steps": 20, + "seed": 42 + }' +``` + +To replay the batching benefit with matched warmup: + +```bash +python benchmarks/diffusion/diffusion_benchmark_serving.py \ + --endpoint /v1/chat/completions \ + --dataset vbench \ + --task t2i \ + --model Qwen/Qwen-Image-2512 \ + --num-prompts 10 \ + --height 1024 \ + --width 1024 \ + --num-inference-steps 50 \ + --seed 42 \ + --port 8089 \ + --max-concurrency 8 \ + --warmup-requests 8 \ + --warmup-concurrency 8 \ + --warmup-num-inference-steps 3 \ + --disable-tqdm +``` + +Run that once against `--max-num-seqs 1`, then rerun it against `--max-num-seqs 8` +and compare the output JSON or terminal metrics. + +#### Benchmark + +The numbers below are loaded-once request-level measurements, not an HTTP +concurrency benchmark. The model is loaded once, warmup requests are run, +then 100 sequential requests are measured. + +```text +Hardware: 1x A800-SXM4-80GB +Requests: 100 +Concurrency: 1 +Resolution: 1024x1024 +Denoising steps: 20 +``` + +| Config | Mean | P95 | P99 | Peak VRAM | +| --- | ---: | ---: | ---: | ---: | +| BF16 | 7.669s | 7.736s | 7.748s | 58,979 MiB | + +#### Notes + +- Memory usage: BF16 peak VRAM ~59 GiB on A800, leaving ~22 GiB headroom for + compile overhead and image decode. +- Key flags: `--step-execution` enables the step-wise runtime; `--max-num-seqs` + controls how many compatible requests may stay active together. +- Keep `--max-num-seqs 1` when you want the more conservative path, when traffic + is mostly single-request, or when you are debugging correctness before + measuring throughput. +- Current batching is still shape-sensitive: different step progress can + co-batch, but different resolutions do not yet co-batch. + +### 2x B200 ModelOpt mixed FP8/NVFP4 + +#### Environment + +- OS: Linux +- Python: 3.12 +- Driver / runtime: NVIDIA CUDA environment with two B200 GPUs +- vLLM version: Match the vLLM-Omni checkout used for deployment +- vLLM-Omni version or commit: Use the commit that contains ModelOpt mixed + FP8/NVFP4 checkpoint loading for diffusion models + +#### Command + +This configuration uses a pre-quantized ModelOpt mixed checkpoint for +Qwen-Image-2512. The validated checkpoint keeps boundary, modulation, +normalization, and output layers in BF16, uses FP8 for sensitive attention +projections, and uses NVFP4 for heavier MLP / attention-output linear layers. + +Validated checkpoint: + +```text +feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy +``` + +Start the server with the CUTLASS quantized linear backend: + +```bash +CUDA_VISIBLE_DEVICES=2,3 \ +vllm serve feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy \ + --omni \ + --host 0.0.0.0 \ + --port 8091 \ + --trust-remote-code \ + --tensor-parallel-size 2 \ + --linear-backend cutlass \ + --force-cutlass-fp8 +``` + +#### Verification + +Run a direct image generation request after the server is ready: + +```bash +curl http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "feizhai123/qwen-image-2512-modelopt-mixed-fp8-sensitive-nvfp4-heavy", + "prompt": "A cinematic Warhammer 40,000 Ultramarine warrior in blue power armor, armor stained with alien ichor, swinging a chainsword at a Tyranid creature, gothic battlefield, dramatic lighting, high detail, no text, no watermark", + "negative_prompt": "blurry, low quality, text, watermark, logo, malformed armor, bad hands", + "size": "1024x1024", + "num_inference_steps": 20, + "true_cfg_scale": 4.0, + "guidance_scale": 4.0, + "seed": 48 + }' +``` + +Check that: + +- The server responds on `http://localhost:8091/health`. +- The response contains one generated image. +- Logs show CUTLASS selection for ModelOpt FP8 / mixed-precision linear + layers instead of falling back to the default automatic backend choice. + +#### Benchmark + +The numbers below are loaded-once request-level measurements, not an HTTP +concurrency benchmark. Each case loads the model once, runs warmup requests, +then measures 100 sequential requests. + +```text +Hardware: 2x B200 +Tensor parallel size: 2 +Requests: 100 +Concurrency: 1 +Resolution: 1024x1024 +Denoising steps: 20 +``` + +| Config | Mean | P90 | P99 | Peak VRAM | +| --- | ---: | ---: | ---: | ---: | +| BF16 | 2.696s | 2.700s | 2.707s | 99,000 MiB | +| ModelOpt FP8 CUTLASS | 2.663s | 2.671s | 2.676s | 85,760 MiB | +| Mixed FP8/NVFP4 CUTLASS | 2.584s | 2.590s | 2.594s | 84,412 MiB | + +The same 2-GPU setup was also used to sweep NVFP4 linear backends. Direct +CUTLASS was the fastest validated backend for this checkpoint. + +| Backend | Mean | P90 | P99 | Peak VRAM | +| --- | ---: | ---: | ---: | ---: | +| auto | 2.927s | 2.948s | 3.535s | 85,340 MiB | +| flashinfer-cudnn | 3.523s | 3.528s | 3.766s | 84,500 MiB | +| flashinfer-trtllm | 2.711s | 2.727s | 3.010s | 84,996 MiB | +| flashinfer-cutlass | 2.935s | 2.950s | 3.011s | 85,340 MiB | +| cutlass | 2.578s | 2.584s | 2.592s | 84,412 MiB | + +#### Notes + +- Use `--linear-backend cutlass` to select the validated quantized linear + backend from the CLI. +- `--force-cutlass-fp8` keeps ModelOpt FP8 layers on the CUTLASS FP8 path. +- The benchmark table intentionally reports Qwen-Image-2512 as + concurrency-1 request-level data; do not compare it directly with an online + HTTP concurrency benchmark. diff --git a/recipes/Qwen/Qwen3-TTS.md b/recipes/Qwen/Qwen3-TTS.md index 73a14a6305a..b6d2c929304 100644 --- a/recipes/Qwen/Qwen3-TTS.md +++ b/recipes/Qwen/Qwen3-TTS.md @@ -28,11 +28,9 @@ Smaller 0.6B variants are also available for `CustomVoice` and `Base`. ## References -- Upstream or canonical docs: - [`docs/user_guide/examples/online_serving/qwen3_tts.md`](../../docs/user_guide/examples/online_serving/qwen3_tts.md) - Related examples under `examples/`: - [`examples/online_serving/qwen3_tts/`](../../examples/online_serving/qwen3_tts/), - [`examples/offline_inference/qwen3_tts/`](../../examples/offline_inference/qwen3_tts/) + [`examples/online_serving/text_to_speech/qwen3_tts/`](../../examples/online_serving/text_to_speech/qwen3_tts/), + [`examples/offline_inference/text_to_speech/qwen3_tts/`](../../examples/offline_inference/text_to_speech/qwen3_tts/) - Related issue or discussion: [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645) @@ -67,9 +65,9 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \ Alternatively, use the convenience script: ```bash -./examples/online_serving/qwen3_tts/run_server.sh # Default: CustomVoice -./examples/online_serving/qwen3_tts/run_server.sh VoiceDesign # VoiceDesign -./examples/online_serving/qwen3_tts/run_server.sh Base # Base (voice clone) +./examples/online_serving/text_to_speech/qwen3_tts/run_server.sh # Default: CustomVoice +./examples/online_serving/text_to_speech/qwen3_tts/run_server.sh VoiceDesign # VoiceDesign +./examples/online_serving/text_to_speech/qwen3_tts/run_server.sh Base # Base (voice clone) ``` The bundled deploy config (`vllm_omni/deploy/qwen3_tts.yaml`) enables async @@ -117,7 +115,7 @@ curl http://localhost:8091/v1/audio/voices **Using the Python client:** ```bash -cd examples/online_serving/qwen3_tts +cd examples/online_serving/text_to_speech/qwen3_tts # CustomVoice python openai_speech_client.py \ @@ -150,6 +148,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \ "voice": "vivian", "language": "English", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 - ``` @@ -157,14 +156,99 @@ curl -X POST http://localhost:8091/v1/audio/speech \ **Offline inference (no server needed):** ```bash -python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice -python examples/offline_inference/qwen3_tts/end2end.py --query-type CustomVoice --streaming +python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type CustomVoice +python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type CustomVoice --streaming ``` ## Notes - Memory usage: The deploy config allocates `gpu_memory_utilization: 0.3` per stage (talker + code2wav share a single GPU). For the 0.6B variants or constrained GPUs, adjust via `--gpu-memory-utilization`. - Key flags: `--omni` is required. `--deploy-config` points to the bundled two-stage pipeline config. -- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Streaming requires `stream=true` with `response_format="pcm"`. +- Async chunking: Enabled by default in `qwen3_tts.yaml` for streaming-friendly first-audio latency. Raw audio streaming requires `stream=true`, `stream_format="audio"`, and `response_format="pcm"`. - Task/model matching: Each task type requires its matching model checkpoint. Using a CustomVoice model for a Base (voice clone) request will fail. - Known limitations: The server serves one model variant at a time. To switch task types (e.g., CustomVoice to Base), restart the server with the corresponding model. + +## Hardware Support + +## GPU + +### 1x RTX 4090 24GB (0.6B CustomVoice) + +#### Environment + +- OS: Ubuntu 24.04 LTS +- Python: 3.12.3 +- PyTorch: 2.11.0+cu130 +- Driver / runtime: NVIDIA 580.126.09 / CUDA 13.0 +- GPU: NVIDIA GeForce RTX 4090, 24 GB +- vLLM version: 0.21.0 +- vLLM-Omni version or commit: 0.21.0rc3.dev91+gd4c13950 + +#### Command + +```bash +vllm serve Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \ + --deploy-config vllm_omni/deploy/qwen3_tts.yaml \ + --omni --port 8091 +``` + +The default deploy config (`qwen3_tts.yaml`) works without modification on the +RTX 4090. Both stages (talker + code2wav) share GPU 0 with +`gpu_memory_utilization: 0.3` each. + +#### Verification + +**English synthesis:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "Hello, this is Qwen3-TTS running on RTX 4090.", + "voice": "vivian", + "language": "English" + }' --output test_english.wav +``` + +**Chinese synthesis:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "你好，这是在RTX 4090上运行的语音合成测试。", + "voice": "vivian", + "language": "Chinese" + }' --output test_chinese.wav +``` + +**With emotion instruction:** + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "input": "I am so excited about this!", + "voice": "vivian", + "language": "English", + "instructions": "Speak with great enthusiasm" + }' --output test_emotion.wav +``` + +#### Notes + +- Memory usage: **~13.5 GiB / 24 GiB** at idle with the default deploy config + (`gpu_memory_utilization: 0.3` per stage). The 0.6B model weights occupy only + ~2.4 GiB (Stage 0: 1.91 GiB, Stage 1: 0.45 GiB); the remainder is KV cache + pre-allocated at startup. To reduce idle footprint to ~5 GiB, use a custom + deploy config with lower utilization (inference peak ~10 GiB): + + ```yaml + # Custom deploy config for Qwen3-TTS-12Hz-0.6B on RTX 4090 + # Copy vllm_omni/deploy/qwen3_tts.yaml and override gpu_memory_utilization: + stages: + - stage_id: 0 + gpu_memory_utilization: 0.15 + - stage_id: 1 + gpu_memory_utilization: 0.15 + ``` diff --git a/recipes/README.md b/recipes/README.md index 94cbcdbd577..d335183caf2 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -12,6 +12,13 @@ that is aligned with that external repository as a reference for structure rather than the place to add files for this repo. Use one Markdown file per model family by default. +> **Note:** The vLLM-Omni TTS recipes (12 models across 8 providers) are also +> published on the rendered recipes site at +> [recipes.vllm.ai](https://recipes.vllm.ai) (see +> [vllm-project/recipes#554](https://github.com/vllm-project/recipes/pull/554)). +> Individual pages resolve at `https://recipes.vllm.ai/<provider>/<model-id>`, +> e.g. [`bosonai/higgs-audio-v3-tts-4b`](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b). + Example layout: ```text @@ -29,31 +36,40 @@ recipes/ |--------|------|----------| | [`audiox/AudioX.md`](./audiox/AudioX.md) | Offline + online unified text/video→audio diffusion | 1x L4 24GB | | [`Baidu/ERNIE-Image.md`](./Baidu/ERNIE-Image.md) | Text-to-image online serving (ERNIE-Image 8B) | 1x or 2x RTX 4090 24GB | +| [`Bagel/BAGEL-7B-MoT.md`](./Bagel/BAGEL-7B-MoT.md) | Text-to-image with shared online/offline examples | 1x A100 80GB / 2x CUDA GPUs | +| [`BosonAI/Higgs-Audio-V3-TTS.md`](./BosonAI/Higgs-Audio-V3-TTS.md) | Online + offline multilingual TTS with voice cloning | 1x H100 80GB | | [`ByteDance/Lance.md`](./ByteDance/Lance.md) | Unified AR+diffusion: text/img/video gen + understanding (Lance 3B) | 1x B300 / A100 80GB | | [`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) | +| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) | +| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU | +| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB | | [`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) | +| [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU | | [`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, 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 | +| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB | | [`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 | -| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB | +| [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 | | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB | | [`Qwen/Qwen3-Omni.md`](./Qwen/Qwen3-Omni.md) | Online serving for multimodal chat | 1x A100 80GB | | [`Qwen/Qwen3-TTS.md`](./Qwen/Qwen3-TTS.md) | Text-to-speech serving (CustomVoice / VoiceDesign / Base) | 1x H100/A100 80GB | | [`SenseNova/SenseNova-U1.md`](./SenseNova/SenseNova-U1.md) | Unified image generation and understanding | 1x H200 (144GB) | +| [`StabilityAI/Stable-Audio-Open.md`](./StabilityAI/Stable-Audio-Open.md) | Offline + online text-to-audio generation (Stable Audio Open) | 1x RTX 4090 24GB | | [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB | | [`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 | +| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) | | [`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 | +| [`JD/JoyAI-VL-Interaction.md`](./JD/JoyAI-VL-Interaction.md) | Real-time streaming video-language interaction (proactive speak/silence/delegate) | 1x GPU 24GB+ | Within a single recipe file, include different hardware support sections such as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like diff --git a/recipes/SenseNova/SenseNova-U1.md b/recipes/SenseNova/SenseNova-U1.md index b1a14eaa2db..5cb30a21d56 100644 --- a/recipes/SenseNova/SenseNova-U1.md +++ b/recipes/SenseNova/SenseNova-U1.md @@ -72,6 +72,31 @@ pytest -s -v tests/e2e/offline_inference/test_sensenova_u1_text2img.py \ - No deploy YAML needed — the engine auto-generates a single-stage diffusion config. - Think mode (`--think`) is recommended for higher image quality. +#### Text-to-Image via the shared example + +SenseNova-U1 is also registered with the shared text-to-image example, so you +can run it through `text_to_image.py` instead of `end2end.py`. Forward the +SenseNova-specific generation parameters as a JSON object through `--extra-body`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model SenseNova/SenseNova-U1-8B-MoT \ + --prompt "A beautiful sunset over mountains" \ + --width 2048 --height 2048 \ + --num-inference-steps 50 \ + --seed 42 \ + --extra-body '{"think": true, "cfg_scale": 4.0, "cfg_norm": "none", "timestep_shift": 3.0, "t_eps": 0.02}' \ + --output sensenova_text2img.png +``` + +The `--extra-body` keys are filtered against the model's declared +`extra_body_params` (see +[`vllm_omni/model_extras/sensenova_u1.py`](../../vllm_omni/model_extras/sensenova_u1.py)): +`think`, `cfg_scale`, `cfg_norm`, `timestep_shift`, `t_eps`, `img_cfg_scale`, +and `max_tokens`. No deploy YAML is needed — the engine auto-generates a +single-stage diffusion config. For richer task modalities (img2text, text2text) +and CPU offload, use the dedicated `end2end.py` script described above. + #### Image-to-Image Editing (img2img) ```bash diff --git a/recipes/StabilityAI/Stable-Audio-Open.md b/recipes/StabilityAI/Stable-Audio-Open.md new file mode 100644 index 00000000000..a7f18128d4a --- /dev/null +++ b/recipes/StabilityAI/Stable-Audio-Open.md @@ -0,0 +1,170 @@ +# Stable-Audio-Open Text-To-Audio Generation on 1x GPU + +> Text-to-audio recipe for Stable Audio Open with offline inference and +> OpenAI-compatible online serving on 1x RTX 4090 24GB. + +## Summary + +- Vendor: Stability AI +- Model: `stabilityai/stable-audio-open-1.0` +- Task: Text-to-audio generation +- Mode: Offline inference and online serving +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want to run Stable Audio Open on a single RTX 4090 +24GB GPU for music or sound-effect generation. The recipe covers a 10-second +offline validation sample with TeaCache and online serving through the +`/v1/audio/generate` endpoint. + +## References + +- Model: <https://huggingface.co/stabilityai/stable-audio-open-1.0> +- Offline example: + [`examples/offline_inference/text_to_audio`](../../examples/offline_inference/text_to_audio) +- Online example: + [`examples/online_serving/stable_audio`](../../examples/online_serving/stable_audio) +- Audio generation API: + [`docs/serving/audio_generate_api.md`](../../docs/serving/audio_generate_api.md) +- Related issue: + [#2645](https://github.com/vllm-project/vllm-omni/issues/2645) + +## Hardware Support + +### GPU + +#### 1x RTX 4090 24GB + +##### Environment + +- OS: Ubuntu 22.04.5 +- Python: 3.12 +- GPU: NVIDIA GeForce RTX 4090, 24564 MiB VRAM +- Driver / runtime: NVIDIA driver 595.80, CUDA-capable runtime matching the + repository build +- vLLM version: 0.22.0 +- vLLM-Omni version: source checkout +- PyTorch: 2.11.0+cu130 +- Model path used in the commands below: `/path/to/stable-audio-open-1.0` + +Stable Audio Open is a gated Hugging Face model. Accept the model license on +the Hugging Face model card before downloading the checkpoint. + +```bash +hf auth login + +hf download stabilityai/stable-audio-open-1.0 \ + --local-dir /path/to/stable-audio-open-1.0 +``` + +##### Commands + +Run a 10-second offline validation sample from the repository root: + +```bash +python examples/offline_inference/text_to_audio/text_to_audio.py \ + --model /path/to/stable-audio-open-1.0 \ + --prompt "A gentle piano melody with soft room ambience" \ + --negative-prompt "Low quality, distorted, noisy" \ + --seed 42 \ + --guidance-scale 7.0 \ + --audio-length 10.0 \ + --num-inference-steps 50 \ + --cache-backend tea_cache \ + --output examples/offline_inference/text_to_audio/stable_audio_10s.wav +``` + +Start the online serving endpoint: + +```bash +vllm-omni serve /path/to/stable-audio-open-1.0 \ + --host 0.0.0.0 \ + --port 8091 \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --enforce-eager \ + --omni +``` + +Generate a 10-second WAV file from the repository root in another terminal: + +```bash +curl http://localhost:8091/health + +curl -X POST http://localhost:8091/v1/audio/generate \ + -H "Content-Type: application/json" \ + -d '{ + "input": "A gentle piano melody with soft room ambience", + "audio_length": 10.0, + "num_inference_steps": 50, + "guidance_scale": 7.0, + "negative_prompt": "Low quality, distorted, noisy", + "seed": 42, + "response_format": "wav" + }' \ + --output examples/online_serving/stable_audio/piano_10s.wav +``` + +##### Verification + +Check that: + +- The offline command writes a valid WAV file. +- The server responds on `http://localhost:8091/health`. +- The online request writes a valid WAV file. +- The generated audio sample rate is 44.1 kHz. +- The generated duration is approximately 10 seconds. +- Peak sampled GPU memory is within the 24GB RTX 4090 budget. In the validated + run, offline and online generation each peaked at about 12.6 GiB. + +Validate the offline outputs: + +```bash +ls -lh examples/offline_inference/text_to_audio/stable_audio_10s.wav + +python - <<'PY' +import soundfile as sf + +path = "examples/offline_inference/text_to_audio/stable_audio_10s.wav" +audio, sample_rate = sf.read(path) +print("sample_rate:", sample_rate) +print("shape:", audio.shape) +print("duration:", len(audio) / sample_rate) +PY +``` + +Validate the online output: + +```bash +ls -lh examples/online_serving/stable_audio/piano_10s.wav + +python - <<'PY' +import soundfile as sf + +path = "examples/online_serving/stable_audio/piano_10s.wav" +audio, sample_rate = sf.read(path) +print("sample_rate:", sample_rate) +print("shape:", audio.shape) +print("duration:", len(audio) / sample_rate) +PY +``` + +##### Notes + +- `stable-audio-open-1.0` can generate up to about 47 seconds of 44.1 kHz + stereo audio. This recipe validates 10-second WAV outputs. +- `--cache-backend tea_cache` is supported and was used for the 10-second + offline validation command. +- The model is gated on Hugging Face and requires license acceptance before + download. +- If online serving fails while importing `torchaudio`, make sure the + `torchaudio` wheel matches the installed PyTorch and CUDA build. The + validated environment used `torch==2.11.0+cu130` and + `torchaudio==2.11.0+cu130`. +- The `NIXL is not available`, `GLOO_SOCKET_IFNAME`, and `torchsde` boundary + warnings observed during validation did not prevent successful generation. +- This recipe was validated on one RTX 4090 24GB GPU only. Other GPU counts, + ROCm, XPU, and NPU setups are not covered here. +- Long generations, higher inference-step counts, and non-WAV response formats + were not benchmarked in this recipe. diff --git a/recipes/StabilityAI/Stable-Diffusion-XL.md b/recipes/StabilityAI/Stable-Diffusion-XL.md new file mode 100644 index 00000000000..94c7e96462b --- /dev/null +++ b/recipes/StabilityAI/Stable-Diffusion-XL.md @@ -0,0 +1,123 @@ +# Stable Diffusion XL + +> Text-to-image generation with Stability AI's SDXL Base 1.0 + +## Summary + +- Vendor: Stability AI +- Model: `stabilityai/stable-diffusion-xl-base-1.0` +- Task: Text-to-image generation +- Mode: Offline inference / Online serving with the OpenAI-compatible images API +- Maintainer: Community + +## When to use this recipe + +Use this recipe when you want a high-quality, open-weight text-to-image model +that generates 1024x1024 images. SDXL Base 1.0 is one of the most widely used +diffusion models and produces photorealistic outputs with strong prompt +adherence. It requires ~6.5 GiB for model weights and runs inference at +~7.5 seconds per image (30 steps) on a single Intel Data Center GPU. + +## References + +- Model: <https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0> +- Architecture reference: <https://github.com/Stability-AI/generative-models> +- Related example under `examples/`: + [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py) + +## Hardware Support + +## GPU + +### 1x Intel Arc BMG GPU (XPU) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: Intel XPU environment with Intel Arc BMG GPU +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +**Offline inference:** + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model stabilityai/stable-diffusion-xl-base-1.0 \ + --prompt "a beautiful sunset over the ocean, photorealistic" \ + --height 1024 --width 1024 \ + --num-inference-steps 30 \ + --guidance-scale 7.5 \ + --seed 42 \ + --output sdxl_output.png +``` + +**Online serving:** + +```bash +vllm serve stabilityai/stable-diffusion-xl-base-1.0 --omni --port 8091 +``` + +#### Verification + +**For offline inference:** Check the output image file exists and has valid content. + +**For online serving:** + +```bash +curl -X POST http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "stabilityai/stable-diffusion-xl-base-1.0", + "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus", + "negative_prompt": "blurry, out of focus, low quality", + "guidance_scale": 7.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: Model weights ~6.5 GiB, peak memory ~10.9 GiB at 1024x1024 with CFG. +- Generation time: ~7.5 seconds for 30 inference steps at 1024x1024. +- `negative_prompt` requires `guidance_scale > 1` to take effect. Recommended range: 5.0-9.0. +- Default resolution is 1024x1024. Other supported sizes: 512x512, 768x768. +- The model uses dual CLIP text encoders (ViT-L/14 + OpenCLIP ViT-bigG/14) for rich text understanding. + +### 1x NVIDIA GPU (24+ GB VRAM) + +#### Environment + +- OS: Linux +- Python: 3.10+ +- Driver / runtime: NVIDIA CUDA environment with an A100 +- vLLM version: Match the repository requirements for your checkout +- vLLM-Omni version or commit: Use the commit you are deploying from + +#### Command + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model stabilityai/stable-diffusion-xl-base-1.0 \ + --prompt "a beautiful sunset over the ocean, photorealistic" \ + --height 1024 --width 1024 \ + --num-inference-steps 30 \ + --guidance-scale 7.5 \ + --seed 42 \ + --output sdxl_output.png +``` + +#### Notes + +- Memory usage: ~6.5 GiB model weights, ~10 GiB peak with CFG at 1024x1024. +- Fits on a single GPU with 24 GB VRAM (e.g., RTX 3090/4090, A5000). +- For GPUs with less VRAM, use `--enable-cpu-offload` to offload text encoders during diffusion. diff --git a/recipes/Wan-AI/Wan2.1-VACE.md b/recipes/Wan-AI/Wan2.1-VACE.md new file mode 100644 index 00000000000..c187e54551b --- /dev/null +++ b/recipes/Wan-AI/Wan2.1-VACE.md @@ -0,0 +1,147 @@ +# Wan2.1 VACE Unified Video Generation + +## Summary + +- Vendor: Wan-AI +- Models: `Wan2.1-VACE-1.3B-diffusers`, `Wan2.1-VACE-14B-diffusers` +- Tasks: T2V, I2V, V2LF, FLF2V, inpainting, and R2V +- Mode: offline inference through the shared task examples + +The shared examples construct VACE's pipeline-native conditioning data from the +provided media inputs. They do not require a mode parameter or a model-specific +example script. + +## References + +- [VACE project](https://github.com/ali-vilab/VACE) +- [Hugging Face 14B checkpoint](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) +- [ModelScope 14B checkpoint](https://modelscope.cn/models/Wan-AI/Wan2.1-VACE-14B-diffusers) +- [Shared text-to-video example](../../examples/offline_inference/text_to_video/text_to_video.md) +- [Shared image-to-video example](../../examples/offline_inference/image_to_video/README.md) +- [Migration RFC #4539](https://github.com/vllm-project/vllm-omni/issues/4539) + +## Hardware Support + +The 1.3B checkpoint was validated on an RTX 5090; the 81-frame T2V command below +used approximately 20 GiB of peak reserved GPU memory with VAE tiling. The 14B +checkpoint requires approximately 70 GiB of disk space and was validated on a +single 48 GiB L40S with VAE tiling and layerwise offload. Start with the smoke +configuration before attempting an 81-frame run. + +## Setup + +The public examples use the Hugging Face model ID by default: + +```bash +MODEL="${MODEL:-Wan-AI/Wan2.1-VACE-14B-diffusers}" +ASSET_DIR="${ASSET_DIR:-$HOME/datasets/vace}" +OUT="${OUT:-$PWD/vace-results/14b}" +mkdir -p "$ASSET_DIR" "$OUT" +``` + +For environments where Hugging Face downloads are unavailable, download the +same checkpoint from ModelScope and point `MODEL` at the local directory: + +```bash +MODEL="$HOME/models/Wan2.1-VACE-14B-diffusers" +modelscope download --model Wan-AI/Wan2.1-VACE-14B-diffusers --local_dir "$MODEL" +``` + +Download the public conditioning images and create a mask matching the output +size. White pixels are regenerated and black pixels are preserved. + +```bash +wget -O "$ASSET_DIR/astronaut.jpg" \ + https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg +wget -O "$ASSET_DIR/vace_first_frame.png" \ + https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png +wget -O "$ASSET_DIR/vace_last_frame.png" \ + https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png + +python - "$ASSET_DIR/vace_center_mask.png" <<'PY' +import sys +from PIL import Image + +mask = Image.new("L", (832, 480), 0) +mask.paste(255, (336, 0, 496, 480)) +mask.save(sys.argv[1]) +PY + +ASTRONAUT="$ASSET_DIR/astronaut.jpg" +FIRST="$ASSET_DIR/vace_first_frame.png" +LAST="$ASSET_DIR/vace_last_frame.png" +MASK="$ASSET_DIR/vace_center_mask.png" +``` + +## Commands + +The commands below reproduce the original VACE examples with `81 frames / 30 +steps`. Most tasks use `480x832`; FLF2V uses `512x512`. For a smoke test, +replace the frame and step counts with `5 frames / 2 steps`. + +```bash +# T2V +python examples/offline_inference/text_to_video/text_to_video.py \ + --model "$MODEL" \ + --prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \ + --seed 0 \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \ + --guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \ + --output "$OUT/t2v.mp4" + +# I2V +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "$MODEL" --image "$ASTRONAUT" \ + --prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \ + --seed 42 \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \ + --guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \ + --output "$OUT/i2v.mp4" + +# V2LF +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "$MODEL" --last-image "$ASTRONAUT" \ + --prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \ + --seed 42 \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \ + --guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \ + --output "$OUT/v2lf.mp4" + +# FLF2V +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "$MODEL" --image "$FIRST" --last-image "$LAST" \ + --prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \ + --seed 42 --height 512 --width 512 --num-frames 81 --num-inference-steps 30 \ + --guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \ + --output "$OUT/flf2v.mp4" + +# Inpaint +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "$MODEL" --image "$ASTRONAUT" --mask-image "$MASK" \ + --prompt "Shrek, the ogre, walks out of a building in a happy mood" \ + --seed 42 \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \ + --guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \ + --output "$OUT/inpaint.mp4" + +# R2V; repeat --reference-image to add references +python examples/offline_inference/image_to_video/image_to_video.py \ + --model "$MODEL" --reference-image "$ASTRONAUT" \ + --prompt "Camera slowly zooms out from the character walking in a garden" \ + --seed 42 \ + --height 480 --width 832 --num-frames 81 --num-inference-steps 30 \ + --guidance-scale 5.0 --flow-shift 5.0 --vae-use-tiling --enable-layerwise-offload \ + --output "$OUT/r2v.mp4" +``` + +## Verification + +```bash +for video in "$OUT"/*.mp4; do + ffprobe -v error -select_streams v:0 \ + -show_entries stream=codec_name,width,height,nb_frames,r_frame_rate,duration \ + -of default=noprint_wrappers=1 "$video" +done +``` + +For the 1.3B checkpoint, `--enable-layerwise-offload` is generally unnecessary. diff --git a/recipes/audiox/AudioX.md b/recipes/audiox/AudioX.md index 6365bd02136..a2e26073666 100644 --- a/recipes/audiox/AudioX.md +++ b/recipes/audiox/AudioX.md @@ -22,8 +22,9 @@ that produces stereo 44.1 kHz audio up to ~10 s per call. - vLLM-Omni weight bundle: <https://huggingface.co/zhangj1an/AudioX> - Pipeline: `vllm_omni.diffusion.models.audiox.pipeline_audiox.AudioXPipeline` - Input transforms: `vllm_omni.transformers_utils.processors.audiox` -- Offline example: [`examples/offline_inference/audiox/`](../../examples/offline_inference/audiox/) -- Online example: [`examples/online_serving/audiox/`](../../examples/online_serving/audiox/) +- Param contract: `vllm_omni/model_extras/audiox.py` (declared `extra_body` knobs) +- Offline example: [`examples/offline_inference/text_to_audio/text_to_audio.py`](../../examples/offline_inference/text_to_audio/text_to_audio.py) +- Online: standard OpenAI chat-completions endpoint (see commands below) ## Hardware Support @@ -41,27 +42,63 @@ that produces stereo 44.1 kHz audio up to ~10 s per call. #### Command -Offline (text-to-audio): +AudioX uses the **standard** `text_to_audio` example offline and the standard +OpenAI chat-completions endpoint online. All model-specific knobs +(`audiox_task`, `seconds_start`, `seconds_total`, `sigma_min`, `sigma_max`, +`cfg_rescale`, `video_path`, `audio_path`) are declared in +`vllm_omni/model_extras/audiox.py` and routed through `extra_body` (online) / +`--extra-body` (offline). The offline example also exposes `--task`, `--video`, +`--audio-start`, and `--audio-length` shortcuts for the most common knobs. + +Offline — text tasks (`t2a` / `t2m`): ```bash huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights -python examples/offline_inference/audiox/end2end.py \ - --model ./audiox_weights \ - --tasks t2a \ - --num-inference-steps 250 \ - --seconds-total 10 +python examples/offline_inference/text_to_audio/text_to_audio.py \ + --model ./audiox_weights --task t2a \ + --prompt "Fireworks burst twice, followed by a clock ticking." \ + --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ + --extra-body '{"sigma_min": 0.03, "sigma_max": 1000.0}' \ + --output t2a.wav +``` + +Offline — video-conditioned tasks (`v2a` / `v2m` / `tv2a` / `tv2m`) need `--video`: + +```bash +python examples/offline_inference/text_to_audio/text_to_audio.py \ + --model ./audiox_weights --task tv2a \ + --prompt "drum beating sound and human talking" \ + --video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4 \ + --num-inference-steps 250 --guidance-scale 6.0 --audio-length 10.0 \ + --output tv2a.wav ``` Online: ```bash -bash examples/online_serving/audiox/run_server.sh -python examples/online_serving/audiox/openai_chat_client.py \ - --task t2a \ - --prompt "Fireworks burst twice, followed by a clock ticking." \ - --output t2a.wav +DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN \ + vllm serve zhangj1an/AudioX --omni --model-class-name AudioXPipeline --port 8099 + +curl -sS -X POST http://localhost:8099/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "zhangj1an/AudioX", + "messages": [{"role": "user", "content": [{"type": "text", "text": "Uplifting ukulele"}]}], + "extra_body": { + "num_inference_steps": 250, + "guidance_scale": 7.0, + "seed": 42, + "audiox_task": "t2m", + "seconds_total": 10.0, + "sigma_min": 0.3, + "sigma_max": 500.0 + } + }' > t2m.json ``` +For `v2*` / `tv2*` online, attach the video as a `video_url` content item +(data URI for local files) in the message `content`. + #### Verification ```bash @@ -75,13 +112,11 @@ ffprobe t2a.wav #### Notes - Memory usage: ~10 GB peak with `num_inference_steps=250`, 10 s of audio. -- Output rate: 44.1 kHz stereo, regardless of `--sample-rate` (resampled in the example - script if requested). -- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via - `extra_args["audiox_task"]` (offline) or the `extra_args` field in the OpenAI - chat-completions body (online). -- Video conditioning: `v2*` and `tv2*` require a video file; the online client - attaches it as an OpenAI `video_url` content item (data URI for local files). +- Output rate: 44.1 kHz stereo regardless of requested rate. +- Supported tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`. Pass via the + `audiox_task` knob (in `extra_body` online, or `--task` offline). +- Video conditioning: `v2*` and `tv2*` require a video; offline via `--video` + (→ `video_path` knob), online via a `video_url` content item. - Cache acceleration is **not** supported (AudioXPipeline is in `_NO_CACHE_ACCELERATION`). - Tensor parallelism is supported via `--tensor-parallel-size` (DiT QKV is sharded with `QKVParallelLinear`); cross-attention K/V is also TP-sharded. diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index 062009fe6b4..c3282ba5174 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -6,7 +6,7 @@ - Vendor: NVIDIA - Model: `nvidia/Cosmos3-Nano` -- 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 +- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional transfer controls, 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 @@ -25,6 +25,9 @@ mode is selected per request: 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. +- **Transfer V2V** — pass one or more transfer hints in `extra_params` + (`edge`, `blur`, `depth`, `seg`, `wsm`) to guide generation with control + frames. Transfer mode is video-only and cannot be combined with sound or action. - **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. @@ -190,6 +193,18 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ -F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \ -o cosmos3_v2v_from_url.mp4 +# Transfer V2V with a precomputed depth control video. `control_path` can point +# to a local image/video; edge and blur can also be computed from `input_reference` +# by passing `"edge":true` or `"blur":true`. +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=Generate a realistic scene following the provided control video." \ + -F "size=1280x720" -F "num_frames=121" \ + -F "num_inference_steps=50" -F "seed=125" \ + -F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \ + -o cosmos3_transfer_depth.mp4 + # Text-to-video-with-sound curl -sS -X POST http://localhost:8000/v1/videos/sync \ -H "Accept: video/mp4" \ @@ -312,6 +327,17 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ 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). +- **Transfer controls:** `extra_params` may include `edge`, `blur`, `depth`, + `seg`, or `wsm`. Each hint accepts `true`, a path string, or an object such as + `{"control_path": "/path/to/control.mp4"}`; `edge` also accepts + `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. + Transfer-level options include `control_guidance`, + `control_guidance_interval`, `num_video_frames_per_chunk` (default `93`, + `101` for WSM), `num_conditional_frames` (default `1`), + `num_first_chunk_conditional_frames`, `max_frames`, + `show_control_condition`, `show_input`, and + `share_vision_temporal_positions`. Non-WSM transfer preserves the input video + fps when available; WSM defaults to 10 fps unless `fps` is supplied. - **DROID OpenPI observations:** include a string `prompt`, either `observation/image` or the three-view DROID camera keys (`observation/wrist_image_left`, `observation/exterior_image_1_left`, @@ -341,60 +367,240 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ #### Command -```python -# cosmos3_offline.py — run with: python cosmos3_offline.py -import torch -from vllm_omni.entrypoints.omni import Omni -from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - -def main(): - omni = Omni( - model="nvidia/Cosmos3-Nano", - model_class_name="Cosmos3OmniDiffusersPipeline", - trust_remote_code=True, - enforce_eager=True, - # Guardrails are disabled here for a quick local run; install - # cosmos-guardrail + gated-repo access and drop this to enable them. - model_config={"guardrails": False}, - ) - gen = torch.Generator(device="cpu").manual_seed(42) - - # Text-to-image (modalities=["image"]). For T2V use modalities=["video"] - # plus num_frames/fps; for I2V add multi_modal_data={"image": <PIL.Image>}. - outputs = omni.generate( - { - "prompt": "A photorealistic red sports car at golden hour, cinematic lighting.", - "negative_prompt": "blurry, distorted, low quality", - "modalities": ["image"], - }, - OmniDiffusionSamplingParams( - height=1024, width=1024, generator=gen, - guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1, - ), - ) - outputs[0].request_output.images[0].save("cosmos3_t2i.png") - omni.close() - - -if __name__ == "__main__": - main() +Cosmos3 runs through the standard task examples; pass model-specific knobs via +`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a +quick local run (install `cosmos-guardrail` + accept the gated repo to enable them). + +```bash +# Text-to-image -> examples/offline_inference/text_to_image +python examples/offline_inference/text_to_image/text_to_image.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A photorealistic red sports car at golden hour, cinematic lighting." \ + --negative-prompt "blurry, distorted, low quality" \ + --height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \ + --extra-body '{"flow_shift": 3.0, "guardrails": false}' \ + --output cosmos3_t2i.png + +# Text-to-video -> examples/offline_inference/text_to_video +python examples/offline_inference/text_to_video/text_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A robot arm is cleaning a plate in the kitchen." \ + --negative-prompt "blurry, distorted, low quality, jittery, deformed" \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, + "use_resolution_template": false, "use_duration_template": false}' \ + --output cosmos3_t2v.mp4 + +# Image-to-video -> examples/offline_inference/image_to_video +# (Cosmos3 bundles example frames under assets/; any RGB image works too.) +python examples/offline_inference/image_to_video/image_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \ + --prompt "The scene comes to life with smooth, natural motion." \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \ + --output cosmos3_i2v.mp4 +``` + +#### Verification + +```bash +python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)" +ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4 +``` + +#### Notes + +- A single `Cosmos3OmniDiffusersPipeline` serves every mode; the standard examples + select it automatically from `model_index.json`. T2I is chosen by the + `text_to_image` prompt builder (which marks `modalities=["image"]`); `text_to_video` + defaults to T2V; `image_to_video` adds `multi_modal_data={"image": ...}` (I2V). + V2V is served online (`/v1/videos/sync`). +- Model-specific knobs (`flow_shift`, `max_sequence_length`, `condition_*`, + `generate_sound`/`sound_duration`, `guardrails`, `action_*`, ...) are declared + once in `vllm_omni/model_extras/cosmos3.py` and forwarded through `--extra-body`; + unknown keys for the model are dropped. + +## NPU + +### 1x Ascend 910B / 910C (Atlas A2 / A3) — Online serving + +#### Environment + +- OS: Linux (aarch64) +- Python: 3.12+ +- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C +- vLLM version: match the repository requirements from your current checkout +- vLLM-Ascend version: match the repository requirements from your current checkout +- vLLM-Omni version or commit: use the commit you are deploying from + +#### Command + +Requires the `vllm-omni` package (or the `quay.io/atlas-ci/vllm-ascend` A2 / A3 container), +which provides the `vllm serve … --omni` entrypoint used below. + +Safety guardrails are **on by default** (NVIDIA Open Model License). They load +the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must: + +1. `pip install cosmos-guardrail` +2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail> +3. Export a token with access: `export HF_TOKEN=hf_...` + +Then launch the recommended server: + +```bash +vllm serve nvidia/Cosmos3-Nano \ + --omni \ + --host 0.0.0.0 --port 8000 \ + --init-timeout 1800 +``` + +To run **without** guardrails (you are responsible for license compliance), +add `--no-guardrails` (no token/`cosmos-guardrail` needed). For tensor parallel +add `--tensor-parallel-size 8`. `--quantization fp8` and +`--enable-layerwise-offload` are not supported on NPU. +The pipeline auto-resolves from `model_index.json`; pass +`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly. + +#### Verification + +Best quality uses the JSON-upsampled prompts from `assets/` (download with +`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal +self-contained examples: + +```bash +curl http://localhost:8000/v1/models + +# Text-to-image -> /v1/images/generations (1024x1024, 10 steps; base64 PNG) +curl -sS -X POST http://localhost:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "model": "nvidia/Cosmos3-Nano", + "prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.", + "negative_prompt": "blurry, distorted, low quality", + "size": "1024x1024", "n": 1, "response_format": "b64_json", + "num_inference_steps": 10, "guidance_scale": 7.0, "seed": 42 + }' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))" + +# Text-to-video -> /v1/videos/sync (720p, 49 frames @ 24fps) +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=A robot arm is cleaning a plate in the kitchen" \ + -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \ + -F "size=1280x720" -F "num_frames=49" -F "fps=24" \ + -F "num_inference_steps=20" -F "guidance_scale=6.0" \ + -F "max_sequence_length=4096" -F "flow_shift=10.0" \ + -F "seed=123" \ + -o cosmos3_t2v.mp4 + +# Image-to-video -> /v1/videos/sync with an uploaded reference image +curl -sS -X POST http://localhost:8000/v1/videos/sync \ + -H "Accept: video/mp4" \ + -F "model=nvidia/Cosmos3-Nano" \ + -F "prompt=The scene comes to life with smooth, natural motion." \ + -F "size=1280x720" -F "num_frames=25" -F "fps=8" \ + -F "num_inference_steps=10" -F "guidance_scale=6.0" \ + -F "seed=42" \ + -F "input_reference=@reference.jpg;type=image/jpeg" \ + -o cosmos3_i2v.mp4 + +# Video-to-video -> /v1/videos/sync with an uploaded reference video +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=17" -F "fps=5" \ + -F "num_inference_steps=10" -F "guidance_scale=6.0" \ + -F "seed=42" \ + -F "input_reference=@reference.mp4;type=video/mp4" \ + -o cosmos3_v2v.mp4 +``` + +#### Notes + +- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):** + - T2I 1024² — 10 steps → ~8 s + - T2V 1280×720 @ 20 steps — 49 frames → ~55 s + - I2V 1280×720 @ 10 steps — 25 frames → ~25 s + - V2V 480×320 @ 10 steps — 17 frames → ~12 s +- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU; + full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer) + ~33 GB on disk. +- **Determinism:** identical seed reproduces identical output on the same + hardware; outputs are not bit-identical across different GPU/NPU 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/V2V + 1280×720, 35 steps, guidance 6.0, `flow_shift=10.0`. +- **Key flags / params:** `--no-guardrails` (optional, to disable guardrails), `--init-timeout 1800` + (for model loading), `--tensor-parallel-size 8` for multi-NPU, and + `--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class. +- **Known limitations:** + - Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a + resolution-parsing bug; basic V2V without transfer hints works. + - FP8 online quantization and layerwise offload are not supported on NPU. + +### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation + +#### Environment + +- OS: Linux (aarch64) +- Python: 3.12+ +- Driver / runtime: CANN 8.5.1 + NNAL + Ascend 910B / 910C +- vLLM-Omni version or commit: use the commit you are deploying from + +#### Command + +The same offline task examples run on NPU; pass model-specific knobs via +`--extra-body`. Guardrails are on by default — pass `"guardrails": false` for a +quick local run (install `cosmos-guardrail` + accept the gated repo to enable them). + +```bash +# Text-to-image -> examples/offline_inference/text_to_image +python examples/offline_inference/text_to_image/text_to_image.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A photorealistic red sports car at golden hour, cinematic lighting." \ + --negative-prompt "blurry, distorted, low quality" \ + --height 1024 --width 1024 --num-inference-steps 50 --guidance-scale 7.0 \ + --extra-body '{"flow_shift": 3.0, "guardrails": false}' \ + --output cosmos3_t2i.png + +# Text-to-video -> examples/offline_inference/text_to_video +python examples/offline_inference/text_to_video/text_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --prompt "A robot arm is cleaning a plate in the kitchen." \ + --negative-prompt "blurry, distorted, low quality, jittery, deformed" \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, + "use_resolution_template": false, "use_duration_template": false}' \ + --output cosmos3_t2v.mp4 + +# Image-to-video -> examples/offline_inference/image_to_video +# (Cosmos3 bundles example frames under assets/; any RGB image works too.) +python examples/offline_inference/image_to_video/image_to_video.py \ + --model nvidia/Cosmos3-Nano \ + --image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \ + --prompt "The scene comes to life with smooth, natural motion." \ + --height 720 --width 1280 --num-frames 189 --fps 24 \ + --num-inference-steps 35 --guidance-scale 6.0 \ + --extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \ + --output cosmos3_i2v.mp4 ``` #### Verification ```bash -python cosmos3_offline.py python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)" +ffprobe -v error -show_entries stream=codec_type,nb_frames,width,height cosmos3_t2v.mp4 ``` #### Notes -- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by - `prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus - `num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and - `multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V. - For video, frames are returned in `outputs[0].request_output.images` as an - `(B, F, H, W, 3)` array. -- The offline entry must be guarded by `if __name__ == "__main__":` — the engine - spawns workers with the `spawn` start method. +- Guardrails are on by default on NPU with `cosmos-guardrail` installed. Pass + `"guardrails": false` in `--extra-body` to disable them (there is no + `--no-guardrails` flag for the offline scripts). +- Video at the 189-frame default takes ~15 min/clip on 1 NPU; reduce + `--num-frames` for faster iteration. diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index 7e63ba26dba..c692b100cae 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -57,7 +57,10 @@ vllm serve nvidia/Cosmos3-Super \ Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to -disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs. +disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs; +`--quantization fp8` (online, no calibration) cuts peak VRAM for 720p video +generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V +composition can shift at the same seed). #### Verification diff --git a/recipes/inclusionAI/Ming-flash-omni-2.0.md b/recipes/inclusionAI/Ming-flash-omni-2.0.md index 4a3a8666928..b9c05606d9a 100644 --- a/recipes/inclusionAI/Ming-flash-omni-2.0.md +++ b/recipes/inclusionAI/Ming-flash-omni-2.0.md @@ -18,6 +18,7 @@ Use this recipe when you want a known-good starting point for serving - **Thinker only** — multimodal understanding with text output. - **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output. +- **Thinker + Imagegen (DiT)** — text-to-image / img2img - **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint. ## References @@ -160,6 +161,129 @@ deltas appear at `choices[0].delta.content`. - Memory usage: size depends on output modalities and multimodal input; leave headroom for video frames and audio caches. +## Image generation (text-to-image / img2img) + +Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage. Launch with the image deploy YAML, which adds an image-gen stage behind the thinker. + +The image-generation stage is a standard vLLM-Omni diffusion pipeline (`MingImagePipeline`); its request knobs are declared in `vllm_omni/model_extras/ming_flash_omni.py` and routed through `extra_body`, so they no longer need a bespoke `sampling_params_list` recipe (that form is still available for per-stage thinker sampling — see below). + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \ + --deploy-config vllm_omni/deploy/ming_flash_omni_image.yaml \ + --stage-init-timeout 1800 \ + --init-timeout 1800 \ + --port 8091 +``` + +With fewer GPUs, copy the YAML and drop the thinker to TP=2 with the DiT on a +free card. + +### Online (text-to-image) + +Request image output with `"modalities": ["image"]`: + +```bash +curl -s http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + { + "role": "user", + "content": "Please draw a cute cat." + } + ], + "modalities": ["image"] + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_imagegen.png +``` + +**Quick form — `extra_body`** (keys are filtered against the declared set and routed into every stage's `extra_args`). Convenient, but it does **not** let you set the thinker (stage-0) sampling params. Wrap the knobs in a literal `"extra_body"` object for raw curl; the OpenAI Python client's `extra_body=` kwarg produces the same request: + +```bash +curl http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "modalities": ["image"], + "messages": [ + { + "role": "user", + "content": "Draw a poster." + } + ], + "extra_body": { + "steps": 6, + "cfg": 1.5, + "height": 512, + "width": 512, + "seed": 123, + "byte5_text": ["理解与生成统一"], + "negative_prompt": "ugly, blurry, distorted" + } + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_imagegen_extra_body.png +``` + +**Full control — `sampling_params_list`** (one entry per stage: `[thinker, imagegen]`). +Use this when you need to tune the thinker's own sampling (`temperature` / `top_p` / `top_k` / `max_tokens`), or to place knobs explicitly per stage. +Note `negative_prompt` must sit on the **stage-0 thinker** `extra_args`; the imagegen-stage knobs (`steps` / `cfg` / `height` / `width` / `seed` / `byte5_text`) go on the **stage-1** entry: + +```bash +curl http://127.0.0.1:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "modalities": ["image"], + "sampling_params_list": [ + { + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "max_tokens": 1, + "seed": 42, + "extra_args": { + "negative_prompt": "ugly, blurry, distorted" + } + }, + { + "seed": 42, + "extra_args": { + "steps": 6, + "cfg": 1.5, + "height": 512, + "width": 512, + "seed": 123, + "byte5_text": ["理解与生成统一"] + } + } + ], + "messages": [ + { + "role": "user", + "content": "Draw a poster." + } + ] + }' \ + | jq -r '.choices[0].message.content[0].image_url.url | split(",")[1]' \ + | base64 -d > ming_imagegen_knobs.png +``` + +### Knobs (declared `extra_body` params) + +| Key | Default | Description | +| --- | --- | --- | +| `height` / `width` | 1024 | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). | +| `steps` | 30 | Number of FlowMatchEuler denoise steps. | +| `cfg` | 2.0 | Classifier-free guidance scale. | +| `seed` | 42 | Per-request RNG seed. | +| `byte5_text` | (auto) | Glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "…". ` format. Auto-extracted from quoted spans in the prompt when omitted. | +| `negative_prompt` | (empty) | Real CFG negative conditioning. Spawns a CFG-text companion via `expand_cfg_prompts`; **online / text-to-image only** (offline uses Ming's default zero-negative). | + +For img2img, add an `image_url` content part to the user message (online) or pass `--image` (offline); the reference image is routed into the DiT stage as `extra[reference_image]`. + ### 1x H100 80GB — standalone TTS (talker only) The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint. diff --git a/recipes/krea/Krea-2.md b/recipes/krea/Krea-2.md new file mode 100644 index 00000000000..80e5105dd5c --- /dev/null +++ b/recipes/krea/Krea-2.md @@ -0,0 +1,143 @@ +# Krea 2 for text-to-image (base + LoRA) + +## Summary + +- Vendor: Krea +- Model: `krea/Krea-2-Turbo` (few-step distilled) and `krea/Krea-2-Raw` +- Task: Text-to-image (T2I) +- Mode: Offline engine and online OpenAI-compatible serving +- Maintainer: Community + +## When to use this recipe + +Use this recipe for a known-good starting point to run the Krea 2 single-stream +MMDiT with vLLM-Omni on a **single 80 GB NVIDIA H100** — both the offline +`text_to_image.py` entrypoint and the online `/v1/images/generations` route, +including PEFT LoRA adapters. `Krea2Pipeline` is auto-detected from +`model_index.json`; guidance is checkpoint-aware (distilled Turbo runs no-CFG, +Raw runs CFG). + +## References + +- Related example under `examples/`: + [`examples/offline_inference/text_to_image/README.md`](../../examples/offline_inference/text_to_image/README.md) + (see the "Krea 2" section) +- Feature support: + [`docs/user_guide/diffusion_features.md`](../../docs/user_guide/diffusion_features.md) +- Related PR: [#4730](https://github.com/vllm-project/vllm-omni/pull/4730) + +## Hardware Support + +This recipe documents a **single-GPU** CUDA layout on H100 80 GB. Add more +platforms (ROCm / NPU) or multi-GPU layouts as community validation lands. + +## GPU + +### 1× H100 80GB + +#### Environment + +Versions from a working editable install (activate `vllm-omni/.venv` or your equivalent): + +- OS: Linux +- Python: 3.12 +- Driver / runtime: NVIDIA driver **595.71.05** (CUDA 13.x) +- torch: **2.11.0+cu130** +- vLLM: **0.24.0** +- vLLM-Omni: editable install from this repo (Git **`960ce8f6`**) +- Transformers: **5.12.1** + +Krea 2 weights are ~30 GB and fit comfortably on a single 80 GB H100. + +#### Command + +Offline — few-step distilled (Turbo) checkpoint (2048×2048, no CFG): + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model krea/Krea-2-Turbo \ + --prompt "a fox in the snow" \ + --num-inference-steps 8 \ + --guidance-scale 0.0 \ + --height 2048 --width 2048 \ + --output krea2_turbo.png +``` + +Offline — Raw checkpoint (1024×1024, CFG enabled): + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model krea/Krea-2-Raw \ + --prompt "a fox in the snow" \ + --num-inference-steps 28 \ + --guidance-scale 4.5 \ + --height 1024 --width 1024 \ + --output krea2_raw.png +``` + +LoRA (offline) — pass a PEFT adapter with `--lora-path` / `--lora-scale`. +`NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA` is a vLLM-Omni-compatible repackaging +of `krea/Krea-2-LoRA-darkbrush`; its trigger word is `monochrome ink wash style`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model krea/Krea-2-Turbo \ + --prompt "a fox in the snow, monochrome ink wash style" \ + --num-inference-steps 8 --guidance-scale 0.0 \ + --lora-path NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA --lora-scale 1.0 \ + --output krea2_darkbrush.png +``` + +Online — serve and query the OpenAI-compatible image route: + +```bash +vllm serve krea/Krea-2-Turbo --omni --enforce-eager --port 8091 +``` + +```bash +curl -s http://localhost:8091/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{ + "prompt": "a fox in the snow, monochrome ink wash style", + "size": "1024x1024", + "num_inference_steps": 8, + "guidance_scale": 0.0, + "seed": 42, + "lora": {"name": "darkbrush", "path": "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA", "scale": 1.0} + }' | jq -r '.data[0].b64_json' | base64 -d > krea2_online.png +``` + +Cache-DiT acceleration (offline) — add `--cache-backend cache_dit`: + +```bash +python examples/offline_inference/text_to_image/text_to_image.py \ + --model krea/Krea-2-Turbo --prompt "a fox in the snow" \ + --num-inference-steps 8 --guidance-scale 0.0 \ + --cache-backend cache_dit --output krea2_cachedit.png +``` + +#### Verification + +```bash +# Offline: the script prints "Saved generated image to <path>"; confirm the file exists. +ls -lh krea2_turbo.png +``` + +For the online route, decode the returned `b64_json` (see the `curl` above) and +confirm a valid 1024×1024 PNG is written. + +#### Notes + +- Key flags: distilled Turbo → few steps (~8) with `guidance_scale=0.0`; Raw → + more steps (~28) with `guidance_scale>0` (CFG). `guidance-scale` follows the + Krea 2 convention `velocity = cond + guidance_scale * (cond - uncond)`. +- LoRA: adapters must be PEFT format (`adapter_config.json` + + `adapter_model.safetensors`); the projections are `ReplicatedLinear` + (264 modules). Offline uses `--lora-path` (which sets + `OmniDiffusionSamplingParams.lora_request` / `lora_scale`); online uses the + top-level `lora` object. Style LoRAs like darkbrush need their trigger word in + the prompt. +- Cache-DiT is supported; `has_separate_cfg` is checkpoint-aware (False for Turbo + no-CFG, True for Raw CFG). +- Known limitations: TP / SP / CFG-Parallel are not yet wired for Krea 2. HSDP, + layerwise CPU offload, and VAE-patch-parallel (decode) are supported. diff --git a/recipes/mistralai/Voxtral-TTS.md b/recipes/mistralai/Voxtral-TTS.md index 595d5c29e5f..8868ed3c333 100644 --- a/recipes/mistralai/Voxtral-TTS.md +++ b/recipes/mistralai/Voxtral-TTS.md @@ -27,6 +27,79 @@ examples in this repository. ## GPU +### 1x NVIDIA GeForce RTX 5090 32GB + +#### Environment + +- OS: Ubuntu 24.04.4 LTS +- Python: 3.12.3 +- PyTorch: 2.11.0+cu130 +- Driver / runtime: NVIDIA driver 580.159.03, CUDA 13.0 +- GPU: NVIDIA GeForce RTX 5090, 32 GB (`32607 MiB`) +- vLLM version: 0.23.0 +- vLLM-Omni version or commit: 0.23.0rc2.dev140+gbe60d7c7c (`be60d7c7`) +- Tested with: `mistral-common==1.11.5` + +#### Command + +```bash +vllm serve mistralai/Voxtral-4B-TTS-2603 \ + --omni \ + --port 8091 \ + --trust-remote-code +``` + +The deploy config is auto-loaded from `vllm_omni/deploy/voxtral_tts.yaml`. + +#### Verification + +Health check and voices: + +```bash +curl -s -o /tmp/voxtral_health.txt -w '%{http_code}' http://127.0.0.1:8091/health +curl -s http://127.0.0.1:8091/v1/audio/voices +``` + +Non-streaming WAV: + +```bash +curl -X POST http://127.0.0.1:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "mistralai/Voxtral-4B-TTS-2603", + "input": "Hello, this is Voxtral TTS running on an NVIDIA GeForce RTX 5090 with vLLM-Omni.", + "voice": "casual_female", + "language": "English", + "response_format": "wav" + }' \ + --output /tmp/voxtral_5090.wav +``` + +Streaming PCM: + +```bash +curl -X POST http://127.0.0.1:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "mistralai/Voxtral-4B-TTS-2603", + "input": "Hello, this is Voxtral TTS streaming PCM on an NVIDIA GeForce RTX 5090.", + "voice": "casual_female", + "language": "English", + "stream": true, + "response_format": "pcm" + }' \ + --output /tmp/voxtral_5090_stream.pcm + +ffmpeg -f s16le -ar 24000 -ac 1 -i /tmp/voxtral_5090_stream.pcm /tmp/voxtral_5090_stream.wav -y +ffprobe -hide_banner /tmp/voxtral_5090_stream.wav +``` + +#### Notes + +- Output: 24 kHz mono WAV/PCM; both non-streaming and streaming requests passed. +- Memory usage: ~27.3 GiB after warmup (Stage 0 ~25.5 GiB, Stage 1 ~1.7 GiB). +- Key flags: `--omni` is required; `--trust-remote-code` is recommended. + ### 1x RTX 4090 24GB #### Environment @@ -80,6 +153,7 @@ curl -X POST http://127.0.0.1:8091/v1/audio/speech \ "voice": "casual_female", "language": "English", "stream": true, + "stream_format": "audio", "response_format": "pcm" }' \ --output voxtral_stream.pcm @@ -105,3 +179,45 @@ voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../exam ) or adding a `ref_audio` request parameter will fail because the public checkpoint does not provide the encoder weights needed to convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.` + +### 1x NVIDIA RTX A6000 48GB + +#### Environment + +- OS: Ubuntu 22.04 +- Python: 3.12 +- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0 +- GPU: NVIDIA RTX A6000, 48 GB +- vLLM version: 0.22.0 +- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d +- Required package: `pip install mistral-common` + +#### Command + +```bash +vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091 +``` + +#### Verification + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "mistralai/Voxtral-4B-TTS-2603", + "input": "Artificial intelligence is transforming the way we interact with technology.", + "voice": "neutral_female", + "response_format": "wav" + }' -o voxtral_test.wav +``` + +Supported voice presets: `ar_male`, `casual_female`, `casual_male`, `cheerful_female`, +`de_female`, `de_male`, `es_female`, `es_male`, `fr_female`, `fr_male`, `hi_female`, +`hi_male`, `it_female`, `it_male`, `neutral_female`, `neutral_male`, `nl_female`, +`nl_male`, `pt_female`, `pt_male`. + +#### Notes + +- Memory usage: Peak VRAM ~40.5 GiB during inference. +- Generation time: ~8 seconds for ~200 words of English text. +- Audio quality is high; mono 24 kHz output. diff --git a/requirements/common.txt b/requirements/common.txt index 469563d3ab0..243a6ce7ed4 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -9,11 +9,11 @@ # That breakage is now handled in code instead of by pinning: # MossTTSNanoForGeneration wraps weight loading in # transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py). -transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0 +transformers >= 5.5.3 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.38.0 -safetensors>=0.8.0rc0 +safetensors>=0.8.0 accelerate==1.12.0 soundfile>=0.13.1 cache-dit==1.3.0 @@ -28,3 +28,9 @@ aenum==3.1.16 pyzmq>=25.0.0 janus>=1.0.0 msgpack>=1.0.0 +pydantic>=2.1.0 +pydantic-settings>=2.0.0 +# Required for GGUF-quantized diffusion model loading (gguf_adapters). The CI +# base image does not ship it, so declare it explicitly. The loader also guards +# this import so a missing gguf no longer breaks all diffusion model loading. +gguf>=0.10.0 diff --git a/tests/ar_diffusion/__init__.py b/tests/ar_diffusion/__init__.py new file mode 100644 index 00000000000..9881313609a --- /dev/null +++ b/tests/ar_diffusion/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/ar_diffusion/test_engine_routing.py b/tests/ar_diffusion/test_engine_routing.py new file mode 100644 index 00000000000..fb55e01e9dc --- /dev/null +++ b/tests/ar_diffusion/test_engine_routing.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the generic engine-backend dispatcher + AR-Diffusion runner wiring. + +``DiffusionEngine.resolve_engine_class`` is a generic dispatcher (``"default"`` / a +``DiffusionEngine`` subclass / an import-path string). DreamZero selects the +AR-Diffusion engine via its deploy config's ``engine_backend`` qualname — no +DreamZero/ar_diffusion knowledge lives in the public base, so the routing check +here is simply that the ``ARDiffusionEngine`` qualname resolves correctly. +""" + +from dataclasses import fields +from types import SimpleNamespace + +import pytest +from vllm.utils.import_utils import resolve_obj_by_qualname + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.experimental.ar_diffusion.engine import ( + AR_DIFFUSION_MODEL_RUNNER_CLS, + ARDiffusionEngine, +) + + +def _cfg(backend): + return SimpleNamespace(engine_backend=backend) + + +def test_default_resolves_to_diffusion_engine(): + assert DiffusionEngine.resolve_engine_class(_cfg("default")) is DiffusionEngine + + +def test_subclass_type_is_returned(): + assert DiffusionEngine.resolve_engine_class(_cfg(ARDiffusionEngine)) is ARDiffusionEngine + + +def test_ar_diffusion_qualname_resolves(): + # How DreamZero's deploy config selects the engine: a full import-path string. + cls = DiffusionEngine.resolve_engine_class(_cfg("vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine")) + assert cls is ARDiffusionEngine + + +def test_non_engine_type_raises(): + with pytest.raises(TypeError): + DiffusionEngine.resolve_engine_class(_cfg(dict)) + + +def test_qualname_not_an_engine_raises(): + with pytest.raises(TypeError): + DiffusionEngine.resolve_engine_class(_cfg("builtins.dict")) + + +def test_bad_qualname_raises(): + with pytest.raises(ValueError): + DiffusionEngine.resolve_engine_class(_cfg("not.a.real.module.NoSuchEngine")) + + +def test_ar_diffusion_engine_is_diffusion_engine_subclass(): + assert issubclass(ARDiffusionEngine, DiffusionEngine) + + +def test_config_engine_backend_field_defaults_to_default(): + field = {f.name: f for f in fields(OmniDiffusionConfig)}["engine_backend"] + assert field.default == "default" + + +# --- AR-Diffusion worker/runner wiring ----------------------------------------------- + + +def test_ar_diffusion_model_runner_cls_resolves(): + cls = resolve_obj_by_qualname(AR_DIFFUSION_MODEL_RUNNER_CLS) + assert cls.__name__ == "ARDiffusionModelRunner" + + +def test_engine_class_declares_its_runner(): + """Routing lives on the engine class (review: hsliuustc0106) — engines + never mutate od_config; the worker resolves the runner off the class.""" + assert ARDiffusionEngine.default_diffusion_model_runner_cls == AR_DIFFUSION_MODEL_RUNNER_CLS + assert DiffusionEngine.default_diffusion_model_runner_cls is None + + +def _select_runner(od_config, platform_default="PLATFORM"): + """Mirror of the worker's runner-selection chain (override > engine > platform).""" + runner_override = getattr(od_config, "diffusion_model_runner_cls", None) + if isinstance(runner_override, str) and runner_override: + return runner_override + engine_cls = DiffusionEngine.resolve_engine_class(od_config) + engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None) + if isinstance(engine_runner, str) and engine_runner: + return engine_runner + return platform_default + + +def test_runner_selection_chain(): + ar = "vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine" + # engine-declared runner wins over platform default + assert ( + _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend=ar)) + == AR_DIFFUSION_MODEL_RUNNER_CLS + ) + # explicit override wins over the engine's declaration + assert ( + _select_runner(SimpleNamespace(diffusion_model_runner_cls="my.custom.Runner", engine_backend=ar)) + == "my.custom.Runner" + ) + # default engine -> platform default + assert _select_runner(SimpleNamespace(diffusion_model_runner_cls=None, engine_backend="default")) == "PLATFORM" + + +def test_dreamzero_pipeline_rejects_non_ar_diffusion_engine(): + """Review (hsliuustc0106): a stale config with engine_backend='default' + must fail at init, not mid-forward on the first KV access.""" + from types import SimpleNamespace + + import pytest + + from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline + + with pytest.raises(ValueError, match="requires the AR-Diffusion engine"): + DreamZeroPipeline(od_config=SimpleNamespace(engine_backend="default")) + with pytest.raises(ValueError, match="requires the AR-Diffusion engine"): + DreamZeroPipeline(od_config=SimpleNamespace()) + + +def test_runner_rejects_multi_sequence_config(): + """Review (hsliuustc0106): the paged path writes batch index 0 only, so + max_num_seqs > 1 must be rejected at init.""" + from types import SimpleNamespace + + import pytest + + from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner + + fake = SimpleNamespace(od_config=SimpleNamespace(max_num_seqs=4)) + with pytest.raises(ValueError, match="max_num_seqs=1"): + ARDiffusionModelRunner._preallocate_kv_cache(fake) + + +def test_execute_model_accepts_kv_prefetch_jobs_kwarg(): + """Regression: the worker passes ``kv_prefetch_jobs`` to every runner + (added with the KV prefetch path, #4448); the AR-Diffusion override must + accept and forward it or DreamZero crashes on the warm-up dummy run.""" + import inspect + + from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner + + for cls in (DiffusionModelRunner, ARDiffusionModelRunner): + params = inspect.signature(cls.execute_model).parameters + assert "kv_prefetch_jobs" in params, cls diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/ar_diffusion/test_kv_cache.py new file mode 100644 index 00000000000..beb575a0b8f --- /dev/null +++ b/tests/ar_diffusion/test_kv_cache.py @@ -0,0 +1,254 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the AR-Diffusion KV cache helpers (Phase 1, PR-2). + +Covers the request adapter, the chunk-window spec/manager (registration + the +eviction policy), and the pool builder — exercised against the installed vLLM +KV stack on CPU (block bookkeeping only, no GPU tensors). +""" + +import pytest +import torch +from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry +from vllm.v1.request import RequestStatus + +from vllm_omni.experimental.ar_diffusion.kv_cache import ( + ARDiffusionKVCache, + ARDiffusionKVConfig, + ARDiffusionRequestAdapter, + ChunkWindowManager, + ChunkWindowSpec, + allocate_kv_pool_with_views, + build_kv_manager, + compute_num_blocks, +) +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens + +BLOCK = 16 + + +def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False): + return ChunkWindowSpec( + block_size=BLOCK, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=window_chunks * chunk_size, + chunk_size=chunk_size, + window_chunks=window_chunks, + sink_chunks=sink_chunks, + reset_at_boundary=reset_at_boundary, + ) + + +# --- ChunkWindowSpec registration ------------------------------------------- + + +def test_spec_registration_resolves_to_chunk_window_manager(): + # Without explicit registration the MRO walk would fall back to the parent + # SlidingWindowManager; assert the subclass manager wins. + spec = make_spec() + assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager + + +def test_spec_kind_is_sliding_window(): + assert get_kv_cache_spec_kind(make_spec()) == KVCacheSpecKind.SLIDING_WINDOW + + +def test_spec_rejects_inconsistent_window(): + with pytest.raises(ValueError): + ChunkWindowSpec( + block_size=BLOCK, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=99, # != window_chunks * chunk_size + chunk_size=BLOCK, + window_chunks=2, + ) + + +# --- eviction policy (pure) ------------------------------------------------- + + +def test_sliding_replace_keeps_window(): + # window = 2 chunks * 16 = 32. Base sliding formula keeps `window` tokens; + # the snap is to chunk boundaries. + def skip(n): + return chunk_window_skipped_tokens(n, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False) + + assert skip(32) == 0 # nothing past the window yet + assert skip(48) == 16 # one chunk fell out of the window + assert skip(64) == 32 + + +def test_sliding_replace_snaps_to_chunk_boundary(): + # A non-chunk-aligned overflow must snap down so a chunk is never half-evicted. + skip = chunk_window_skipped_tokens(50, chunk_size=16, sliding_window=32, sink_chunks=0, reset_at_boundary=False) + assert skip % 16 == 0 and skip == 16 + + +def test_sink_chunks_protected(): + # sink = 1 chunk (16 tokens) is never skipped. + skip = chunk_window_skipped_tokens(80, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=False) + # raw overflow = 80 - 32 - 16 = 32 -> snapped 32 + assert skip == 32 + + +def test_reset_at_boundary_drops_completed_past_sink(): + skip = chunk_window_skipped_tokens(48, chunk_size=16, sliding_window=32, sink_chunks=1, reset_at_boundary=True) + # completed = 48; sink = 16 -> drop 32 + assert skip == 32 + + +# --- ARDiffusionKVConfig ------------------------------------------------------------ + + +def test_kv_config_sliding_window_property(): + assert ARDiffusionKVConfig(chunk_size=16, window_chunks=3).sliding_window == 48 + assert ARDiffusionKVConfig(chunk_size=16, window_chunks=None).sliding_window is None + + +# --- ARDiffusionRequestAdapter ------------------------------------------------------ + + +def test_adapter_advances_per_chunk_not_per_step(): + a = ARDiffusionRequestAdapter("r0", chunk_size=16) + assert a.num_computed_tokens == 0 + assert a.num_tokens == 16 # in-flight chunk + a.on_chunk_committed() + assert a.num_computed_tokens == 16 + assert a.num_tokens == 32 + + +def test_adapter_accounts_for_prefill_prefix(): + a = ARDiffusionRequestAdapter("r0", chunk_size=16, prefill_prefix_tokens=4) + assert a.num_computed_tokens == 4 + assert a.num_prompt_tokens == 4 + assert a.num_tokens == 20 + + +def test_adapter_status_is_vllm_enum(): + assert isinstance(ARDiffusionRequestAdapter("r0", chunk_size=16).status, RequestStatus) + + +# --- pool / manager --------------------------------------------------------- + + +def test_compute_num_blocks(): + # 1 MiB budget at fraction 0.5 with 16 KiB pages -> 32 blocks. + assert compute_num_blocks(1 << 20, 0.5, 16 << 10) == 32 + + +def test_paged_pool_layout_exposes_flat_slot_views(): + kv_pools, k_pools, v_pools = allocate_kv_pool_with_views( + num_blocks=4, + block_size=BLOCK, + num_layers=1, + num_kv_heads=4, + head_dim=64, + dtype=torch.float32, + device=torch.device("cpu"), + ) + assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64) + assert k_pools[0].shape == (4 * BLOCK, 4, 64) + assert v_pools[0].shape == (4 * BLOCK, 4, 64) + + k_pools[0][BLOCK + 3].fill_(7) + v_pools[0][2 * BLOCK + 5].fill_(11) + assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3]) + assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5]) + + +def test_build_manager_allocate_free_roundtrip(): + """End-to-end: a ARDiffusionRequestAdapter drives a real KVCacheManager. + + This is the adapter conformance check — if the adapter were missing an + attribute the manager reads, allocate_slots/free would raise here. + """ + spec = make_spec() + mgr = build_kv_manager(spec, ["layer0"], num_blocks=16, max_model_len=1024) + free_before = mgr.block_pool.get_num_free_blocks() + + adapter = ARDiffusionRequestAdapter("req-0", chunk_size=BLOCK) + blocks = mgr.allocate_slots(adapter, num_new_tokens=BLOCK, full_sequence_must_fit=True) + assert blocks is not None + assert mgr.block_pool.get_num_free_blocks() < free_before + + mgr.free(adapter) + assert mgr.block_pool.get_num_free_blocks() == free_before + + +def test_cross_attn_pool_deducted_from_self_attn_budget(): + """The cross-attn pool is allocated directly; its bytes are subtracted from + the self-attn paged-pool budget so the two together stay within the free + memory budget (review: zwhzzz0821).""" + L = 512 + avail = 1 << 30 # 1 GiB + kv = ARDiffusionKVCache( + ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=2, gpu_memory_fraction=0.5), + num_layers=2, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + block_size=BLOCK, + max_model_len=4096, + available_bytes=avail, + cross_attn_length=L, + device=torch.device("cpu"), + ) + cross_bytes = 2 * 2 * L * 4 * 64 * torch.float16.itemsize * 2 # K+V, pos+neg, layers + expected = compute_num_blocks(avail - cross_bytes, 0.5, kv.spec.page_size_bytes * 2) + assert kv.num_blocks == expected + assert expected > 2 * (2 + 1) + 2 # above the min-blocks floor, so the deduction is what's tested + + +def _make_kv(*, local_branches, num_frame_per_block=2, window_chunks=9): + return ARDiffusionKVCache( + ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks), + num_layers=1, + num_kv_heads=4, + head_size=64, + dtype=torch.float32, + block_size=BLOCK, + max_model_len=4096, + available_bytes=1 << 16, # tiny -> the floor binds + device=torch.device("cpu"), + local_branches=local_branches, + num_frame_per_block=num_frame_per_block, + ) + + +def test_pool_floor_is_branch_aware(): + """CFG-parallel rank (one local branch) sizes for one window + in-flight chunk; + a single-process run (both branches) sizes for two. Scratch scales the same way.""" + one = _make_kv(local_branches=1) + two = _make_kv(local_branches=2) + + assert one.managed_num_blocks == 1 * (9 + 2) + 2 # 13 + assert two.managed_num_blocks == 2 * (9 + 2) + 2 # 24 + assert one.scratch_num_blocks == one.scratch_blocks_per_branch + assert two.scratch_num_blocks == 2 * two.scratch_blocks_per_branch + assert one.num_blocks_total == 13 + one.scratch_blocks_per_branch + + +def test_scratch_maps_to_slot_zero_with_one_local_branch(): + """A CFG-parallel rank runs exactly one branch: whichever CFG side it is, + its scratch lands in the rank's single slot (no dead second slot).""" + one = _make_kv(local_branches=1) + assert one.scratch_block_ids(True, 0, 2) == one.scratch_block_ids(False, 0, 2) + + two = _make_kv(local_branches=2) + assert two.scratch_block_ids(True, 0, 2) != two.scratch_block_ids(False, 0, 2) + + +def test_scratch_exhaustion_still_raises(): + one = _make_kv(local_branches=1) + cap = one.scratch_blocks_per_branch + with pytest.raises(RuntimeError, match="scratch blocks exhausted"): + one.scratch_block_ids(False, 0, cap + 1) + + +def test_invalid_local_branches_rejected(): + with pytest.raises(ValueError, match="local_branches"): + _make_kv(local_branches=3) diff --git a/tests/ar_diffusion/test_kv_orchestration.py b/tests/ar_diffusion/test_kv_orchestration.py new file mode 100644 index 00000000000..4ba5cb49bf1 --- /dev/null +++ b/tests/ar_diffusion/test_kv_orchestration.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for ARDiffusionKVCache — the engine-level KV orchestration body (Phase 1).""" + +import pytest +import torch + +from vllm_omni.experimental.ar_diffusion.kv_cache import ARDiffusionKVCache, ARDiffusionKVConfig + +DIMS = dict(num_layers=2, num_kv_heads=4, head_size=64, dtype=torch.float16, block_size=16) + + +def make_cache(*, chunk_size=16, window_chunks=2, available_bytes=1 << 24): + cfg = ARDiffusionKVConfig(enable=True, chunk_size=chunk_size, window_chunks=window_chunks) + return ARDiffusionKVCache(cfg, max_model_len=4096, available_bytes=available_bytes, **DIMS) + + +def test_requires_enabled_config(): + with pytest.raises(ValueError): + ARDiffusionKVCache(ARDiffusionKVConfig(enable=False), max_model_len=256, available_bytes=1 << 20, **DIMS) + + +def test_requires_bounded_window(): + cfg = ARDiffusionKVConfig(enable=True, chunk_size=16, window_chunks=None) + with pytest.raises(ValueError): + ARDiffusionKVCache(cfg, max_model_len=256, available_bytes=1 << 20, **DIMS) + + +def test_full_request_lifecycle_and_eviction(): + """begin -> per-chunk allocate/slots/commit over a long rollout -> free. + + Exercises the orchestrator end-to-end and asserts the chunk window bounds + memory (pool plateaus) and frees cleanly. + """ + kv = make_cache(chunk_size=16, window_chunks=2) + free_total = kv.manager.block_pool.get_num_free_blocks() + + adapter = kv.begin_request("req-0") + free_after = [] + for k in range(10): + block_table = kv.allocate_chunk(adapter) + slots = kv.chunk_write_slots(adapter) + # Slots target real blocks in this chunk's table; length == chunk_size. + assert len(slots) == kv.spec.chunk_size + used = {int(s) // kv.block_size for s in slots} + assert kv.null_block_id not in used + assert used <= set(block_table) + # Resident window stays bounded by window + the in-flight chunk. + assert len(kv.window_block_ids(adapter)) <= kv.spec.window_chunks + 1 + free_after.append(kv.manager.block_pool.get_num_free_blocks()) + kv.commit_chunk(adapter) + + # Pool memory plateaus once the window is full (eviction recycles blocks). + assert free_after[-1] == free_after[kv.spec.window_chunks] + + kv.end_request(adapter) + assert kv.manager.block_pool.get_num_free_blocks() == free_total + + +def test_num_computed_advances_per_chunk(): + kv = make_cache(chunk_size=16, window_chunks=2) + a = kv.begin_request("r") + assert a.num_computed_tokens == 0 + kv.allocate_chunk(a) + kv.commit_chunk(a) + assert a.num_computed_tokens == 16 + kv.end_request(a) + + +def test_state_close_frees_both_branch_blocks(): + """ARDiffusionKVState.close() returns both CFG branches' pool blocks to the pool. + + This is the primitive the runner's LRU eviction relies on: when a session is + evicted, close() must free the blocks both adapters hold, or session churn + leaks pool ownership (review P1). + """ + from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + + kv = make_cache(chunk_size=16, window_chunks=2) + free_total = kv.manager.block_pool.get_num_free_blocks() + + pos = kv.begin_request("bde__s") + neg = kv.begin_request("bde__s__neg") + state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers) + for _ in range(3): + for adapter in (pos, neg): + kv.allocate_chunk(adapter) + kv.commit_chunk(adapter) + # Both branches hold resident blocks now. + assert kv.manager.block_pool.get_num_free_blocks() < free_total + + state.close() + # All blocks returned — no leak across an evicted session. + assert kv.manager.block_pool.get_num_free_blocks() == free_total diff --git a/tests/ar_diffusion/test_paged_attention.py b/tests/ar_diffusion/test_paged_attention.py new file mode 100644 index 00000000000..1836656f940 --- /dev/null +++ b/tests/ar_diffusion/test_paged_attention.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for AR-Diffusion paged self-attention contexts.""" + +from __future__ import annotations + +import subprocess +from importlib.util import find_spec +from pathlib import Path +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm_omni.experimental.ar_diffusion.kv_cache import ( + ARDiffusionKVCache, + ARDiffusionKVConfig, + ARDiffusionPagedLayerContext, + ARDiffusionPagedLayerInputs, + ar_diffusion_paged_attention, + paged_write_attn, +) +from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + +BLOCK = 16 +N_HEADS = 4 +HEAD_DIM = 64 + + +def make_state(*, num_layers=1, window_chunks=2, dtype=torch.float32, device=torch.device("cpu")): + cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks) + kv = ARDiffusionKVCache( + cfg, + num_layers=num_layers, + num_kv_heads=N_HEADS, + head_size=HEAD_DIM, + dtype=dtype, + block_size=BLOCK, + max_model_len=4096, + available_bytes=1 << 26, + device=device, + ) + pos = kv.begin_request("r-pos") + neg = kv.begin_request("r-neg") + return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers) + + +def _dense_attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor: + scores = torch.einsum("bqhd,bkhd->bhqk", query.float(), key.float()) * (HEAD_DIM**-0.5) + probs = torch.softmax(scores, dim=-1).to(value.dtype) + return torch.einsum("bhqk,bkhd->bqhd", probs, value) + + +def _cuda_flash_attn_usable() -> bool: + if not torch.cuda.is_available(): + return False + try: + spec = find_spec("vllm.vllm_flash_attn") + if spec is None or spec.origin is None: + return True + fa2_so = Path(spec.origin).parent / "_vllm_fa2_C.abi3.so" + linked = subprocess.check_output(["ldd", str(fa2_so)], text=True, timeout=5) + except Exception: + return True + if "libcudart.so.13" not in linked: + return True + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], + text=True, + timeout=5, + ) + driver_major = int(out.splitlines()[0].split(".")[0]) + except Exception: + return True + return driver_major >= 580 + + +def _commit_video_span( + kv: ARDiffusionKVCache, + st: ARDiffusionKVState, + *, + is_negative: bool, + n_chunks: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(device) + k = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + v = torch.randn(1, n_chunks * BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][ctx.current_video_slot_mapping] = k[0] + kv._v_pools[0][ctx.current_video_slot_mapping] = v[0] + st.commit_paged_context(is_negative) + return k, v + + +def test_paged_context_allocates_lazily_and_commits_after_forward(): + _, st = make_state() + + contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True) + ctx = contexts[0].forward_ctx + assert isinstance(contexts[0], ARDiffusionPagedLayerContext) + assert st.pos.completed_chunks == 0 + assert ctx.current_video_slot_mapping is None + + ctx.ensure_video_slots(torch.device("cpu")) + assert st.pos.completed_chunks == 0 + assert len(ctx.current_video_block_ids) == 1 + + st.commit_paged_context(False) + assert st.pos.completed_chunks == 1 + assert st._committed[False] == BLOCK + + +def test_scratch_video_and_action_blocks_do_not_commit(): + kv, st = make_state() + + ctx = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=False)[0].forward_ctx + ctx.ensure_video_slots(torch.device("cpu")) + ctx.ensure_action_slots(3, torch.device("cpu")) + + assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 2) + assert ctx.action_scratch_block_ids == kv.scratch_block_ids(False, 2, 1) + st.commit_paged_context(False) + assert st.pos.completed_chunks == 0 + assert st._committed[False] == 0 + + +def test_pipeline_kv_get_paged_path_has_no_gather_backend(): + kv, st = make_state() + assert not hasattr(kv, "gather_window_all_layers") + + from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline + + pipeline = DreamZeroPipeline.__new__(DreamZeroPipeline) + pipeline._ar_diffusion_kv_state = st + contexts = pipeline._kv_get(MagicMock(), False, seq_len=BLOCK, update_kv_cache=False) + + assert len(contexts) == 1 + assert isinstance(contexts[0], ARDiffusionPagedLayerContext) + + +@pytest.mark.parametrize("history_chunks", [0, 1, 3]) +@pytest.mark.parametrize("action_len", [0, 3]) +@pytest.mark.parametrize("commit_current", [False, True]) +def test_paged_attention_matches_dense_reference_cpu(history_chunks, action_len, commit_current): + torch.manual_seed(0) + device = torch.device("cpu") + dtype = torch.float32 + kv, st = make_state(dtype=dtype, device=device, window_chunks=2) + + history_k_parts: list[torch.Tensor] = [] + history_v_parts: list[torch.Tensor] = [] + if history_chunks: + k, v = _commit_video_span( + kv, + st, + is_negative=False, + n_chunks=history_chunks, + dtype=dtype, + device=device, + ) + history_k_parts.append(k) + history_v_parts.append(v) + + ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0].forward_ctx + ctx.ensure_video_slots(device) + current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][ctx.current_video_slot_mapping] = current_k[0] + kv._v_pools[0][ctx.current_video_slot_mapping] = current_v[0] + + action_k = action_v = None + if action_len: + ctx.ensure_action_slots(action_len, device) + action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + kv._k_pools[0][ctx.action_slot_mapping] = action_k[0] + kv._v_pools[0][ctx.action_slot_mapping] = action_v[0] + + query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + block_table, query_start_loc, seq_lens, max_query_len, max_seq_len = ctx.build_block_table( + action_len=action_len, + query_len=query.shape[1], + device=device, + ) + paged = ar_diffusion_paged_attention( + query, + kv.key_cache(0), + kv.value_cache(0), + block_table=block_table, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + softmax_scale=HEAD_DIM**-0.5, + causal=False, + ) + + if history_k_parts: + history_k = torch.cat(history_k_parts, dim=1) + history_v = torch.cat(history_v_parts, dim=1) + else: + history_k = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + history_v = torch.empty(1, 0, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :] + new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :] + if action_len: + new_k = torch.cat([new_k, action_k], dim=1) + new_v = torch.cat([new_v, action_v], dim=1) + ref = _dense_attention(query, new_k, new_v) + + torch.testing.assert_close(paged, ref, rtol=1e-5, atol=1e-5) + + before = st.pos.completed_chunks + st.commit_paged_context(False) + assert st.pos.completed_chunks == before + (1 if commit_current else 0) + + +@pytest.mark.skipif(not _cuda_flash_attn_usable(), reason="usable CUDA FlashAttention is required") +@pytest.mark.parametrize("history_chunks", [1, 3]) +@pytest.mark.parametrize("action_len", [0, 3]) +@pytest.mark.parametrize("commit_current", [False, True]) +def test_paged_attention_matches_dense_reference_gpu(history_chunks, action_len, commit_current): + pytest.importorskip("vllm.vllm_flash_attn") + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.float16 + kv, st = make_state(dtype=dtype, device=device, window_chunks=2) + + history_k, history_v = _commit_video_span( + kv, + st, + is_negative=False, + n_chunks=history_chunks, + dtype=dtype, + device=device, + ) + + layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit_current)[0] + ctx = layer_ctx.forward_ctx + current_k = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + current_v = torch.randn(1, BLOCK, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + action_k = action_v = None + if action_len: + action_k = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + action_v = torch.randn(1, action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + + query = torch.randn(1, BLOCK + action_len, N_HEADS, HEAD_DIM, dtype=dtype, device=device) + + # The production path: once-per-forward host prep, then the fused + # write+attend custom op consuming the NamedTuple payload. + ctx.prepare(device=device, action_len=action_len, query_len=query.shape[1]) + inputs = layer_ctx.to_layer_inputs() + assert isinstance(inputs, ARDiffusionPagedLayerInputs) + paged = paged_write_attn( + inputs, + query[0], + current_k[0], + current_v[0], + action_k[0] if action_len else None, + action_v[0] if action_len else None, + HEAD_DIM**-0.5, + ).unsqueeze(0) + + # Direct python-fn call on the same (already written) pools must be + # bit-exact: identical kernel, identical inputs. + direct = ar_diffusion_paged_attention( + query, + kv.key_cache(0), + kv.value_cache(0), + block_table=ctx.block_table, + query_start_loc=ctx.query_start_loc, + seq_lens=ctx.seq_lens, + max_query_len=ctx.max_query_len, + max_seq_len=ctx.max_seq_len, + softmax_scale=HEAD_DIM**-0.5, + causal=False, + ) + assert torch.equal(paged, direct) + + new_k = torch.cat([history_k, current_k], dim=1)[:, -kv.spec.sliding_window :] + new_v = torch.cat([history_v, current_v], dim=1)[:, -kv.spec.sliding_window :] + if action_len: + new_k = torch.cat([new_k, action_k], dim=1) + new_v = torch.cat([new_v, action_v], dim=1) + ref = _dense_attention(query, new_k, new_v) + + torch.testing.assert_close(paged, ref, rtol=2e-2, atol=2e-2) + + +def test_block_table_padded_to_fixed_width(): + """Shapes must be constant across window growth: only values change.""" + device = torch.device("cpu") + kv, st = make_state(window_chunks=2) + + ctx1 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx + ctx1.prepare(device=device, action_len=0, query_len=BLOCK) + st.commit_paged_context(False) + + ctx2 = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx + ctx2.prepare(device=device, action_len=0, query_len=BLOCK) + st.commit_paged_context(False) + + # 1-block vs 2-block visible history: same table width, same max_seq_len. + assert ctx1.block_table.shape == ctx2.block_table.shape + assert ctx1.max_seq_len == ctx2.max_seq_len + expected_width = kv.spec.sliding_window // kv.block_size + 1 + assert ctx1.block_table.shape == (1, expected_width) + # Real lengths live in seq_lens, not the padded table. + assert int(ctx1.seq_lens[0]) == BLOCK + assert int(ctx2.seq_lens[0]) == 2 * BLOCK + + +def test_prepare_is_idempotent_and_layers_share_metadata(): + device = torch.device("cpu") + _, st = make_state(num_layers=2) + contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False) + fctx = contexts[0].forward_ctx + fctx.prepare(device=device, action_len=0, query_len=BLOCK) + table = fctx.block_table + fctx.prepare(device=device, action_len=0, query_len=BLOCK) + assert fctx.block_table is table # memoized, not rebuilt + + i0, i1 = contexts[0].to_layer_inputs(), contexts[1].to_layer_inputs() + # 0-dim tensors (NOT python ints) so dynamo doesn't install per-layer + # value guards on the shared block code object. + assert isinstance(i0.layer_idx, torch.Tensor) and int(i0.layer_idx) == 0 + assert isinstance(i1.layer_idx, torch.Tensor) and int(i1.layer_idx) == 1 + # All layers share the same metadata tensor objects. + assert i0.block_table is i1.block_table + assert i0.seq_lens is i1.seq_lens + assert i0.video_slots is i1.video_slots + + +def test_layer_inputs_before_prepare_raises(): + _, st = make_state() + layer_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0] + with pytest.raises(RuntimeError, match="before prepare"): + layer_ctx.to_layer_inputs() + + +def test_custom_op_registration_idempotent(): + import importlib + import sys + + assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn") + mod = "vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention" + saved = sys.modules.pop(mod) + try: + importlib.import_module(mod) # re-registration must not raise + finally: + sys.modules[mod] = saved + assert hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn") + + +def test_custom_op_compiles_fullgraph_without_recompile_on_value_change(): + """The op must trace as one opaque node: fullgraph OK, and changed tensor + VALUES (new slots / block ids) must not trigger recompilation.""" + import torch._dynamo + + device = torch.device("cpu") + kv, st = make_state(num_layers=2, window_chunks=2) + + def run_one_forward(commit): + contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=commit) + fctx = contexts[0].forward_ctx + fctx.prepare(device=device, action_len=0, query_len=BLOCK) + q = torch.randn(BLOCK, N_HEADS, HEAD_DIM) + k = torch.randn(BLOCK, N_HEADS, HEAD_DIM) + v = torch.randn(BLOCK, N_HEADS, HEAD_DIM) + # Both layers through ONE compiled fn: layer_idx is a tensor, so a + # different layer must NOT recompile (all 40 DiT blocks share the + # block-forward code object in production). + for layer_ctx in contexts: + out = compiled(layer_ctx.to_layer_inputs(), q, k, v) + st.commit_paged_context(False) + return out + + torch._dynamo.reset() + from torch._dynamo.testing import CompileCounter + + counter = CompileCounter() + + def fn(inputs, q, k, v): + return paged_write_attn(inputs, q, k, v, None, None, HEAD_DIM**-0.5) * 1.0 + + compiled = torch.compile(fn, backend=counter, fullgraph=True) + + run_one_forward(commit=True) # history grows between calls -> + run_one_forward(commit=True) # block-table VALUES change, shapes don't + run_one_forward(commit=False) + + assert counter.frame_count == 1, f"recompiled: frame_count={counter.frame_count}" diff --git a/tests/ar_diffusion/test_pipeline_bridge.py b/tests/ar_diffusion/test_pipeline_bridge.py new file mode 100644 index 00000000000..28c3c959dc6 --- /dev/null +++ b/tests/ar_diffusion/test_pipeline_bridge.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the ARDiffusionKVState paged-attention pipeline bridge.""" + +import pytest +import torch + +from vllm_omni.experimental.ar_diffusion.kv_cache import ( + ARDiffusionKVCache, + ARDiffusionKVConfig, + ARDiffusionPagedLayerContext, +) +from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + +BLOCK = 16 +N_HEADS = 4 +HEAD_DIM = 64 + + +def make_state(num_layers=1, window_chunks=4, cross_attn_length=0): + cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks) + kv = ARDiffusionKVCache( + cfg, + num_layers=num_layers, + num_kv_heads=N_HEADS, + head_size=HEAD_DIM, + dtype=torch.float32, + block_size=BLOCK, + max_model_len=4096, + available_bytes=1 << 24, + cross_attn_length=cross_attn_length, + device=torch.device("cpu"), + ) + pos = kv.begin_request("r-pos") + neg = kv.begin_request("r-neg") + return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers) + + +def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None: + ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx + ctx.ensure_video_slots(torch.device("cpu")) + st.commit_paged_context(is_negative) + + +def test_get_returns_paged_layer_contexts_when_nothing_committed(): + _, st = make_state(num_layers=2) + for neg in (False, True): + contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False) + assert len(contexts) == 2 + assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts) + assert contexts[0].history_block_ids == [] + + +def test_managed_paged_context_commits_only_after_forward(): + kv, st = make_state(num_layers=1) + contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True) + ctx = contexts[0].forward_ctx + + assert st.pos.completed_chunks == 0 + assert kv.window_block_ids(st.pos) == [] + ctx.ensure_video_slots(torch.device("cpu")) + assert st.pos.completed_chunks == 0 + assert len(ctx.current_video_block_ids) == 2 + + st.commit_paged_context(False) + + assert st.pos.completed_chunks == 2 + assert st._committed[False] == 2 * BLOCK + assert len(kv.window_block_ids(st.pos)) == 2 + + +def test_scratch_paged_context_does_not_commit(): + kv, st = make_state(num_layers=1) + contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False) + ctx = contexts[0].forward_ctx + ctx.ensure_video_slots(torch.device("cpu")) + + assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1) + st.commit_paged_context(False) + + assert st.pos.completed_chunks == 0 + assert st._committed[False] == 0 + assert kv.window_block_ids(st.pos) == [] + + +def test_eviction_bounds_resident_window(): + kv, st = make_state(num_layers=1, window_chunks=3) + _prepare_and_commit(st, False, 3) + assert len(kv.window_block_ids(st.pos)) == 3 + + _prepare_and_commit(st, False, 2) + ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx + visible_blocks, video_len = ctx.video_block_table(torch.device("cpu")) + assert len(visible_blocks) == 3 + assert video_len == 3 * BLOCK + + +def test_branches_are_independent(): + kv, st = make_state() + _prepare_and_commit(st, False, 1) + assert len(kv.window_block_ids(st.pos)) == 1 + assert kv.window_block_ids(st.neg) == [] + + +def test_reset_clears_session_window(): + kv, st = make_state(num_layers=1, window_chunks=4) + _prepare_and_commit(st, False, 2) + _prepare_and_commit(st, True, 1) + assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK + free_before = kv.manager.block_pool.get_num_free_blocks() + + st.reset() + + assert st._committed == {False: 0, True: 0} + assert st._paged_pending == {False: None, True: None} + assert kv.window_block_ids(st.pos) == [] + assert kv.window_block_ids(st.neg) == [] + assert kv.manager.block_pool.get_num_free_blocks() > free_before + _prepare_and_commit(st, False, 1) + assert len(kv.window_block_ids(st.pos)) == 1 + + +def test_get_cross_kv_caches_raises_before_populate(): + """Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud.""" + _, st = make_state() # cross_attn_length == 0, nothing populated + for neg in (False, True): + with pytest.raises(RuntimeError, match="cross-attn read before"): + st.get_cross_kv_caches(neg) + # Guards the full AND: the populated flag alone is not enough without a pool. + st._cross_text_populated[False] = True + with pytest.raises(RuntimeError, match="cross-attn read before"): + st.get_cross_kv_caches(False) + + +def test_get_cross_kv_caches_returns_pool_dicts_when_populated(): + L = 8 + kv, st = make_state(num_layers=2, cross_attn_length=L) + written = [] + for i in range(2): + k = torch.randn(1, L, N_HEADS, HEAD_DIM) + v = torch.randn(1, L, N_HEADS, HEAD_DIM) + kv.write_cross_kv(i, False, k, v) + written.append((k, v)) + st._cross_text_populated[False] = True + + out = st.get_cross_kv_caches(False) + assert len(out) == 2 + for i, (k, v) in enumerate(written): + assert out[i]["is_init"] is True + assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM) + assert torch.equal(out[i]["k"], k) + assert torch.equal(out[i]["v"], v) + + +def test_kv_create_is_noop_engine_owns_allocation(): + from unittest.mock import MagicMock + + from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline + + p = DreamZeroPipeline.__new__(DreamZeroPipeline) + p._ar_diffusion_kv_state = object() + state = MagicMock() + p._kv_create(state, 1, "float32", "cpu", 24, 4, 64) + state.create_kv_caches.assert_not_called() + + +def test_get_kv_cache_requires_frame_aligned_seqlen(): + _, st = make_state() + with pytest.raises(AssertionError, match="frame-aligned"): + st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True) + + +def test_prepare_one_branch_allocates_nothing_for_the_other(): + """CFG-parallel laziness: a rank prepares only the branch it runs.""" + kv, st = make_state(num_layers=1, window_chunks=4) + pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx + neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx + + pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK) + + assert pos_ctx._allocated_video + assert not neg_ctx._allocated_video + assert kv.window_block_ids(st.neg) == [] + + +def test_failed_forward_tears_down_session(monkeypatch): + """Review follow-up (hsliuustc0106): a forward that dies partway must not + leave allocated-but-uncommitted KV behind — the session is torn down (pool + blocks freed, model-local state dropped) and the error propagates.""" + from types import SimpleNamespace + + from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner + from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner + + kv, st = make_state(num_layers=1, window_chunks=4) + _prepare_and_commit(st, False, 2) # session owns pool blocks + free_before_failure = kv.manager.block_pool.get_num_free_blocks() + + runner = object.__new__(ARDiffusionModelRunner) + runner.kv_cache = kv + runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st}) + runner._max_ar_diffusion_states = 4 + runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None) + runner.device = None + runner._perf_e2e_times = [] + + def boom(self, req, kv_prefetch_jobs=None): + raise RuntimeError("layer 17 exploded") + + monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom) + req = SimpleNamespace( + request_id="r0", + sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}), + ) + + with pytest.raises(RuntimeError, match="layer 17 exploded"): + ARDiffusionModelRunner.execute_model(runner, req) + + assert "s1" not in runner._ar_diffusion_states + assert "s1" not in runner.pipeline._states + assert runner.pipeline._ar_diffusion_kv_state is None + # close() freed the session's resident blocks. + assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure diff --git a/tests/ar_diffusion/test_slot_mapping.py b/tests/ar_diffusion/test_slot_mapping.py new file mode 100644 index 00000000000..061df4d9c64 --- /dev/null +++ b/tests/ar_diffusion/test_slot_mapping.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Slot-mapping / block-table tests for the AR-Diffusion engine (Phase 1, Step 3).""" + +import pytest +import torch + +from vllm_omni.experimental.ar_diffusion.kv_cache import ( + ARDiffusionRequestAdapter, + ChunkWindowSpec, + build_kv_manager, + chunk_slot_mapping, + compute_slot_mapping, + resident_block_ids, +) + +BLOCK = 16 + + +def make_spec(*, window_chunks=2): + return ChunkWindowSpec( + block_size=BLOCK, + num_kv_heads=4, + head_size=64, + dtype=torch.float16, + sliding_window=window_chunks * BLOCK, + chunk_size=BLOCK, + window_chunks=window_chunks, + ) + + +def test_slot_mapping_matches_block_offsets(): + # block table [5, 2, 9], block_size 4. slot = block*4 + (pos % 4). + slots = compute_slot_mapping([5, 2, 9], [0, 1, 4, 5, 8, 9], block_size=4) + assert slots.tolist() == [20, 21, 8, 9, 36, 37] + + +def test_slot_mapping_rejects_bad_block_size(): + with pytest.raises(ValueError): + compute_slot_mapping([1], [0], block_size=0) + + +def test_chunk_slot_mapping_targets_current_chunk(): + # After 2 committed chunks (num_computed=32), the in-flight chunk maps to the + # block covering positions 32..47 — i.e. block table index 2. + block_ids = [7, 4, 3] # block index 2 -> physical block 3 + slots = chunk_slot_mapping(block_ids, num_computed_tokens=32, chunk_size=4, block_size=16) + # positions 32..35 -> block_index 2 -> physical 3 -> slots 48..51 + assert slots.tolist() == [48, 49, 50, 51] + + +def test_resident_block_ids_excludes_null(): + assert resident_block_ids([0, 1, 0, 2, 3], null_block_id=0) == [1, 2, 3] + + +def test_blocktable_build_single_request(): + """Integration: a real KVCacheManager block table -> a valid slot mapping. + + Every slot for the in-flight chunk must land inside a real (non-null) block + that the manager actually allocated for the request. + """ + spec = make_spec() + mgr = build_kv_manager(spec, ["l0"], num_blocks=16, max_model_len=1024) + null_id = mgr.block_pool.null_block.block_id + adapter = ARDiffusionRequestAdapter("req", chunk_size=BLOCK) + + mgr.allocate_slots(adapter, num_new_tokens=BLOCK) + block_ids = mgr.get_block_ids(adapter.request_id)[0] + real = set(resident_block_ids(block_ids, null_id)) + assert real, "expected at least one real block after allocation" + + slots = chunk_slot_mapping(block_ids, adapter.num_computed_tokens, spec.chunk_size, spec.block_size) + # Each slot resolves to a real allocated block, none to null_block. + blocks_used = {int(s) // spec.block_size for s in slots} + assert blocks_used <= real + assert null_id not in blocks_used + assert len(slots) == spec.chunk_size diff --git a/tests/assets/indextts2/ref_audio.wav b/tests/assets/indextts2/ref_audio.wav new file mode 100644 index 00000000000..d13a4c84749 Binary files /dev/null and b/tests/assets/indextts2/ref_audio.wav differ diff --git a/tests/assets/soulxsinger/music.json b/tests/assets/soulxsinger/music.json new file mode 100644 index 00000000000..d9e0d8882ce --- /dev/null +++ b/tests/assets/soulxsinger/music.json @@ -0,0 +1,16 @@ +[ + { + "index": "vocal_240_51240", + "language": "Mandarin", + "time": [ + 240, + 51240 + ], + "duration": "0.21 0.18 0.34 0.38 0.26 0.22 0.50 0.20 0.33 0.13 0.60 0.22 0.38 0.52 0.30 0.12 1.82 0.98 0.20 0.26 0.38 0.36 0.58 0.54 0.26 0.50 1.46 3.03 0.24 0.24 0.29 0.25 0.20 0.14 0.80 0.22 0.54 0.30 0.24 0.16 0.58 0.28 0.38 1.73 0.79 0.24 0.30 0.34 0.30 0.30 0.34 0.21 0.21 0.36 0.34 0.23 1.85 1.90 0.23 0.39 0.68 0.50 0.31 0.43 0.76 0.38 2.00 1.87 0.68 0.72 0.56 0.62 0.80 0.40 0.42 1.68 1.79 0.70 0.66 0.54 0.24 0.48 0.68 0.40 2.34 0.14", + "text": "<SP> 只 是 因 为 为 在 人 群 群 中 多 看 了 你 一 眼 <SP> 再 也 没 能 忘 掉 你 容 颜 <SP> 梦 想 着 着 偶 偶 然 然 有 <SP> 一 一 天 再 相 见 <SP> 从 此 我 开 始 始 孤 孤 单 思 念 念 <SP> 想 想 你 时 你 你 在 天 边 <SP> 想 你 时 你 在 眼 前 前 <SP> 想 你 时 你 你 在 脑 海 <SP>", + "phoneme": "<SP> zh_zhi3 zh_shi4 zh_yin1 zh_wei4 zh_wei2 zh_zai4 zh_ren2 zh_qun2 zh_qun2 zh_zhong1 zh_duo1 zh_kan4 zh_le5 zh_ni3 zh_yi1 zh_yan3 <SP> zh_zai4 zh_ye3 zh_mei2 zh_neng2 zh_wang4 zh_diao4 zh_ni3 zh_rong2 zh_yan2 <SP> zh_meng4 zh_xiang3 zh_zhe5 zh_zhe5 zh_ou3 zh_ou3 zh_ran2 zh_ran2 zh_you3 <SP> zh_yi1 zh_yi1 zh_tian1 zh_zai4 zh_xiang1 zh_jian4 <SP> zh_cong2 zh_ci3 zh_wo3 zh_kai1 zh_shi3 zh_shi3 zh_gu1 zh_gu1 zh_dan1 zh_si1 zh_nian4 zh_nian4 <SP> zh_xiang3 zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_tian1 zh_bian1 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_zai4 zh_yan3 zh_qian2 zh_qian2 <SP> zh_xiang3 zh_ni3 zh_shi2 zh_ni3 zh_ni3 zh_zai4 zh_nao3 zh_hai3 <SP>", + "note_pitch": "0 64 64 64 66 68 66 66 66 64 64 64 64 66 64 60 61 0 63 63 63 64 66 63 61 59 56 0 68 66 66 68 68 66 66 64 64 0 64 66 61 61 61 64 0 62 63 63 64 64 66 65 66 61 58 59 56 0 69 71 66 68 68 71 66 64 61 0 66 61 68 66 64 63 61 59 0 71 66 68 68 71 66 64 61 0", + "note_type": "1 2 2 2 2 3 2 2 2 3 2 2 2 2 2 2 2 1 2 2 2 2 2 2 2 2 2 1 2 2 2 3 2 3 2 3 2 1 2 3 2 2 2 2 1 2 2 2 2 2 3 2 3 2 2 2 3 1 2 3 2 2 2 3 2 2 2 1 2 2 2 2 2 2 2 3 1 2 2 2 2 3 2 2 2 1", + "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 318.9 332.2 331.2 323.9 312.3 0.0 0.0 0.0 0.0 0.0 0.0 340.9 342.3 340.2 337.4 333.5 329.3 327.8 328.3 329.1 330.5 331.9 331.5 329.2 327.6 327.4 327.7 329.6 330.5 329.5 328.1 328.5 330.1 330.4 329.5 328.9 331.0 333.6 332.8 331.7 330.2 330.6 330.9 329.5 327.6 332.4 343.5 363.9 368.7 369.3 368.2 366.7 371.3 382.0 392.9 402.3 409.8 413.8 416.3 416.4 415.4 414.1 413.4 415.3 417.4 418.3 413.6 416.0 0.0 0.0 0.0 0.0 368.2 362.5 361.3 362.0 362.9 364.3 366.4 367.5 368.4 368.9 368.1 367.6 369.7 371.1 371.1 369.5 368.0 367.3 366.6 367.1 369.5 372.0 371.3 368.2 368.1 369.6 369.9 373.4 376.5 360.2 285.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 384.0 376.7 372.0 371.7 373.9 376.3 375.5 361.2 333.0 317.8 319.0 329.1 333.5 333.2 333.2 333.3 328.0 314.5 0.0 0.0 0.0 320.6 326.5 332.9 335.7 334.1 329.9 327.4 326.3 326.6 329.1 330.5 330.0 328.0 326.7 328.6 331.1 332.4 332.1 332.8 333.2 333.1 331.8 328.6 318.9 293.2 323.6 327.5 330.0 332.8 331.1 319.4 272.2 0.0 0.0 0.0 0.0 278.5 294.5 310.0 318.5 322.7 327.4 334.1 337.2 336.3 328.5 321.0 324.2 341.5 362.2 374.7 375.8 370.5 366.0 365.2 368.7 371.7 374.4 373.0 368.7 370.4 374.8 375.1 372.5 368.7 363.0 357.1 359.0 366.8 377.1 379.5 371.5 359.9 351.0 358.4 371.7 377.9 375.8 367.8 359.7 363.6 375.3 379.9 377.5 372.9 359.8 345.2 337.1 334.5 333.4 332.8 329.6 319.2 292.8 261.7 253.4 262.8 273.0 278.1 279.3 278.6 277.6 277.9 278.3 278.0 277.3 276.7 275.6 275.3 276.1 277.8 278.2 277.7 277.6 277.2 276.4 275.5 273.9 271.8 270.7 274.7 282.7 285.5 281.4 273.1 264.4 262.2 268.9 278.4 284.3 283.6 277.3 268.5 262.2 263.2 270.7 278.8 285.0 287.3 282.0 272.1 265.2 266.0 272.1 278.2 283.8 285.1 281.4 272.6 267.2 269.1 276.1 282.7 289.3 290.5 285.6 273.9 267.5 267.4 271.8 277.7 282.6 283.1 277.8 269.8 265.6 268.6 273.8 281.1 286.2 287.2 282.3 270.4 264.5 263.3 263.0 268.5 268.3 268.7 270.9 277.2 277.4 284.3 285.7 282.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 311.5 306.4 306.0 306.8 305.8 305.0 307.2 310.8 312.4 313.5 314.4 314.6 313.7 312.2 311.1 311.1 312.8 314.6 312.7 306.5 299.7 302.7 308.7 315.2 314.2 310.9 309.4 309.8 310.6 310.7 310.8 311.4 312.2 313.0 312.5 312.2 312.2 311.9 312.0 312.2 311.4 306.1 303.3 309.3 315.0 325.3 324.4 323.9 325.0 326.4 327.7 327.7 328.1 328.6 327.8 328.3 328.3 327.6 328.6 336.1 347.2 359.4 371.0 374.9 371.3 366.5 369.5 375.5 376.2 370.2 365.7 367.3 372.9 376.9 375.3 368.4 360.1 358.0 365.4 380.8 383.5 380.4 374.4 367.6 365.9 372.0 376.3 377.9 375.7 372.6 368.6 361.2 353.7 0.0 309.4 302.3 299.2 300.6 305.7 308.5 308.6 309.0 310.9 310.9 309.5 308.6 309.0 311.3 313.0 314.4 314.5 313.0 312.1 311.2 308.0 299.5 295.8 295.0 285.7 278.0 280.1 281.2 279.9 278.3 278.3 279.5 279.7 272.2 259.1 0.0 0.0 0.0 240.7 236.0 233.8 233.5 234.6 235.7 235.6 236.2 239.8 245.4 248.8 250.1 250.4 249.9 249.1 247.2 241.7 231.7 216.0 197.6 192.1 197.2 204.3 206.5 205.4 201.7 200.8 203.5 208.2 210.9 210.3 207.1 203.5 202.3 202.1 202.8 205.1 207.3 208.1 207.3 205.4 202.1 199.3 200.5 206.1 212.9 212.9 206.5 197.1 191.1 191.2 196.8 203.6 207.8 208.7 206.4 201.7 197.3 197.4 200.7 206.5 209.9 209.9 206.9 203.3 201.2 203.4 207.7 210.5 208.9 207.6 206.5 204.2 202.0 203.6 205.7 210.8 213.1 214.3 210.6 204.1 199.7 202.3 211.9 217.7 215.1 215.0 215.3 213.3 0.0 0.0 216.3 0.0 0.0 195.2 209.0 205.3 201.3 196.8 195.7 195.8 195.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 431.9 425.8 421.0 417.7 419.3 422.7 423.4 421.1 417.9 411.9 408.6 411.5 414.9 419.6 426.7 424.2 0.0 0.0 0.0 280.8 279.4 369.7 371.8 370.6 369.8 370.5 370.5 371.4 374.6 374.0 361.0 332.7 328.0 0.0 359.0 365.0 371.1 373.8 371.9 369.7 369.6 371.0 378.3 390.4 403.4 414.1 417.9 417.5 417.0 416.6 415.7 414.9 414.5 413.8 413.5 413.6 413.8 415.2 416.9 417.3 415.4 415.0 414.1 411.4 409.2 403.5 392.7 375.6 367.5 365.8 368.8 370.7 370.8 369.0 367.0 366.0 366.2 367.8 369.2 368.8 367.2 366.7 367.3 367.3 368.8 368.5 366.8 365.3 364.0 363.0 365.0 367.3 367.7 365.5 364.0 364.9 367.4 370.3 371.3 369.9 366.5 364.1 368.3 385.8 406.7 409.4 399.7 376.0 357.2 359.9 372.4 377.5 366.9 345.1 320.6 315.8 321.8 329.0 329.9 326.4 324.5 324.8 324.9 325.3 325.1 324.7 325.9 326.8 326.6 326.6 326.4 325.4 323.2 322.5 323.9 326.2 328.6 329.9 329.5 329.0 328.2 327.7 327.9 329.6 331.2 331.4 334.2 340.1 335.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 309.0 320.1 329.6 334.4 334.4 334.5 338.2 346.0 355.1 361.4 366.1 369.5 372.7 373.3 372.9 375.2 375.8 374.7 367.6 360.1 0.0 0.0 0.0 0.0 0.0 0.0 283.6 280.4 277.2 274.3 271.8 271.0 273.3 277.2 278.8 278.1 277.2 277.4 278.8 279.3 280.2 281.4 279.2 274.9 268.3 254.1 239.9 0.0 0.0 267.7 267.7 270.1 271.8 273.1 276.3 278.3 269.1 256.1 0.0 0.0 0.0 0.0 286.0 281.1 275.7 272.7 271.7 271.7 272.5 274.0 275.5 277.9 281.1 287.1 314.0 342.2 363.0 376.8 375.4 365.8 351.8 323.3 0.0 0.0 0.0 0.0 299.6 322.7 336.5 336.9 335.4 332.9 330.0 326.6 323.4 322.3 324.4 326.7 328.5 328.5 326.4 324.0 322.5 323.2 325.5 327.7 328.6 328.1 325.4 320.9 316.5 316.1 320.8 327.7 333.7 333.2 325.6 315.5 306.9 305.5 315.1 331.0 343.3 341.6 331.4 319.5 308.3 307.6 318.1 329.5 338.2 342.0 337.4 329.4 321.0 316.6 319.8 330.8 340.6 344.5 342.3 331.4 319.3 315.7 322.9 329.3 336.5 344.0 341.1 328.3 318.7 320.0 328.2 333.5 337.9 339.7 338.5 327.3 323.3 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 285.3 284.6 284.3 288.2 295.1 305.6 312.2 315.5 287.8 0.0 0.0 0.0 0.0 0.0 309.8 312.4 313.8 312.8 310.8 310.7 312.5 313.2 310.7 305.7 305.9 308.1 309.9 311.3 311.0 309.8 309.7 310.6 311.1 311.1 311.0 311.9 313.7 316.4 319.5 311.6 281.3 283.4 306.6 0.0 0.0 0.0 334.4 326.7 322.6 321.6 321.7 324.5 332.0 336.7 333.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 332.5 331.7 332.7 334.8 334.7 335.9 339.7 345.6 353.9 362.3 370.2 374.0 373.8 371.4 370.0 368.2 367.0 368.0 370.4 371.2 371.3 370.0 360.3 0.0 0.0 0.0 319.1 326.8 331.5 332.0 331.6 327.2 328.2 334.1 343.1 350.7 358.0 363.5 368.4 370.8 371.8 371.1 370.2 368.5 368.0 371.1 375.5 376.3 369.5 0.0 0.0 0.0 280.4 267.2 263.6 265.8 271.2 275.0 276.4 276.9 277.1 280.9 287.5 284.7 280.7 0.0 0.0 0.0 0.0 0.0 0.0 231.6 226.0 222.9 224.1 226.2 228.5 234.3 240.9 246.6 249.8 249.6 246.3 245.2 245.8 247.2 248.3 249.3 249.0 245.9 242.7 235.8 226.3 217.8 213.1 209.3 208.0 207.0 205.9 205.8 203.4 201.5 205.3 207.5 209.1 210.4 208.9 203.2 199.2 199.3 201.2 205.6 206.3 204.0 202.7 202.2 203.3 206.1 205.6 201.9 198.8 195.5 195.6 198.7 204.2 214.5 218.3 214.3 207.0 199.1 192.4 189.9 193.6 201.0 210.6 212.4 209.2 202.0 195.2 190.0 190.4 196.0 204.2 209.2 208.9 203.0 195.8 190.7 189.8 194.6 200.8 208.7 213.5 210.5 202.0 193.9 187.7 187.8 193.1 199.8 202.8 204.1 203.8 200.4 197.0 193.9 191.6 192.5 198.4 204.8 203.9 203.5 201.3 200.2 198.4 198.5 201.0 204.0 204.6 205.4 202.0 199.2 194.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 426.6 432.4 436.6 437.2 434.6 433.0 434.0 453.1 473.8 492.8 503.2 501.8 491.3 479.6 478.1 489.7 503.6 508.3 503.2 489.5 482.3 484.2 495.7 505.4 506.6 501.6 495.4 492.7 497.1 500.4 496.1 490.1 480.4 453.8 412.6 373.6 357.7 354.5 356.8 361.0 365.0 367.8 369.5 370.1 371.1 371.3 370.8 370.1 369.7 369.0 369.5 370.4 371.1 370.9 370.2 370.8 371.3 372.7 380.2 386.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.0 427.8 417.0 411.6 412.6 411.2 410.7 412.5 413.7 414.9 414.8 415.7 415.4 414.0 413.8 414.6 415.2 416.3 415.2 413.6 413.4 414.9 416.8 416.4 414.7 413.8 413.2 414.7 417.5 420.7 428.2 443.1 459.8 481.7 502.3 510.8 509.1 497.8 490.6 493.5 500.8 508.3 511.8 504.6 493.5 487.3 492.3 507.2 518.8 514.8 496.4 489.6 490.4 491.8 493.3 0.0 0.0 0.0 0.0 0.0 391.9 368.8 360.2 358.3 360.0 362.0 361.8 360.9 360.3 361.8 365.7 368.7 369.3 367.7 364.7 365.6 368.4 370.4 371.7 369.0 366.8 367.5 367.9 369.4 374.6 384.6 396.8 388.2 382.8 0.0 0.0 0.0 0.0 0.0 352.6 339.5 328.2 326.0 326.5 326.9 328.2 329.4 329.4 329.3 337.3 349.1 350.9 343.6 333.4 326.5 327.1 331.5 333.4 326.8 318.4 0.0 0.0 0.0 282.0 283.6 281.6 279.5 273.2 267.9 268.8 273.2 277.5 279.8 278.7 276.8 273.3 268.7 265.0 266.1 272.4 282.3 284.9 279.4 268.5 255.3 253.0 259.9 270.2 281.5 284.3 279.7 270.1 260.2 256.4 260.1 266.6 273.0 279.6 283.6 282.3 272.3 262.1 255.8 259.8 269.6 274.9 278.7 278.4 271.5 261.7 255.8 256.8 263.9 272.1 279.7 278.4 269.9 259.0 255.9 262.1 269.2 274.5 279.6 282.1 280.7 274.3 269.9 270.9 270.7 273.8 276.9 274.9 271.1 266.5 266.7 269.4 276.9 282.9 282.0 279.5 274.9 271.2 267.9 263.1 270.2 281.3 285.5 283.9 280.6 271.1 262.8 263.9 263.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 362.9 364.0 366.8 366.6 363.6 359.8 357.7 359.9 367.6 371.6 365.4 359.3 358.6 359.5 365.5 375.7 378.1 371.5 359.5 350.8 358.6 374.7 381.7 376.8 363.6 353.3 356.9 367.8 378.3 380.8 375.7 372.3 372.6 374.7 372.9 365.2 346.4 314.8 285.6 277.3 275.4 275.3 277.5 279.2 279.2 278.8 277.6 276.4 276.5 277.6 278.8 279.9 281.8 282.7 281.0 279.2 278.6 279.4 279.9 279.2 278.7 281.7 285.8 289.9 288.7 284.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 428.0 417.9 410.0 409.5 411.2 414.0 417.5 417.9 418.1 419.3 419.0 417.8 417.1 417.2 416.6 415.8 415.0 415.1 414.5 410.9 401.5 388.2 371.5 359.2 363.0 371.7 375.7 370.6 359.6 353.4 358.7 367.0 374.5 378.2 373.7 362.8 358.0 361.6 368.5 373.7 377.7 375.8 368.8 361.7 360.5 366.7 372.5 375.6 380.1 384.9 383.5 376.6 0.0 0.0 0.0 0.0 0.0 331.7 320.0 316.8 320.7 325.1 327.0 327.4 326.6 326.4 325.5 325.1 326.2 327.6 328.5 328.4 327.7 328.7 329.9 329.8 329.7 327.7 326.5 327.8 328.9 329.3 329.8 329.2 328.8 329.4 330.4 331.0 331.0 330.5 329.2 328.2 328.3 328.7 329.4 327.5 322.9 314.8 299.3 290.8 292.3 299.0 305.4 309.2 312.1 313.9 316.8 320.4 327.6 333.1 338.0 341.2 342.1 338.4 331.0 0.0 0.0 0.0 0.0 0.0 0.0 311.7 292.5 281.8 276.1 275.1 276.6 277.6 277.4 274.0 264.8 249.6 239.1 238.8 243.0 245.7 245.0 241.9 237.9 234.6 237.1 242.7 250.4 252.5 248.3 240.7 231.6 228.3 234.0 240.6 247.0 250.6 248.5 242.1 233.6 226.8 227.2 234.0 240.7 246.9 248.9 244.4 237.6 230.8 228.8 236.4 246.9 250.2 249.2 245.4 239.9 232.5 224.7 226.5 238.5 252.0 257.4 255.9 248.4 237.4 230.5 231.6 238.8 247.6 252.6 253.9 253.0 248.9 241.8 236.4 234.1 236.5 246.2 257.7 255.4 248.2 238.5 234.8 237.8 244.4 250.5 253.8 251.7 246.7 240.0 238.3 244.8 251.6 256.4 256.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 484.0 484.5 487.5 490.6 487.2 479.7 478.9 483.8 494.1 505.7 510.3 501.1 485.4 472.4 472.3 495.6 515.8 517.1 503.8 486.5 475.1 479.9 491.4 507.9 509.9 505.0 496.2 485.7 484.5 490.8 495.4 497.7 496.6 489.9 469.1 424.3 387.2 373.0 367.3 364.4 364.6 366.2 366.6 368.8 371.9 373.1 371.4 368.9 367.1 367.0 367.7 368.8 370.1 369.8 369.2 369.8 370.0 369.0 369.0 371.0 375.2 385.3 391.6 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 425.2 418.5 409.4 409.5 408.1 406.3 408.4 410.3 412.8 413.8 412.6 411.2 410.8 410.4 410.2 410.5 412.2 413.6 413.4 412.6 412.6 413.4 415.6 417.9 418.5 417.0 415.0 415.0 416.8 421.5 437.9 455.0 472.6 493.0 502.2 496.7 484.6 480.6 487.2 498.6 506.3 506.2 499.4 489.7 489.5 496.2 503.8 507.2 501.8 491.1 489.9 501.1 520.5 533.7 525.8 0.0 0.0 0.0 0.0 0.0 421.8 365.2 351.0 352.0 358.0 362.6 365.6 365.4 363.2 362.5 362.2 362.4 364.1 365.2 364.9 363.5 363.4 366.3 368.5 370.0 369.7 368.5 366.9 365.7 366.3 368.6 370.8 372.2 370.2 373.1 377.0 376.5 371.0 346.8 326.1 325.2 328.4 331.4 334.0 332.2 327.2 324.4 331.3 349.7 363.2 362.0 346.1 327.4 321.8 329.8 340.7 348.4 349.9 346.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 269.5 275.7 282.6 282.9 278.8 275.2 274.3 273.5 273.9 275.3 276.4 277.3 276.1 272.3 268.0 268.3 272.5 279.1 281.7 280.1 271.3 261.0 257.3 262.7 273.2 280.8 283.3 279.3 269.1 258.3 258.0 266.3 278.2 286.7 287.8 283.4 274.9 264.7 257.7 260.0 272.0 286.3 294.7 289.8 274.5 263.8 263.4 268.8 277.3 284.1 286.5 285.2 281.7 272.7 264.3 260.3 267.7 281.7 289.2 289.0 281.2 266.2 256.0 254.4 261.5 276.4 286.8 288.8 286.7 273.8 260.9 260.8 270.2 283.4 291.3 292.8 286.1 273.8 265.2 264.0 271.3 283.8 290.3 289.7 277.7 266.9 260.5 263.0 267.8 281.5 286.7 286.1 285.5 280.1 279.0 283.0 284.1 284.5 285.5 282.2 0.0 0.0 280.0 276.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" + } +] diff --git a/tests/assets/soulxsinger/music.mp3 b/tests/assets/soulxsinger/music.mp3 new file mode 100644 index 00000000000..f3126e79e38 Binary files /dev/null and b/tests/assets/soulxsinger/music.mp3 differ diff --git a/tests/assets/soulxsinger/zh_prompt.json b/tests/assets/soulxsinger/zh_prompt.json new file mode 100644 index 00000000000..86ff3614ea2 --- /dev/null +++ b/tests/assets/soulxsinger/zh_prompt.json @@ -0,0 +1,16 @@ +[ + { + "index": "vocal_320_10687", + "language": "Mandarin", + "time": [ + 320, + 10687 + ], + "duration": "0.23 0.34 0.26 0.70 0.52 0.46 0.36 0.44 0.14 0.24 0.64 0.47 0.51 1.10 0.28 0.38 0.32 0.32 0.38 0.32 0.31 0.19 1.45", + "text": "<SP> 除 了 想 你 你 <SP> 除 了 了 爱 你 你 <SP> 我 什 么 什 么 都 愿 愿 意", + "phoneme": "<SP> zh_chu2 zh_le5 zh_xiang3 zh_ni3 zh_ni3 <SP> zh_chu2 zh_le5 zh_le5 zh_ai4 zh_ni3 zh_ni3 <SP> zh_wo3 zh_shen2 zh_me5 zh_shen2 zh_me5 zh_dou1 zh_yuan4 zh_yuan4 zh_yi4", + "note_pitch": "0 62 65 67 67 69 0 67 69 67 65 67 69 67 67 66 64 64 60 60 65 67 0", + "note_type": "1 2 2 2 2 3 1 2 2 3 2 2 3 1 2 2 2 2 2 2 2 3 2", + "f0": "0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 294.1 288.2 290.6 294.6 295.7 292.8 291.5 294.4 295.4 294.7 293.5 292.2 294.4 295.8 293.2 297.7 320.1 338.1 348.3 348.7 344.4 342.6 346.2 354.8 356.9 353.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 401.1 403.8 405.0 402.1 398.4 395.6 393.3 392.0 391.4 390.3 390.2 390.3 390.3 391.3 392.8 393.6 391.7 390.5 391.4 391.6 391.5 393.5 393.8 390.8 387.4 387.8 389.3 390.8 392.2 391.6 390.2 389.8 389.1 388.4 390.0 395.5 397.2 396.7 395.5 395.1 394.6 394.9 395.6 395.2 394.6 395.4 395.9 394.0 391.7 390.7 391.7 392.6 391.6 395.7 405.8 441.7 462.5 463.8 450.2 430.9 414.5 415.2 426.7 439.8 454.4 462.9 447.8 422.5 400.6 403.2 423.5 451.3 482.3 492.8 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 435.9 414.3 406.2 402.8 398.2 396.7 396.7 395.2 398.3 416.9 441.6 446.9 441.5 437.3 435.3 432.9 431.6 432.3 434.8 436.7 433.8 422.8 406.7 390.9 382.0 381.7 384.7 382.7 368.1 357.4 355.6 355.1 352.4 348.1 346.4 348.7 351.6 355.0 354.4 351.7 349.9 349.2 348.1 346.0 345.4 344.2 344.4 345.5 346.6 349.0 349.7 349.1 349.5 349.6 349.7 349.4 349.5 352.2 354.5 355.7 355.6 356.9 359.4 361.5 363.8 360.4 354.3 357.2 363.9 372.4 382.6 399.1 402.9 400.6 395.5 390.5 388.9 390.2 391.1 391.9 391.5 390.4 390.2 391.3 391.6 391.0 388.6 386.2 389.1 403.8 430.2 441.8 449.5 448.1 443.2 438.3 432.9 430.7 434.6 442.0 447.6 446.0 440.3 434.7 431.1 435.7 442.3 445.9 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 423.8 402.1 398.5 397.6 393.4 390.5 391.4 402.5 427.6 442.5 442.5 435.1 430.1 430.8 439.9 447.4 442.2 426.1 412.3 399.8 391.8 389.5 388.5 387.8 386.4 384.7 384.5 387.9 391.2 391.8 392.9 393.8 392.0 392.0 395.4 398.1 398.2 396.3 393.1 391.1 388.9 386.6 383.1 381.5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 357.8 367.6 370.9 365.4 359.3 355.7 358.1 372.7 396.8 404.0 398.4 392.6 389.2 388.8 383.6 362.3 341.1 325.1 326.3 327.9 331.3 333.3 326.0 319.4 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 356.4 353.5 349.0 346.5 338.4 328.4 323.7 326.2 334.2 338.4 331.6 309.8 280.6 256.4 252.8 256.1 258.3 256.9 257.6 261.1 259.7 258.7 258.6 260.2 262.2 262.4 262.9 264.2 263.4 260.2 256.4 253.2 238.7 223.3 231.6 257.7 258.1 258.4 258.8 258.0 256.8 255.8 254.0 255.1 258.1 261.9 263.7 262.6 256.7 253.1 250.2 246.7 258.7 294.4 327.1 342.6 346.0 344.1 341.2 342.3 345.2 350.1 364.7 382.5 396.1 396.1 389.2 381.8 381.1 387.2 397.0 399.3 390.7 374.3 360.1 350.6 346.6 347.4 350.7 354.4 354.3 351.7 349.8 348.4 346.9 347.0 348.4 349.5 351.0 352.3 353.6 353.3 350.5 348.3 345.5 344.3 344.4 347.0 350.6 352.0 351.0 350.8 350.1 347.6 345.7 347.0 350.3 351.7 350.7 348.5 346.9 347.7 349.0 349.1 348.5 346.8 346.3 348.4 349.0 349.2 351.1 349.6 348.3 350.5 351.1 348.0 347.6 349.1 351.3 356.0 361.3 360.6 354.0 341.0 316.2 302.9 302.1 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0" + } +] diff --git a/tests/assets/soulxsinger/zh_prompt.mp3 b/tests/assets/soulxsinger/zh_prompt.mp3 new file mode 100644 index 00000000000..bc306410a44 Binary files /dev/null and b/tests/assets/soulxsinger/zh_prompt.mp3 differ diff --git a/tests/attention/test_fish_kvcache_attn.py b/tests/attention/test_fish_kvcache_attn.py index 89b43e6279d..c82fd185882 100644 --- a/tests/attention/test_fish_kvcache_attn.py +++ b/tests/attention/test_fish_kvcache_attn.py @@ -92,6 +92,18 @@ def _metadata( )() +def _decode_kv_cache( + *, + num_blocks: int = 8, + block_size: int = 16, + num_kv_heads: int = 8, + head_size: int = 128, + dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + # vLLM >=0.23.0 KV cache layout: key/value live on dim 1. + return torch.zeros((num_blocks, 2, block_size, num_kv_heads, head_size), dtype=dtype) + + def test_fish_kvcache_enabled_by_default(monkeypatch): monkeypatch.delenv("VLLM_OMNI_FISH_KVCACHE_ATTN", raising=False) @@ -182,7 +194,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc metadata = _metadata() query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -201,6 +213,110 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc assert fish_kvcache_backend.get_fish_kvcache_attn_stats()["small_hit_count"] == 1 +def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size): + # Source of truth for the paged KV cache layout the backend must unpack. + from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend + + return FlashAttentionBackend.get_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + + +def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch): + # vLLM >=0.23.0 lays the KV cache out as (num_blocks, 2, block_size, + # num_kv_heads, head_size). The backend must unbind on dim 1 so each of + # key/value comes back as the 4-D (num_blocks, block_size, num_kv_heads, + # head_size) tensor the decode kernel expects. + monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") + monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) + monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None) + monkeypatch.setattr(fish_kvcache_backend, "can_use_fish_kvcache_attn", lambda **_: True) + + num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128 + cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + assert cache_shape == (num_blocks, 2, block_size, num_kv_heads, head_size) + + captured = {} + + def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, scale, max_seq_len): + del query, block_table, seq_lens, scale, max_seq_len + captured["key"] = tuple(key_cache.shape) + captured["value"] = tuple(value_cache.shape) + out.fill_(1) + return out + + monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode) + + model = _FakeModel() + assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1 + + query = torch.zeros((2, 32, head_size), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros(cache_shape, dtype=torch.float16) + + model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output) + + expected = (num_blocks, block_size, num_kv_heads, head_size) + assert captured["key"] == expected + assert captured["value"] == expected + + +def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkeypatch): + # End-to-end with the real vLLM cache layout and the real (unmocked) guard. + # + # The point of the unbind(1) fix: unbind(0) raised "too many values to + # unpack" on the 5-D vLLM 0.23.0 layout, so the request crashed outright. + # unbind(1) makes the request *work* again -- it unpacks cleanly and the + # backend gracefully falls back to the original attention with correct + # output. This test guards that fix and must NOT depend on the fast path + # firing. + # + # It also documents a known limitation (tracked separately): unbind(1) on + # the interleaved (num_blocks, 2, ...) layout yields *non-contiguous* + # key/value views, which the guard rejects on the is_contiguous() check. + # So the dimensions are compatible (4-D, block_size=16, head_size=128) -- + # the only thing keeping the fast path from engaging is contiguity, not + # shape. Making the fast path actually fire would require a .contiguous() + # copy and is out of scope here. + fish_kvcache_backend.reset_fish_kvcache_attn_stats() + monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") + monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) + monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None) + monkeypatch.setattr(fish_kvcache_attn, "is_available", lambda: True) + + num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128 + cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size) + + decode_calls = [] + + def fake_decode(*args, **kwargs): + del args, kwargs + decode_calls.append(True) + + monkeypatch.setattr(fish_kvcache_backend, "fish_decode_kvcache_attn", fake_decode) + + model = _FakeModel() + assert fish_kvcache_backend.install_fish_kvcache_attn_backend(model) == 1 + + query = torch.zeros((2, 32, head_size), dtype=torch.float16) + output = torch.zeros_like(query) + kv_cache = torch.zeros(cache_shape, dtype=torch.float16) + + # Root cause record: unbind(1) gives the right 4-D shape but a non-contiguous view. + key_cache, value_cache = kv_cache.unbind(1) + assert tuple(key_cache.shape) == (num_blocks, block_size, num_kv_heads, head_size) + assert not key_cache.is_contiguous() + + # unbind(1) no longer crashes (unbind(0) would have raised here) and the + # request completes via the correct fallback path. + result = model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output) + + assert result is output + assert decode_calls == [] + assert model.layers[0].self_attn.attn.impl.original_calls == 1 + assert torch.equal(output, torch.full_like(output, 2)) + stats = fish_kvcache_backend.get_fish_kvcache_attn_stats() + assert stats["fallback_count_by_reason"] == {"guard_miss": 1} + + def test_fish_kvcache_backend_install_is_idempotent(monkeypatch): monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1") monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True) @@ -215,7 +331,7 @@ def test_fish_kvcache_backend_install_is_idempotent(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -250,7 +366,7 @@ def fake_decode(*args, **kwargs): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -297,7 +413,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -340,7 +456,7 @@ def fake_can_use(**kwargs): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -383,7 +499,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() model.layers[0].self_attn.attn.impl.forward( None, @@ -428,7 +544,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -475,7 +591,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -509,7 +625,7 @@ def test_fish_kvcache_backend_required_rejects_non_cpu_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="must be a CPU tensor"): model.layers[0].self_attn.attn.impl.forward( @@ -539,7 +655,7 @@ def test_fish_kvcache_backend_required_rejects_underestimated_upper_bound(monkey query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="underestimates"): model.layers[0].self_attn.attn.impl.forward( @@ -569,7 +685,7 @@ def test_fish_kvcache_backend_required_rejects_missing_upper_bound(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="requires seq_lens_cpu_upper_bound"): model.layers[0].self_attn.attn.impl.forward( @@ -603,7 +719,7 @@ def test_fish_kvcache_backend_required_raises_on_guard_miss(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() with pytest.raises(RuntimeError, match="required"): model.layers[0].self_attn.attn.impl.forward( @@ -629,7 +745,7 @@ def test_fish_kvcache_backend_required_allows_prefill_fallback(monkeypatch): query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, @@ -675,7 +791,7 @@ def fake_decode(query, key_cache, value_cache, block_table, seq_lens, out, *, sc query = torch.zeros((2, 32, 128), dtype=torch.float16) output = torch.zeros_like(query) - kv_cache = torch.zeros((2, 8, 16, 8, 128), dtype=torch.float16) + kv_cache = _decode_kv_cache() result = model.layers[0].self_attn.attn.impl.forward( None, diff --git a/tests/benchmarks/metrics/test_metrics.py b/tests/benchmarks/metrics/test_metrics.py index 2688b224f5b..32d9d7958f5 100644 --- a/tests/benchmarks/metrics/test_metrics.py +++ b/tests/benchmarks/metrics/test_metrics.py @@ -14,6 +14,37 @@ pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu] +def test_tpot_matches_mean_itl_per_request(): + """TPOT and pooled ITL should agree when output_len tracks ITL samples.""" + output = MixRequestFuncOutput() + output.success = True + output.prompt_len = 100 + # Simulate server reporting more tokens than SSE chunks (bundled tokens). + # len(itl)=2 → 2 inter-chunk intervals, but server generated 10 tokens. + output.output_tokens = 10 + output.generated_text = "hello world" + output.ttft = 0.05 + output.text_latency = 0.25 + output.latency = 0.30 + output.itl = [0.10, 0.10] + + metrics, _ = calculate_metrics( + input_requests=[], + outputs=[output], + dur_s=1.0, + tokenizer=None, + selected_percentiles=[50.0], + goodput_config_dict={}, + task_type=TaskType.GENERATION, + selected_percentile_metrics=["tpot", "itl"], + max_concurrency=None, + request_rate=float("inf"), + benchmark_duration=1.0, + ) + + assert metrics.mean_tpot_ms == pytest.approx(metrics.mean_itl_ms, rel=1e-6, abs=1e-6) + + def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutput: """Build a minimal successful MixRequestFuncOutput for metrics aggregation.""" output = MixRequestFuncOutput() diff --git a/tests/benchmarks/patch/test_patch.py b/tests/benchmarks/patch/test_patch.py index 35a18aea33c..ac3aef31e17 100644 --- a/tests/benchmarks/patch/test_patch.py +++ b/tests/benchmarks/patch/test_patch.py @@ -133,7 +133,9 @@ async def test_output_tokens_not_assigned_without_metrics(mocker: MockerFixture) # Assert assert output.success is True - assert output.output_tokens == 0, "output_tokens should default to 0 when no metrics" + # Without server metrics we no longer infer count from SSE chunks because + # one chunk may carry multiple tokens. output_tokens stays 0 here. + assert output.output_tokens == 0, "output_tokens should be 0 when no server metrics are present" assert output.generated_text == "Hello world" @@ -277,7 +279,7 @@ async def test_output_tokens_with_missing_num_tokens_out(mocker: MockerFixture): # Assert assert output.success is True - assert output.output_tokens == 0, "output_tokens should default to 0 when num_tokens_out is missing" + assert output.output_tokens == 0, "output_tokens should be 0 when num_tokens_out is missing" @pytest.mark.asyncio @@ -579,6 +581,167 @@ async def test_text_latency_value_consistency(self, mocker: MockerFixture): # ============================================================================ +@pytest.mark.asyncio +async def test_skips_empty_role_chunk_for_ttft_and_itl(mocker: MockerFixture): + """Role/empty text chunks must not count toward TTFT or ITL.""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"role": "assistant", "content": ""}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": " world"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 2, "total_tokens": 12}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks, delay_between_chunks=0.01) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + assert output.success is True + assert output.output_tokens == 2 + assert len(output.itl) == 1 + assert output.ttft > 0 + assert output.text_latency - output.ttft == pytest.approx(sum(output.itl), rel=1e-6, abs=1e-6) + + +@pytest.mark.asyncio +async def test_bundled_tokens_split_itl_from_usage(mocker: MockerFixture): + """Multiple tokens in one SSE chunk should expand ITL via usage deltas.""" + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "A"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + } + ), + create_sse_chunk( + { + "choices": [{"delta": {"content": "BCD"}}], + "modality": "text", + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks, delay_between_chunks=0.02) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + assert output.success is True + assert output.output_tokens == 4 + assert len(output.itl) == 3 + assert output.text_latency - output.ttft == pytest.approx(sum(output.itl), rel=1e-6, abs=1e-6) + + +@pytest.mark.asyncio +async def test_output_tokens_prefers_usage_over_metrics(mocker: MockerFixture): + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Hello"}}], + "modality": "text", + "metrics": {"num_tokens_out": 5}, + "usage": {"prompt_tokens": 10, "completion_tokens": 8, "total_tokens": 18}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + assert output.success is True + assert output.output_tokens == 8 + + +@pytest.mark.asyncio +async def test_metrics_only_chunk_updates_output_tokens(mocker: MockerFixture): + request_input = RequestFuncInput( + model="test-model", + model_name="test-model", + prompt="test prompt", + api_url="http://test.com/v1/chat/completions", + prompt_len=10, + output_len=20, + ) + + chunks = [ + create_sse_chunk( + { + "choices": [{"delta": {"content": "Text response"}}], + "modality": "text", + } + ), + create_sse_chunk( + { + "modality": "text", + "metrics": {"num_tokens_out": 25, "num_tokens_in": 10}, + } + ), + b"data: [DONE]\n\n", + ] + + mock_response = MockResponse(200, chunks) + mock_session = mocker.AsyncMock() + mock_session.post = mocker.MagicMock(return_value=mock_response) + + output = await async_request_openai_chat_omni_completions(request_input, mock_session) + + assert output.success is True + assert output.output_tokens == 25 + + @pytest.mark.asyncio async def test_prompt_len_assigned_from_usage(mocker: MockerFixture): # Arrange: request claims prompt_len=100, but server reports 4992 (multimodal). diff --git a/tests/config/composable_parallel/test_apply.py b/tests/config/composable_parallel/test_apply.py new file mode 100644 index 00000000000..044b2cb024b --- /dev/null +++ b/tests/config/composable_parallel/test_apply.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for applying strategy specs onto merged stage configs.""" + +from __future__ import annotations + +import pytest + +from vllm_omni.config.composable_parallel import ( + Broadcast, + FanInByStage, + MeshAxisSpec, + RouteByStage, + StrategyApplyError, + StrategySpec, + TakeRank, + apply_strategy_specs, +) +from vllm_omni.config.stage_config import StageConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _tp(size: int) -> StrategySpec: + return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank()) + + +def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec: + return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage()) + + +def _stage(stage_id: int, model_stage: str, engine_args=None, runtime=None) -> StageConfig: + return StageConfig( + stage_id=stage_id, + model_stage=model_stage, + yaml_engine_args=dict(engine_args or {}), + yaml_runtime=dict(runtime or {"num_replicas": 1}), + ) + + +def _qwen_stages() -> list[StageConfig]: + return [ + _stage(0, "thinker"), + _stage(1, "talker"), + _stage(2, "code2wav"), + ] + + +def test_apply_tp_by_role(): + stages = _qwen_stages() + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2 + # untouched roles keep their config + assert "tensor_parallel_size" not in stages[1].yaml_engine_args + + +def test_apply_by_model_stage(): + stages = _qwen_stages() + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2 + + +def test_apply_stage_replica_sets_num_replicas_and_surfaces_lb(): + stages = _qwen_stages() + result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]}) + assert stages[1].yaml_runtime["num_replicas"] == 2 + assert result.omni_lb_policy == "round-robin" + + +def test_only_declared_axes_are_written(): + stages = _qwen_stages() + # strategy declares only stage_replica -> tp must not be forced. + apply_strategy_specs(stages, {"talker": [_stage_replica(2)]}) + assert "tensor_parallel_size" not in stages[1].yaml_engine_args + + +def test_conflict_on_explicit_tp(): + stages = _qwen_stages() + stages[0].yaml_engine_args["tensor_parallel_size"] = 4 + with pytest.raises(StrategyApplyError): + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + + +def test_equal_explicit_value_is_noop(): + stages = _qwen_stages() + stages[0].yaml_engine_args["tensor_parallel_size"] = 2 + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2 + + +def test_explicit_none_conflicts_with_derived_value(): + # An explicit YAML ``null`` (``tensor_parallel_size: null``) is a *present* + # value, not a missing key, so a strategy deriving a non-None size must raise + # rather than silently clobber the explicit None. + stages = _qwen_stages() + stages[0].yaml_engine_args["tensor_parallel_size"] = None + with pytest.raises(StrategyApplyError): + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + + +def test_missing_key_is_filled_without_conflict(): + # A genuinely absent key (never set in the YAML) is filled by the strategy + # and must NOT raise — the contrast case to an explicit None. + stages = _qwen_stages() + assert "tensor_parallel_size" not in stages[0].yaml_engine_args + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2 + + +def test_num_replicas_conflict(): + stages = _qwen_stages() + stages[1].yaml_runtime["num_replicas"] = 3 + with pytest.raises(StrategyApplyError): + apply_strategy_specs(stages, {"talker": [_stage_replica(2)]}) + + +def test_device_count_ok_template(): + stages = _qwen_stages() + stages[0].yaml_runtime["devices"] = "0,1" + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + assert stages[0].yaml_engine_args["tensor_parallel_size"] == 2 + + +def test_device_count_ok_pool(): + # tp=2 -> world=2; 2 replicas -> pool of 4 device ids is valid. + stages = _qwen_stages() + stages[1].yaml_runtime["devices"] = "0,1,2,3" + apply_strategy_specs(stages, {"talker": [_tp(2), _stage_replica(2)]}) + assert stages[1].yaml_runtime["num_replicas"] == 2 + + +def test_device_count_mismatch(): + stages = _qwen_stages() + stages[0].yaml_runtime["devices"] = "0,1,2" + with pytest.raises(StrategyApplyError): + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + + +def test_unknown_role_raises(): + stages = _qwen_stages() + with pytest.raises(StrategyApplyError): + apply_strategy_specs(stages, {"nonexistent": [_tp(2)]}) + + +def test_conflicting_lb_policy_across_roles(): + stages = _qwen_stages() + with pytest.raises(StrategyApplyError): + apply_strategy_specs( + stages, + { + "talker": [_stage_replica(2, "round_robin")], + "code2wav": [_stage_replica(2, "least_queue")], + }, + ) diff --git a/tests/config/composable_parallel/test_factory_integration.py b/tests/config/composable_parallel/test_factory_integration.py new file mode 100644 index 00000000000..468b3db2c4b --- /dev/null +++ b/tests/config/composable_parallel/test_factory_integration.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration tests: strategy overlay on the real merge_pipeline_deploy seam.""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import pytest + +from vllm_omni.config.composable_parallel import ( + Broadcast, + FanInByStage, + MeshAxisSpec, + RouteByStage, + StrategyApplyError, + StrategySpec, + TakeRank, + apply_strategy_specs, +) +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES +from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEPLOY = Path(__file__).parents[3] / "vllm_omni" / "deploy" / "qwen2_5_omni.yaml" + + +def _tp(size: int) -> StrategySpec: + return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank()) + + +def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec: + return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage()) + + +def _qwen_stages(): + pipeline = OMNI_PIPELINES["qwen2_5_omni"] + deploy = load_deploy_config(_DEPLOY) + return merge_pipeline_deploy(pipeline, deploy) + + +def _stage(stages, role): + return next(s for s in stages if s.model_stage == role) + + +def test_overlay_tp_on_thinker(): + stages = _qwen_stages() + # The bundled deploy pins the thinker to one GPU; a TP=2 strategy needs a + # matching 2-GPU layout (mirrors what a TP2 deploy would declare). + _stage(stages, "thinker").yaml_runtime["devices"] = "0,1" + result = apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + assert _stage(result.stages, "thinker").yaml_engine_args["tensor_parallel_size"] == 2 + + +def test_device_guard_rejects_tp_on_single_gpu_deploy(): + # The bundled deploy pins the thinker to a single GPU, so the pre-spawn + # device check must refuse a TP=2 strategy on it. + stages = _qwen_stages() + with pytest.raises(StrategyApplyError): + apply_strategy_specs(stages, {"thinker": [_tp(2)]}) + + +def test_overlay_stage_replica_on_talker(): + stages = _qwen_stages() + result = apply_strategy_specs(stages, {"talker": [_stage_replica(2, "round_robin")]}) + assert _stage(result.stages, "talker").yaml_runtime["num_replicas"] == 2 + assert result.omni_lb_policy == "round-robin" + + +def test_overlay_mixed_roles(): + stages = _qwen_stages() + _stage(stages, "thinker").yaml_runtime["devices"] = "0,1" + result = apply_strategy_specs( + stages, + { + "thinker": [_tp(2)], + "talker": [_stage_replica(2, "round_robin")], + "code2wav": [_stage_replica(2, "round_robin")], + }, + ) + by_role = {s.model_stage: s for s in result.stages} + assert by_role["thinker"].yaml_engine_args["tensor_parallel_size"] == 2 + assert by_role["talker"].yaml_runtime["num_replicas"] == 2 + assert by_role["code2wav"].yaml_runtime["num_replicas"] == 2 + assert result.omni_lb_policy == "round-robin" + + +def _resolved(stages, role): + """Return a role's fully-resolved (post-CLI) OmegaConf stage config.""" + stage = next(s for s in stages if s.model_stage == role) + return stage.to_omegaconf() + + +def test_device_check_survives_cli_override(): + # Strategy replicates the talker (1-GPU template -> valid at apply time), + # but a CLI --stage_1_devices with 3 ids must NOT slip past the device + # guard: effective world=1, replicas=2 admits only 1 or 2 device ids. + pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"] + with pytest.raises(StrategyApplyError): + StageConfigFactory._create_legacy_from_registry( + pipeline_cfg, + cli_overrides={"stage_1_devices": "0,1,2"}, + strategy_specs={"talker": [_stage_replica(2, "round_robin")]}, + ) + + +def test_cli_overrides_strategy_with_warning(): + # Strategy derives num_replicas=2 for the talker; a CLI override to 3 wins + # (it is the most explicit user action) but must be surfaced loudly rather + # than silently. The resulting layout (1 template device) stays valid. + # + # vLLM's logger sets propagate=False, so attach a handler directly to it + # rather than relying on pytest's caplog (which listens on the root logger). + messages: list[str] = [] + + class _Capture(logging.Handler): + def emit(self, record: logging.LogRecord) -> None: + messages.append(record.getMessage()) + + log = logging.getLogger("vllm_omni.config.config_factory") + handler = _Capture(level=logging.WARNING) + log.addHandler(handler) + try: + pipeline_cfg = OMNI_PIPELINES["qwen2_5_omni"] + stages, _ = StageConfigFactory._create_legacy_from_registry( + pipeline_cfg, + cli_overrides={"stage_1_num_replicas": 3}, + strategy_specs={"talker": [_stage_replica(2, "round_robin")]}, + ) + finally: + log.removeHandler(handler) + + # CLI value wins in the resolved config. + assert _resolved(stages, "talker").runtime.num_replicas == 3 + # ...and the override was warned about, naming the conflicting field. + assert any("num_replicas" in m and "overrides the strategy-derived" in m for m in messages) diff --git a/tests/config/composable_parallel/test_spec_types.py b/tests/config/composable_parallel/test_spec_types.py new file mode 100644 index 00000000000..90b00cf9506 --- /dev/null +++ b/tests/config/composable_parallel/test_spec_types.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the vendored spec/pattern data types.""" + +from __future__ import annotations + +import pytest + +from vllm_omni.config.composable_parallel.aggregation import AggregationPattern, TakeRank +from vllm_omni.config.composable_parallel.routing import RoutingPattern +from vllm_omni.config.composable_parallel.spec import MeshAxisSpec + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_mesh_axis_rejects_bad_kind(): + with pytest.raises(ValueError): + MeshAxisSpec(kind="not_a_kind", size=2) # type: ignore[arg-type] + + +def test_mesh_axis_rejects_nonpositive_size(): + with pytest.raises(ValueError): + MeshAxisSpec(kind="tp", size=0) + + +def test_routing_pattern_is_closed(): + with pytest.raises(TypeError): + + class RogueRouting(RoutingPattern): # noqa: F811 + pass + + +def test_aggregation_pattern_is_closed(): + with pytest.raises(TypeError): + + class RogueAggregation(AggregationPattern): # noqa: F811 + pass + + +def test_take_rank_type_guard(): + with pytest.raises(TypeError): + TakeRank(rank=True) # bool is not a valid rank diff --git a/tests/config/composable_parallel/test_strategy_loader.py b/tests/config/composable_parallel/test_strategy_loader.py new file mode 100644 index 00000000000..01e202d0518 --- /dev/null +++ b/tests/config/composable_parallel/test_strategy_loader.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for parsing/loading strategy files.""" + +from __future__ import annotations + +import pytest + +from vllm_omni.config.composable_parallel.routing import Broadcast, RouteByStage +from vllm_omni.config.composable_parallel.strategy_loader import ( + StrategyLoadError, + load_strategy_specs, + parse_strategy_specs, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_parse_basic(): + data = { + "strategies": { + "thinker": [{"axis": "tp", "size": 2}], + "talker": [{"axis": "stage_replica", "size": 2, "routing": "round_robin"}], + } + } + specs = parse_strategy_specs(data) + assert set(specs) == {"thinker", "talker"} + + (tp_spec,) = specs["thinker"] + assert tp_spec.mesh_axis.kind == "tp" + assert tp_spec.mesh_axis.size == 2 + assert isinstance(tp_spec.routing, Broadcast) + + (sr_spec,) = specs["talker"] + assert sr_spec.mesh_axis.kind == "stage_replica" + assert isinstance(sr_spec.routing, RouteByStage) + assert sr_spec.routing.routing_policy == "round_robin" + + +def test_parse_without_strategies_key(): + # A bare role mapping (no top-level "strategies") is also accepted. + specs = parse_strategy_specs({"thinker": [{"axis": "tp", "size": 4}]}) + assert specs["thinker"][0].mesh_axis.size == 4 + + +def test_l1_owner_goes_to_shard_extension(): + specs = parse_strategy_specs({"talker": [{"axis": "stage_replica", "size": 2, "l1_owner": "delegated"}]}) + assert specs["talker"][0].shard_extension["l1_owner"] == "delegated" + + +def test_missing_axis_raises(): + with pytest.raises(StrategyLoadError): + parse_strategy_specs({"thinker": [{"size": 2}]}) + + +def test_missing_size_raises(): + with pytest.raises(StrategyLoadError): + parse_strategy_specs({"thinker": [{"axis": "tp"}]}) + + +def test_routing_on_non_policy_axis_raises(): + with pytest.raises(StrategyLoadError): + parse_strategy_specs({"thinker": [{"axis": "tp", "size": 2, "routing": "random"}]}) + + +def test_bad_size_raises(): + with pytest.raises(StrategyLoadError): + parse_strategy_specs({"thinker": [{"axis": "tp", "size": "two"}]}) + + +def test_entries_must_be_list(): + with pytest.raises(StrategyLoadError): + parse_strategy_specs({"thinker": {"axis": "tp", "size": 2}}) + + +def test_non_mapping_entry_raises(): + # A list whose elements are not mappings (e.g. a bare string) must raise a + # StrategyLoadError, not an opaque TypeError from dict(entry). + with pytest.raises(StrategyLoadError): + parse_strategy_specs({"thinker": ["tp"]}) + + +def test_load_from_file(tmp_path): + path = tmp_path / "strategy.yaml" + path.write_text( + "strategies:\n" + " thinker:\n" + " - axis: tp\n" + " size: 2\n" + " talker:\n" + " - axis: stage_replica\n" + " size: 2\n" + " routing: least_queue\n" + ) + specs = load_strategy_specs(str(path)) + assert specs["thinker"][0].mesh_axis.size == 2 + assert specs["talker"][0].routing.routing_policy == "least_queue" diff --git a/tests/config/composable_parallel/test_translator.py b/tests/config/composable_parallel/test_translator.py new file mode 100644 index 00000000000..b55c10e9ec1 --- /dev/null +++ b/tests/config/composable_parallel/test_translator.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for StrategySpec -> OmniParallelConfig translation.""" + +from __future__ import annotations + +import dataclasses + +import pytest + +from vllm_omni.config.composable_parallel import ( + AxisTranslationError, + Broadcast, + FanInByStage, + MeshAxisSpec, + PartitionByHash, + PipelineMicrobatch, + RouteByStage, + StrategySpec, + TakeRank, + Union, + translate_strategy_stack, +) +from vllm_omni.config.composable_parallel.aggregation import StitchPipeline + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _tp(size: int) -> StrategySpec: + return StrategySpec("tp", MeshAxisSpec("tp", size), Broadcast(), TakeRank()) + + +def _dp(size: int, policy: str = "round_robin") -> StrategySpec: + return StrategySpec("dp", MeshAxisSpec("dp", size), RouteByStage(policy), Union()) + + +def _pp(size: int) -> StrategySpec: + return StrategySpec("pp", MeshAxisSpec("pp", size), PipelineMicrobatch(), StitchPipeline()) + + +def _ep(size: int) -> StrategySpec: + return StrategySpec("ep", MeshAxisSpec("ep", size), Broadcast(), Union()) + + +def _stage_replica(size: int, policy: str = "round_robin") -> StrategySpec: + return StrategySpec("stage_replica", MeshAxisSpec("stage_replica", size), RouteByStage(policy), FanInByStage()) + + +def test_tp_only(): + cfg = translate_strategy_stack([_tp(2)]) + assert cfg.tensor_parallel_size == 2 + assert cfg.data_parallel_size == 1 + assert cfg.world_size == 2 + assert cfg.stage_replica_size == 1 + assert cfg.l1_owners == {"tp": "engine"} + + +def test_dp_is_engine_owned(): + cfg = translate_strategy_stack([_dp(2)]) + assert cfg.data_parallel_size == 2 + assert cfg.world_size == 2 + # engine DP does not emit an omni LB policy (vLLM's internal DP LB is used). + assert cfg.omni_lb_policy is None + assert cfg.as_engine_kwargs().get("omni_lb_policy") is None + assert cfg.l1_owners["dp"] == "engine" + + +def test_tp_times_dp_world_size(): + cfg = translate_strategy_stack([_tp(2), _dp(2)]) + assert cfg.world_size == 4 + assert cfg.tensor_parallel_size == 2 + assert cfg.data_parallel_size == 2 + + +def test_stage_replica_is_delegated(): + cfg = translate_strategy_stack([_stage_replica(3, "least_queue")]) + # stage_replica is NOT a world dimension. + assert cfg.world_size == 1 + assert cfg.stage_replica_size == 3 + assert cfg.omni_lb_policy == "least-queue-length" + assert cfg.l1_owners["stage_replica"] == "delegated" + # stage_replica_size is not an engine kwarg; it's a per-stage deploy knob. + assert "stage_replica_size" not in cfg.as_engine_kwargs() + + +def test_stage_replica_lb_policy_only_when_replicated(): + cfg = translate_strategy_stack([_stage_replica(1)]) + assert cfg.stage_replica_size == 1 + assert cfg.as_engine_kwargs().get("omni_lb_policy") is None + + +def test_tp_with_stage_replica(): + cfg = translate_strategy_stack([_tp(2), _stage_replica(2)]) + assert cfg.world_size == 2 # tp only + assert cfg.tensor_parallel_size == 2 + assert cfg.stage_replica_size == 2 + assert cfg.omni_lb_policy == "round-robin" + + +def test_ep_must_match_tp_dp_product(): + cfg = translate_strategy_stack([_tp(2), _ep(2)]) + assert cfg.enable_expert_parallel is True + assert cfg.as_engine_kwargs()["enable_expert_parallel"] is True + + +def test_ep_mismatch_raises(): + with pytest.raises(Exception) as exc: + translate_strategy_stack([_tp(2), _ep(4)]) + assert "must equal" in str(exc.value) + + +def test_duplicate_kind_rejected(): + with pytest.raises(AxisTranslationError): + translate_strategy_stack([_tp(2), _tp(2)]) + + +def test_unsupported_kind_rejected(): + spec = StrategySpec("sp", MeshAxisSpec("sp_ulysses", 2), Broadcast(), TakeRank()) + with pytest.raises(AxisTranslationError): + translate_strategy_stack([spec]) + + +def test_dp_hash_routing_rejected(): + spec = StrategySpec("dp", MeshAxisSpec("dp", 2), PartitionByHash(), Union()) + with pytest.raises(NotImplementedError): + translate_strategy_stack([spec]) + + +def test_dp_invalid_routing_policy_rejected(): + spec = StrategySpec("dp", MeshAxisSpec("dp", 2), RouteByStage("bogus"), Union()) + with pytest.raises(AxisTranslationError): + translate_strategy_stack([spec]) + + +def test_lb_policy_never_in_engine_kwargs_even_when_replicated(): + # omni_lb_policy is pipeline-wide (applied at engine construction), not a + # per-stage engine arg, so it must never leak into as_engine_kwargs(). + cfg = translate_strategy_stack([_tp(2), _stage_replica(2, "least_queue")]) + assert cfg.omni_lb_policy == "least-queue-length" + assert "omni_lb_policy" not in cfg.as_engine_kwargs() + + +def test_stage_replica_hash_routing_rejected(): + spec = StrategySpec("sr", MeshAxisSpec("stage_replica", 2), RouteByStage("hash"), FanInByStage()) + with pytest.raises(NotImplementedError): + translate_strategy_stack([spec]) + + +def test_dp_wrong_owner_rejected(): + spec = StrategySpec( + "dp", + MeshAxisSpec("dp", 2), + RouteByStage("round_robin"), + Union(), + shard_extension={"l1_owner": "delegated"}, + ) + with pytest.raises(AxisTranslationError): + translate_strategy_stack([spec]) + + +def test_stage_replica_wrong_owner_rejected(): + spec = StrategySpec( + "sr", + MeshAxisSpec("stage_replica", 2), + RouteByStage("round_robin"), + FanInByStage(), + shard_extension={"l1_owner": "engine"}, + ) + with pytest.raises(AxisTranslationError): + translate_strategy_stack([spec]) + + +@pytest.mark.parametrize( + "specs", + [ + [_tp(2)], + [_dp(2)], + [_tp(2), _dp(2)], + [_tp(2), _ep(2)], + [_tp(2), _stage_replica(2)], + [_pp(2)], + ], +) +def test_as_engine_kwargs_are_valid_omni_engine_args(specs): + # Guard against desync: every key emitted by as_engine_kwargs() must be a + # real OmniEngineArgs field, otherwise constructing the per-stage engine args + # would explode at runtime. Keep this CPU-only by checking field names rather + # than instantiating the engine. + from vllm_omni.engine.arg_utils import OmniEngineArgs + + valid_fields = {f.name for f in dataclasses.fields(OmniEngineArgs)} + cfg = translate_strategy_stack(specs) + kwargs = cfg.as_engine_kwargs() + unknown = set(kwargs) - valid_fields + assert not unknown, f"as_engine_kwargs emitted non-OmniEngineArgs fields: {unknown}" diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py new file mode 100644 index 00000000000..23e53ea3d27 --- /dev/null +++ b/tests/config/test_omni_config.py @@ -0,0 +1,859 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the additive structured Omni config.""" + +from __future__ import annotations + +from dataclasses import fields +from pathlib import Path + +import pytest +from pydantic import ValidationError +from transformers import Qwen3OmniMoeConfig + +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.config import omni_config as omni_config_module +from vllm_omni.config.omni_config import ( + BaseVllmOmniStageConfig, + OmniStageCacheConfig, + OmniStageConnectorConfig, + OmniStageDiffusionParallelConfig, + OmniStageLoadConfig, + OmniStageModelConfig, + OmniStageParallelConfig, + OmniStageRuntimeConfig, + OmniStageSchedulerConfig, + VllmOmniARStageConfig, + VllmOmniConfig, + VllmOmniDiffusionStageConfig, + VllmOmniGenerationStageConfig, +) +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config +from vllm_omni.config.stage_config import ( + _STAGE_DEPLOY_FIELDS, + PIPELINE_WIDE_ENGINE_FIELDS, + DeployConfig, + PipelineConfig, + StageDeployConfig, + StageExecutionType, + load_deploy_config, + merge_pipeline_deploy, +) +from vllm_omni.engine.stage_init_utils import build_engine_args_dict + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy" + + +@pytest.fixture(autouse=True) +def _stable_test_platform(monkeypatch): + from vllm_omni import platforms + + platform = platforms.current_omni_platform + monkeypatch.setattr(platform, "device_name", "cpu", raising=False) + monkeypatch.setattr(platform, "device_type", "cpu", raising=False) + + +def _load_default_deploy(pipeline: PipelineConfig) -> DeployConfig: + if pipeline.default_deploy_config_name is not None: + return load_deploy_config(_DEPLOY_DIR / pipeline.default_deploy_config_name) + return DeployConfig() + + +def _resolve_pipeline_or_skip(model_type: str, hf_config=None) -> PipelineConfig: + pipeline = resolve_pipeline_config(model_type, hf_config) + if pipeline is None: + pytest.skip(f"Pipeline {model_type!r} requires an HF config to resolve") + return pipeline + + +def _from_pipeline_key( + model_type: str, + hf_config=None, + deploy_config_path: str | None = None, + cli_overrides: dict | None = None, +) -> VllmOmniConfig: + return VllmOmniConfig.from_pipeline_config( + _resolve_pipeline_or_skip(model_type, hf_config), + deploy_config_path=deploy_config_path, + cli_overrides=cli_overrides, + ) + + +@pytest.mark.parametrize("model_type", sorted(OMNI_PIPELINES)) +def test_vllm_omni_config_from_pipeline_config_matches_merge_pipeline_deploy(model_type: str): + pipeline = _resolve_pipeline_or_skip(model_type) + legacy_deploy = _load_default_deploy(pipeline) + + legacy_stages = merge_pipeline_deploy(pipeline, legacy_deploy) + omni_config = VllmOmniConfig.from_pipeline_config(pipeline) + + assert omni_config.pipeline_config is pipeline + assert len(omni_config.stage_configs) == len(legacy_stages) + + for legacy_stage, omni_stage in zip(legacy_stages, omni_config.stage_configs, strict=True): + assert omni_config.stage_by_id(legacy_stage.stage_id) is omni_stage + + assert omni_stage.stage_pipeline_config is pipeline.get_stage(legacy_stage.stage_id) + assert omni_stage.model_config.default_sampling_params == legacy_stage.yaml_extras.get( + "default_sampling_params" + ) + assert omni_stage.connector_config.output_connectors == legacy_stage.yaml_extras.get("output_connectors") + assert omni_stage.connector_config.input_connectors == legacy_stage.yaml_extras.get("input_connectors") + assert omni_stage.runtime_config.devices == legacy_stage.yaml_runtime.get("devices") + assert omni_stage.runtime_config.num_replicas == legacy_stage.yaml_runtime.get("num_replicas", 1) + + engine_args = legacy_stage.yaml_engine_args + assert omni_stage.model_config.enforce_eager == engine_args.get("enforce_eager", False) + assert omni_stage.load_config.load_format == engine_args.get("load_format", "auto") + assert omni_stage.load_config.tokenizer_mode == engine_args.get("tokenizer_mode", "auto") + assert omni_stage.cache_config.gpu_memory_utilization == engine_args.get("gpu_memory_utilization", 0.90) + assert omni_stage.scheduler_config.max_num_seqs == engine_args.get("max_num_seqs", 128) + assert omni_stage.scheduler_config.max_num_batched_tokens == engine_args.get("max_num_batched_tokens") + assert omni_stage.scheduler_config.async_scheduling == engine_args.get("async_scheduling", True) + legacy_parallel_config = engine_args.get("parallel_config") or {} + assert omni_stage.parallel_config.tensor_parallel_size == legacy_parallel_config.get( + "tensor_parallel_size", + engine_args.get("tensor_parallel_size", 1), + ) + + if omni_stage.stage_pipeline_config.execution_type == StageExecutionType.DIFFUSION: + assert isinstance(omni_stage, VllmOmniDiffusionStageConfig) + assert omni_stage.diffusion_config is not None + assert omni_stage.diffusion_config.stage_id == legacy_stage.stage_id + assert omni_stage.diffusion_config.model_arch == engine_args.get("model_arch") + elif omni_stage.stage_pipeline_config.execution_type == StageExecutionType.LLM_AR: + assert isinstance(omni_stage, VllmOmniARStageConfig) + assert not hasattr(omni_stage, "diffusion_config") + else: + assert isinstance(omni_stage, VllmOmniGenerationStageConfig) + assert not hasattr(omni_stage, "diffusion_config") + + +def test_stage_by_id_raises_for_unknown_stage(): + omni_config = _from_pipeline_key("qwen3_tts") + + with pytest.raises(KeyError, match="no stage 99"): + omni_config.stage_by_id(99) + + +def test_resolve_execution_mode_rejects_unknown_execution_type(): + with pytest.raises(ValueError, match="Unsupported stage execution type"): + omni_config_module._resolve_execution_mode("unknown_execution_type") + + +def test_from_pipeline_config_preserves_current_pipeline_config_object(): + omni_config = _from_pipeline_key("minicpmo_4_5") + pipeline = _resolve_pipeline_or_skip("minicpmo_4_5") + + assert omni_config.pipeline_config is pipeline + assert not hasattr(omni_config, "pipeline") + assert "hf_config_predicate" in {f.name for f in fields(PipelineConfig)} + assert omni_config.pipeline_config.hf_config_predicate is pipeline.hf_config_predicate + + +def test_from_pipeline_config_normalizes_stage_engine_extras_without_expanding_stage_deploy_config(): + assert not hasattr(StageDeployConfig, "model_config") + assert not hasattr(StageDeployConfig, "parallel_config") + + stage = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2").stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.parallel_config.tensor_parallel_size == 1 + assert stage.parallel_config.cfg_parallel_size == 2 + assert stage.diffusion_config.model_config["default_robot_embodiment"] == "roboarena" + + +def test_from_pipeline_config_applies_cli_overrides_without_stage_config_runtime_bridge(): + omni_config = _from_pipeline_key( + "qwen3_tts", + cli_overrides={ + "stage_0_max_num_seqs": 7, + "stage_1_tensor_parallel_size": 2, + }, + ) + + stage0 = omni_config.stage_by_id(0) + stage1 = omni_config.stage_by_id(1) + + assert stage0.scheduler_config.max_num_seqs == 7 + assert stage1.parallel_config.tensor_parallel_size == 2 + assert stage1.runtime_config.num_gpus == stage1.parallel_config.world_size + + +def test_runtime_num_gpus_is_derived_from_parallel_world_size(): + omni_config = _from_pipeline_key("hunyuan_image3_dit") + stage = omni_config.stage_by_id(0) + + assert stage.parallel_config.tensor_parallel_size == 4 + assert stage.parallel_config.world_size == 4 + assert stage.runtime_config.num_gpus == 4 + + +def test_runtime_num_gpus_ignores_stale_runtime_override(): + omni_config = _from_pipeline_key( + "hunyuan_image3_dit", + cli_overrides={ + "stage_0_num_gpus": 1, + }, + ) + stage = omni_config.stage_by_id(0) + + assert stage.parallel_config.world_size == 4 + assert stage.runtime_config.num_gpus == 4 + + +def test_from_pipeline_config_does_not_route_server_cli_keys_to_diffusion_stage(): + omni_config = _from_pipeline_key( + "dreamzero", + deploy_config_path="dreamzero_tp1_cfg2", + cli_overrides={ + "host": "0.0.0.0", + "port": 8000, + "api_key": "secret", + "stage_0_host": "127.0.0.1", + "stage_0_port": 23456, + }, + ) + + stage = omni_config.stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.host == "127.0.0.1" + assert stage.diffusion_config.port == 23456 + assert not hasattr(stage.diffusion_config, "api_key") + + +def test_pipeline_deploy_cli_fields_reuse_legacy_pipeline_wide_engine_fields(): + assert omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS is PIPELINE_WIDE_ENGINE_FIELDS + assert "active_stream_window" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS + assert "custom_voice_dir" in omni_config_module._PIPELINE_DEPLOY_CLI_FIELDS + + +def test_pipeline_wide_model_fields_are_retained_on_structured_stage_configs(tmp_path): + custom_voice_dir = tmp_path / "voices" + omni_config = _from_pipeline_key( + "qwen3_tts", + cli_overrides={ + "active_stream_window": 2, + "custom_voice_dir": str(custom_voice_dir), + }, + ) + + assert {stage.model_config.active_stream_window for stage in omni_config.stage_configs} == {2} + assert {stage.model_config.custom_voice_dir for stage in omni_config.stage_configs} == {str(custom_voice_dir)} + + +def test_stage_deploy_engine_fields_reuse_legacy_stage_deploy_fields(): + assert omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS == tuple(_STAGE_DEPLOY_FIELDS) + assert "tensor_parallel_size" in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS + assert "stage_id" not in omni_config_module._STAGE_DEPLOY_ENGINE_FIELDS + + +def test_public_config_exports_use_stage_specific_sub_config_names(): + import vllm_omni.config as config_pkg + + generic_names = { + "CacheConfig", + "ConnectorConfig", + "LoadConfig", + "ModelConfig", + "OrchestratorConfig", + "ParallelConfig", + "RuntimeConfig", + "SchedulerConfig", + } + + assert generic_names.isdisjoint(config_pkg.__all__) + assert { + "OmniStageCacheConfig", + "OmniStageConnectorConfig", + "OmniStageDiffusionParallelConfig", + "OmniStageLoadConfig", + "OmniStageModelConfig", + "VllmOmniOrchestratorConfig", + "OmniStageParallelConfig", + "OmniStageRuntimeConfig", + "OmniStageSchedulerConfig", + "StageConfigType", + }.issubset(config_pkg.__all__) + + +def test_from_pipeline_config_keeps_worker_backend_separate_from_distributed_executor_backend(): + omni_config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2") + + stage = omni_config.stage_by_id(0) + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.distributed_executor_backend == "mp" + assert omni_config.orchestrator_config.worker_backend == "multi_process" + + +def test_from_pipeline_config_maps_orchestrator_cli_overrides(): + omni_config = _from_pipeline_key( + "qwen3_tts", + cli_overrides={ + "stage_init_timeout": 1200, + "init_timeout": 1800, + "worker_backend": "ray", + "ray_address": "ray://127.0.0.1:10001", + "omni_master_address": "127.0.0.1", + "omni_master_port": 12345, + "omni_dp_size_local": 2, + "omni_lb_policy": "round_robin", + "omni_heartbeat_timeout": 9.5, + "shm_threshold_bytes": 4096, + "batch_timeout": 3, + }, + ) + + orchestrator_config = omni_config.orchestrator_config + assert orchestrator_config.stage_init_timeout == 1200 + assert orchestrator_config.init_timeout == 1800 + assert orchestrator_config.worker_backend == "ray" + assert orchestrator_config.ray_address == "ray://127.0.0.1:10001" + assert orchestrator_config.omni_master_address == "127.0.0.1" + assert orchestrator_config.omni_master_port == 12345 + assert orchestrator_config.omni_dp_size_local == 2 + assert orchestrator_config.omni_lb_policy == "round_robin" + assert orchestrator_config.omni_heartbeat_timeout == 9.5 + assert orchestrator_config.shm_threshold_bytes == 4096 + assert orchestrator_config.batch_timeout == 3 + + +def test_from_pipeline_config_records_loaded_deploy_path_on_orchestrator_config(): + omni_config = _from_pipeline_key("dreamzero", deploy_config_path="dreamzero_tp1_cfg2") + + assert omni_config.pipeline_config.model_type == "dreamzero" + assert omni_config.orchestrator_config.deploy_config_path == str(_DEPLOY_DIR / "dreamzero_tp1_cfg2.yaml") + + +def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating_topology(): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + + async_config = _from_pipeline_key("qwen3_tts") + assert async_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_async_chunk") + assert async_config.stage_by_id(1).custom_process_input_func is None + + sync_config = _from_pipeline_key("qwen3_tts", cli_overrides={"async_chunk": False}) + assert sync_config.stage_by_id(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload") + assert sync_config.stage_by_id(1).custom_process_input_func.endswith("talker2code2wav_token_only") + + assert pipeline.get_stage(0).custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload") + assert pipeline.get_stage(1).custom_process_input_func is None + + +def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations(): + assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config") + assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config") + + public_fields = {f.name for f in fields(BaseVllmOmniStageConfig)} + + assert public_fields == { + "stage_pipeline_config", + "model_config", + "load_config", + "cache_config", + "scheduler_config", + "connector_config", + "runtime_config", + "parallel_config", + "quantization_config", + } + assert "diffusion_config" not in public_fields + assert {f.name for f in fields(VllmOmniDiffusionStageConfig)} == public_fields | {"diffusion_config"} + assert {f.name for f in fields(VllmOmniARStageConfig)} == public_fields + assert {f.name for f in fields(VllmOmniGenerationStageConfig)} == public_fields + + +def test_runtime_config_fields_match_rfc_runtime_scope(): + assert {f.name for f in fields(OmniStageRuntimeConfig)} == { + "devices", + "num_replicas", + "env", + "num_gpus", + "log_level", + "log_stats", + "profiler_config", + } + + +def test_sub_config_fields_match_rfc_scopes(): + assert {f.name for f in fields(OmniStageModelConfig)} == { + "active_stream_window", + "enable_sleep_mode", + "default_sampling_params", + "subtalker_sampling_params", + "has_sampling_extra_args", + "custom_voice_dir", + "task_type", + "codec_frame_rate_hz", + "enforce_eager", + "enable_flashinfer_autotune", + "compilation_config", + "enable_multithread_weight_load", + "num_weight_load_threads", + "disable_autocast", + } + assert {f.name for f in fields(OmniStageLoadConfig)} == { + "load_format", + "tokenizer_mode", + "config_format", + "skip_mm_profiling", + } + assert {f.name for f in fields(OmniStageCacheConfig)} == { + "gpu_memory_utilization", + "enable_prefix_caching", + "disable_hybrid_kv_cache_manager", + "mm_processor_cache_gb", + } + assert {f.name for f in fields(OmniStageSchedulerConfig)} == { + "max_num_seqs", + "max_num_batched_tokens", + "max_model_len", + "enable_chunked_prefill", + "async_scheduling", + } + assert {f.name for f in fields(OmniStageConnectorConfig)} == { + "stage_connector", + "output_connectors", + "input_connectors", + } + assert {f.name for f in fields(OmniStageParallelConfig)} == { + "pipeline_parallel_size", + "data_parallel_size", + "tensor_parallel_size", + "enable_expert_parallel", + "world_size", + } + assert {f.name for f in fields(OmniStageDiffusionParallelConfig)} == { + "pipeline_parallel_size", + "data_parallel_size", + "tensor_parallel_size", + "sequence_parallel_size", + "ulysses_degree", + "ring_degree", + "ulysses_mode", + "cfg_parallel_size", + "vae_patch_parallel_size", + "vae_parallel_mode", + "use_hsdp", + "mask_sp_padding", + "hsdp_shard_size", + "hsdp_replicate_size", + "enable_expert_parallel", + "world_size", + } + + +def test_diffusion_parallel_config_fields_cover_legacy_surface(): + from vllm_omni.diffusion.data import DiffusionParallelConfig + + legacy_fields = {f.name for f in fields(DiffusionParallelConfig)} + structured_fields = {f.name for f in fields(OmniStageDiffusionParallelConfig)} + expected_upstream_fields = {"mask_sp_padding"} + + assert legacy_fields | expected_upstream_fields <= structured_fields + assert structured_fields - legacy_fields - expected_upstream_fields == {"world_size"} + + +def test_diffusion_parallel_config_keeps_current_diffusion_parallel_surface(): + cfg = OmniStageDiffusionParallelConfig( + pipeline_parallel_size=2, + data_parallel_size=3, + tensor_parallel_size=4, + cfg_parallel_size=3, + mask_sp_padding=True, + ) + + assert cfg.pipeline_parallel_size == 2 + assert cfg.data_parallel_size == 3 + assert cfg.cfg_parallel_size == 3 + assert cfg.mask_sp_padding is True + assert cfg.world_size == 72 + + +def test_parallel_config_derived_fields_are_not_init_inputs(): + with pytest.raises(ValidationError): + OmniStageParallelConfig(world_size=4) + + with pytest.raises(ValidationError): + OmniStageDiffusionParallelConfig(world_size=4) + + with pytest.raises(ValidationError): + OmniStageDiffusionParallelConfig(sequence_parallel_size=2) + + +def test_diffusion_parallel_config_matches_diffusion_parallel_world_size_for_vae_patch_parallel(): + cfg = OmniStageDiffusionParallelConfig( + tensor_parallel_size=2, + cfg_parallel_size=2, + vae_patch_parallel_size=4, + ) + + assert cfg.vae_patch_parallel_size == 4 + assert cfg.world_size == 4 + + +def test_diffusion_parallel_config_supports_diffusion_hsdp_auto_sharding(): + cfg = OmniStageDiffusionParallelConfig( + pipeline_parallel_size=2, + ulysses_degree=2, + use_hsdp=True, + hsdp_shard_size=-1, + hsdp_replicate_size=2, + ) + + assert cfg.hsdp_shard_size == 2 + assert cfg.world_size == 4 + + +def test_diffusion_parallel_config_rejects_hsdp_with_tp_or_dp(): + with pytest.raises(ValueError, match="cannot be used with TP or DP"): + OmniStageDiffusionParallelConfig(tensor_parallel_size=2, use_hsdp=True, hsdp_shard_size=2) + + with pytest.raises(ValueError, match="cannot be used with TP or DP"): + OmniStageDiffusionParallelConfig(data_parallel_size=2, use_hsdp=True, hsdp_shard_size=2) + + +def test_from_pipeline_config_preserves_legacy_pp_dp_for_world_size(): + cfg = _from_pipeline_key("hunyuan_image3_dit").stage_by_id(0).parallel_config + + assert cfg.pipeline_parallel_size == 1 + assert cfg.data_parallel_size == 1 + assert cfg.tensor_parallel_size == 4 + assert cfg.world_size == 4 + + +def test_from_pipeline_config_derives_sequence_parallel_size_from_degrees(tmp_path): + deploy_path = tmp_path / "dreamzero_derived_parallel.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " parallel_config:", + " sequence_parallel_size: 99", + " ulysses_degree: 2", + " ring_degree: 3", + ] + ) + ) + + stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.parallel_config.sequence_parallel_size == 6 + assert stage.parallel_config.world_size == 6 + + +def test_diffusion_parallel_config_rejects_cfg_parallel_size_outside_current_bound(): + with pytest.raises(ValidationError): + OmniStageDiffusionParallelConfig(cfg_parallel_size=4) + + +def test_stage_realizations_use_stage_specific_parallel_config_types(): + qwen_config = _from_pipeline_key("qwen3_tts") + ar_stage = qwen_config.stage_by_id(0) + generation_stage = qwen_config.stage_by_id(1) + diffusion_stage = _from_pipeline_key("hunyuan_image3_dit").stage_by_id(0) + + assert isinstance(ar_stage, VllmOmniARStageConfig) + assert type(ar_stage.parallel_config) is OmniStageParallelConfig + assert not hasattr(ar_stage.parallel_config, "cfg_parallel_size") + assert not hasattr(ar_stage.parallel_config, "sequence_parallel_size") + assert not hasattr(ar_stage.parallel_config, "ulysses_degree") + + assert isinstance(generation_stage, VllmOmniGenerationStageConfig) + assert type(generation_stage.parallel_config) is OmniStageParallelConfig + assert not hasattr(generation_stage.parallel_config, "cfg_parallel_size") + assert not hasattr(generation_stage.parallel_config, "sequence_parallel_size") + assert not hasattr(generation_stage.parallel_config, "ulysses_degree") + + assert isinstance(diffusion_stage, VllmOmniDiffusionStageConfig) + assert isinstance(diffusion_stage.parallel_config, OmniStageDiffusionParallelConfig) + assert diffusion_stage.parallel_config.cfg_parallel_size == 1 + assert diffusion_stage.parallel_config.sequence_parallel_size == 1 + assert diffusion_stage.parallel_config.ulysses_degree == 1 + + +def test_from_pipeline_config_preserves_diffusion_parallel_mask_sp_padding(tmp_path): + deploy_path = tmp_path / "dreamzero_mask_sp_padding.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " parallel_config:", + " mask_sp_padding: true", + ] + ) + ) + + stage = _from_pipeline_key("dreamzero", deploy_config_path=str(deploy_path)).stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.parallel_config.mask_sp_padding is True + + +def test_from_pipeline_config_matches_stage_config_to_omegaconf_behavior_for_representative_stage(): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] + omega_stage = legacy_stage.to_omegaconf() + omni_stage = _from_pipeline_key("qwen3_tts").stage_by_id(legacy_stage.stage_id) + + assert omega_stage.stage_id == omni_stage.stage_id + assert omega_stage.stage_type == omni_stage.stage_type.value + assert omega_stage.engine_input_source == omni_stage.input_sources + assert omega_stage.final_output == omni_stage.final_output + assert omega_stage.final_output_type == omni_stage.final_output_type + assert omega_stage.is_comprehension == omni_stage.is_comprehension + assert omega_stage.engine_args.model_stage == omni_stage.model_stage + assert omega_stage.engine_args.worker_type == omni_stage.worker_type + assert omega_stage.engine_args.scheduler_cls == omni_stage.scheduler_cls + assert omega_stage.runtime.process is True + assert omega_stage.runtime.requires_multimodal_data == omni_stage.requires_multimodal_data + + +def test_from_pipeline_config_uses_hf_config_for_callable_resolver(): + hf_config = Qwen3OmniMoeConfig() + hf_config.enable_audio_output = False + + omni_config = _from_pipeline_key("qwen3_omni_moe", hf_config=hf_config) + + assert omni_config.pipeline_config.model_type == "qwen3_omni_moe_thinker_only" + assert len(omni_config.stage_configs) == 1 + assert omni_config.orchestrator_config.deploy_config_path is None + + thinker = omni_config.stage_configs[0] + assert thinker.model_stage == "thinker" + assert thinker.model_config.default_sampling_params == {"detokenize": True} + + +def test_from_pipeline_config_accepts_pre_resolved_pipeline(): + resolved_pipeline = PipelineConfig(model_type="callable_resolved_variant") + + omni_config = VllmOmniConfig.from_pipeline_config(resolved_pipeline) + + assert omni_config.pipeline_config is resolved_pipeline + + +def test_from_pipeline_config_prefers_loaded_user_deploy_config(monkeypatch): + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + user_deploy_config = DeployConfig( + stages=[StageDeployConfig(stage_id=0, max_num_seqs=7)], + ) + monkeypatch.setattr( + omni_config_module, + "load_deploy_config", + lambda _path: pytest.fail("default deploy config should not be loaded"), + ) + + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + user_deploy_config=user_deploy_config, + ) + + assert omni_config.stage_by_id(0).scheduler_config.max_num_seqs == 7 + + +def test_from_pipeline_config_default_deploy_name_ignores_cwd(monkeypatch, tmp_path): + default_name = "pipeline_default.yaml" + (tmp_path / default_name).write_text("stages: []\n", encoding="utf-8") + monkeypatch.chdir(tmp_path) + pipeline = PipelineConfig( + model_type="pipeline_with_default", + default_deploy_config_name=default_name, + ) + loaded_paths = [] + + def _load_deploy_config(path): + loaded_paths.append(Path(path)) + return DeployConfig() + + monkeypatch.setattr(omni_config_module, "load_deploy_config", _load_deploy_config) + + omni_config = VllmOmniConfig.from_pipeline_config(pipeline) + + assert omni_config.orchestrator_config.deploy_config_path == str(_DEPLOY_DIR / default_name) + assert loaded_paths == [_DEPLOY_DIR / default_name] + + +def test_from_pipeline_config_uses_resolved_deploy_pipeline(): + deploy_path = get_deploy_config_path("aura_omni.yaml") + pipeline = _resolve_pipeline_or_skip("aura_omni") + + omni_config = VllmOmniConfig.from_pipeline_config( + pipeline, + deploy_config_path=str(deploy_path), + ) + + assert omni_config.pipeline_config.model_type == "aura_omni" + assert [stage.model_stage for stage in omni_config.stage_configs] == [ + "asr", + "aura", + "qwen3_tts", + "code2wav", + ] + + +def test_from_pipeline_config_matches_to_omegaconf_diffusion_parallel_config(): + pipeline = _resolve_pipeline_or_skip("hunyuan_image3_dit") + legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] + omega_stage = legacy_stage.to_omegaconf() + omni_stage = _from_pipeline_key("hunyuan_image3_dit").stage_by_id(legacy_stage.stage_id) + + assert ( + omega_stage.engine_args.parallel_config.pipeline_parallel_size + == omni_stage.parallel_config.pipeline_parallel_size + ) + assert omega_stage.engine_args.parallel_config.data_parallel_size == omni_stage.parallel_config.data_parallel_size + assert ( + omega_stage.engine_args.parallel_config.tensor_parallel_size == omni_stage.parallel_config.tensor_parallel_size + ) + assert ( + omega_stage.engine_args.parallel_config.sequence_parallel_size + == omni_stage.parallel_config.sequence_parallel_size + ) + assert omega_stage.engine_args.parallel_config.cfg_parallel_size == omni_stage.parallel_config.cfg_parallel_size + assert ( + omega_stage.engine_args.parallel_config.vae_patch_parallel_size + == omni_stage.parallel_config.vae_patch_parallel_size + ) + + +def test_from_pipeline_config_matches_build_engine_args_dict_behavior_for_representative_stage(monkeypatch): + from vllm_omni.engine import stage_init_utils + + monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None) + pipeline = _resolve_pipeline_or_skip("qwen3_tts") + legacy_stage = merge_pipeline_deploy(pipeline, _load_default_deploy(pipeline))[0] + omega_stage = legacy_stage.to_omegaconf() + legacy_engine_args = build_engine_args_dict( + omega_stage, + model="/tmp/qwen3-tts", + stage_connector_spec={"name": "SharedMemoryConnector", "extra": {}}, + ) + omni_stage = _from_pipeline_key("qwen3_tts").stage_by_id(legacy_stage.stage_id) + + assert legacy_engine_args["model"] == "/tmp/qwen3-tts" + assert legacy_engine_args["stage_id"] == omni_stage.stage_id + assert legacy_engine_args["model_stage"] == omni_stage.model_stage + assert legacy_engine_args["worker_type"] == omni_stage.worker_type + assert legacy_engine_args["scheduler_cls"] == omni_stage.scheduler_cls + assert legacy_engine_args["stage_connector_spec"] == {"name": "SharedMemoryConnector", "extra": {}} + assert legacy_engine_args["has_sampling_extra_args"] == bool( + (omni_stage.model_config.default_sampling_params or {}).get("extra_args") + ) + assert omni_stage.model_config.has_sampling_extra_args == legacy_engine_args["has_sampling_extra_args"] + + +def test_from_pipeline_config_derives_has_sampling_extra_args_from_stage_defaults(): + stage = _from_pipeline_key("voxtral_tts").stage_by_id(0) + + assert (stage.model_config.default_sampling_params or {}).get("extra_args") + assert stage.model_config.has_sampling_extra_args is True + + +def test_diffusion_config_preserves_existing_coercion_hooks(): + import torch + + from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig + + cfg = omni_config_module._DiffusionConfigProjection( + dtype="float32", + cache_config={"rel_l1_thresh": 0.3}, + diffusion_attention_config={"default": "flash_attn"}, + diffusion_kv_cache_skip_steps="0-2,4", + diffusion_kv_cache_skip_layers=[1, 3], + ) + + assert cfg.dtype is torch.float32 + assert isinstance(cfg.cache_config, DiffusionCacheConfig) + assert isinstance(cfg.diffusion_attention_config, AttentionConfig) + assert cfg.diffusion_attention_config.default.backend == "flash_attn" + assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1, 2, 4} + assert cfg.diffusion_kv_cache_skip_layer_indices == {1, 3} + assert cfg.max_cpu_loras == 1 + + +def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): + monkeypatch.setenv("DIFFUSION_CACHE_BACKEND", "TEA_CACHE") + + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( + diffusion_attention_backend="flash_attn", + kv_cache_dtype="fp8", + kv_cache_skip_steps="0-1", + kv_cache_skip_layers=[2], + static_lora_scale=0.25, + diffusers_load_kwargs=None, + diffusers_call_kwargs=None, + ) + + assert cfg.diffusion_attention_config.default.backend == "flash_attn" + assert cfg.diffusion_kv_cache_dtype == "fp8" + assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1} + assert cfg.diffusion_kv_cache_skip_layer_indices == {2} + assert cfg.lora_scale == 0.25 + assert cfg.cache_backend == "tea_cache" + assert cfg.diffusers_load_kwargs == {} + assert cfg.diffusers_call_kwargs == {} + + +def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path): + deploy_path = tmp_path / "dreamzero_diffusion_aliases.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " diffusion_attention_backend: flash_attn", + ] + ) + ) + + stage = _from_pipeline_key( + "dreamzero", + deploy_config_path=str(deploy_path), + ).stage_by_id(0) + + assert isinstance(stage, VllmOmniDiffusionStageConfig) + assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn" + + +def test_diffusion_config_field_classification_covers_current_fields(): + classified_fields = ( + omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS + | omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS + | omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS + ) + + assert classified_fields == {f.name for f in fields(omni_config_module._DiffusionConfigProjection)} + assert { + "enable_prompt_embed_cache", + "prompt_embed_cache_size", + "diffusion_kv_cache_dtype", + } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS + assert { + "revision", + "trust_remote_code", + "distributed_executor_backend", + } <= omni_config_module._DIFFUSION_SHARED_CONFIG_FIELDS + assert "prompt_file_path" in omni_config_module._DIFFUSION_RUNTIME_CONFIG_FIELDS + + +def test_diffusion_config_projection_keeps_mapping_quantization_config_serializable(): + quantization_config = { + "method": "example_quant", + "weights": "weights.bin", + } + + cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(quantization_config=quantization_config) + + assert cfg.quantization_config == quantization_config diff --git a/tests/config/test_pipeline_registry.py b/tests/config/test_pipeline_registry.py index 336f3805412..4bc19b11632 100644 --- a/tests/config/test_pipeline_registry.py +++ b/tests/config/test_pipeline_registry.py @@ -1,130 +1,68 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for the central pipeline registry (2.5/N).""" +"""Tests for out of tree registration to OMNI_PIPELINES.""" -from __future__ import annotations +import pytest +from transformers import PretrainedConfig -import logging +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline +from vllm_omni.config.stage_config import PipelineConfig, pipeline_cfg_resolver + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def custom_resolver(): + """Build a reusable custom resolver for PipelineConfigs.""" + + class CustomConfigType(PretrainedConfig): + pass + + @pipeline_cfg_resolver(config_type=CustomConfigType) + def custom_resolver( + hf_config: CustomConfigType, + ) -> PipelineConfig: + return PipelineConfig(model_type="resolved_type") + + return custom_resolver + + +def test_register_pipeline_config(clean_pipeline_registry): + """Ensure that we can register a custom pipeline config to OMNI_PIPELINES.""" + new_model_type = "new_model_type" + pipe_cfg = PipelineConfig(model_type=new_model_type) + assert new_model_type not in OMNI_PIPELINES + register_pipeline(pipe_cfg) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is pipe_cfg + + +def test_register_pipeline_config_with_model_type(clean_pipeline_registry): + """Ensure that we can register a custom pipeline config with an explicit model_type to OMNI_PIPELINES.""" + new_model_type = "new_model_type" + unused_model_type = "foo" + pipe_cfg = PipelineConfig(model_type=unused_model_type) + assert new_model_type not in OMNI_PIPELINES + assert unused_model_type not in OMNI_PIPELINES + + # Registering with an explicitly provided model_type uses + # the passed value instead of the pipeline_cfg.model_type + register_pipeline(pipe_cfg, new_model_type) + assert new_model_type in OMNI_PIPELINES + assert unused_model_type not in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is pipe_cfg + + +def test_register_resolver(custom_resolver, clean_pipeline_registry): + """Ensure that we can register a custom resolver to OMNI_PIPELINES.""" + new_model_type = "new_model_type" + assert new_model_type not in OMNI_PIPELINES + register_pipeline(custom_resolver, new_model_type) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is custom_resolver -import pytest -from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES -from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, - PipelineConfig, - StageExecutionType, - StagePipelineConfig, - register_pipeline, -) - - -class TestCentralRegistryDeclarations: - """Every in-tree pipeline must be declared exactly once in the central registry.""" - - def test_omni_entries_visible_in_registry(self): - for key in _OMNI_PIPELINES: - assert key in _PIPELINE_REGISTRY - - def test_expected_omni_pipelines_present(self): - # Guard against accidental removal during future refactors. - assert "qwen2_5_omni" in _OMNI_PIPELINES - assert "qwen2_5_omni_thinker_only" in _OMNI_PIPELINES - assert "qwen3_omni_moe" in _OMNI_PIPELINES - assert "qwen3_tts" in _OMNI_PIPELINES - - -class TestLazyLoading: - """Pipelines are imported only on first access.""" - - def test_contains_without_import(self): - # ``in`` hits the lazy map, not the loaded cache. - assert "qwen3_omni_moe" in _PIPELINE_REGISTRY - - def test_getitem_loads_correct_pipeline(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] - assert pipeline.model_type == "qwen3_omni_moe" - assert pipeline.model_arch == "Qwen3OmniMoeForConditionalGeneration" - - def test_unknown_model_type_returns_none_via_get(self): - assert _PIPELINE_REGISTRY.get("not_a_real_pipeline") is None - - def test_unknown_model_type_raises_keyerror_via_getitem(self): - with pytest.raises(KeyError): - _PIPELINE_REGISTRY["not_a_real_pipeline"] - - def test_iteration_yields_registered_pipelines(self): - keys = set(_PIPELINE_REGISTRY) - assert "qwen2_5_omni" in keys - assert "qwen3_omni_moe" in keys - - def test_iteration_resilience(self, monkeypatch, caplog): - # Synthetic failure to verify iteration robustness against load errors. - target_logger = logging.getLogger("vllm_omni.config.stage_config") - target_logger.addHandler(caplog.handler) - prev_level = target_logger.level - target_logger.setLevel(logging.ERROR) - try: - monkeypatch.setattr( - _PIPELINE_REGISTRY, - "_lazy_map", - { - "healthy": ("vllm_omni.model_executor.models.qwen2_5_omni.pipeline", "QWEN2_5_OMNI_PIPELINE"), - "broken": ("vllm_omni.non_existent_module", "SomeConfig"), - "attr_fail": ("vllm_omni.config.stage_config", "NON_EXISTENT_VAR"), - }, - ) - - # Assert that list(_PIPELINE_REGISTRY.values()) and dict(_PIPELINE_REGISTRY.items()) - # both yield only healthy entries. - healthy_instance = _PIPELINE_REGISTRY["healthy"] - assert list(_PIPELINE_REGISTRY.values()) == [healthy_instance] - assert dict(_PIPELINE_REGISTRY.items()) == {"healthy": healthy_instance} - assert "Failed to import pipeline module" in caplog.text - assert "Pipeline variable 'NON_EXISTENT_VAR' not found" in caplog.text - - finally: - target_logger.removeHandler(caplog.handler) - target_logger.setLevel(prev_level) - - -class TestDynamicRegistration: - """``register_pipeline()`` still works for plugins and tests.""" - - def test_register_adds_to_registry(self): - custom = PipelineConfig( - model_type="_test_dynamic_registration", - model_arch="DynamicTestModel", - stages=( - StagePipelineConfig( - stage_id=0, - model_stage="test", - execution_type=StageExecutionType.LLM_AR, - input_sources=(), - final_output=True, - ), - ), - ) - register_pipeline(custom) - try: - assert "_test_dynamic_registration" in _PIPELINE_REGISTRY - assert _PIPELINE_REGISTRY["_test_dynamic_registration"] is custom - finally: - # Don't leak the test registration into other tests. - if "_test_dynamic_registration" in _PIPELINE_REGISTRY: - del _PIPELINE_REGISTRY["_test_dynamic_registration"] - - def test_dynamic_registration_overrides_lazy_entry(self): - # Build a substitute for qwen3_omni_moe that we can distinguish. - original = _PIPELINE_REGISTRY["qwen3_omni_moe"] - override = PipelineConfig( - model_type="qwen3_omni_moe", - model_arch="OverriddenArch", - stages=original.stages, - ) - register_pipeline(override) - try: - assert _PIPELINE_REGISTRY["qwen3_omni_moe"].model_arch == "OverriddenArch" - finally: - # Remove the dynamic override so later tests see the original. - if "qwen3_omni_moe" in _PIPELINE_REGISTRY._loaded: - del _PIPELINE_REGISTRY["qwen3_omni_moe"] +def test_register_resolver_requires_model_type(custom_resolver, clean_pipeline_registry): + """Ensure that registering a custom resolver to OMNI_PIPELINES requires an explicit model_type.""" + with pytest.raises(ValueError): + register_pipeline(custom_resolver) diff --git a/tests/config/test_server_settings.py b/tests/config/test_server_settings.py new file mode 100644 index 00000000000..53c15524ce6 --- /dev/null +++ b/tests/config/test_server_settings.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for ServerSettings env-var parameterization.""" + +import pytest + +from vllm_omni.config.server_settings import ServerSettings + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +ENV_VARS = [ + "VLLM_OMNI_SERVER_STORAGE__PATH", + "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", + "VLLM_OMNI_SERVER_STORAGE__FILE_TTL", + "VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", + "VLLM_OMNI_STORAGE_PATH", + "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", +] + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + for name in ENV_VARS: + monkeypatch.delenv(name, raising=False) + + +def test_defaults_when_no_env_set(): + storage = ServerSettings().storage + + assert storage.path == "/tmp/storage" + assert storage.file_concurrency == 4 + assert storage.file_ttl is None + assert storage.ttl_sweep_interval is None + + +def test_nested_env_vars_populate_storage(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__PATH", "/var/storage") + monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", "8") + monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_TTL", "120") + monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL", "30") + + storage = ServerSettings().storage + + assert storage.path == "/var/storage" + assert storage.file_concurrency == 8 + assert storage.file_ttl == 120 + assert storage.ttl_sweep_interval == 30 + + +def test_ttl_sweep_interval_defaults_when_ttl_set(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__FILE_TTL", "60") + + storage = ServerSettings().storage + + assert storage.file_ttl == 60 + assert storage.ttl_sweep_interval == 300 + + +def test_deprecated_path_alias_populates_and_warns(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_STORAGE_PATH", "/legacy/path") + + with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_PATH"): + storage = ServerSettings().storage + + assert storage.path == "/legacy/path" + + +def test_deprecated_concurrency_alias_populates_and_warns(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "16") + + with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_MAX_CONCURRENCY"): + storage = ServerSettings().storage + + assert storage.file_concurrency == 16 + + +def test_new_env_var_wins_over_deprecated_alias(monkeypatch): + monkeypatch.setenv("VLLM_OMNI_SERVER_STORAGE__PATH", "/new/path") + monkeypatch.setenv("VLLM_OMNI_STORAGE_PATH", "/legacy/path") + + with pytest.warns(DeprecationWarning, match="VLLM_OMNI_STORAGE_PATH"): + storage = ServerSettings().storage + + assert storage.path == "/new/path" diff --git a/tests/conftest.py b/tests/conftest.py index 163309b8caa..11619b782ea 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,6 +17,7 @@ bootstrap_vllm_layer_custom_op_modules() pytest_plugins = ( + "tests.helpers.fixtures.config", "tests.helpers.fixtures.env", "tests.helpers.fixtures.log", "tests.helpers.fixtures.run_args", diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py new file mode 100644 index 00000000000..27087d253eb --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -0,0 +1,164 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod +import vllm_omni.model_executor.models.voxcpm2.scheduler as voxcpm2_scheduler_mod +from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _MockQueue: + def __init__(self, items: list | None = None) -> None: + self._items = list(items or []) + + def __bool__(self) -> bool: + return bool(self._items) + + def __iter__(self): + return iter(self._items) + + def __len__(self) -> int: + return len(self._items) + + def add_request(self, request) -> None: + self._items.append(request) + + def prepend_requests(self, requests) -> None: + self._items = list(requests) + self._items + + +class _MockRequest: + def __init__( + self, + request_id: str, + *, + status: RequestStatus = RequestStatus.RUNNING, + num_prompt_tokens: int = 4, + num_computed_tokens: int = 4, + num_output_placeholders: int = 0, + ) -> None: + self.request_id = request_id + self.status = status + self.num_prompt_tokens = num_prompt_tokens + self.num_computed_tokens = num_computed_tokens + self.num_output_placeholders = num_output_placeholders + + def is_finished(self) -> bool: + return RequestStatus.is_finished(self.status) + + +@pytest.fixture(autouse=True) +def _mock_cuda_graph_platform(monkeypatch) -> None: + monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: True) + + +def _make_scheduler( + *, + enable_unified_decode_graph: bool | None = True, + deterministic_cfm_noise: bool = False, + runtime_config_on_model_config: bool = False, +) -> VoxCPM2OmniARAsyncScheduler: + sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler) + hf_config = SimpleNamespace() + model_config = SimpleNamespace(hf_config=hf_config) + if enable_unified_decode_graph is not None or deterministic_cfm_noise: + runtime_config = SimpleNamespace( + enable_unified_decode_graph=True if enable_unified_decode_graph is None else enable_unified_decode_graph, + deterministic_cfm_noise=deterministic_cfm_noise, + ) + if runtime_config_on_model_config: + model_config.voxcpm2_runtime_config = runtime_config + else: + hf_config.voxcpm2_runtime_config = runtime_config + sched.vllm_config = SimpleNamespace(model_config=model_config) + return sched + + +def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None: + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_does_not_defer_without_decode_ready() -> None: + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("prefill-running", num_prompt_tokens=8, num_computed_tokens=4)] + scheduler.waiting = _MockQueue([_MockRequest("waiting", status=RequestStatus.WAITING)]) + + assert not scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None: + scheduler = _make_scheduler(enable_unified_decode_graph=False) + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert not scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_uses_model_runtime_defaults() -> None: + scheduler = _make_scheduler(enable_unified_decode_graph=None) + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_reads_model_config_runtime_config() -> None: + scheduler = _make_scheduler(enable_unified_decode_graph=True, runtime_config_on_model_config=True) + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_does_not_defer_with_deterministic_noise() -> None: + scheduler = _make_scheduler(deterministic_cfm_noise=True) + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert not scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypatch) -> None: + monkeypatch.setattr(voxcpm2_scheduler_mod.current_omni_platform, "is_cuda", lambda: False) + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert not scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None: + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("decode")] + original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING) + deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING) + original_waiting = _MockQueue([original_waiting_req]) + scheduler.waiting = original_waiting + scheduler.policy = "fcfs" + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler._consume_pending_connector_output = lambda model_mode: None + + monkeypatch.setattr(scheduler_mod, "create_request_queue", lambda _policy: _MockQueue()) + + def fake_upstream_schedule(self, throttle_prefills: bool = False): + assert self.waiting is not original_waiting + assert not self.waiting + self.waiting.add_request(deferred_by_upstream) + raise RuntimeError("stop before output wrapping") + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule) + + with pytest.raises(RuntimeError, match="stop before output wrapping"): + scheduler.schedule() + + assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req] diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 1b36cd784d8..068cbd36203 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -115,6 +115,7 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase): ("MiMoAudioModel", "code2wav"), ("Qwen3TTSCode2Wav", "code2wav"), ("CosyVoice3Model", "cosyvoice3_code2wav"), + ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"), ("DyninOmniForConditionalGeneration", "token2image"), ("DyninOmniForConditionalGeneration", "token2audio"), } diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index a22f760b255..2eb88bbd62c 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -11,11 +11,180 @@ import pytest from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw +from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import modify_stage_config from vllm_omni.platforms import current_omni_platform +def _named_pytest_marks(names: list[str]) -> list[pytest.MarkDecorator]: + marks: list[pytest.MarkDecorator] = [] + for name in names: + name = name.strip() + if not name: + raise ValueError("mark name must be a non-empty string") + marks.append(getattr(pytest.mark, name)) + return marks + + +def _hardware_marks_from_dict(hw: Any) -> list[pytest.MarkDecorator]: + if not isinstance(hw, dict): + raise ValueError(f"mark.hardware_marks must be a dict, got {type(hw).__name__}") + res = hw.get("res") + if not isinstance(res, dict): + raise ValueError(f"mark.hardware_marks.res must be a dict, got {type(res).__name__}") + num_cards = hw.get("num_cards", 1) + return list(hardware_marks(res=res, num_cards=num_cards)) + + +def resolve_pytest_marks(mark_field: Any) -> list[pytest.MarkDecorator]: + """Convert a JSON ``mark`` field into pytest mark decorators. + + Supported form (per test-case object in perf/stability JSON):: + + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, + "full_model", + "diffusion" + ] + + Exactly one ``hardware_marks`` object is required when ``mark`` is present. + ``hardware_marks`` delegates to :func:`tests.helpers.mark.hardware_marks` + (same shape as ``@hardware_test``). Additional array entries are registered + pytest marker names (strings). + """ + if mark_field is None: + return [] + + if isinstance(mark_field, list): + marks: list[pytest.MarkDecorator] = [] + hw_seen = False + for item in mark_field: + if isinstance(item, dict) and "hardware_marks" in item: + if hw_seen: + raise ValueError("mark array must contain at most one hardware_marks object") + hw_seen = True + marks.extend(_hardware_marks_from_dict(item["hardware_marks"])) + unknown_keys = set(item) - {"hardware_marks"} + if unknown_keys: + raise ValueError( + f"mark hardware_marks object only allows hardware_marks; unknown keys: {sorted(unknown_keys)}" + ) + elif isinstance(item, str): + item = item.strip() + if not item: + raise ValueError("mark name must be a non-empty string") + marks.extend(_named_pytest_marks([item])) + else: + raise ValueError( + f"mark array entries must be hardware_marks objects or marker name strings; got {type(item).__name__}" + ) + if not hw_seen: + raise ValueError("mark array must contain a hardware_marks object") + return marks + + raise ValueError(f"mark must be a list; got {type(mark_field).__name__}") + + +def _mark_names(mark_field: Any) -> set[str]: + if isinstance(mark_field, list): + return {str(item) for item in mark_field if isinstance(item, str)} + return set() + + +def is_diffusion_perf_config(cfg: dict[str, Any]) -> bool: + """True for perf JSON cases intended for ``run_diffusion_benchmark.py``.""" + if cfg.get("server_type") is not None: + return True + return "diffusion" in _mark_names(cfg.get("mark")) + + +def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]: + return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs} + + +def create_unique_server_pytest_params( + configs: list[dict[str, Any]], + stage_configs_dir: Path, +) -> list[Any]: + """Like :func:`create_unique_server_params`, but wrap each row in ``pytest.param`` with JSON marks.""" + marks_by_name = _marks_by_test_name(configs) + return [ + pytest.param( + row, + marks=marks_by_name.get(row[0], []), + id=row[0], + ) + for row in create_unique_server_params(configs, stage_configs_dir) + ] + + +def create_benchmark_pytest_params( + benchmark_configs: list[dict[str, Any]], + server_to_benchmark_mapping: dict[str, dict], +) -> list[Any]: + """Like :func:`create_benchmark_indices`, but wrap each index in ``pytest.param`` with JSON marks.""" + marks_by_name = _marks_by_test_name(benchmark_configs) + params: list[Any] = [] + seen: set[str] = set() + for config in benchmark_configs: + test_name = config["test_name"] + if test_name in seen: + continue + seen.add(test_name) + params_list = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping) + id_suffixes = _unique_benchmark_param_id_suffixes(params_list) + for idx, id_suffix in enumerate(id_suffixes): + params.append( + pytest.param( + (test_name, idx), + marks=marks_by_name.get(test_name, []), + id=f"{test_name}-{id_suffix}", + ) + ) + return params + + +def create_paired_benchmark_pytest_params( + server_entries: list[tuple[Any, str]], + params_by_test_name: dict[str, list[Any]], + marks_by_name: dict[str, list[pytest.MarkDecorator]], +) -> list[Any]: + """One ``pytest.param`` per ``(server entry, benchmark index)`` pair. + + Pass ``server_entry`` and ``(test_name, idx)`` as separate ``pytest.param`` + arguments for ``@pytest.mark.parametrize("server_fixture,benchmark_params", ...)``. + """ + pairs: list[Any] = [] + for server_entry, test_name in server_entries: + params_list = params_by_test_name.get(test_name, []) + id_suffixes = _unique_benchmark_param_id_suffixes(params_list) + for idx, id_suffix in enumerate(id_suffixes): + pairs.append( + pytest.param( + server_entry, + (test_name, idx), + marks=marks_by_name.get(test_name, []), + id=f"{test_name}-{id_suffix}", + ) + ) + return pairs + + +def create_paired_omni_benchmark_pytest_params( + configs: list[dict[str, Any]], + stage_configs_dir: Path, +) -> list[Any]: + """Paired params for ``run_benchmark.py`` (omni/tts).""" + mapping = create_test_parameter_mapping(configs) + marks_by_name = _marks_by_test_name(configs) + server_entries = [(row, row[0]) for row in create_unique_server_params(configs, stage_configs_dir)] + params_by_test_name = { + test_name: get_benchmark_params_for_server(test_name, mapping) for _, test_name in server_entries + } + return create_paired_benchmark_pytest_params(server_entries, params_by_test_name, marks_by_name) + + def load_configs(config_path: str) -> list[dict[str, Any]]: try: abs_path = Path(config_path).resolve() @@ -32,6 +201,29 @@ def load_configs(config_path: str) -> list[dict[str, Any]]: raise RuntimeError(f"Failed to load configuration file: {str(e)}") +def load_benchmark_configs( + config_path: str | None = None, + *, + config_dir: Path | None = None, +) -> list[dict[str, Any]]: + """Load one benchmark JSON file, or merge all ``*.json`` under *config_dir*. + + When *config_path* is omitted, every ``*.json`` in *config_dir* is loaded and + concatenated. Pytest ``-m`` expressions then select cases via each entry's + ``mark`` field (e.g. ``-m tts`` after bulk load). + """ + if config_path is not None: + return load_configs(config_path) + if config_dir is None: + raise ValueError("load_benchmark_configs requires config_path or config_dir") + configs: list[dict[str, Any]] = [] + for path in sorted(config_dir.glob("*.json")): + configs.extend(load_configs(str(path))) + if not configs: + raise ValueError(f"No benchmark JSON files found under {config_dir}") + return configs + + def modify_stage(default_path: str, updates: dict[str, Any] | None, deletes: dict[str, Any] | None) -> str: kwargs: dict[str, Any] = {} if updates is not None: @@ -293,6 +485,158 @@ def _safe_filename_token(value: Any | None, *, default: str = "na") -> str: return s if s else default +def _benchmark_param_id_suffix(param: dict[str, Any], *, idx: int) -> str: + """Derive a readable pytest id suffix from one ``benchmark_params`` entry.""" + name = param.get("name") + if isinstance(name, str) and name.strip(): + return _safe_filename_token(name.strip()) + + parts: list[str] = [] + for key in ("task", "eval_phase", "dataset_name", "dataset"): + value = param.get(key) + if value is None or value == "": + continue + token = _safe_filename_token(str(value)) + if token != "na": + parts.append(token) + if parts: + return "_".join(parts) + + return f"case{idx}" + + +def _unique_benchmark_param_id_suffixes(params_list: list[dict[str, Any]]) -> list[str]: + """Return unique pytest id suffixes for a server's benchmark param list.""" + raw = [_benchmark_param_id_suffix(param, idx=idx) for idx, param in enumerate(params_list)] + seen: dict[str, int] = {} + unique: list[str] = [] + for suffix in raw: + count = seen.get(suffix, 0) + seen[suffix] = count + 1 + if count == 0: + unique.append(suffix) + else: + unique.append(f"{suffix}_{count}") + return unique + + +def extract_mark_resource_label(mark_field: Any) -> str: + """Return a filename-safe hardware label from ``mark.hardware_marks.res`` values. + + Example: ``{"cuda": "H100"}`` -> ``"H100"``; multiple platforms join with ``-``. + + Prefer :func:`get_runtime_resource_label` for perf result filenames so labels + reflect the machine that actually ran the benchmark. + """ + if isinstance(mark_field, list): + for item in mark_field: + if isinstance(item, dict) and "hardware_marks" in item: + return extract_mark_resource_label(item) + return "na" + if not isinstance(mark_field, dict): + return "na" + hw = mark_field.get("hardware_marks") + if not isinstance(hw, dict): + return "na" + res = hw.get("res") + if not isinstance(res, dict) or not res: + return "na" + labels = [_safe_filename_token(value) for value in res.values()] + return "-".join(labels) if labels else "na" + + +_KNOWN_RUNTIME_RESOURCE_TOKENS: tuple[str, ...] = ( + "H100", + "H800", + "H200", + "H20", + "L40S", + "L40", + "L4", + "A100", + "A800", + "A10G", + "A10", + "A30", + "MI325", + "MI300", + "MI250", + "B60", + "S5000", + "910B4", + "910B", + "910", + "310P", + "A2", + "A3", +) +_RUNTIME_RESOURCE_LABEL: str | None = None + + +def _normalize_runtime_device_label(raw: str) -> str: + """Map a platform device name to a short filename-safe resource token.""" + if not raw or not str(raw).strip(): + return "na" + upper = str(raw).upper() + for token in _KNOWN_RUNTIME_RESOURCE_TOKENS: + if token.upper() in upper: + return _safe_filename_token(token) + compact = re.sub(r"[^a-zA-Z0-9]+", "", str(raw)) + for prefix in ("NVIDIA", "AMD", "ASCEND", "HUAWEI"): + if compact.upper().startswith(prefix): + compact = compact[len(prefix) :] + break + return _safe_filename_token(compact[:48]) if compact else "na" + + +def _read_runtime_device_name(*, device_id: int = 0) -> str | None: + """Device name from the active Omni platform.""" + if current_omni_platform.device_count() <= device_id: + return None + get_name = getattr(current_omni_platform, "get_device_name", None) + if not callable(get_name): + return None + raw = get_name(device_id) + if isinstance(raw, str) and raw.strip(): + return raw.strip() + return None + + +def get_runtime_resource_label(*, device_id: int = 0, refresh: bool = False) -> str: + """Return a filename-safe hardware label detected on the running machine.""" + global _RUNTIME_RESOURCE_LABEL + if not refresh and _RUNTIME_RESOURCE_LABEL is not None: + return _RUNTIME_RESOURCE_LABEL + raw = _read_runtime_device_name(device_id=device_id) + label = _normalize_runtime_device_label(raw) if raw else "na" + if not refresh: + _RUNTIME_RESOURCE_LABEL = label + return label + + +_FILENAME_OMIT_RESOURCE_LABELS = frozenset({"H100"}) + + +def hardware_json_value(resource_label: str | None) -> str: + """Hardware token stored in perf result JSON (empty when unknown).""" + token = _safe_filename_token(resource_label) + return "" if token == "na" else token + + +def resource_label_for_filename(resource_label: str | None) -> str: + """Hardware token embedded in result filenames (H100 omitted on default CI pool).""" + token = _safe_filename_token(resource_label) + if token in _FILENAME_OMIT_RESOURCE_LABELS: + return "" + return token + + +def extract_configs_resource_label(configs: list[dict[str, Any]]) -> str: + """Return runtime hardware label for perf result filenames.""" + del configs + return get_runtime_resource_label() + + def resolve_baseline_value( baseline_raw: Any, *, @@ -357,6 +701,7 @@ def run_benchmark( max_concurrency: Any | None = None, random_input_len: Any | None = None, random_output_len: Any | None = None, + resource_label: str | None = None, ) -> dict[str, Any]: """Run one ``vllm bench serve --omni`` iteration and return parsed metrics. @@ -367,7 +712,11 @@ def run_benchmark( current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") ri = _safe_filename_token(random_input_len) ro = _safe_filename_token(random_output_len) - result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json" + hw = resource_label_for_filename(resource_label) + if hw: + result_filename = f"result_{test_name}_{hw}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json" + else: + result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json" if "--result-filename" in args: print(f"The result file will be overwritten by {result_filename}") command = ( @@ -437,6 +786,7 @@ def _forward_stream(stream) -> None: result["random_input_len"] = random_input_len if random_output_len is not None: result["random_output_len"] = random_output_len + result["Hardware"] = hardware_json_value(resource_label) with open(result_path, "w", encoding="utf-8") as f: json.dump(result, f, ensure_ascii=False, indent=2) return result diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 7ca37632364..21167cb0772 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -7,21 +7,24 @@ import pytest from tests.dfx.conftest import ( - create_benchmark_indices, + create_paired_omni_benchmark_pytest_params, create_test_parameter_mapping, - create_unique_server_params, get_benchmark_params_for_server, - load_configs, + get_runtime_resource_label, + is_diffusion_perf_config, + load_benchmark_configs, resolve_baseline_value, run_benchmark, ) from tests.helpers.runtime import OmniServer -pytestmark = [pytest.mark.full_model] - # Compare metrics to each test JSON ``baseline`` block only when pytest is run with ``--assert-baseline`` # (registered in ``tests/dfx/conftest.py``; default: off). ``run_benchmark`` and ``_resolve_baseline_value`` are # defined in the same module. +# +# Optional JSON field ``mark`` is applied as pytest marks via +# ``create_paired_omni_benchmark_pytest_params`` (e.g. ``"mark": [{"hardware_marks": +# {"res": {"cuda": "H100"}, "num_cards": 2}}, "full_model", "omni"]``). os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -40,22 +43,22 @@ def _get_config_file_from_argv() -> str | None: _PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests" -_DEFAULT_CONFIG_FILE = str(_PERF_TESTS_DIR / "test_qwen_omni.json") CONFIG_FILE_PATH = _get_config_file_from_argv() if CONFIG_FILE_PATH is None: + _all_configs = load_benchmark_configs(config_dir=_PERF_TESTS_DIR) + BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if not is_diffusion_perf_config(cfg)] print( - "No --test-config-file in argv, using default: tests/dfx/perf/tests/test_qwen_omni.json " - "(override with e.g. --test-config-file tests/dfx/perf/tests/test_tts.json)" + f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} omni/tts case(s) from " + f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} diffusion; " + f"use -m to filter, e.g. -m tts)" ) - CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE - -BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) - +else: + BENCHMARK_CONFIGS = load_benchmark_configs(CONFIG_FILE_PATH) DEPLOY_CONFIGS_DIR = Path(__file__).parent.parent / "deploy" -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) +paired_benchmark_params = create_paired_omni_benchmark_pytest_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) _omni_server_lock = threading.Lock() @@ -91,17 +94,11 @@ def omni_server(request): print("OmniServer stopped") -benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) - - -@pytest.fixture(params=benchmark_indices) -def benchmark_params(request, omni_server): - """Benchmark parameters fixture with proper parametrization""" +@pytest.fixture +def benchmark_params(request): + """Benchmark parameters fixture; paired with ``omni_server`` via parametrization.""" test_name, param_index = request.param - if test_name != omni_server.test_name: - pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}") - all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping) if not all_params: @@ -153,8 +150,11 @@ def assert_result( @pytest.mark.benchmark -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True) +@pytest.mark.parametrize( + "omni_server,benchmark_params", + paired_benchmark_params, + indirect=["omni_server", "benchmark_params"], +) def test_performance_benchmark(omni_server, benchmark_params, request): test_name = benchmark_params["test_name"] params = benchmark_params["params"] @@ -168,6 +168,7 @@ def test_performance_benchmark(omni_server, benchmark_params, request): print(f"Benchmark parameters: {benchmark_params}") assert_baseline = request.config.getoption("--assert-baseline", default=False) + resource_label = get_runtime_resource_label() def to_list(value, default=None): if value is None: @@ -239,6 +240,7 @@ def to_list(value, default=None): max_concurrency=None, random_input_len=params.get("random_input_len"), random_output_len=params.get("random_output_len"), + resource_label=resource_label, ) assert_result( result, @@ -264,6 +266,7 @@ def to_list(value, default=None): max_concurrency=concurrency, random_input_len=params.get("random_input_len"), random_output_len=params.get("random_output_len"), + resource_label=resource_label, ) assert_result( result, diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 62ce5960614..d19c0900ed0 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -8,15 +8,21 @@ 2. ``benchmark_endpoint``: which serving API the benchmark client calls. Examples: ``/v1/chat/completions`` and ``/v1/videos``. -A config JSON file is REQUIRED via --test-config-file: +A config JSON file may be passed via --test-config-file. If omitted, every ``*.json`` under +``tests/dfx/perf/tests/`` is loaded and pytest ``-m`` filters by each case's ``mark``: + pytest run_diffusion_benchmark.py -m "diffusion" pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json Optional: ``--assert-baseline`` compares metrics to the ``baseline`` block in each benchmark entry (default: off). -All benchmark results for a session are consolidated into a single JSON file under -BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable). -Each entry in the file contains the test metadata (test_name, endpoint, benchmark_params, -timestamp) together with the raw metrics returned by the benchmark script. +Optional JSON field ``mark`` is applied as pytest marks on that case via +``pytest.param`` (e.g. ``"mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, "full_model", "diffusion"]``). + +All benchmark results are written under BENCHMARK_RESULT_DIR (override via the +DIFFUSION_BENCHMARK_DIR environment variable). Each source JSON file gets one +aggregated ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` (JSON array +of all runs from cases in that file). Bulk load without ``--test-config-file`` uses +the same per-file aggregation; ``-m`` only selects which cases run. """ import json @@ -35,8 +41,15 @@ import pytest from benchmarks.diffusion.backends import endpoint_filename_token, normalize_endpoint - -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +from tests.dfx.conftest import ( + create_paired_benchmark_pytest_params, + get_runtime_resource_label, + hardware_json_value, + is_diffusion_perf_config, + resolve_pytest_marks, + resource_label_for_filename, +) +from tests.helpers.runtime import get_open_port os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN") @@ -126,6 +139,10 @@ def _process_inline_fields(obj: Any, parent_key: str = "") -> None: DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json" +_DIFFUSION_SOURCE_CONFIG_KEY = "_source_config_file" +_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests" + + def _get_config_file_from_argv() -> str | None: """Read --test-config-file from sys.argv at import time so pytest parametrize can use it. @@ -142,9 +159,6 @@ def _get_config_file_from_argv() -> str | None: CONFIG_FILE_PATH = _get_config_file_from_argv() -if CONFIG_FILE_PATH is None: - print("No config file provided, using default config file: tests/dfx/perf/tests/test_qwen_image_vllm_omni.json") - CONFIG_FILE_PATH = "tests/dfx/perf/tests/test_qwen_image_vllm_omni.json" # --------------------------------------------------------------------------- # Config loading @@ -184,28 +198,85 @@ def load_configs(config_path: str) -> list[dict[str, Any]]: raise RuntimeError(f"Failed to load configuration file: {str(e)}") -BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +def load_diffusion_benchmark_configs( + config_path: str | None = None, + *, + config_dir: Path | None = None, +) -> list[dict[str, Any]]: + """Load one diffusion benchmark JSON, or merge all ``*.json`` under *config_dir*.""" + if config_path is not None: + configs = load_configs(config_path) + source = str(Path(config_path).resolve()) + for cfg in configs: + cfg.setdefault(_DIFFUSION_SOURCE_CONFIG_KEY, source) + return configs + if config_dir is None: + raise ValueError("load_diffusion_benchmark_configs requires config_path or config_dir") + configs: list[dict[str, Any]] = [] + for path in sorted(config_dir.glob("*.json")): + source = str(path.resolve()) + for cfg in load_configs(str(path)): + cfg[_DIFFUSION_SOURCE_CONFIG_KEY] = source + configs.append(cfg) + if not configs: + raise ValueError(f"No benchmark JSON files found under {config_dir}") + return configs -_config_stem = Path(CONFIG_FILE_PATH).stem # e.g. "test_qwen_image_vllm_omni" -AGGREGATED_RESULT_FILE = BENCHMARK_RESULT_DIR / f"diffusion_result_{_config_stem}_{_SESSION_TIMESTAMP}.json" +if CONFIG_FILE_PATH is None: + _all_configs = load_diffusion_benchmark_configs(config_dir=_PERF_TESTS_DIR) + BENCHMARK_CONFIGS = [cfg for cfg in _all_configs if is_diffusion_perf_config(cfg)] + print( + f"No --test-config-file: loaded {len(BENCHMARK_CONFIGS)} diffusion case(s) from " + f"{_PERF_TESTS_DIR}/*.json (skipped {len(_all_configs) - len(BENCHMARK_CONFIGS)} omni/tts; " + f"use -m to filter, e.g. -m diffusion)" + ) +else: + BENCHMARK_CONFIGS = load_diffusion_benchmark_configs(CONFIG_FILE_PATH) + +_AGGREGATED_RESULT_FILES_BY_SOURCE: dict[str, Path] = {} + + +def _normalized_source_path(source_file: str) -> str: + return str(Path(source_file).resolve()) + + +def _aggregated_result_file_for_source(source_file: str) -> Path: + """One session aggregate per source JSON (same naming as single ``--test-config-file``).""" + key = _normalized_source_path(source_file) + if key not in _AGGREGATED_RESULT_FILES_BY_SOURCE: + stem = Path(key).stem + resource = resource_label_for_filename(get_runtime_resource_label()) + if resource: + result_name = f"diffusion_result_{stem}_{resource}_{_SESSION_TIMESTAMP}.json" + else: + result_name = f"diffusion_result_{stem}_{_SESSION_TIMESTAMP}.json" + _AGGREGATED_RESULT_FILES_BY_SOURCE[key] = BENCHMARK_RESULT_DIR / result_name + return _AGGREGATED_RESULT_FILES_BY_SOURCE[key] -def _append_to_aggregated_file(record: dict[str, Any]) -> None: - """Thread-safe append of *record* to the session-level aggregated JSON file. - The file contains a JSON array; each call loads the existing array (or - starts a new one), appends the record, and writes the file back atomically. - """ +def _write_result_record(record: dict[str, Any]) -> Path: + """Append one benchmark record to the aggregate file for its source JSON.""" + source_file = record.get("source_file") or CONFIG_FILE_PATH + if not source_file: + raise ValueError("benchmark record missing source_file") + target = _aggregated_result_file_for_source(str(source_file)) with _RESULT_LOCK: BENCHMARK_RESULT_DIR.mkdir(parents=True, exist_ok=True) - if AGGREGATED_RESULT_FILE.exists(): - with open(AGGREGATED_RESULT_FILE, encoding="utf-8") as f: + if target.exists(): + with open(target, encoding="utf-8") as f: records: list[dict] = json.load(f) else: records = [] records.append(record) - with open(AGGREGATED_RESULT_FILE, "w", encoding="utf-8") as f: + with open(target, "w", encoding="utf-8") as f: json.dump(records, f, indent=2, ensure_ascii=False) + return target + + +def _append_to_aggregated_file(record: dict[str, Any]) -> None: + """Backward-compatible wrapper; prefer :func:`_write_result_record`.""" + _write_result_record(record) _server_lock = threading.Lock() @@ -215,14 +286,6 @@ def _append_to_aggregated_file(record: dict[str, Any]) -> None: # --------------------------------------------------------------------------- -def _get_open_port() -> int: - """Return an available TCP port on localhost.""" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("", 0)) - s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - return s.getsockname()[1] - - def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None: """Block until the given host:port accepts connections or timeout expires. @@ -312,7 +375,7 @@ def __init__( self.model = server_cfg["model"] self.serve_args = server_cfg["serve_args"] self.host = "127.0.0.1" - self.port = port if port is not None else _get_open_port() + self.port = port if port is not None else get_open_port(self.host) self.proc: subprocess.Popen | None = None self.test_name: str = "" @@ -482,6 +545,7 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]] "serve_args": _build_serve_args(serve_args_dict), "benchmark_endpoint": cfg.get("benchmark_endpoint", cfg.get("benchmark_backend")), "server_params": cfg["server_params"], + "mark": cfg.get("mark"), } ) return result @@ -496,19 +560,29 @@ def _test_param_mapping(configs: list[dict[str, Any]]) -> dict[str, list[dict]]: return mapping -def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer: - """Factory: return a vLLM-Omni diffusion server instance for the config.""" - return DiffusionServer(server_cfg=server_cfg) +def _marks_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[pytest.MarkDecorator]]: + return {str(cfg["test_name"]): resolve_pytest_marks(cfg.get("mark")) for cfg in configs} + + +def _paired_diffusion_benchmark_pytest_params(configs: list[dict[str, Any]]) -> list[Any]: + """Paired params for ``run_diffusion_benchmark.py``; same shape as omni runner.""" + test_param_map = _test_param_mapping(configs) + server_entries = [(cfg, cfg["test_name"]) for cfg in _unique_server_params(configs)] + return create_paired_benchmark_pytest_params(server_entries, test_param_map, _marks_by_test_name(configs)) # --------------------------------------------------------------------------- # Parametrize data # --------------------------------------------------------------------------- -server_params = _unique_server_params(BENCHMARK_CONFIGS) test_param_map = _test_param_mapping(BENCHMARK_CONFIGS) +paired_benchmark_params = _paired_diffusion_benchmark_pytest_params(BENCHMARK_CONFIGS) + + +def _make_server(server_cfg: dict[str, Any]) -> DiffusionServer: + """Factory: return a vLLM-Omni diffusion server instance for the config.""" + return DiffusionServer(server_cfg=server_cfg) -benchmark_indices: list[int] = list(range(max(len(v) for v in test_param_map.values()))) # --------------------------------------------------------------------------- # Pytest fixtures @@ -534,16 +608,13 @@ def diffusion_server(request): @pytest.fixture -def benchmark_params(request, diffusion_server): - """Yield the benchmark params dict for the current (server, index) pair.""" - param_index: int = request.param - test_name = diffusion_server.test_name +def benchmark_params(request): + """Benchmark params for the paired server/index parametrization.""" + test_name, param_index = request.param params_list = test_param_map.get(test_name, []) if not params_list: raise ValueError(f"No benchmark params for test: {test_name}") - if param_index >= len(params_list): - pytest.skip(f"Param index {param_index} out of range for {test_name} (has {len(params_list)} params)") current = param_index + 1 total = len(params_list) @@ -556,6 +627,10 @@ def benchmark_params(request, diffusion_server): # --------------------------------------------------------------------------- +_STAGE_METRICS_ENDPOINTS = {"/v1/chat/completions"} +_DIFFUSION_PIPELINE_PROFILER_ARG = "enable-diffusion-pipeline-profiler" + + def run_benchmark( host: str, port: int, @@ -571,8 +646,8 @@ def run_benchmark( The raw metrics are written to a temporary file by the subprocess. After the run completes the metrics are merged with full metadata (test_name, endpoint, benchmark_params, timestamp, flat reporting fields) and appended - to the session-wide aggregated JSON file (AGGREGATED_RESULT_FILE). The - temporary file is removed afterwards. Subprocess stdout/stderr are tee'd + to ``diffusion_result_{config_stem}_{hardware}_{timestamp}.json`` for the + source JSON file. The temporary file is removed afterwards. Subprocess stdout/stderr are tee'd to a .log file under BENCHMARK_RESULT_DIR/logs/; its path is stored in the record. """ @@ -582,7 +657,12 @@ def run_benchmark( log_dir = BENCHMARK_RESULT_DIR / "logs" log_dir.mkdir(parents=True, exist_ok=True) endpoint_label = endpoint_filename_token(endpoint) - log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log" + resource_label = get_runtime_resource_label() + hw_for_filename = resource_label_for_filename(resource_label) + if hw_for_filename: + log_file = log_dir / f"{test_name}_{hw_for_filename}_{endpoint_label}_{timestamp}.log" + else: + log_file = log_dir / f"{test_name}_{endpoint_label}_{timestamp}.log" with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp: tmp_result_file = Path(tmp.name) @@ -608,6 +688,11 @@ def run_benchmark( str(tmp_result_file), ] + serve_args_dict = (server_cfg or {}).get("serve_args_dict") + profiler_enabled = isinstance(serve_args_dict, dict) and bool(serve_args_dict.get(_DIFFUSION_PIPELINE_PROFILER_ARG)) + if endpoint in _STAGE_METRICS_ENDPOINTS and profiler_enabled: + cmd.append("--return-stage-metrics") + for key, value in params.items(): if key in exclude_keys or value is None: continue @@ -687,7 +772,7 @@ def run_benchmark( "Model": model, "Framework": server_type, "API Endpoint": endpoint, - "Hardware": "", + "Hardware": hardware_json_value(resource_label), "Deployment": "", "Task": params.get("task", "t2i"), "Dataset": params.get("dataset", "random"), @@ -719,8 +804,8 @@ def run_benchmark( "build_url": os.environ.get("BUILDKITE_BUILD_URL", ""), "source_file": source_file, } - _append_to_aggregated_file(record) - print(f"\n Result appended to: {AGGREGATED_RESULT_FILE}") + result_path = _write_result_record(record) + print(f"\n Result saved to: {result_path}") print(f" Log saved to: {log_file}") return metrics @@ -798,7 +883,7 @@ def assert_result( def _default_benchmark_endpoint_for_task(task: str) -> str: """Return the default client-side benchmark endpoint for a diffusion task.""" - if task in {"t2v", "i2v", "ti2v"}: + if task in {"t2v", "i2v", "ti2v", "v2v"}: return "/v1/videos" if task in {"t2i", "i2i", "ti2i"}: return "/v1/chat/completions" @@ -923,12 +1008,10 @@ def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]: # --------------------------------------------------------------------------- @pytest.mark.benchmark @pytest.mark.parametrize( - "diffusion_server", - server_params, - ids=[p["test_name"] for p in server_params], - indirect=True, + "diffusion_server,benchmark_params", + paired_benchmark_params, + indirect=["diffusion_server", "benchmark_params"], ) -@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True) def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, request): """Run the diffusion performance benchmark and verify request completion. @@ -953,7 +1036,10 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req test_name=test_name, endpoint=endpoint, server_cfg=server_cfg, - source_file=cast(str, CONFIG_FILE_PATH), + source_file=server_cfg.get( + _DIFFUSION_SOURCE_CONFIG_KEY, + CONFIG_FILE_PATH or f"{_PERF_TESTS_DIR}/*.json", + ), ) print(f"\n{'=' * 60}") @@ -971,7 +1057,9 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req if key in result: print(f" {key}: {result[key]:.4f}") - print(f"\n Aggregated results: {AGGREGATED_RESULT_FILE}") + source = server_cfg.get(_DIFFUSION_SOURCE_CONFIG_KEY) or CONFIG_FILE_PATH + if source: + print(f"\n Aggregated results: {_aggregated_result_file_for_source(str(source))}") print("=" * 60) assert_baseline = request.config.getoption("--assert-baseline", default=False) diff --git a/tests/dfx/perf/tests/test_bagel_vllm_omni.json b/tests/dfx/perf/tests/test_bagel_vllm_omni.json index 09fc96030b8..3038393872f 100644 --- a/tests/dfx/perf/tests/test_bagel_vllm_omni.json +++ b/tests/dfx/perf/tests/test_bagel_vllm_omni.json @@ -1,6 +1,18 @@ [ { "test_name": "test_bagel_single_device_single_stage_t2i", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], "description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task", "server_type": "vllm-omni", "server_params": { @@ -38,6 +50,18 @@ }, { "test_name": "test_bagel_single_device_single_stage_i2i", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], "description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task", "server_type": "vllm-omni", "server_params": { @@ -76,6 +100,18 @@ }, { "test_name": "test_bagel_single_device_multi_stage_t2i", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task", "server_type": "vllm-omni", "server_params": { @@ -113,6 +149,18 @@ }, { "test_name": "test_bagel_single_device_multi_stage_i2i", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], "description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task", "server_type": "vllm-omni", "server_params": { diff --git a/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json new file mode 100644 index 00000000000..9bf31658d15 --- /dev/null +++ b/tests/dfx/perf/tests/test_cosmos3_vllm_omni.json @@ -0,0 +1,261 @@ +[ + { + "test_name": "test_cosmos3_t2i_official_demo_2gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Cosmos3-Nano text-to-image official demo workload on 2 GPUs", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/images/generations", + "server_params": { + "model": "nvidia/Cosmos3-Nano", + "serve_args": { + "cfg-parallel-size": 2, + "use-hsdp": true, + "hsdp-shard-size": 2, + "vae-patch-parallel-size": 2, + "vae-use-tiling": true, + "model-class-name": "Cosmos3OmniDiffusersPipeline", + "no-guardrails": true, + "stage-init-timeout": 900, + "init-timeout": 1200, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1024x1024_steps4", + "dataset": "random", + "task": "t2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 4, + "num-prompts": 3, + "max-concurrency": 1, + "seed": 42, + "extra-body": { + "guidance_scale": 7.0, + "flow_shift": 3.0, + "guardrails": false + }, + "baseline": { + "throughput_qps": 1.107, + "latency_mean": 0.8929, + "peak_memory_mb_mean": 80000.0 + } + } + ] + }, + { + "test_name": "test_cosmos3_t2v_official_demo_2gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Cosmos3-Nano text-to-video official demo workload on 2 GPUs", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "nvidia/Cosmos3-Nano", + "serve_args": { + "cfg-parallel-size": 2, + "use-hsdp": true, + "hsdp-shard-size": 2, + "vae-patch-parallel-size": 2, + "vae-use-tiling": true, + "model-class-name": "Cosmos3OmniDiffusersPipeline", + "no-guardrails": true, + "stage-init-timeout": 900, + "init-timeout": 1200, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1280x720_frames189_steps4", + "dataset": "random", + "task": "t2v", + "num-prompts": 3, + "max-concurrency": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1280, + "height": 720, + "prompt": "A robot arm is cleaning a plate in the kitchen.", + "num_inference_steps": 4, + "num_frames": 189, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 10.0, + "max_sequence_length": 4096, + "guardrails": false, + "use_resolution_template": false, + "use_duration_template": false + }, + "baseline": { + "throughput_qps": 0.027, + "latency_mean": 36.6824, + "peak_memory_mb_mean": 30155.4 + } + } + ] + }, + { + "test_name": "test_cosmos3_i2v_official_demo_2gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Cosmos3-Nano image-to-video official demo workload on 2 GPUs", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "nvidia/Cosmos3-Nano", + "serve_args": { + "cfg-parallel-size": 2, + "use-hsdp": true, + "hsdp-shard-size": 2, + "vae-patch-parallel-size": 2, + "vae-use-tiling": true, + "model-class-name": "Cosmos3OmniDiffusersPipeline", + "no-guardrails": true, + "stage-init-timeout": 900, + "init-timeout": 1200, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1280x720_frames189_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 3, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1280, + "height": 720, + "prompt": "The scene comes to life with smooth, natural motion.", + "num_inference_steps": 4, + "num_frames": 189, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 10.0, + "max_sequence_length": 4096, + "guardrails": false + }, + "baseline": { + "throughput_qps": 0.027, + "latency_mean": 36.968, + "peak_memory_mb_mean": 30443.6 + } + } + ] + }, + { + "test_name": "test_cosmos3_v2v_official_demo_2gpu", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "Cosmos3-Nano video-to-video official demo workload on 2 GPUs", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/videos", + "server_params": { + "model": "nvidia/Cosmos3-Nano", + "serve_args": { + "cfg-parallel-size": 2, + "use-hsdp": true, + "hsdp-shard-size": 2, + "vae-patch-parallel-size": 2, + "vae-use-tiling": true, + "model-class-name": "Cosmos3OmniDiffusersPipeline", + "no-guardrails": true, + "stage-init-timeout": 900, + "init-timeout": 1200, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1280x720_frames189_steps4", + "dataset": "random", + "task": "v2v", + "num-prompts": 3, + "max-concurrency": 1, + "seed": 42, + "random-request-config": [ + { + "width": 1280, + "height": 720, + "prompt": "Continue the same scene with smooth natural motion.", + "num_inference_steps": 4, + "num_frames": 189, + "fps": 24, + "weight": 1 + } + ], + "extra-body": { + "guidance_scale": 6.0, + "flow_shift": 10.0, + "max_sequence_length": 4096, + "guardrails": false, + "condition_frame_indexes_vision": [ + 0, + 1 + ], + "condition_video_keep": "first" + }, + "baseline": { + "throughput_qps": 0.027, + "latency_mean": 37.2117, + "peak_memory_mb_mean": 30135.6 + } + } + ] + } +] diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json index 6fdaec5807e..9b7b2eb5cf1 100644 --- a/tests/dfx/perf/tests/test_higgs_audio_v3.json +++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json @@ -1,6 +1,18 @@ [ { "test_name": "test_higgs_audio_v3", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "tts" + ], "server_params": { "model": "bosonai/higgs-audio-v3-tts-4b", "stage_overrides": { diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json index b95aacd054e..1b6902cba07 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json +++ b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json @@ -1,132 +1,152 @@ [ - { - "test_name": "test_hunyuan_image3_ti2i_inline", - "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config", - "server_type": "vllm-omni", - "benchmark_endpoint": "/v1/images/edits", - "server_params": { - "model": "tencent/HunyuanImage-3.0-Instruct", - "serve_args": { - "deploy-config-inline": { - "type": "yaml", - "content": { - "pipeline": "hunyuan_image_3_moe", - "async_chunk": false, - "trust_remote_code": true, - "connectors": { - "rdma_connector": { - "name": "MooncakeTransferEngineConnector", - "extra": { - "host": "auto", - "zmq_port": 50051, - "protocol": "rdma", - "device_name": "", - "memory_pool_size": 17179869184, - "memory_pool_device": "cpu" - } - } - }, - "stages": [ - { - "stage_id": 0, - "is_comprehension": true, - "final_output": true, - "final_output_type": "text", - "max_num_seqs": 16, - "gpu_memory_utilization": 0.95, - "enforce_eager": false, - "enable_prefix_caching": false, - "max_num_batched_tokens": 32768, - "devices": "0,1", - "tensor_parallel_size": 2, - "hf_overrides": { - "rope_parameters": { - "mrope_section": [0, 32, 32], - "rope_type": "default" - } - }, - "omni_kv_config": { - "need_send_cache": true - }, - "output_connectors": { - "to_stage_1": "rdma_connector" - }, - "default_sampling_params": { - "temperature": 0.0, - "top_p": 1, - "top_k": -1, - "max_tokens": 8192, - "detokenize": true, - "skip_special_tokens": false, - "stop_token_ids": [128025] - } - }, - { - "stage_id": 1, - "max_num_seqs": 1, - "enforce_eager": true, - "devices": "2,3", - "distributed_executor_backend": "mp", - "omni_kv_config": { - "need_recv_cache": true - }, - "parallel_config": { - "tensor_parallel_size": 2, - "vae_patch_parallel_size": 2, - "enable_expert_parallel": true - }, - "input_connectors": { - "from_stage_0": "rdma_connector" - }, - "default_sampling_params": { - "num_inference_steps": 50, - "guidance_scale": 0 - } - } - ], - "edges": [ - { - "from": 0, - "to": 1, - "window_size": -1, - "max_inflight": 1 - } - ] - } + { + "test_name": "test_hunyuan_image3_ti2i_inline", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "diffusion", + "local_model" + ], + "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config", + "server_type": "vllm-omni", + "benchmark_endpoint": "/v1/images/edits", + "server_params": { + "model": "tencent/HunyuanImage-3.0-Instruct", + "serve_args": { + "trust-remote-code": true, + "deploy-config-inline": { + "type": "yaml", + "content": { + "pipeline": "hunyuan_image_3_moe", + "async_chunk": false, + "trust_remote_code": true, + "connectors": { + "rdma_connector": { + "name": "MooncakeTransferEngineConnector", + "extra": { + "host": "auto", + "zmq_port": 50051, + "protocol": "rdma", + "device_name": "", + "memory_pool_size": 17179869184, + "memory_pool_device": "cpu" + } + } + }, + "stages": [ + { + "stage_id": 0, + "is_comprehension": true, + "final_output": true, + "final_output_type": "text", + "max_num_seqs": 16, + "gpu_memory_utilization": 0.95, + "enforce_eager": false, + "enable_prefix_caching": false, + "max_num_batched_tokens": 32768, + "devices": "0,1", + "tensor_parallel_size": 2, + "hf_overrides": { + "rope_parameters": { + "mrope_section": [ + 0, + 32, + 32 + ], + "rope_type": "default" + } }, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "inline_dataset_5_requests", - "dataset": "custom", - "dataset-path-inline": { - "type": "jsonl", - "content": [ - { - "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。", - "image_urls": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"], - "num_inference_steps": 8, - "seed": 42, - "guidance_scale": 1.0 - } - ] + "omni_kv_config": { + "need_send_cache": true + }, + "output_connectors": { + "to_stage_1": "rdma_connector" + }, + "default_sampling_params": { + "temperature": 0.0, + "top_p": 1, + "top_k": -1, + "max_tokens": 8192, + "detokenize": true, + "skip_special_tokens": false, + "stop_token_ids": [ + 128025 + ] + } + }, + { + "stage_id": 1, + "max_num_seqs": 1, + "enforce_eager": true, + "devices": "2,3", + "distributed_executor_backend": "mp", + "omni_kv_config": { + "need_recv_cache": true + }, + "parallel_config": { + "tensor_parallel_size": 2, + "vae_patch_parallel_size": 2, + "enable_expert_parallel": true + }, + "input_connectors": { + "from_stage_0": "rdma_connector" }, - "task": "ti2i", - "bot-task": "think_recaption", - "max-concurrency": 16, - "num-prompts": 96, - "skip-performance-assertion": true, - "baseline": { - "throughput_qps": 0.1, - "latency_p99": 44.46, - "latency_mean": 35.57, - "latency_std_dev": 0.5, - "peak_memory_mb_max": 70000, - "peak_memory_mb_mean": 70000 + "default_sampling_params": { + "num_inference_steps": 50, + "guidance_scale": 0 } + } + ], + "edges": [ + { + "from": 0, + "to": 1, + "window_size": -1, + "max_inflight": 1 + } + ] + } + }, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "inline_dataset_5_requests", + "dataset": "custom", + "dataset-path-inline": { + "type": "jsonl", + "content": [ + { + "prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。", + "image_urls": [ + "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png" + ], + "num_inference_steps": 8, + "seed": 42, + "guidance_scale": 1.0 } - ] - } + ] + }, + "task": "ti2i", + "bot-task": "think_recaption", + "max-concurrency": 16, + "num-prompts": 96, + "skip-performance-assertion": true, + "baseline": { + "throughput_qps": 0.1874, + "latency_p99": 101.852, + "latency_mean": 79.0385, + "peak_memory_mb_max": 70000, + "peak_memory_mb_mean": 70000 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json index e6d7ce65948..7a699f53a41 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json @@ -1,35 +1,47 @@ [ - { - "test_name": "test_hunyuan_image_tp2_cfgp2", - "description": "TP=2 CfgP=2 baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "tencent/HunyuanImage-3.0-Instruct", - "serve_args": { - "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", - "tensor-parallel-size": 2, - "cfg-parallel-size": 2, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1024x1024_steps8", - "dataset": "random", - "task": "t2i", - "width": 1024, - "height": 1024, - "num-inference-steps": 8, - "num-prompts": 10, - "max-concurrency": 1, - "skip-performance-assertion": true, - "baseline": { - "throughput_qps": 0.21, - "latency_p99": 4.7469, - "peak_memory_mb_max": 101100, - "peak_memory_mb_mean": 101100 - } - } - ] - } + { + "test_name": "test_hunyuan_image_tp2_cfgp2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "diffusion", + "local_model" + ], + "description": "TP=2 CfgP=2 baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "tencent/HunyuanImage-3.0-Instruct", + "serve_args": { + "deploy_config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", + "tensor-parallel-size": 2, + "cfg-parallel-size": 2, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1024x1024_steps8", + "dataset": "random", + "task": "t2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 8, + "num-prompts": 10, + "max-concurrency": 1, + "skip-performance-assertion": true, + "baseline": { + "throughput_qps": 0.21, + "latency_p99": 4.7469, + "peak_memory_mb_max": 101100, + "peak_memory_mb_mean": 101100 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json index fbbe1643e71..521febd2ded 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json @@ -1,35 +1,47 @@ [ - { - "test_name": "test_hunyuan_image_tp2_sp2", - "description": "TP=2 SP=2 baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "tencent/HunyuanImage-3.0-Instruct", - "serve_args": { - "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", - "tensor-parallel-size": 2, - "ulysses-degree": 2, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1024x1024_steps8", - "dataset": "random", - "task": "t2i", - "width": 1024, - "height": 1024, - "num-inference-steps": 8, - "num-prompts": 10, - "max-concurrency": 1, - "skip-performance-assertion": true, - "baseline": { - "throughput_qps": 0.20, - "latency_p99": 5.1025, - "peak_memory_mb_max": 97402, - "peak_memory_mb_mean": 97402 - } - } - ] - } + { + "test_name": "test_hunyuan_image_tp2_sp2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "diffusion", + "local_model" + ], + "description": "TP=2 SP=2 baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "tencent/HunyuanImage-3.0-Instruct", + "serve_args": { + "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", + "tensor-parallel-size": 2, + "ulysses-degree": 2, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1024x1024_steps8", + "dataset": "random", + "task": "t2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 8, + "num-prompts": 10, + "max-concurrency": 1, + "skip-performance-assertion": true, + "baseline": { + "throughput_qps": 0.20, + "latency_p99": 5.1025, + "peak_memory_mb_max": 97402, + "peak_memory_mb_mean": 97402 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json index 8f4c0ec5000..f2d05564b0f 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image_tp4.json +++ b/tests/dfx/perf/tests/test_hunyuan_image_tp4.json @@ -1,34 +1,47 @@ [ - { - "test_name": "test_hunyuan_image_tp4", - "description": "TP=4 baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "tencent/HunyuanImage-3.0-Instruct", - "serve_args": { - "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", - "tensor-parallel-size": 4, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1024x1024_steps8", - "dataset": "random", - "task": "t2i", - "width": 1024, - "height": 1024, - "num-inference-steps": 8, - "num-prompts": 10, - "max-concurrency": 1, - "skip-performance-assertion": true, - "baseline": { - "throughput_qps": 0.21, - "latency_p99": 4.9235, - "peak_memory_mb_max": 57576, - "peak_memory_mb_mean": 57576 - } - } - ] - } + { + "test_name": "test_hunyuan_image_tp4", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "diffusion", + "local_model" + ], + "description": "TP=4 baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "tencent/HunyuanImage-3.0-Instruct", + "serve_args": { + "deploy-config": "../vllm_omni/deploy/hunyuan_image3_dit.yaml", + "tensor-parallel-size": 4, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "1024x1024_steps8", + "dataset": "random", + "task": "t2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 8, + "num-prompts": 10, + "max-concurrency": 1, + "skip-performance-assertion": true, + "baseline": { + "throughput_qps": 0.213, + "latency_mean": 4.6953, + "latency_p99": 4.8601, + "peak_memory_mb_max": 56912.0, + "peak_memory_mb_mean": 56912.0 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json index 4a6f9e3501f..1a0f31e9809 100644 --- a/tests/dfx/perf/tests/test_ltx2_vllm_omni.json +++ b/tests/dfx/perf/tests/test_ltx2_vllm_omni.json @@ -1,217 +1,273 @@ [ - { - "test_name": "test_ltx2_baseline_eager", - "description": "Single-device baseline with enforce-eager (no torch.compile)", - "server_type": "vllm-omni", - "server_params": { - "model": "Lightricks/LTX-2", - "serve_args": { - "enforce-eager": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "256x256_145f_steps6", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 256, - "height": 256, - "num-frames": 145, - "fps": 24, - "num-inference-steps": 6, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - }, - { - "name": "480x768_41f_steps20", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 768, - "height": 480, - "num-frames": 41, - "fps": 24, - "num-inference-steps": 20, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - } - ] + { + "test_name": "test_ltx2_baseline_eager", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline with enforce-eager (no torch.compile)", + "server_type": "vllm-omni", + "server_params": { + "model": "Lightricks/LTX-2", + "serve_args": { + "enforce-eager": true, + "enable-diffusion-pipeline-profiler": true + } }, - - { - "test_name": "test_ltx2_torch_compile", - "description": "Single-device with torch.compile (default, no enforce-eager)", - "server_type": "vllm-omni", - "server_params": { - "model": "Lightricks/LTX-2", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "256x256_145f_steps6", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 256, - "height": 256, - "num-frames": 145, - "fps": 24, - "num-inference-steps": 6, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - }, - { - "name": "480x768_41f_steps20", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 768, - "height": 480, - "num-frames": 41, - "fps": 24, - "num-inference-steps": 20, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - } - ] + "benchmark_params": [ + { + "name": "256x256_145f_steps6", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 256, + "height": 256, + "num-frames": 145, + "fps": 24, + "num-inference-steps": 6, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + }, + { + "name": "480x768_41f_steps20", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 768, + "height": 480, + "num-frames": 41, + "fps": 24, + "num-inference-steps": 20, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + } + ] + }, + { + "test_name": "test_ltx2_torch_compile", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device with torch.compile (default, no enforce-eager)", + "server_type": "vllm-omni", + "server_params": { + "model": "Lightricks/LTX-2", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } }, - - { - "test_name": "test_ltx2_cfg2_eager", - "description": "CFG-parallel=2 with enforce-eager", - "server_type": "vllm-omni", - "server_params": { - "model": "Lightricks/LTX-2", - "serve_args": { - "cfg-parallel-size": 2, - "enforce-eager": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "256x256_145f_steps6", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 256, - "height": 256, - "num-frames": 145, - "fps": 24, - "num-inference-steps": 6, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - }, - { - "name": "480x768_41f_steps20", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 768, - "height": 480, - "num-frames": 41, - "fps": 24, - "num-inference-steps": 20, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - } - ] + "benchmark_params": [ + { + "name": "256x256_145f_steps6", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 256, + "height": 256, + "num-frames": 145, + "fps": 24, + "num-inference-steps": 6, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + }, + { + "name": "480x768_41f_steps20", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 768, + "height": 480, + "num-frames": 41, + "fps": 24, + "num-inference-steps": 20, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + } + ] + }, + { + "test_name": "test_ltx2_cfg2_eager", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "CFG-parallel=2 with enforce-eager", + "server_type": "vllm-omni", + "server_params": { + "model": "Lightricks/LTX-2", + "serve_args": { + "cfg-parallel-size": 2, + "enforce-eager": true, + "enable-diffusion-pipeline-profiler": true + } }, - - { - "test_name": "test_ltx2_cfg2_compile", - "description": "CFG-parallel=2 with torch.compile", - "server_type": "vllm-omni", - "server_params": { - "model": "Lightricks/LTX-2", - "serve_args": { - "cfg-parallel-size": 2, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "256x256_145f_steps6", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 256, - "height": 256, - "num-frames": 145, - "fps": 24, - "num-inference-steps": 6, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - }, - { - "name": "480x768_41f_steps20", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 768, - "height": 480, - "num-frames": 41, - "fps": 24, - "num-inference-steps": 20, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - } - ] + "benchmark_params": [ + { + "name": "256x256_145f_steps6", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 256, + "height": 256, + "num-frames": 145, + "fps": 24, + "num-inference-steps": 6, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + }, + { + "name": "480x768_41f_steps20", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 768, + "height": 480, + "num-frames": 41, + "fps": 24, + "num-inference-steps": 20, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + } + ] + }, + { + "test_name": "test_ltx2_cfg2_compile", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "CFG-parallel=2 with torch.compile", + "server_type": "vllm-omni", + "server_params": { + "model": "Lightricks/LTX-2", + "serve_args": { + "cfg-parallel-size": 2, + "enable-diffusion-pipeline-profiler": true + } }, - - { - "test_name": "test_ltx2_cache_dit_eager", - "description": "CacheDiT with enforce-eager", - "server_type": "vllm-omni", - "server_params": { - "model": "Lightricks/LTX-2", - "serve_args": { - "cache-backend": "cache_dit", - "enforce-eager": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "256x256_145f_steps6", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 256, - "height": 256, - "num-frames": 145, - "fps": 24, - "num-inference-steps": 6, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - }, - { - "name": "480x768_41f_steps20", - "dataset": "random", - "task": "t2v", - "backend": "v1/videos", - "width": 768, - "height": 480, - "num-frames": 41, - "fps": 24, - "num-inference-steps": 20, - "num-prompts": 3, - "max-concurrency": 1, - "enable-negative-prompt": true - } - ] - } + "benchmark_params": [ + { + "name": "256x256_145f_steps6", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 256, + "height": 256, + "num-frames": 145, + "fps": 24, + "num-inference-steps": 6, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + }, + { + "name": "480x768_41f_steps20", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 768, + "height": 480, + "num-frames": 41, + "fps": 24, + "num-inference-steps": 20, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + } + ] + }, + { + "test_name": "test_ltx2_cache_dit_eager", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "CacheDiT with enforce-eager", + "server_type": "vllm-omni", + "server_params": { + "model": "Lightricks/LTX-2", + "serve_args": { + "cache-backend": "cache_dit", + "enforce-eager": true, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "256x256_145f_steps6", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 256, + "height": 256, + "num-frames": 145, + "fps": 24, + "num-inference-steps": 6, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + }, + { + "name": "480x768_41f_steps20", + "dataset": "random", + "task": "t2v", + "backend": "v1/videos", + "width": 768, + "height": 480, + "num-frames": 41, + "fps": 24, + "num-inference-steps": 20, + "num-prompts": 3, + "max-concurrency": 1, + "enable-negative-prompt": true + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json index bc954cee59f..66ffa231f71 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json @@ -1,331 +1,343 @@ [ - { - "test_name": "test_qwen3_omni_chunk", - "server_params": { - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "extra_cli_args": [ - "--async-chunk" - ] + { + "test_name": "test_qwen3_omni_chunk", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "omni" + ], + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "extra_cli_args": [ + "--async-chunk" + ] + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], + "random_input_len": 2500, + "random_output_len": 900, + "ignore_eos": true, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 135.4379, + 187.0896, + 400, + 435.675, + 724.6836 + ], + "mean_e2el_ms": [ + 17619.9851, + 34811.8335, + 43656.9501, + 56657.1548, + 85097.6016 + ], + "mean_audio_rtf": [ + 0.1279, + 0.2238, + 0.2735, + 0.3712, + 0.5634 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 1, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "audio": 1 }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 4, - 16, - 32, - 64, - 128 - ], - "max_concurrency": [ - 1, - 4, - 8, - 16, - 32 - ], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 400, - 300, - 400, - 600, - 900 - ], - "mean_e2el_ms": [ - 26000, - 43000, - 52000, - 72000, - 109000 - ], - "mean_audio_rtf": [ - 0.15, - 0.30, - 0.35, - 0.50, - 0.75 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 10 - ], - "request_rate": [ - 0.1 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 1, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "audio": 1 - }, - "random_mm_bucket_config": { - "(0, 60, 3)": 1.0 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 600 - ], - "mean_e2el_ms": [ - 6000 - ], - "mean_audio_rtf": [ - 0.20 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 40 - ], - "request_rate": [ - 0.5 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 2, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.5, - "(720, 1280, 2)": 0.5 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 300 - ], - "mean_e2el_ms": [ - 7000 - ], - "mean_audio_rtf": [ - 0.25 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 100 - ], - "request_rate": [ - 1.0 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 3, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1, - "audio": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.34, - "(720, 1280, 2)": 0.33, - "(0, 60, 3)": 0.33 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 700 - ], - "mean_e2el_ms": [ - 11000 - ], - "mean_audio_rtf": [ - 0.35 - ] - } - }, - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 4, - 16, - 32, - 64, - 128 - ], - "max_concurrency": [ - 1, - 4, - 8, - 16, - 32 - ], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 100, - 300, - 500, - 800, - 1500 - ], - "mean_e2el_ms": [ - 5000, - 8000, - 10000, - 14000, - 20000 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 10 - ], - "request_rate": [ - 0.1 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 1, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "audio": 1 - }, - "random_mm_bucket_config": { - "(0, 60, 3)": 1.0 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 500 - ], - "mean_e2el_ms": [ - 1000 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 40 - ], - "request_rate": [ - 0.5 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 2, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.5, - "(720, 1280, 2)": 0.5 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 200 - ], - "mean_e2el_ms": [ - 800 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 100 - ], - "request_rate": [ - 1.0 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 3, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1, - "audio": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.34, - "(720, 1280, 2)": 0.33, - "(0, 60, 3)": 0.33 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 600 - ], - "mean_e2el_ms": [ - 1300 - ] - } - } - ] - } + "random_mm_bucket_config": { + "(0, 60, 3)": 1.0 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 340.329 + ], + "mean_e2el_ms": [ + 4305.477 + ], + "mean_audio_rtf": [ + 0.1305 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 2, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.5, + "(720, 1280, 2)": 0.5 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 220.0909 + ], + "mean_e2el_ms": [ + 5475.9565 + ], + "mean_audio_rtf": [ + 0.1818 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 3, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1, + "audio": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.34, + "(720, 1280, 2)": 0.33, + "(0, 60, 3)": 0.33 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 450.0155 + ], + "mean_e2el_ms": [ + 8516.8316 + ], + "mean_audio_rtf": [ + 0.2717 + ] + } + }, + { + "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], + "random_input_len": 2500, + "random_output_len": 900, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 78.4278, + 195.8037, + 333.7318, + 561.805, + 1038.043 + ], + "mean_e2el_ms": [ + 4288.3029, + 7429.9429, + 9451.6184, + 12607.5292, + 18153.2547 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "random_mm_base_items_per_request": 1, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "audio": 1 + }, + "random_mm_bucket_config": { + "(0, 60, 3)": 1.0 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 275.3875 + ], + "mean_e2el_ms": [ + 730.6995 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "random_mm_base_items_per_request": 2, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.5, + "(720, 1280, 2)": 0.5 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 166.2201 + ], + "mean_e2el_ms": [ + 673.5789 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "random_mm_base_items_per_request": 3, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1, + "audio": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.34, + "(720, 1280, 2)": 0.33, + "(0, 60, 3)": 0.33 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 357.7781 + ], + "mean_e2el_ms": [ + 953.3493 + ] + } + } + ] + } ] 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 bee799a776f..ac6f679c16e 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json @@ -1,115 +1,127 @@ [ - { - "test_name": "test_qwen3_omni_3gpu_replica2", - "server_params": { - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_overrides": { - "1": {"num_replicas": 2, "devices": "1,2"}, - "2": {"num_replicas": 2, "devices": "1,2"} - }, - "extra_cli_args": [ - "--async-chunk" - ] + { + "test_name": "test_qwen3_omni_3gpu_replica2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 3 + } + }, + "full_model", + "omni" + ], + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "stage_overrides": { + "1": {"num_replicas": 2, "devices": "1,2"}, + "2": {"num_replicas": 2, "devices": "1,2"} + }, + "extra_cli_args": [ + "--async-chunk" + ] + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [128, 128, 128, 128], + "max_concurrency": [8, 16, 24, 32], + "random_input_len": 2500, + "random_output_len": 900, + "ignore_eos": true, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "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] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [10], + "request_rate": [0.1], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 1, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "audio": 1 }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [128, 128, 128, 128], - "max_concurrency": [8, 16, 24, 32], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "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] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [10], - "request_rate": [0.1], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 1, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "audio": 1 - }, - "random_mm_bucket_config": { - "(0, 60, 3)": 1.0 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [377], - "mean_audio_ttfp_ms": [684], - "mean_audio_rtf": [0.14] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [40], - "request_rate": [0.5], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 2, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.5, - "(720, 1280, 2)": 0.5 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [228], - "mean_audio_ttfp_ms": [686], - "mean_audio_rtf": [0.16] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [100], - "request_rate": [1.0], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 3, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1, - "audio": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.34, - "(720, 1280, 2)": 0.33, - "(0, 60, 3)": 0.33 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [522], - "mean_audio_ttfp_ms": [1119], - "mean_audio_rtf": [0.21] - } - } - ] - } + "random_mm_bucket_config": { + "(0, 60, 3)": 1.0 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [377], + "mean_audio_ttfp_ms": [684], + "mean_audio_rtf": [0.14] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [40], + "request_rate": [0.5], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 2, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.5, + "(720, 1280, 2)": 0.5 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [228], + "mean_audio_ttfp_ms": [686], + "mean_audio_rtf": [0.16] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [100], + "request_rate": [1.0], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 3, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1, + "audio": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.34, + "(720, 1280, 2)": 0.33, + "(0, 60, 3)": 0.33 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [522], + "mean_audio_ttfp_ms": [1119], + "mean_audio_rtf": [0.21] + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json index e57acecba4a..78daf691b9e 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json @@ -1,170 +1,182 @@ [ - { - "test_name": "test_qwen3_omni", - "server_params": { - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "extra_cli_args": [ - "--no-async-chunk" - ] + { + "test_name": "test_qwen3_omni", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "omni" + ], + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "extra_cli_args": [ + "--no-async-chunk" + ] + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], + "random_input_len": 2500, + "random_output_len": 900, + "ignore_eos": true, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 109.1218, + 247.4411, + 354.9384, + 447.7885, + 694.7487 + ], + "mean_e2el_ms": [ + 22413.667, + 30808.9715, + 35677.844, + 46345.8978, + 68350.4853 + ], + "mean_audio_rtf": [ + 0.1695, + 0.2033, + 0.2467, + 0.3234, + 0.4438 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 1, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "audio": 1 }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 4, - 16, - 32, - 64, - 128 - ], - "max_concurrency": [ - 1, - 4, - 8, - 16, - 32 - ], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 400, - 300, - 400, - 600, - 900 - ], - "mean_e2el_ms": [ - 32000, - 34000, - 45000, - 61000, - 90000 - ], - "mean_audio_rtf": [ - 0.20, - 0.25, - 0.30, - 0.45, - 0.60 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 10 - ], - "request_rate": [ - 0.1 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 1, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "audio": 1 - }, - "random_mm_bucket_config": { - "(0, 60, 3)": 1.0 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 500 - ], - "mean_e2el_ms": [ - 6000 - ], - "mean_audio_rtf": [ - 0.20 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 40 - ], - "request_rate": [ - 0.5 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 2, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.5, - "(720, 1280, 2)": 0.5 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 300 - ], - "mean_e2el_ms": [ - 7000 - ], - "mean_audio_rtf": [ - 0.20 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "num_prompts": [ - 100 - ], - "request_rate": [ - 1.0 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "random_mm_base_items_per_request": 3, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1, - "audio": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.34, - "(720, 1280, 2)": 0.33, - "(0, 60, 3)": 0.33 - }, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_ttft_ms": [ - 700 - ], - "mean_e2el_ms": [ - 9000 - ], - "mean_audio_rtf": [ - 0.30 - ] - } - } - ] - } + "random_mm_bucket_config": { + "(0, 60, 3)": 1.0 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 323.3573 + ], + "mean_e2el_ms": [ + 4623.4173 + ], + "mean_audio_rtf": [ + 0.1406 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 2, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.5, + "(720, 1280, 2)": 0.5 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 215.4502 + ], + "mean_e2el_ms": [ + 4924.149 + ], + "mean_audio_rtf": [ + 0.1637 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "random_mm_base_items_per_request": 3, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1, + "audio": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.34, + "(720, 1280, 2)": 0.33, + "(0, 60, 3)": 0.33 + }, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_ttft_ms": [ + 442.5815 + ], + "mean_e2el_ms": [ + 6441.0986 + ], + "mean_audio_rtf": [ + 0.2091 + ] + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json index a973bb0e64b..b9d5530987a 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json +++ b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json @@ -1,180 +1,192 @@ [ - { - "test_name": "test_qwen3_omni_vllm_text", - "server_params": { - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "use_omni": false, - "extra_cli_args": [ - "--no-enable-prefix-caching", - "--skip-mm-profiling" - ] + { + "test_name": "test_qwen3_omni_vllm_text", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "omni" + ], + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "use_omni": false, + "extra_cli_args": [ + "--no-enable-prefix-caching", + "--skip-mm-profiling" + ] + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [ + 4, + 16, + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8, + 16, + 32 + ], + "random_input_len": 2500, + "random_output_len": 900, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "vllm", - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "endpoint": "/v1/completions", - "num_prompts": [ - 4, - 16, - 32, - 64, - 128 - ], - "max_concurrency": [ - 1, - 4, - 8, - 16, - 32 - ], - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 300, - 300, - 1000, - 600, - 800 - ], - "mean_e2el_ms": [ - 6000, - 8000, - 11000, - 13000, - 17000 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "vllm", - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "endpoint": "/v1/completions", - "num_prompts": [ - 10 - ], - "request_rate": [ - 0.1 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 1, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "audio": 1 - }, - "random_mm_bucket_config": { - "(0, 60, 3)": 1.0 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 30 - ], - "mean_e2el_ms": [ - 540 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "vllm", - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "endpoint": "/v1/completions", - "num_prompts": [ - 40 - ], - "request_rate": [ - 0.5 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 2, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.5, - "(720, 1280, 2)": 0.5 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 30 - ], - "mean_e2el_ms": [ - 600 - ] - } - }, - { - "dataset_name": "random-mm", - "backend": "vllm", - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "endpoint": "/v1/completions", - "num_prompts": [ - 100 - ], - "request_rate": [ - 1.0 - ], - "random_input_len": 100, - "random_output_len": 100, - "random_range_ratio": 0.0, - "ignore_eos": true, - "extra_body": { - "modalities": [ - "text" - ] - }, - "random_mm_base_items_per_request": 3, - "random_mm_num_mm_items_range_ratio": 0.5, - "random_mm_limit_mm_per_prompt": { - "image": 1, - "video": 1, - "audio": 1 - }, - "random_mm_bucket_config": { - "(256, 256, 1)": 0.34, - "(720, 1280, 2)": 0.33, - "(0, 60, 3)": 0.33 - }, - "percentile-metrics": "ttft,tpot,itl,e2el", - "baseline": { - "mean_ttft_ms": [ - 40 - ], - "mean_e2el_ms": [ - 700 - ] - } - } - ] - } + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 69.1276, + 255.1882, + 310.639, + 422.0337, + 532.0582 + ], + "mean_e2el_ms": [ + 4382.7211, + 7510.3815, + 9440.4854, + 12553.0186, + 16932.6065 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [ + 10 + ], + "request_rate": [ + 0.1 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "random_mm_base_items_per_request": 1, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "audio": 1 + }, + "random_mm_bucket_config": { + "(0, 60, 3)": 1.0 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 26.3381 + ], + "mean_e2el_ms": [ + 485.2436 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [ + 40 + ], + "request_rate": [ + 0.5 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "random_mm_base_items_per_request": 2, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.5, + "(720, 1280, 2)": 0.5 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 26.4517 + ], + "mean_e2el_ms": [ + 517.2017 + ] + } + }, + { + "dataset_name": "random-mm", + "backend": "vllm", + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "endpoint": "/v1/completions", + "num_prompts": [ + 100 + ], + "request_rate": [ + 1.0 + ], + "random_input_len": 100, + "random_output_len": 100, + "random_range_ratio": 0.0, + "ignore_eos": true, + "extra_body": { + "modalities": [ + "text" + ] + }, + "random_mm_base_items_per_request": 3, + "random_mm_num_mm_items_range_ratio": 0.5, + "random_mm_limit_mm_per_prompt": { + "image": 1, + "video": 1, + "audio": 1 + }, + "random_mm_bucket_config": { + "(256, 256, 1)": 0.34, + "(720, 1280, 2)": 0.33, + "(0, 60, 3)": 0.33 + }, + "percentile-metrics": "ttft,tpot,itl,e2el", + "baseline": { + "mean_ttft_ms": [ + 27.6972 + ], + "mean_e2el_ms": [ + 561.3684 + ] + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json index 1903bc6708f..53e13daedf6 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json @@ -1,149 +1,185 @@ [ - { - "test_name": "test_qwen_image_edit_2511_single_device", - "description": "Single-device baseline (two input images)", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit-2511", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20_i2i_2img", - "dataset": "random", - "task": "i2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 2, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.06, - "latency_mean": 20, - "peak_memory_mb_max": 57000, - "peak_memory_mb_mean": 57000 - } - }, - { - "name": "1536x1536_steps35_i2i_2img", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 2, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.01, - "latency_mean": 60, - "peak_memory_mb_max": 67000, - "peak_memory_mb_mean": 67000 - } - } - ] + { + "test_name": "test_qwen_image_edit_2511_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline (two input images)", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Edit-2511", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4", - "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit-2511", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "vae-patch-parallel-size": 4, - "vae-use-tiling": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1536x1536_steps35_i2i_2img", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 2, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.05, - "latency_mean": 20, - "peak_memory_mb_max": 55000, - "peak_memory_mb_mean": 55000 - } - } - ] + "benchmark_params": [ + { + "name": "512x512_steps20_i2i_2img", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 2, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0721, + "latency_mean": 13.8398, + "peak_memory_mb_max": 54731.4, + "peak_memory_mb_mean": 54731.4 + } + }, + { + "name": "1536x1536_steps35_i2i_2img", + "dataset": "random", + "task": "i2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 2, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0186, + "latency_mean": 53.6942, + "peak_memory_mb_max": 64402.4, + "peak_memory_mb_mean": 64402.4 + } + } + ] + }, + { + "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Edit-2511", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit", - "description": "Ulysses SP=2 + CFG=2 + CacheDiT", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit-2511", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "cache-backend": "cache_dit", - "cache-config": { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": false, - "taylorseer_order": 1, - "scm_steps_mask_policy": null, - "scm_steps_policy": "dynamic" - }, - "enable-diffusion-pipeline-profiler": true - } + "benchmark_params": [ + { + "name": "1536x1536_steps35_i2i_2img", + "dataset": "random", + "task": "i2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 2, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0538, + "latency_mean": 18.5591, + "peak_memory_mb_max": 53485.0, + "peak_memory_mb_mean": 53485.0 + } + } + ] + }, + { + "test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "Ulysses SP=2 + CFG=2 + CacheDiT", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Edit-2511", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "cache-backend": "cache_dit", + "cache-config": { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": false, + "taylorseer_order": 1, + "scm_steps_mask_policy": null, + "scm_steps_policy": "dynamic" }, - "benchmark_params": [ - { - "name": "512x512_steps20_i2i_2img", - "dataset": "random", - "task": "i2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 2, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.21, - "latency_mean": 10, - "peak_memory_mb_max": 57000, - "peak_memory_mb_mean": 57000 - } - }, - { - "name": "1536x1536_steps35_i2i_2img", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 2, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.08, - "latency_mean": 20, - "peak_memory_mb_max": 66000, - "peak_memory_mb_mean": 66000 - } - } - ] - } + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps20_i2i_2img", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 2, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.2776, + "latency_mean": 3.5938, + "peak_memory_mb_max": 55214.0, + "peak_memory_mb_mean": 55214.0 + } + }, + { + "name": "1536x1536_steps35_i2i_2img", + "dataset": "random", + "task": "i2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 2, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.1051, + "latency_mean": 9.4918, + "peak_memory_mb_max": 64265.6, + "peak_memory_mb_mean": 64265.6 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json index 21695b7528c..bfda3073172 100644 --- a/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json @@ -1,144 +1,180 @@ [ - { - "test_name": "test_qwen_image_edit_single_device", - "description": "Single-device baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20_i2i", - "dataset": "random", - "task": "i2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.1, - "latency_mean": 10, - "peak_memory_mb_max": 57000, - "peak_memory_mb_mean": 57000 - } - }, - { - "name": "1536x1536_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.02, - "latency_mean": 50, - "peak_memory_mb_max": 65000, - "peak_memory_mb_mean": 65000 - } - } - ] + { + "test_name": "test_qwen_image_edit_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Edit", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4", - "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "vae-patch-parallel-size": 4, - "vae-use-tiling": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1536x1536_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.06, - "latency_mean": 20, - "peak_memory_mb_max": 56000, - "peak_memory_mb_mean": 56000 - } - } - ] + "benchmark_params": [ + { + "name": "512x512_steps20_i2i", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.1129, + "latency_mean": 8.8352, + "peak_memory_mb_max": 55356.5, + "peak_memory_mb_mean": 55356.5 + } + }, + { + "name": "1536x1536_steps35_i2i", + "dataset": "random", + "task": "i2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0240, + "latency_mean": 41.5686, + "peak_memory_mb_max": 62762.7, + "peak_memory_mb_mean": 62762.7 + } + } + ] + }, + { + "test_name": "test_qwen_image_edit_ulysses2_cfg2_vae_patch4", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Edit", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit", - "description": "Ulysses SP=2 + CFG=2 + CacheDiT", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Edit", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "cache-backend": "cache_dit", - "cache-config": { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": false, - "taylorseer_order": 1, - "scm_steps_mask_policy": null, - "scm_steps_policy": "dynamic" - }, - "enable-diffusion-pipeline-profiler": true - } + "benchmark_params": [ + { + "name": "1536x1536_steps35_i2i", + "dataset": "random", + "task": "i2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0672, + "latency_mean": 14.8517, + "peak_memory_mb_max": 53931.5, + "peak_memory_mb_mean": 53931.5 + } + } + ] + }, + { + "test_name": "test_qwen_image_edit_ulysses2_cfg2_cache_dit", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "Ulysses SP=2 + CFG=2 + CacheDiT", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Edit", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "cache-backend": "cache_dit", + "cache-config": { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": false, + "taylorseer_order": 1, + "scm_steps_mask_policy": null, + "scm_steps_policy": "dynamic" }, - "benchmark_params": [ - { - "name": "512x512_steps20_i2i", - "dataset": "random", - "task": "i2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.26, - "latency_mean": 10, - "peak_memory_mb_max": 58000, - "peak_memory_mb_mean": 58000 - } - }, - { - "name": "1536x1536_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.1, - "latency_mean": 10, - "peak_memory_mb_max": 67000, - "peak_memory_mb_mean": 67000 - } - } - ] - } + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps20_i2i", + "dataset": "random", + "task": "i2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.2879, + "latency_mean": 3.4701, + "peak_memory_mb_max": 55525.6, + "peak_memory_mb_mean": 55525.6 + } + }, + { + "name": "1536x1536_steps35_i2i", + "dataset": "random", + "task": "i2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.1113, + "latency_mean": 8.9656, + "peak_memory_mb_max": 63758.3, + "peak_memory_mb_mean": 63758.3 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json index c0cd3b7c015..3445dd94ca4 100644 --- a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -1,49 +1,61 @@ [ - { - "test_name": "test_qwen_image_layered_single_device", - "description": "Single-device baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image-Layered", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "640x640_steps20_i2i", - "dataset": "random", - "task": "i2i", - "width": 640, - "height": 640, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.05, - "latency_mean": 20, - "peak_memory_mb_max": 63000, - "peak_memory_mb_mean": 63000 - } - }, - { - "name": "1024x1024_steps35_i2i", - "dataset": "random", - "task": "i2i", - "width": 1024, - "height": 1024, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.03, - "latency_mean": 30, - "peak_memory_mb_max": 63000, - "peak_memory_mb_mean": 63000 - } - } - ] - } + { + "test_name": "test_qwen_image_layered_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Layered", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "640x640_steps20_i2i", + "dataset": "random", + "task": "i2i", + "width": 640, + "height": 640, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0703, + "latency_mean": 14.208, + "peak_memory_mb_max": 60543.5, + "peak_memory_mb_mean": 60543.5 + } + }, + { + "name": "1024x1024_steps35_i2i", + "dataset": "random", + "task": "i2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0417, + "latency_mean": 23.914, + "peak_memory_mb_max": 60543.5, + "peak_memory_mb_mean": 60543.5 + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index a6eacf88e59..b3daf758732 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -1,219 +1,279 @@ [ - { - "test_name": "test_qwen_image_single_device", - "description": "Single-device baseline (no parallelism)", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20", - "dataset": "random", - "task": "t2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.35, - "latency_mean": 3, - "peak_memory_mb_mean": 56000 - } - }, - { - "name": "1536x1536_steps35", - "dataset": "random", - "task": "t2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.03, - "latency_mean": 26, - "peak_memory_mb_mean": 65000 - } - } - ] + { + "test_name": "test_qwen_image_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline (no parallelism)", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_single_device_step_execution", - "description": "Single-device baseline (no parallelism) with step execution", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image", - "serve_args": { - "enable-diffusion-pipeline-profiler": true, - "step-execution": true - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20", - "dataset": "random", - "task": "t2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.35, - "latency_mean": 3, - "peak_memory_mb_mean": 56000 - } - }, - { - "name": "1536x1536_steps35", - "dataset": "random", - "task": "t2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.03, - "latency_mean": 26, - "peak_memory_mb_mean": 65000 - } - } - ] + "benchmark_params": [ + { + "name": "512x512_steps20", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.4153, + "latency_mean": 2.4065, + "peak_memory_mb_mean": 52289.9 + } + }, + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0423, + "latency_mean": 23.5991, + "peak_memory_mb_mean": 61360.5 + } + } + ] + }, + { + "test_name": "test_qwen_image_single_device_step_execution", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline (no parallelism) with step execution", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "enable-diffusion-pipeline-profiler": true, + "step-execution": true + } }, - { - "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4", - "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "vae-patch-parallel-size": 4, - "vae-use-tiling": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "1536x1536_steps35", - "dataset": "random", - "task": "t2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.11, - "latency_mean": 9, - "peak_memory_mb_mean": 55000 - } - } - ] + "benchmark_params": [ + { + "name": "512x512_steps20", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.4175, + "latency_mean": 2.3953, + "peak_memory_mb_mean": 56000 + } + }, + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.0423, + "latency_mean": 23.5768, + "peak_memory_mb_mean": 65000 + } + } + ] + }, + { + "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "vae-patch-parallel-size": 4, + "vae-use-tiling": true, + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit", - "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image", - "serve_args": { - "ulysses-degree": 2, - "cfg-parallel-size": 2, - "cache-backend": "cache_dit", - "cache-config": { - "Fn_compute_blocks": 1, - "Bn_compute_blocks": 0, - "max_warmup_steps": 4, - "residual_diff_threshold": 0.24, - "max_continuous_cached_steps": 3, - "enable_taylorseer": false, - "taylorseer_order": 1, - "scm_steps_mask_policy": null, - "scm_steps_policy": "dynamic" - }, - "enable-diffusion-pipeline-profiler": true - } + "benchmark_params": [ + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.1214, + "latency_mean": 8.2211, + "peak_memory_mb_mean": 51704.7 + } + } + ] + }, + { + "test_name": "test_qwen_image_ulysses2_cfg2_cache_dit", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 4 + } + }, + "full_model", + "diffusion" + ], + "description": "Ulysses SP=2 + CFG-parallel=2 + CacheDiT acceleration", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "ulysses-degree": 2, + "cfg-parallel-size": 2, + "cache-backend": "cache_dit", + "cache-config": { + "Fn_compute_blocks": 1, + "Bn_compute_blocks": 0, + "max_warmup_steps": 4, + "residual_diff_threshold": 0.24, + "max_continuous_cached_steps": 3, + "enable_taylorseer": false, + "taylorseer_order": 1, + "scm_steps_mask_policy": null, + "scm_steps_policy": "dynamic" }, - "benchmark_params": [ - { - "name": "512x512_steps20", - "dataset": "random", - "task": "t2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.46, - "latency_mean": 3, - "peak_memory_mb_mean": 56000 - } - }, - { - "name": "1536x1536_steps35", - "dataset": "random", - "task": "t2i", - "width": 1536, - "height": 1536, - "num-inference-steps": 35, - "num-prompts": 10, - "max-concurrency": 1, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": 0.18, - "latency_mean": 6, - "peak_memory_mb_mean": 67000 - } - } - ] + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_qwen_image_single_device_step_execution_high_concurrency", - "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep", - "server_type": "vllm-omni", - "server_params": { - "model": "Qwen/Qwen-Image", - "serve_args": { - "enable-diffusion-pipeline-profiler": true, - "step-execution": true, - "max-num-seqs": 8 - } - }, - "benchmark_params": [ - { - "name": "512x512_steps20_high_concurrency", - "dataset": "random", - "task": "t2i", - "width": 512, - "height": 512, - "num-inference-steps": 20, - "num-prompts": [20, 40, 80, 160], - "max-concurrency": [1, 2, 4, 8], - "warmup-requests": 8, - "warmup-concurrency": 8, - "warmup-num-inference-steps": 20, - "enable-negative-prompt": true, - "baseline": { - "throughput_qps": [0.30, 0.50, 0.60, 0.70], - "latency_mean": [3.50, 4.00, 7.00, 12.00], - "peak_memory_mb_mean": [67000, 67000, 67000, 67000] - } - } - ] - } + "benchmark_params": [ + { + "name": "512x512_steps20", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.559, + "latency_mean": 1.794, + "peak_memory_mb_mean": 52397.6775 + } + }, + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.1918, + "latency_mean": 5.2022, + "peak_memory_mb_mean": 61468.8 + } + } + ] + }, + { + "test_name": "test_qwen_image_single_device_step_execution_high_concurrency", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "enable-diffusion-pipeline-profiler": true, + "step-execution": true, + "max-num-seqs": 8 + } + }, + "benchmark_params": [ + { + "name": "512x512_steps20_high_concurrency", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": [20, 40, 80, 160], + "max-concurrency": [1, 2, 4, 8], + "warmup-requests": 8, + "warmup-concurrency": 8, + "warmup-num-inference-steps": 20, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": [0.4125, 0.6714, 0.6674, 0.8159], + "latency_mean": [2.4215, 2.992, 5.9293, 9.656], + "peak_memory_mb_mean": [67000, 67000, 67000, 67000] + } + } + ] + } ] diff --git a/tests/dfx/perf/tests/test_runner_metadata.py b/tests/dfx/perf/tests/test_runner_metadata.py index a195c4546a9..71d4aa95bf0 100644 --- a/tests/dfx/perf/tests/test_runner_metadata.py +++ b/tests/dfx/perf/tests/test_runner_metadata.py @@ -86,3 +86,202 @@ def test_enabled_false_entry_is_skipped(): # Only the enabled=True entry should appear assert len(params) == 1 assert params[0].get("task") == "voice_clone" + + +def test_resolve_pytest_marks_hardware_dict_with_extra(): + from tests.dfx.conftest import resolve_pytest_marks + + marks = resolve_pytest_marks( + [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 2}}, + "full_model", + "diffusion", + "local_model", + ] + ) + names = {m.name for m in marks} + assert "H100" in names + assert "cuda" in names + assert "gpu" in names + assert "distributed_cuda" in names + assert "full_model" in names + assert "diffusion" in names + assert "local_model" in names + + +def test_resolve_pytest_marks_rejects_legacy_object_format(): + from tests.dfx.conftest import resolve_pytest_marks + + with pytest.raises(ValueError, match="mark must be a list"): + resolve_pytest_marks( + { + "hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}, + "marks": ["full_model"], + } + ) + + +def test_extract_mark_resource_label(): + from tests.dfx.conftest import extract_mark_resource_label + + assert extract_mark_resource_label(None) == "na" + assert ( + extract_mark_resource_label( + [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, + "full_model", + ] + ) + == "H100" + ) + assert ( + extract_mark_resource_label( + [ + { + "hardware_marks": { + "res": {"cuda": "H100", "rocm": "MI325"}, + "num_cards": 2, + } + }, + "diffusion", + ] + ) + == "H100-MI325" + ) + + +def test_resource_label_for_filename(): + from tests.dfx.conftest import resource_label_for_filename + + assert resource_label_for_filename("H100") == "" + assert resource_label_for_filename("L4") == "L4" + assert resource_label_for_filename("910B") == "910B" + assert resource_label_for_filename("na") == "na" + + +def test_hardware_json_value(): + from tests.dfx.conftest import hardware_json_value + + assert hardware_json_value("H100") == "H100" + assert hardware_json_value("na") == "" + assert hardware_json_value(None) == "" + + +def test_extract_configs_resource_label(monkeypatch): + from tests.dfx.conftest import extract_configs_resource_label, get_runtime_resource_label + + monkeypatch.setattr( + "tests.dfx.conftest._read_runtime_device_name", + lambda *, device_id=0: "NVIDIA H100 80GB HBM3", + ) + get_runtime_resource_label(refresh=True) + assert extract_configs_resource_label([]) == "H100" + get_runtime_resource_label(refresh=True) + monkeypatch.setattr( + "tests.dfx.conftest._read_runtime_device_name", + lambda *, device_id=0: "Ascend910B2", + ) + assert get_runtime_resource_label(refresh=True) == "910B" + + +def test_load_benchmark_configs_from_dir(tmp_path): + from tests.dfx.conftest import load_benchmark_configs + + (tmp_path / "a.json").write_text( + json.dumps([{"test_name": "test_a", "server_params": {"model": "m/a"}, "benchmark_params": []}]), + encoding="utf-8", + ) + (tmp_path / "b.json").write_text( + json.dumps([{"test_name": "test_b", "server_params": {"model": "m/b"}, "benchmark_params": []}]), + encoding="utf-8", + ) + configs = load_benchmark_configs(config_dir=tmp_path) + assert [c["test_name"] for c in configs] == ["test_a", "test_b"] + + +def test_create_unique_server_pytest_params_applies_marks(tmp_path): + from tests.dfx.conftest import create_unique_server_pytest_params + + configs = [ + { + "test_name": "test_with_mark", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, + "full_model", + ], + "server_params": {"model": "some/model"}, + "benchmark_params": [{"name": "p0", "num_prompts": 1}], + }, + { + "test_name": "test_without_mark", + "server_params": {"model": "other/model"}, + "benchmark_params": [{"name": "p0", "num_prompts": 1}], + }, + ] + params = create_unique_server_pytest_params(configs, tmp_path) + by_id = {p.id: p for p in params} + assert len(by_id["test_with_mark"].values) == 1 + assert isinstance(by_id["test_with_mark"].values[0], tuple) + assert any(m.name == "H100" for m in by_id["test_with_mark"].marks) + assert not any(m.name == "H100" for m in by_id["test_without_mark"].marks) + + +def test_is_diffusion_perf_config(): + from tests.dfx.conftest import is_diffusion_perf_config + + assert not is_diffusion_perf_config( + {"test_name": "omni_a", "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "omni"]} + ) + assert is_diffusion_perf_config( + { + "test_name": "diff_a", + "server_type": "vllm-omni", + "mark": [{"hardware_marks": {"res": {"cuda": "H100"}}}, "diffusion"], + } + ) + + +def test_benchmark_param_id_suffix_from_task_eval_phase(): + from tests.dfx.conftest import _unique_benchmark_param_id_suffixes + + params = [ + {"task": "default_voice", "eval_phase": "latency"}, + {"task": "default_voice", "eval_phase": "throughput"}, + ] + assert _unique_benchmark_param_id_suffixes(params) == [ + "default_voice_latency", + "default_voice_throughput", + ] + + +def test_create_paired_omni_benchmark_pytest_params(tmp_path): + from tests.dfx.conftest import create_paired_omni_benchmark_pytest_params + + configs = [ + { + "test_name": "test_omni", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, + "omni", + ], + "server_params": {"model": "m/omni"}, + "benchmark_params": [{"name": "p0", "num_prompts": 1}, {"name": "p1", "num_prompts": 2}], + }, + { + "test_name": "test_tts", + "mark": [ + {"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 1}}, + "tts", + ], + "server_params": {"model": "m/tts"}, + "benchmark_params": [{"name": "p0", "num_prompts": 1}], + }, + ] + params = create_paired_omni_benchmark_pytest_params(configs, tmp_path) + by_id = {p.id: p for p in params} + assert set(by_id) == {"test_omni-p0", "test_omni-p1", "test_tts-p0"} + omni_row, bench_row = by_id["test_tts-p0"].values + assert omni_row[0] == "test_tts" + assert bench_row == ("test_tts", 0) + assert any(m.name == "tts" for m in by_id["test_tts-p0"].marks) + assert not any(m.name == "tts" for m in by_id["test_omni-p0"].marks) diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json index 9246b502f2f..85236eab735 100644 --- a/tests/dfx/perf/tests/test_tts.json +++ b/tests/dfx/perf/tests/test_tts.json @@ -1,6 +1,19 @@ [ { "test_name": "test_qwen3_tts_base", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100", + "npu": "A3" + }, + "num_cards": 2 + } + }, + "full_model", + "tts" + ], "server_params": { "model": "Qwen/Qwen3-TTS-12Hz-1.7B-Base", "stage_overrides": { @@ -27,10 +40,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 3000 + 139.9434 ], "median_audio_rtf": [ - 0.6 + 0.1683 ] } }, @@ -55,19 +68,19 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 1000, - 3000, - 14000 + 227.5653, + 647.1272, + 4777.2714 ], "median_audio_rtf": [ - 0.70, - 1.15, - 3.90 + 0.3313, + 0.3642, + 1.2282 ], "audio_throughput": [ - 12, - 14, - 14 + 23.501, + 40.634, + 43.0758 ] } }, @@ -98,6 +111,18 @@ }, { "test_name": "test_qwen3_tts_customvoice", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "tts" + ], "server_params": { "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", "stage_overrides": { @@ -129,10 +154,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 2000 + 49.7908 ], "median_audio_rtf": [ - 0.40 + 0.1425 ] } }, @@ -162,19 +187,19 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 300, - 1000, - 5000 + 75.3341, + 713.7727, + 5941.7179 ], "median_audio_rtf": [ - 0.25, - 0.35, - 1.10 + 0.1859, + 0.3148, + 1.1662 ], "audio_throughput": [ - 30, - 40, - 40 + 36.465, + 49.6268, + 48.3091 ] } }, @@ -228,13 +253,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 100 + 56.6977 ], "median_audio_rtf": [ - 0.20 + 0.1555 ], "audio_throughput": [ - 10 + 11.5292 ] } }, @@ -259,10 +284,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 100 + 52.0447 ], "median_audio_rtf": [ - 0.15 + 0.1452 ] } }, @@ -291,19 +316,19 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 100, - 1000, - 5000 + 75.5967, + 587.5877, + 4960.3726 ], "median_audio_rtf": [ - 0.25, - 0.40, - 1.20 + 0.1871, + 0.3143, + 1.1703 ], "audio_throughput": [ - 30, - 40, - 40 + 40.2198, + 48.9046, + 48.0505 ] } }, @@ -328,13 +353,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 100 + 56.3961 ], "median_audio_rtf": [ - 0.20 + 0.1542 ], "audio_throughput": [ - 9 + 9.6672 ] } } @@ -342,6 +367,18 @@ }, { "test_name": "test_voxtral_tts", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "tts" + ], "server_params": { "model": "mistralai/Voxtral-4B-TTS-2603" }, @@ -353,14 +390,24 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [20], - "max_concurrency": [1], + "num_prompts": [ + 20 + ], + "max_concurrency": [ + 1 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "casual_male"}, + "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] + "median_audio_ttfp_ms": [ + 150 + ], + "median_audio_rtf": [ + 0.15 + ] } }, { @@ -370,15 +417,27 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [80], - "max_concurrency": [8], + "num_prompts": [ + 80 + ], + "max_concurrency": [ + 8 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "casual_male"}, + "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] + "median_audio_ttfp_ms": [ + 1500 + ], + "median_audio_rtf": [ + 0.3 + ], + "audio_throughput": [ + 30.0 + ] } }, { @@ -388,15 +447,27 @@ "backend": "openai-audio-speech", "endpoint": "/v1/audio/speech", "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", - "num_prompts": [100], - "request_rate": [2.0], + "num_prompts": [ + 100 + ], + "request_rate": [ + 2.0 + ], "seed_tts_locale": "en", - "extra_body": {"voice": "casual_male"}, + "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] + "median_audio_ttfp_ms": [ + 3000 + ], + "median_audio_rtf": [ + 0.25 + ], + "audio_throughput": [ + 4.0 + ] } } ] diff --git a/tests/dfx/perf/tests/test_voxcpm2.json b/tests/dfx/perf/tests/test_voxcpm2.json index f19f95960c7..ef72b5d88f7 100644 --- a/tests/dfx/perf/tests/test_voxcpm2.json +++ b/tests/dfx/perf/tests/test_voxcpm2.json @@ -1,9 +1,24 @@ [ { "test_name": "test_voxcpm2", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "tts" + ], "server_params": { "model": "openbmb/VoxCPM2", - "trust_remote_code": true + "trust_remote_code": true, + "extra_cli_args": [ + "--trust-remote-code" + ] }, "benchmark_params": [ { @@ -23,10 +38,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 600 + 130.2115 ], "median_audio_rtf": [ - 0.2 + 0.0488 ] } }, @@ -47,13 +62,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 3500 + 475.3357 ], "median_audio_rtf": [ - 0.85 + 0.1671 ], "audio_throughput": [ - 8 + 22.8717 ] } }, @@ -100,10 +115,10 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 350 + 75.7427 ], "median_audio_rtf": [ - 0.15 + 0.0419 ] } }, @@ -127,13 +142,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 3000 + 351.5855 ], "median_audio_rtf": [ - 0.5 + 0.1207 ], "audio_throughput": [ - 10 + 73.7632 ] } }, @@ -157,13 +172,13 @@ "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { "median_audio_ttfp_ms": [ - 500 + 80.1095 ], "median_audio_rtf": [ - 0.3 + 0.0485 ], "audio_throughput": [ - 5 + 18.3715 ] } }, diff --git a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json index d70a28efd80..78f720a6a78 100644 --- a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json +++ b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json @@ -1,107 +1,131 @@ [ - { - "test_name": "test_wan22_i2v_single_device", - "description": "Single-device baseline", - "server_type": "vllm-omni", - "server_params": { - "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", - "serve_args": { - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "832x480_frames81_steps4", - "dataset": "random", - "task": "i2v", - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 1, - "seed": 42, - "enable-negative-prompt": true, - "random-request-config": [ - { - "width": 832, - "height": 480, - "num_inference_steps": 4, - "num_frames": 81, - "fps": 16, - "weight": 1 - } - ], - "baseline": { - "throughput_qps": 0.04, - "latency_mean": 30, - "peak_memory_mb_mean": 75000 - } - } - ] + { + "test_name": "test_wan22_i2v_single_device", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 1 + } + }, + "full_model", + "diffusion" + ], + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } }, - { - "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing", - "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing", - "server_type": "vllm-omni", - "server_params": { - "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", - "serve_args": { - "usp": 2, - "vae-patch-parallel-size": 2, - "use-hsdp": true, - "vae-use-slicing": true, - "enable-diffusion-pipeline-profiler": true - } - }, - "benchmark_params": [ - { - "name": "832x480_frames81_steps4", - "dataset": "random", - "task": "i2v", - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 1, - "seed": 42, - "enable-negative-prompt": true, - "random-request-config": [ - { - "width": 832, - "height": 480, - "num_inference_steps": 4, - "num_frames": 81, - "fps": 16, - "weight": 1 - } - ], - "baseline": { - "throughput_qps": 0.06, - "latency_mean": 20, - "peak_memory_mb_mean": 48000 - } - }, - { - "name": "1280x720_frames121_steps4", - "dataset": "random", - "task": "i2v", - "num-prompts": 10, - "max-concurrency": 1, - "num-input-images": 1, - "seed": 42, - "enable-negative-prompt": true, - "random-request-config": [ - { - "width": 1280, - "height": 720, - "num_inference_steps": 4, - "num_frames": 121, - "fps": 16, - "weight": 1 - } - ], - "baseline": { - "throughput_qps": 0.01, - "latency_mean": 90, - "peak_memory_mb_mean": 56000 - } - } - ] - } + "benchmark_params": [ + { + "name": "832x480_frames81_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 81, + "fps": 16, + "weight": 1 + } + ], + "baseline": { + "throughput_qps": 0.0465, + "latency_mean": 21.4354, + "peak_memory_mb_mean": 72557.58 + } + } + ] + }, + { + "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing", + "mark": [ + { + "hardware_marks": { + "res": { + "cuda": "H100" + }, + "num_cards": 2 + } + }, + "full_model", + "diffusion" + ], + "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing", + "server_type": "vllm-omni", + "server_params": { + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "serve_args": { + "usp": 2, + "vae-patch-parallel-size": 2, + "use-hsdp": true, + "vae-use-slicing": true, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames81_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 81, + "fps": 16, + "weight": 1 + } + ], + "baseline": { + "throughput_qps": 0.0642, + "latency_mean": 15.5589, + "peak_memory_mb_mean": 45742.88 + } + }, + { + "name": "1280x720_frames121_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 1280, + "height": 720, + "num_inference_steps": 4, + "num_frames": 121, + "fps": 16, + "weight": 1 + } + ], + "baseline": { + "throughput_qps": 0.0128, + "latency_mean": 77.8565, + "peak_memory_mb_mean": 53947.84 + } + } + ] + } ] diff --git a/tests/dfx/reliability/README.md b/tests/dfx/reliability/README.md index 7196101bcdd..8dc69dc5522 100644 --- a/tests/dfx/reliability/README.md +++ b/tests/dfx/reliability/README.md @@ -6,9 +6,13 @@ The detailed RFC document is maintained outside this repository (local/internal - Fault helpers: `conftest.py`, `helpers.py` - Qwen3 tests: `test_reliability_qwen3_omni.py` - Wan2.2 tests: `test_reliability_wan22.py` +- HunyuanImage tests: `test_reliability_hunyuan_image.py` +- VoxCPM2 tests: `test_reliability_voxcpm2.py` ```bash pytest --collect-only tests/dfx/reliability pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow +pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow +pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow ``` diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py index 0bb8d2901b6..6531a29bb11 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""HTTP + WebSocket validation for Qwen3-TTS Base: ``/v1/audio/speech``, ``/v1/audio/speech/stream``, batch, voices.""" +"""HTTP + WebSocket validation for Qwen3-TTS Base and Higgs Audio V3 ``/v1/audio/speech`` (and related routes).""" from __future__ import annotations @@ -39,6 +39,7 @@ pytestmark = [pytest.mark.slow, pytest.mark.tts] +_SKIP_ISSUE_3649 = pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3649") _L4_SPEECH_HW = hardware_marks(res={"cuda": "L4"}) _SPEECH_SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"] @@ -54,6 +55,19 @@ ), ] +_HIGGS_AUDIO_V3_SPEECH = [ + pytest.param( + OmniServerParams( + model="bosonai/higgs-audio-v3-tts-4b", + stage_config_path=get_deploy_config_path("higgs_multimodal_qwen3.yaml"), + server_args=_SPEECH_SERVER_ARGS, + env_dict={"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}, + ), + id="higgs_audio_v3", + marks=_L4_SPEECH_HW, + ), +] + def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None: if loc == "batch": @@ -119,9 +133,6 @@ def test_speech_missing_required_fields(omni_server: OmniServer, openai_client: {"response_format": "mpeg"}, ("response_format", "literal_error", "wav"), id="response_format_invalid" ), pytest.param({"stream_format": 1}, ("stream_format", "literal_error", "audio"), id="stream_format_wrong_type"), - pytest.param( - {"stream_format": "sse"}, ("stream_format", "value_error", "audio"), id="stream_format_sse_blocked" - ), pytest.param({"stream": "wrong_type"}, ("stream", "bool_parsing", "validation error"), id="stream_wrong_type"), pytest.param( { @@ -240,6 +251,78 @@ def test_speech_invalid_field_values( ) +# Higgs Audio V3: ``references[]`` cookbook alias and serving-layer 400 rejections. +# Raw HTTP keeps JSON bodies intact (OpenAI SDK strips unknown fields such as ``references``). +@pytest.mark.parametrize( + "extra_json, err_message", + [ + pytest.param({"input": "", "response_format": "wav"}, ("input", "empty"), id="input_empty"), + pytest.param({"input": " ", "response_format": "wav"}, ("input", "empty"), id="input_whitespace_only"), + pytest.param( + {"input": "Hello world.", "max_new_tokens": 0}, ("max_new_tokens", "least 1"), id="max_new_tokens_below_min" + ), + pytest.param( + { + "input": "Hello world.", + "references": [ + {"audio_path": REF_AUDIO_URL, "text": REF_TEXT}, + {"audio_path": REF_AUDIO_URL, "text": REF_TEXT}, + ], + "response_format": "wav", + }, + ("references", "only supports a single reference"), + id="multi_reference_payload", + ), + pytest.param( + {"input": "Hello world.", "references": [], "response_format": "wav"}, + ("references", "non-empty"), + id="references_empty_array", + ), + pytest.param( + {"input": "Hello world.", "references": [{"text": REF_TEXT}], "response_format": "wav"}, + ("audio_path", "required"), + id="references_missing_audio_path", + ), + pytest.param( + {"input": "Hello world.", "references": [1], "response_format": "wav"}, + ("references[0]", "must be an object"), + id="references_bad_element_type", + ), + pytest.param( + { + "input": "Hello world.", + "ref_audio": REF_AUDIO_URL, + "references": [{"audio_path": "https://example.com/other.wav"}], + "response_format": "wav", + }, + ("mutually exclusive", "ref_audio", "references"), + id="ref_audio_and_references_exclusive", + ), + pytest.param( + { + "input": "Hello world.", + "ref_text": "different transcript", + "references": [{"audio_path": REF_AUDIO_URL, "text": REF_TEXT}], + "response_format": "wav", + }, + ("references[0].text", "ref_text", "conflict"), + id="references_text_conflicts_ref_text", + ), + ], +) +@pytest.mark.parametrize("omni_server", _HIGGS_AUDIO_V3_SPEECH, indirect=True) +def test_speech_higgs_invalid_field_values( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, + extra_json: dict[str, object], + err_message: str | tuple[str, ...], +) -> None: + body = {"model": omni_server.model, **extra_json} + openai_client.send_audio_speech_http_request( + {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message} + ) + + # ───────────────────────────────────────────────────────────────────────────── # POST /v1/audio/speech/batch (JSON) # ───────────────────────────────────────────────────────────────────────────── diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py index fa7f0f9059c..5ffd25632f5 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_editing.py @@ -62,8 +62,8 @@ def _finalize_edits_form_data(template: dict[str, Any], server_model: str) -> di ), pytest.param( True, - {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "2"}, - ("Invalid layers", "3", "1"), + {"prompt": "edit", "model": _IMAGE_EDITS_SERVER_MODEL, "layers": "1"}, + ("Invalid layers", "2", "10"), 400, id="invalid_layers", ), diff --git a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py index 05a2a5c2e1c..7efc1e28e59 100644 --- a/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py +++ b/tests/dfx/reliability/invalid_param_test/test_invalid_image_generation.py @@ -85,9 +85,9 @@ def _minimal_images_gen_json(omni_server: OmniServer) -> dict[str, object]: pytest.param( {"true_cfg_scale": 21.0}, ("true_cfg_scale", "less_than_equal", "20"), id="true_cfg_scale_above_max" ), - pytest.param({"layers": 2}, ("layers", "value_error", "3", "10"), id="layers_below_min"), - pytest.param({"layers": 11}, ("layers", "value_error", "3", "10"), id="layers_above_max"), - pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"), + pytest.param({"layers": 1}, ("layers", "value_error", "2", "10"), id="layers_below_min"), + pytest.param({"layers": 11}, ("layers", "value_error", "2", "10"), id="layers_above_max"), + pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative", marks=_SKIP_ISSUE_3649), pytest.param( {"seed": 2**32}, ("seed", "less_than_equal", "4294967295"), 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 6cd4ef4ff43..7d5dce37483 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 @@ -15,7 +15,6 @@ pytestmark = [pytest.mark.slow, pytest.mark.omni] - def _minimal_chat_json(omni_server: OmniServer) -> dict[str, object]: """Minimal valid chat body; individual tests override one offending field.""" return { diff --git a/tests/dfx/reliability/test_reliability_hunyuan_image.py b/tests/dfx/reliability/test_reliability_hunyuan_image.py new file mode 100644 index 00000000000..43a684d1d86 --- /dev/null +++ b/tests/dfx/reliability/test_reliability_hunyuan_image.py @@ -0,0 +1,478 @@ +""" +HunyuanImage-3.0-Instruct reliability integration tests (DiT-only deploy). + +Uses ``hunyuan_image3_dit.yaml`` (single diffusion stage, 4-GPU TP=4, no +AR/Mooncake connector) for simpler local bring-up. Scenarios follow +``docs/user_guide/fault_injection_reliability_matrix.md``. +""" + +from __future__ import annotations + +import concurrent.futures +import os +import time +from pathlib import Path +from typing import Any + +import pytest + +from tests.dfx.conftest import ( + assert_fault_exception, + create_reliability_omni_server_params, + resolve_oom_device_spec, +) +from tests.dfx.reliability.helpers import ( + FaultInjector, + assert_no_server_tree_process_residual_and_gpu_release, + get_health_raw, + inject_gpu_oom, + make_process_kill_fault_injector, + make_server_root_kill_fault_injector, + make_server_tree_kill_fault_injector, + post_json_raw, + run_fault_injection_with_rate_load, + stop_gpu_oom_hogs, + worker_residual_timeout_after_kill_signal, +) +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OpenAIClientHandler + +RELIABILITY_SCENARIOS: list[dict[str, Any]] = [ + { + "test_name": "hunyuan_image_dit_reliability_default", + "server_params": { + "model": "tencent/HunyuanImage-3.0-Instruct", + "stage_config_name": "hunyuan_image3_dit.yaml", + }, + }, +] + +DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy" +OOM_INJECTION_CONFIG = { + "target_mem_ratio": 0.95, + "hold_seconds": 0, + "startup_timeout_sec": 20, + "strict": True, +} +FAULT_ERROR_KEYWORDS = ( + "oom", + "out of memory", + "cuda", + "job failed", + "unknown error", + "internal server error", + "500 server error", +) +PROCESS_KILL_ERROR_KEYWORDS = ( + "timeout", + "did not complete within", + "connection", + "engine", + "orchestrator", + "dead", + "internal", + "500", + "503", +) +# Matrix: worker SIGTERM/SIGKILL; DiT-only HunyuanImage workers are vLLM-Omni::. +WORKER_KILL_PATTERNS = ("vLLM-Omni::",) +SERVE_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGINT", id="sigint"), + pytest.param("SIGKILL", id="sigkill"), +] +TREE_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGKILL", id="sigkill"), +] +TREE_WITH_LOAD_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGKILL", id="sigkill"), +] +WORKER_SIGNAL_FAULT_PARAMS = [ + pytest.param( + make_process_kill_fault_injector( + grep_patterns=WORKER_KILL_PATTERNS, + signal_name="SIGTERM", + limit=1, + post_kill_wait_seconds=2.0, + ), + id="runtime_process_chain_sigterm", + ), + pytest.param( + make_process_kill_fault_injector( + grep_patterns=WORKER_KILL_PATTERNS, + signal_name="SIGKILL", + limit=1, + post_kill_wait_seconds=2.0, + ), + id="runtime_process_chain_sigkill", + ), +] + +HUNYUAN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR) +INFLIGHT_INJECTION_REQUEST_RATE = 0.3 +INFLIGHT_INJECTION_REQUEST_COUNT = 10 + + +def _image_request_config(omni_server: Any) -> dict[str, Any]: + return { + "json": { + "model": omni_server.model, + "prompt": "A simple red apple on a white background.", + "size": "512x512", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": 4, + "guidance_scale": 7.5, + "seed": 42, + "bot_task": "none", + "use_system_prompt": "en_unified", + }, + "timeout": 600, + } + + +def _submit_image_generation(openai_client: OpenAIClientHandler, omni_server: Any) -> None: + responses = openai_client.send_images_generations_http_request(_image_request_config(omni_server)) + resp = responses[0] + if resp.status_code >= 400: + raise RuntimeError( + f"/v1/images/generations failed: status={resp.status_code}, body={getattr(resp, 'error_message', '')!r}" + ) + + +def _assert_post_fault_image_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None: + payload = { + "model": model, + "prompt": "post-fault fast-fail check", + "size": "512x512", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": 4, + "bot_task": "none", + "use_system_prompt": "en_unified", + } + start = time.monotonic() + try: + status, body = post_json_raw(host, port, "/v1/images/generations", payload, timeout_sec=20) + elapsed = time.monotonic() - start + assert elapsed < 15, f"[{scenario} fast_fail] /v1/images/generations did not fail fast: {elapsed:.2f}s" + assert status >= 500, ( + f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}" + ) + except Exception: # noqa: BLE001 + elapsed = time.monotonic() - start + assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s" + + +def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None: + deadline = time.monotonic() + 20.0 + last_observation = "" + while time.monotonic() < deadline: + try: + status, body = get_health_raw(host, port, timeout_sec=5) + last_observation = f"http={status}, body={body[:200]!r}" + if status == 503: + return + except Exception as exc: # noqa: BLE001 + last_observation = f"exception={exc!r}" + return + time.sleep(0.5) + pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}") + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.skip(reason="issue#4285") +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_gpu_oom_image_large_request_failure( + omni_server_function, + openai_client_function, +) -> None: + stage_config_path = getattr(omni_server_function, "stage_config_path", None) + device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) + handle = inject_gpu_oom( + device=device_spec, + target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"], + hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"], + startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"], + strict=OOM_INJECTION_CONFIG["strict"], + ) + try: + try: + _submit_image_generation(openai_client_function, omni_server_function) + except Exception as exc: + assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) + else: + pytest.fail("expected large /v1/images/generations request failure during GPU OOM injection") + finally: + stop_gpu_oom_hogs(handle) + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize( + "fault_injector", + WORKER_SIGNAL_FAULT_PARAMS, + indirect=True, +) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_request_failure( + omni_server_after_fault_function, + openai_client_function, +) -> None: + try: + _submit_image_generation(openai_client_function, omni_server_after_fault_function) + except Exception as exc: + assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) + else: + pytest.fail("expected /v1/images/generations request failure after process-kill injection") + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.skip(reason="issue#4522") +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize( + "fault_injector", + WORKER_SIGNAL_FAULT_PARAMS, + indirect=True, +) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_health_fast_fail_and_concurrent( + omni_server_after_fault_function, + openai_client_function, +) -> None: + host = omni_server_after_fault_function.host + port = omni_server_after_fault_function.port + model = omni_server_after_fault_function.model + + payload = { + "model": model, + "prompt": "fast-fail check", + "size": "512x512", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": 4, + "bot_task": "none", + "use_system_prompt": "en_unified", + } + start = time.monotonic() + try: + status, body = post_json_raw(host, port, "/v1/images/generations", payload, timeout_sec=20) + elapsed = time.monotonic() - start + assert elapsed < 15, ( + f"[process_kill fast_fail] /v1/images/generations did not fail fast after fault: {elapsed:.2f}s" + ) + assert status >= 500, ( + "[process_kill fast_fail] expected server-side error after fatal fault, " + f"got status={status}, body={body[:200]!r}" + ) + except Exception: + elapsed = time.monotonic() - start + assert elapsed < 15, ( + f"[process_kill fast_fail] /v1/images/generations exception was too slow after fault: {elapsed:.2f}s" + ) + + start = time.monotonic() + with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: + futures = [ + executor.submit(_submit_image_generation, openai_client_function, omni_server_after_fault_function) + for _ in range(3) + ] + done, pending = concurrent.futures.wait( + futures, + timeout=40, + return_when=concurrent.futures.ALL_COMPLETED, + ) + + elapsed = time.monotonic() - start + assert not pending, f"[process_kill concurrent] some fault-time image requests hung: pending={len(pending)}" + assert elapsed < 40, f"[process_kill concurrent] fault-time image request convergence is too slow: {elapsed:.2f}s" + + fault_observed = False + for future in done: + try: + future.result() + except Exception as exc: + fault_observed = True + assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) + assert fault_observed, ( + "[process_kill concurrent] expected at least one /v1/images/generations request to fail after fault" + ) + + deadline = time.monotonic() + 20.0 + last_observation = "" + final_health_status: int | None = None + while time.monotonic() < deadline: + try: + status, body = get_health_raw(host, port, timeout_sec=5) + last_observation = f"http={status}, body={body[:200]!r}" + final_health_status = status + if status == 503: + break + except Exception as exc: # noqa: BLE001 + last_observation = f"exception={exc!r}" + time.sleep(0.5) + assert final_health_status == 503, ( + "[process_kill health] expected /health 503 after fatal fault, " + f"got status={final_health_status}, last_observation={last_observation}" + ) + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.skip(reason="issue#4526") +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize( + "fault_injector", + WORKER_SIGNAL_FAULT_PARAMS, + indirect=True, +) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_worker_with_load_request_failure( + omni_server_function, + openai_client_function, + fault_injector: FaultInjector, +) -> None: + scenario = "kill_worker_with_load" + load_result = run_fault_injection_with_rate_load( + submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function), + inject_fault=lambda: fault_injector(omni_server_function), + num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, + request_rate=INFLIGHT_INJECTION_REQUEST_RATE, + completion_timeout_sec=120.0, + ) + assert load_result["failure_observed"], ( + f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" + ) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + + +@pytest.mark.slow +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup( + omni_server_function, + openai_client_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_root_with_load_{signal_name.lower()}" + injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) + load_result = run_fault_injection_with_rate_load( + submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function), + inject_fault=lambda: injector(omni_server_function), + num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, + request_rate=INFLIGHT_INJECTION_REQUEST_RATE, + completion_timeout_sec=120.0, + ) + assert load_result["failure_observed"], ( + f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" + ) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) + + +@pytest.mark.slow +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup( + omni_server_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_root_no_load_{signal_name.lower()}" + injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) + injector(omni_server_function) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup( + omni_server_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_tree_no_load_{signal_name.lower()}" + injector = make_server_tree_kill_fault_injector( + signal_name=signal_name, + post_kill_wait_seconds=2.0, + inter_kill_wait_seconds=0.1, + ) + injector(omni_server_function) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) + + +@pytest.mark.slow +@pytest.mark.diffusion +@pytest.mark.skip(reason="issue#4526") +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", HUNYUAN_PARAMS, indirect=True) +def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( + omni_server_function, + openai_client_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_tree_with_load_{signal_name.lower()}" + injector = make_server_tree_kill_fault_injector( + signal_name=signal_name, + post_kill_wait_seconds=2.0, + inter_kill_wait_seconds=0.1, + ) + load_result = run_fault_injection_with_rate_load( + submit_request=lambda: _submit_image_generation(openai_client_function, omni_server_function), + inject_fault=lambda: injector(omni_server_function), + num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, + request_rate=INFLIGHT_INJECTION_REQUEST_RATE, + completion_timeout_sec=120.0, + ) + assert load_result["failure_observed"], ( + f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" + ) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_image_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) diff --git a/tests/dfx/reliability/test_reliability_voxcpm2.py b/tests/dfx/reliability/test_reliability_voxcpm2.py new file mode 100644 index 00000000000..632c51081c6 --- /dev/null +++ b/tests/dfx/reliability/test_reliability_voxcpm2.py @@ -0,0 +1,508 @@ +""" +VoxCPM2 reliability integration tests. + +Scenarios follow ``docs/user_guide/fault_injection_reliability_matrix.md``. +""" + +from __future__ import annotations + +import concurrent.futures +import os +import time +from pathlib import Path +from typing import Any + +import pytest + +from tests.dfx.conftest import ( + assert_fault_exception, + create_reliability_omni_server_params, + resolve_oom_device_spec, +) +from tests.dfx.reliability.helpers import ( + FaultInjector, + assert_no_server_tree_process_residual_and_gpu_release, + extract_openai_error_contract_from_bytes, + get_health_raw, + inject_gpu_oom, + make_process_kill_fault_injector, + make_server_root_kill_fault_injector, + make_server_tree_kill_fault_injector, + post_json_raw, + run_fault_injection_with_rate_load, + stop_gpu_oom_hogs, + worker_residual_timeout_after_kill_signal, +) +from tests.helpers.mark import hardware_test +from tests.helpers.media import load_test_audio_data_url +from tests.helpers.runtime import OpenAIClientHandler + +RELIABILITY_SCENARIOS: list[dict[str, Any]] = [ + { + "test_name": "voxcpm2_reliability_default", + "server_params": { + "model": "openbmb/VoxCPM2", + "stage_config_name": "voxcpm2.yaml", + # VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False + # would otherwise override deploy YAML settings during server startup. + "server_args": ["--trust-remote-code", "--disable-log-stats"], + }, + }, +] + +DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy" +OOM_INJECTION_CONFIG = { + "target_mem_ratio": 0.96, + "hold_seconds": 0, + "startup_timeout_sec": 20, + "strict": False, +} +FAULT_ERROR_KEYWORDS = ( + "the request failed", + "oom", + "out of memory", + "cuda", + "orchestrator", + "timeout", + "connection", + "500", + "503", +) +PROCESS_KILL_ERROR_KEYWORDS = ( + "timeout", + "did not complete within", + "connection", + "engine", + "orchestrator", + "dead", + "internal", + "500", + "503", +) +RUNTIME_WORKER_PATTERN = "VLLM::" +SERVE_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGINT", id="sigint"), + pytest.param("SIGKILL", id="sigkill"), +] +SERVE_WITH_LOAD_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGINT", id="sigint"), + pytest.param("SIGKILL", id="sigkill"), +] +TREE_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGKILL", id="sigkill"), +] +TREE_WITH_LOAD_SIGNAL_PARAMS = [ + pytest.param("SIGTERM", id="sigterm"), + pytest.param("SIGKILL", id="sigkill"), +] +WORKER_SIGNAL_FAULT_PARAMS = [ + pytest.param( + make_process_kill_fault_injector( + grep_patterns=RUNTIME_WORKER_PATTERN, + signal_name="SIGTERM", + limit=1, + post_kill_wait_seconds=2.0, + ), + id="runtime_process_chain_sigterm", + ), + pytest.param( + make_process_kill_fault_injector( + grep_patterns=RUNTIME_WORKER_PATTERN, + signal_name="SIGKILL", + limit=1, + post_kill_wait_seconds=2.0, + ), + id="runtime_process_chain_sigkill", + ), +] + +VOXCPM2_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR) +INFLIGHT_INJECTION_REQUEST_RATE = 0.3 +INFLIGHT_INJECTION_REQUEST_COUNT = 10 +_SPEECH_FF_HTTP_TIMEOUT_SEC = 30 +_SPEECH_FF_MAX_ELAPSED_SEC = 30 + + +def _speech_request_config(omni_server: Any) -> dict[str, Any]: + return { + "model": omni_server.model, + "input": "The weather is nice today, perfect for a walk in the park.", + "stream": False, + "timeout": 120.0, + "response_format": "wav", + "voice": "default", + } + + +def _submit_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None: + openai_client.send_audio_speech_request(_speech_request_config(omni_server), request_num=1) + + +def _oom_speech_request_config(omni_server: Any) -> dict[str, Any]: + base = "The weather is nice today, perfect for a walk in the park. " + return { + "model": omni_server.model, + "input": base * 80, + "max_new_tokens": 4096, + "ref_audio": load_test_audio_data_url("qwen3_tts/clone_2.wav"), + "stream": False, + "timeout": 300.0, + "response_format": "wav", + "voice": "default", + } + + +def _submit_oom_speech_request(openai_client: OpenAIClientHandler, omni_server: Any) -> None: + openai_client.send_audio_speech_request(_oom_speech_request_config(omni_server), request_num=1) + + +def _assert_post_fault_speech_fast_fail(host: str, port: int, *, model: str, scenario: str) -> None: + payload = { + "model": model, + "input": "post-fault fast-fail check", + "voice": "default", + } + start = time.monotonic() + try: + status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=20) + elapsed = time.monotonic() - start + assert elapsed < 15, f"[{scenario} fast_fail] /v1/audio/speech did not fail fast: {elapsed:.2f}s" + assert status >= 500, ( + f"[{scenario} fast_fail] expected server-side error after fault, got status={status}, body={body[:200]!r}" + ) + except Exception: # noqa: BLE001 + elapsed = time.monotonic() - start + assert elapsed < 15, f"[{scenario} fast_fail] exception was too slow after fault: {elapsed:.2f}s" + + +def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) -> None: + deadline = time.monotonic() + 20.0 + last_observation = "" + while time.monotonic() < deadline: + try: + status, body = get_health_raw(host, port, timeout_sec=5) + last_observation = f"http={status}, body={body[:200]!r}" + if status == 503: + return + except Exception as exc: # noqa: BLE001 + last_observation = f"exception={exc!r}" + return + time.sleep(0.5) + pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}") + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skip(reason="issue#4285") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_gpu_oom_speech_request_failure( + omni_server_function, + openai_client_function, +) -> None: + stage_config_path = getattr(omni_server_function, "stage_config_path", None) + device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) + handle = inject_gpu_oom( + device=device_spec, + target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"], + hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"], + startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"], + strict=OOM_INJECTION_CONFIG["strict"], + ) + try: + try: + _submit_oom_speech_request(openai_client_function, omni_server_function) + except Exception as exc: + assert_fault_exception(exc, FAULT_ERROR_KEYWORDS) + else: + pytest.fail("expected /v1/audio/speech request failure during GPU OOM injection") + finally: + stop_gpu_oom_hogs(handle) + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skip(reason="issue#4285") +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_gpu_oom_speech_error_contract( + omni_server_function, +) -> None: + stage_config_path = getattr(omni_server_function, "stage_config_path", None) + device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path) + handle = inject_gpu_oom( + device=device_spec, + target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"], + hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"], + startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"], + strict=OOM_INJECTION_CONFIG["strict"], + ) + host = omni_server_function.host + port = omni_server_function.port + oom_config = _oom_speech_request_config(omni_server_function) + payload = { + "model": oom_config["model"], + "input": oom_config["input"], + "voice": oom_config["voice"], + "max_new_tokens": oom_config["max_new_tokens"], + "ref_audio": oom_config["ref_audio"], + } + try: + status, body = post_json_raw(host, port, "/v1/audio/speech", payload, timeout_sec=300) + finally: + stop_gpu_oom_hogs(handle) + + assert status >= 500, f"expected speech runtime error under OOM, got status={status}" + speech_error = extract_openai_error_contract_from_bytes(body) + assert speech_error is not None, f"speech error payload not OpenAI-compatible: {body[:300]!r}" + assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}" + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize( + "fault_injector", + WORKER_SIGNAL_FAULT_PARAMS, + indirect=True, +) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_request_failure( + omni_server_after_fault_function, + openai_client_function, +) -> None: + try: + _submit_speech_request(openai_client_function, omni_server_after_fault_function) + except Exception as exc: + assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) + else: + pytest.fail("expected /v1/audio/speech request failure after process-kill injection") + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize( + "fault_injector", + WORKER_SIGNAL_FAULT_PARAMS, + indirect=True, +) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_fast_fail_and_concurrent( + omni_server_after_fault_function, + openai_client_function, +) -> None: + host = omni_server_after_fault_function.host + port = omni_server_after_fault_function.port + model = omni_server_after_fault_function.model + payload = { + "model": model, + "input": "Say hello in one short sentence.", + "voice": "default", + } + + start = time.monotonic() + try: + status, body = post_json_raw( + host, + port, + "/v1/audio/speech", + payload, + timeout_sec=_SPEECH_FF_HTTP_TIMEOUT_SEC, + ) + elapsed = time.monotonic() - start + assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, ( + f"[process_kill fast_fail] /v1/audio/speech did not fail fast after fault: {elapsed:.2f}s" + ) + assert status >= 500, ( + "[process_kill fast_fail] expected server-side failure after fault, " + f"got status={status}, body={body[:200]!r}" + ) + except Exception: + elapsed = time.monotonic() - start + assert elapsed <= _SPEECH_FF_MAX_ELAPSED_SEC + 1.0, ( + f"[process_kill fast_fail] /v1/audio/speech exception was too slow after fault: {elapsed:.2f}s" + ) + + start = time.monotonic() + with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor: + futures = [ + executor.submit(_submit_speech_request, openai_client_function, omni_server_after_fault_function) + for _ in range(3) + ] + done, pending = concurrent.futures.wait( + futures, + timeout=30, + return_when=concurrent.futures.ALL_COMPLETED, + ) + + elapsed = time.monotonic() - start + assert not pending, f"[process_kill concurrent] some fault-time speech requests hung: pending={len(pending)}" + assert elapsed < 30, f"[process_kill concurrent] fault-time speech request convergence is too slow: {elapsed:.2f}s" + + fault_observed = False + for future in done: + try: + future.result() + except Exception as exc: + fault_observed = True + assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS) + assert fault_observed, ( + "[process_kill concurrent] expected at least one /v1/audio/speech request to fail after fault" + ) + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize( + "fault_injector", + WORKER_SIGNAL_FAULT_PARAMS, + indirect=True, +) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_worker_with_load_request_failure( + omni_server_function, + openai_client_function, + fault_injector: FaultInjector, +) -> None: + scenario = "kill_worker_with_load" + load_result = run_fault_injection_with_rate_load( + submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function), + inject_fault=lambda: fault_injector(omni_server_function), + num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, + request_rate=INFLIGHT_INJECTION_REQUEST_RATE, + completion_timeout_sec=120.0, + ) + assert load_result["failure_observed"], ( + f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" + ) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", SERVE_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_serve_root_no_load_fast_fail_and_cleanup( + omni_server_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_root_no_load_{signal_name.lower()}" + injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) + injector(omni_server_function) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) + + +@pytest.mark.slow +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", SERVE_WITH_LOAD_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_serve_root_with_load_fast_fail_and_cleanup( + omni_server_function, + openai_client_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_root_with_load_{signal_name.lower()}" + injector = make_server_root_kill_fault_injector(signal_name=signal_name, post_kill_wait_seconds=2.0) + load_result = run_fault_injection_with_rate_load( + submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function), + inject_fault=lambda: injector(omni_server_function), + num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, + request_rate=INFLIGHT_INJECTION_REQUEST_RATE, + completion_timeout_sec=300.0, + ) + assert load_result["failure_observed"], ( + f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" + ) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", TREE_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_tree_no_load_fast_fail_and_cleanup( + omni_server_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_tree_no_load_{signal_name.lower()}" + injector = make_server_tree_kill_fault_injector( + signal_name=signal_name, + post_kill_wait_seconds=2.0, + inter_kill_wait_seconds=0.1, + ) + injector(omni_server_function) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only") +@pytest.mark.parametrize("signal_name", TREE_WITH_LOAD_SIGNAL_PARAMS) +@pytest.mark.parametrize("omni_server_function", VOXCPM2_PARAMS, indirect=True) +def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( + omni_server_function, + openai_client_function, + signal_name: str, +) -> None: + scenario = f"kill_serve_tree_with_load_{signal_name.lower()}" + injector = make_server_tree_kill_fault_injector( + signal_name=signal_name, + post_kill_wait_seconds=2.0, + inter_kill_wait_seconds=0.1, + ) + load_result = run_fault_injection_with_rate_load( + submit_request=lambda: _submit_speech_request(openai_client_function, omni_server_function), + inject_fault=lambda: injector(omni_server_function), + num_requests=INFLIGHT_INJECTION_REQUEST_COUNT, + request_rate=INFLIGHT_INJECTION_REQUEST_RATE, + completion_timeout_sec=120.0, + ) + assert load_result["failure_observed"], ( + f"[{scenario}] expected at least one load request failure after fault; load_result={load_result}" + ) + host = omni_server_function.host + port = omni_server_function.port + _assert_post_fault_speech_fast_fail(host, port, model=omni_server_function.model, scenario=scenario) + _assert_post_fault_health_terminal(host, port, scenario=scenario) + assert_no_server_tree_process_residual_and_gpu_release( + omni_server_function, + scenario=scenario, + timeout_sec=worker_residual_timeout_after_kill_signal(signal_name), + ) diff --git a/tests/dfx/stability/scripts/test_stability_voxcpm2.py b/tests/dfx/stability/scripts/test_stability_voxcpm2.py new file mode 100644 index 00000000000..e26a2050c54 --- /dev/null +++ b/tests/dfx/stability/scripts/test_stability_voxcpm2.py @@ -0,0 +1,69 @@ +""" +VoxCPM2 stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration. + +Configuration: ``tests/dfx/stability/tests/test_voxcpm2.json``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.dfx.conftest import ( + create_benchmark_indices, + create_test_parameter_mapping, + create_unique_server_params, + load_configs, +) +from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop + +STABILITY_DIR = Path(__file__).resolve().parent.parent +DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" +CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_voxcpm2.json") +DEFAULT_NUM_PROMPTS_PER_BATCH = 20 +STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"] + +try: + BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +except FileNotFoundError: + BENCHMARK_CONFIGS = [] + +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} +benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) + + +@pytest.mark.slow +@pytest.mark.tts +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) +def test_stability_voxcpm2(omni_server, stability_benchmark_params): + test_name = stability_benchmark_params["test_name"] + params = stability_benchmark_params["params"] + duration_sec = params.get("duration_sec", 300) + num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) + request_rate = params.get("request_rate") + max_concurrency = params.get("max_concurrency") + + bench_params = { + k: v + for k, v in params.items() + if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") + } + + result = run_stability_benchmark_loop( + host=omni_server.host, + port=omni_server.port, + model=omni_server.model, + duration_sec=duration_sec, + params=bench_params, + request_rate=request_rate, + max_concurrency=max_concurrency, + result_dir=str(STABILITY_DIR), + num_prompts_per_batch=num_prompts_per_batch, + run_one_batch=_run_one_vllm_bench_batch, + ) + + assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" + assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/tests/test_voxcpm2.json b/tests/dfx/stability/tests/test_voxcpm2.json new file mode 100644 index 00000000000..dd912da36b5 --- /dev/null +++ b/tests/dfx/stability/tests/test_voxcpm2.json @@ -0,0 +1,35 @@ +[ + { + "test_name": "test_voxcpm2_stability", + "server_params": { + "model": "openbmb/VoxCPM2", + "trust_remote_code": true + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "trust_remote_code": true, + "duration_sec": 86400, + "max_concurrency": 2, + "num_prompts_per_batch": 50, + "random_input_len": { + "min": 0, + "max": 1000 + }, + "random_output_len": { + "min": 0, + "max": 1000 + }, + "random_range_ratio": 0.0, + "extra_body": { + "voice": "default", + "language": "English" + }, + "ignore_eos": true, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration" + } + ] + } +] diff --git a/tests/diffusion/attention/test_kernels_hub.py b/tests/diffusion/attention/test_kernels_hub.py new file mode 100644 index 00000000000..0f22a1244ee --- /dev/null +++ b/tests/diffusion/attention/test_kernels_hub.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys +import types + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum +from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl +from vllm_omni.platforms import current_omni_platform + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_kernels_hub_platform_fallback(monkeypatch: pytest.MonkeyPatch): + """Test that when kernels package is not available, the platform fallback logic + + routes to local FLASH_ATTN. + """ + from vllm_omni.platforms.cuda.platform import CudaOmniPlatform + + # Temporarily hide/remove kernels module if present + monkeypatch.setitem(sys.modules, "kernels", None) + + # Use monkeypatch to mock CudaOmniPlatform capability and package checker to allow FLASH_ATTN + from vllm.platforms.interface import DeviceCapability + + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + + monkeypatch.setattr( + CudaOmniPlatform, + "get_device_capability", + classmethod(lambda cls, device_id=0: DeviceCapability(8, 0)), + ) + monkeypatch.setattr(PACKAGES_CHECKER, "get_packages_info", lambda: {"has_flash_attn": True}) + + # Test FLASH_ATTN_HUB falls back to FLASH_ATTN when kernels is unavailable + backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_HUB", head_size=64) + assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + + # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN when kernels is unavailable + backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() + + # Test FLASH_ATTN_3_HUB falls back to FLASH_ATTN_HUB on pre-Hopper GPUs + kernels_module = types.ModuleType("kernels") + monkeypatch.setitem(sys.modules, "kernels", kernels_module) + backend_path = CudaOmniPlatform.get_diffusion_attn_backend_cls("FLASH_ATTN_3_HUB", head_size=64) + assert backend_path == DiffusionAttentionBackendEnum.FLASH_ATTN_HUB.get_path() + + +@pytest.mark.core_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_kernels_hub_execution(): + """Verify basic forward of flash_attn_hub and flash_attn_3_hub, comparing with SDPA reference.""" + device = torch.device(current_omni_platform.device_type) + dtype = torch.bfloat16 + + num_heads = 8 + head_dim = 64 + seq_len = 32 + batch_size = 1 + + torch.manual_seed(42) + q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) + k = q.clone() + v = q.clone() + + # 1. Test PyTorch SDPA reference + sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False) + attn_metadata_sdpa = AttentionMetadata(attn_mask=None) + output_ref = sdpa_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa) + + # 2. Test FlashAttentionHubBackend (FlashAttention 2) + from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttentionHubImpl + + fa_hub_impl = FlashAttentionHubImpl( + num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False + ) + output_fa_hub = fa_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa) + assert output_fa_hub.shape == q.shape + assert not torch.isnan(output_fa_hub).any() + max_diff = torch.max(torch.abs(output_ref - output_fa_hub)).item() + assert max_diff < 1e-2, f"FlashAttentionHub output differs too much from SDPA reference: {max_diff}" + + # 3. Test FlashAttention3HubBackend (FlashAttention 3, Hopper+ only) + major, _minor = torch.cuda.get_device_capability() + if major < 9: + pytest.skip("FLASH_ATTN_3_HUB execution requires Hopper-class GPU (compute capability >= 9.0)") + + from vllm_omni.diffusion.attention.backends.flash_attn_hub import FlashAttention3HubImpl + + fa3_hub_impl = FlashAttention3HubImpl( + num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False + ) + output_fa3_hub = fa3_hub_impl.forward(q.clone(), k.clone(), v.clone(), attn_metadata_sdpa) + assert output_fa3_hub.shape == q.shape + assert not torch.isnan(output_fa3_hub).any() + max_diff = torch.max(torch.abs(output_ref - output_fa3_hub)).item() + assert max_diff < 1e-2, f"FlashAttention3Hub output differs too much from SDPA reference: {max_diff}" diff --git a/tests/diffusion/batching/test_diffusion_batching.py b/tests/diffusion/batching/test_diffusion_batching.py index 01638a16360..fa96b54a45c 100644 --- a/tests/diffusion/batching/test_diffusion_batching.py +++ b/tests/diffusion/batching/test_diffusion_batching.py @@ -220,76 +220,11 @@ async def run_batch( return elapsed -# ------------------------------------------------------------------ -# Explicit batch — single generate() call with list of prompts -# ------------------------------------------------------------------ - - -async def run_batch_explicit( - omni: AsyncOmni, - prompts: list[dict[str, str]], - label: str = "batch_explicit", -) -> float: - """Send all prompts as a single batch via generate(prompt=[...]). - - Passes a *list* of prompts so that all are processed in **one** - ``DiffusionEngine.step()`` call. The orchestrator detects - ``isinstance(prompt, list)`` and routes to the batch path. - - A single ``OmniRequestOutput`` is yielded containing ALL generated - images combined. - """ - print(f"⚡ Running {label.upper()} mode – {len(prompts)} prompts in ONE engine call ...") - sp = _default_sampling_params() - request_id = f"{label}-{uuid.uuid4().hex[:8]}" - start = time.perf_counter() - - result: OmniRequestOutput | None = None - async for output in omni.generate( - prompt=prompts, - request_id=request_id, - sampling_params_list=[sp], - ): - result = output - - elapsed = time.perf_counter() - start - if result is not None: - images = _extract_images(result) - print(f" Got {len(images)} images total from batch, request_id={result.request_id}") - else: - print(" ⚠️ No output received from batch generate()") - - print(f" ✅ Total {label} mode: {elapsed:.2f}s\n") - return elapsed - - # ------------------------------------------------------------------ # Async validation helpers # ------------------------------------------------------------------ -async def validate_batch_explicit(omni: AsyncOmni, prompts: list[dict[str, str]]) -> None: - """Validate generate(prompt=[...]) returns a single result with all images.""" - print(f"🔍 Validating batch generate() correctness with {len(prompts)} prompts ...") - sp = _default_sampling_params() - request_id = f"validate-batch-{uuid.uuid4().hex[:8]}" - - result: OmniRequestOutput | None = None - async for output in omni.generate( - prompt=prompts, - request_id=request_id, - sampling_params_list=[sp], - ): - result = output - - assert result is not None, "No output received from batch generate()" - images = _extract_images(result) - # Batch mode returns ONE output with ALL images combined - assert len(images) == len(prompts), f"Expected {len(prompts)} images (one per prompt), got {len(images)}" - assert result.request_id == request_id, f"Expected request_id={request_id}, got {result.request_id}" - print(f" ✅ Batch returned {len(images)} images with correct request_id.\n") - - async def validate_concurrent(omni: AsyncOmni, prompts: list[dict[str, str]]) -> None: """Validate that every concurrent request receives a distinct result with its own request_id.""" @@ -330,17 +265,14 @@ async def compare_single_vs_parallel( await warmup(omni, WARMUP_PROMPTS) single_time = await run_single(omni, prompts) parallel_time = await run_batch(omni, prompts, label="parallel") - explicit_time = await run_batch_explicit(omni, prompts, label="batch_explicit") finally: omni.shutdown() speedup_parallel = single_time / parallel_time if parallel_time > 0 else float("inf") - speedup_explicit = single_time / explicit_time if explicit_time > 0 else float("inf") print("=" * 60) print(f"📊 Summary ({len(prompts)} prompts)") print(f" Sequential : {single_time:.2f}s") print(f" Parallel (gather) : {parallel_time:.2f}s ({speedup_parallel:.2f}x)") - print(f" Explicit batch : {explicit_time:.2f}s ({speedup_explicit:.2f}x)") print("=" * 60) @@ -362,12 +294,8 @@ async def main(model: str, num_prompts: int, mode: str, batch_size: int = 1) -> if mode == "validate": await validate_concurrent(omni, prompts) - elif mode == "validate_batch": - await validate_batch_explicit(omni, prompts) elif mode == "batch": await run_batch(omni, prompts, label="measurement") - elif mode == "batch_explicit": - await run_batch_explicit(omni, prompts) elif mode == "single": await run_single(omni, prompts) else: @@ -495,13 +423,10 @@ async def _inner(): @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}) @pytest.mark.parametrize("model_name", models) -def test_diffusion_batching_async_explicit_batch(model_name: str): - """Test that AsyncOmni batch mode (generate(prompt=[...])) dispatches - all prompts in a single engine call and returns a single combined result. - - The list-prompt path routes through the orchestrator's - ``add_batch_request_async`` → ``AsyncOmni.generate_batch`` - and yields ONE ``OmniRequestOutput`` with ALL images combined. +def test_diffusion_batching_list_prompt_rejected(model_name: str): + """Test that list-prompt batch requests are rejected at the diffusion + stage boundary. Users should submit multiple independent requests to + leverage scheduler batching instead. """ async def _inner(): @@ -511,25 +436,14 @@ async def _inner(): sp = _default_sampling_params() request_id = f"explicit-batch-{uuid.uuid4().hex[:8]}" - # Batch mode: pass list of prompts → single request_id - result: OmniRequestOutput | None = None - async for output in omni.generate( - prompt=prompts, - request_id=request_id, - sampling_params_list=[sp], - ): - result = output - - assert result is not None, "No output received from batch generate()" - - images = _extract_images(result) - # One image per prompt, all in a single output - assert len(images) == len(prompts), f"Expected {len(prompts)} images in combined output, got {len(images)}" - assert result.request_id == request_id, f"Expected request_id={request_id}, got {result.request_id}" - for i, img in enumerate(images): - assert img.width == 256, f"Image {i} width mismatch" - assert img.height == 256, f"Image {i} height mismatch" - print(f" ✅ Batch returned {len(images)} images, request_id={result.request_id}") + with pytest.raises(ValueError, match="Diffusion stages accept only a single prompt per request"): + async for _output in omni.generate( + prompt=prompts, + request_id=request_id, + sampling_params_list=[sp], + ): + pass + print(" ✅ List-prompt batch correctly rejected") finally: omni.shutdown() @@ -613,12 +527,11 @@ def test_diffusion_batching_distinct_results(model_name: str): parser.add_argument("--batch-size", type=int, default=1, help="Diffusion batch size (1 = no batching)") parser.add_argument( "--mode", - choices=["batch", "batch_explicit", "single", "compare", "validate", "validate_batch"], + choices=["batch", "single", "compare", "validate"], default="compare", help=( - "Run mode: 'batch' (parallel gather), 'batch_explicit' (list-prompt batch API), " - "'single' (sequential), 'compare' (all three), " - "'validate' (concurrent correctness), 'validate_batch' (list-prompt correctness)" + "Run mode: 'batch' (parallel gather), 'single' (sequential), " + "'compare' (single vs parallel), 'validate' (concurrent correctness)" ), ) args = parser.parse_args() diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index 021473ca5cc..f5981b92748 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -13,10 +13,12 @@ from unittest.mock import Mock, patch +import cache_dit import pytest +from cache_dit import ForwardPattern from vllm_omni.diffusion.cache.cache_dit_backend import ( - CUSTOM_DIT_ENABLERS, + CacheDiTAdapterConfig, CacheDiTBackend, ) from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend @@ -45,13 +47,20 @@ def test_init_with_config_object(self): assert backend.config.Fn_compute_blocks == 4 assert backend.enabled is False + @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") - def test_enable_single_transformer(self, mock_cache_dit): + def test_enable_single_transformer(self, mock_cache_dit, mock_block_adapter): """Test enabling cache-dit on single-transformer pipeline.""" # Mock pipeline mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "DiTPipeline" mock_transformer = Mock() + mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_0, + }, + ) + mock_pipeline.transformer = mock_transformer # Mock cache_dit functions @@ -66,14 +75,20 @@ def test_enable_single_transformer(self, mock_cache_dit): assert backend._refresh_func is not None mock_cache_dit.enable_cache.assert_called_once() + @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") - def test_refresh(self, mock_cache_dit): + def test_refresh(self, mock_cache_dit, mock_block_adapter): """Test refreshing cache context with SCM mask policy updates when num_inference_steps changes.""" # Mock pipeline mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "DiTPipeline" mock_transformer = Mock() mock_pipeline.transformer = mock_transformer + mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_0, + }, + ) # Mock cache_dit functions mock_cache_dit.enable_cache = Mock() @@ -124,23 +139,24 @@ def test_refresh(self, mock_cache_dit): assert "cache_config" in call_args[1] assert backend._last_num_inference_steps == 100 - def test_hunyuan_custom_enabler_registered(self): - """Test HunyuanImage3 custom cache-dit enabler is registered.""" - assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS - - def test_dreamid_custom_enabler_registered(self): - """Test DreamIDOmni custom cache-dit enabler is registered.""" - assert "DreamIDOmniPipeline" in CUSTOM_DIT_ENABLERS - @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter): - """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh.""" + """Test HunyuanImage3 uses pipeline.transformer for cache enable/refresh. + + NOTE: HunyuanImage3 no longer has a custom enabler, so this tests against the generic path. + """ mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline" mock_pipeline.model = Mock() mock_pipeline.model.layers = Mock() - + mock_pipeline.model._cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_4, + }, + ) + # NOTE: pipe.transformer is pipe.model in HunyuanImage3Pipelines + mock_pipeline.transformer = mock_pipeline.model mock_cache_dit.enable_cache = Mock() mock_cache_dit.refresh_context = Mock() @@ -152,9 +168,9 @@ def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mo mock_block_adapter.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["transformer"] is mock_pipeline.model - assert adapter_kwargs["blocks"] is mock_pipeline.model.layers - assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4 - assert len(adapter_kwargs["params_modifiers"]) == 1 + assert len(adapter_kwargs["blocks"]) == 1 + assert adapter_kwargs["blocks"][0] == mock_pipeline.model.layers + assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_4 mock_cache_dit.enable_cache.assert_called_once() backend.refresh(mock_pipeline, num_inference_steps=12) @@ -163,24 +179,23 @@ def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mo assert call_args[0][0] is mock_pipeline.model assert call_args[1]["num_inference_steps"] == 12 - def test_enable_hunyuan_pipeline_requires_model_layers(self): - """Test HunyuanImage3 enabler fails with a formatted pipeline class name.""" - mock_pipeline = Mock() - mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline" - mock_pipeline.model = Mock(spec=[]) - - backend = CacheDiTBackend({"Fn_compute_blocks": 2}) - - with pytest.raises(ValueError, match="HunyuanImage3Pipeline"): - backend.enable(mock_pipeline) - @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_block_adapter): + """Test DreamID uses pipeline.transformer for cache enable/refresh. + + NOTE: DreamID no longer has a custom enabler, so this tests against the generic path. + """ mock_pipeline = Mock() mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline" mock_pipeline.transformer = Mock() mock_pipeline.transformer.fused_blocks = Mock() + mock_pipeline.transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "fused_blocks": ForwardPattern.Pattern_0, + }, + has_separate_cfg=True, + ) mock_cache_dit.enable_cache = Mock() mock_cache_dit.refresh_context = Mock() @@ -193,10 +208,10 @@ def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_bl mock_block_adapter.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["transformer"] is mock_pipeline.transformer - assert adapter_kwargs["blocks"] is mock_pipeline.transformer.fused_blocks - assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_0 + assert len(adapter_kwargs["blocks"]) == 1 + assert adapter_kwargs["blocks"][0] == mock_pipeline.transformer.fused_blocks + assert adapter_kwargs["forward_pattern"][0] == ForwardPattern.Pattern_0 assert adapter_kwargs["has_separate_cfg"] is True - assert len(adapter_kwargs["params_modifiers"]) == 1 mock_cache_dit.enable_cache.assert_called_once() backend.refresh(mock_pipeline, num_inference_steps=12) @@ -205,15 +220,43 @@ def test_enable_dreamid_pipeline_uses_fused_blocks(self, mock_cache_dit, mock_bl assert call_args[0][0] is mock_pipeline.transformer assert call_args[1]["num_inference_steps"] == 12 - def test_enable_dreamid_pipeline_requires_fused_blocks(self): + @pytest.mark.parametrize("num_inference_steps", [1, 7]) + @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") + @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") + def test_refresh_scm_bypassed_for_unsupported_step_counts( + self, mock_cache_dit, mock_block_adapter, num_inference_steps + ): + """Ensure SCM is bypassed when num_inference_steps < 8 and not in (4, 6), + because cache_dit.steps_mask() raises for unsupported step count. + For these cases, we fall back to the non-SCM path to avoid crashing. + """ mock_pipeline = Mock() - mock_pipeline.__class__.__name__ = "DreamIDOmniPipeline" - mock_pipeline.transformer = Mock(spec=[]) + mock_pipeline.__class__.__name__ = "DiTPipeline" + mock_transformer = Mock() + mock_pipeline.transformer = mock_transformer + mock_transformer._cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_0, + }, + ) - backend = CacheDiTBackend({"Fn_compute_blocks": 2}) + mock_cache_dit.enable_cache = Mock() + mock_cache_dit.refresh_context = Mock() + mock_cache_dit.steps_mask = cache_dit.steps_mask + + # Create a cache config with an scm policy & enable it + config = DiffusionCacheConfig(scm_steps_mask_policy="fast") + backend = CacheDiTBackend(config) + backend.enable(mock_pipeline) - with pytest.raises(ValueError, match="DreamIDOmniPipeline"): - backend.enable(mock_pipeline) + backend.refresh(mock_pipeline, num_inference_steps=num_inference_steps) + + mock_cache_dit.refresh_context.assert_called_once() + call_args = mock_cache_dit.refresh_context.call_args + # Ensure that we properly guard, i.e., cache config is filtered + assert call_args[0][0] == mock_transformer + assert call_args[1]["num_inference_steps"] == num_inference_steps + assert "cache_config" not in call_args[1] class TestTeaCacheBackend: diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index 146084a508c..22f9b861167 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -5,6 +5,7 @@ Model specific tests for CacheDiT enablement. """ +import sys from unittest.mock import Mock, patch import pytest @@ -12,42 +13,87 @@ from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 import vllm_omni.diffusion.cache.cache_dit_backend as cd_backend -from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, CacheDiTBackend, cache_summary from vllm_omni.diffusion.data import DiffusionCacheConfig +from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer +from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel +from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel from vllm_omni.platforms import current_omni_platform +# NOTE: We patch DreamID Omni's modules here with mocks so that we can import and inspect +# the class even though the dependency may not be set up correctly; this is ok for these +# tests because we just inspect it and never initialize the model. +for mod in ("dreamid_omni", "dreamid_omni.modules", "dreamid_omni.modules.model"): + sys.modules.setdefault(mod, Mock()) +# isort: split +from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel as DreamIdOmniModel # noqa: E402 + pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -SEPARATE_CFG_ENABLERS = [ - cd_backend.enable_cache_for_dreamid_omni, - cd_backend.enable_cache_for_ltx2, - cd_backend.enable_cache_for_helios, - cd_backend.enable_cache_for_wan22, - cd_backend.enable_cache_for_longcat_image, - cd_backend.enable_cache_for_cosmos3, +SEPARATE_CFG_TRANSFORMERS = [ + DreamIdOmniModel, + LTX2VideoTransformer3DModel, + HeliosTransformer3DModel, + LongCatImageTransformer2DModel, + Cosmos3VFMTransformer, ] SAMPLE_CACHE_CONFIG = DiffusionCacheConfig() -@pytest.mark.parametrize("enabler", SEPARATE_CFG_ENABLERS) +def test_wan22_vace_uses_wan22_custom_cache_dit_enabler(): + assert cd_backend.CUSTOM_DIT_ENABLERS["Wan22VACEPipeline"] is cd_backend.enable_cache_for_wan22 + + +@pytest.mark.parametrize("transformer_model", SEPARATE_CFG_TRANSFORMERS) +def test_cache_dit_configs_have_separate_cfg(transformer_model): + """Check models with separate CFG set has_separate_cfg=True in their configs.""" + assert hasattr(transformer_model, "_cache_dit_adapter_config") + assert isinstance(transformer_model._cache_dit_adapter_config, CacheDiTAdapterConfig) + assert transformer_model._cache_dit_adapter_config.has_separate_cfg is True + + @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") -def test_separate_cfg(mock_cache_dit, mock_block_adapter, enabler): - """Ensure that custom enablers for models with separate CFG pass - the param through to cache_dit correctly. - - Regression test for: https://github.com/vllm-project/vllm-omni/pull/2860 - """ +def test_separate_wan22_custom_enabler_has_separate_cfg(mock_cache_dit, mock_block_adapter): + """Ensure that Wan22, which has a custom enabler, setts custom CFG correctly.""" mock_pipeline = Mock() - enabler(mock_pipeline, SAMPLE_CACHE_CONFIG) + cd_backend.enable_cache_for_wan22(mock_pipeline, SAMPLE_CACHE_CONFIG) mock_cache_dit.enable_cache.assert_called_once() adapter_kwargs = mock_block_adapter.call_args.kwargs assert adapter_kwargs["has_separate_cfg"] is True +@pytest.mark.parametrize("has_transformer_2", [False, True]) +@patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") +@patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") +def test_wan22_custom_enabler_passes_taylorseer_calibrator( + mock_cache_dit, + mock_block_adapter, + has_transformer_2, +): + mock_pipeline = Mock() + mock_pipeline.transformer.blocks = [Mock()] + if has_transformer_2: + mock_pipeline.transformer_2.blocks = [Mock()] + else: + mock_pipeline.transformer_2 = None + cache_config = DiffusionCacheConfig(enable_taylorseer=True, taylorseer_order=1) + + cd_backend.enable_cache_for_wan22(mock_pipeline, cache_config) + + enable_cache_kwargs = mock_cache_dit.enable_cache.call_args.kwargs + calibrator_config = enable_cache_kwargs["calibrator_config"] + assert calibrator_config is not None + assert calibrator_config.taylorseer_order == 1 + + adapter_kwargs = mock_block_adapter.call_args.kwargs + for modifier in adapter_kwargs["params_modifiers"]: + assert modifier._context_kwargs["calibrator_config"] is calibrator_config + + @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter): @@ -55,6 +101,7 @@ def test_cosmos3_cache_dit_wraps_gen_layers(mock_cache_dit, mock_block_adapter): mock_pipeline = Mock() gen_layers = object() mock_pipeline.transformer.gen_layers = gen_layers + mock_pipeline.transformer._cache_dit_adapter_config = Cosmos3VFMTransformer._cache_dit_adapter_config cd_backend.enable_cache_for_cosmos3(mock_pipeline, SAMPLE_CACHE_CONFIG) @@ -146,3 +193,12 @@ def call_Fn_blocks_and_capture(self, hidden_states, encoder_hidden_states, *a, * # Pattern_0 maps (hidden_states, encoder_hidden_states) to (video, audio) assert torch.equal(captured["hidden_states"], video_in) assert torch.equal(captured["encoder_hidden_states"], audio_in) + + +def test_summary_with_no_transformer_is_nonfatal(): + """Regression test for https://github.com/vllm-project/vllm-omni/issues/4325.""" + + class FakePipeline: + pass + + cache_summary(pipeline=FakePipeline()) diff --git a/tests/diffusion/cache/test_stepcache.py b/tests/diffusion/cache/test_stepcache.py new file mode 100644 index 00000000000..1fc3bbb5922 --- /dev/null +++ b/tests/diffusion/cache/test_stepcache.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for StepCacheBackend and StepCacheState.""" + +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.cache.stepcache import ( + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheBackend, + StepCacheConfig, + StepCacheState, + is_stepcache_active, +) +from vllm_omni.diffusion.data import DiffusionCacheConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]: + v_prev = torch.randn(1, dim) + if sim >= 0.99: + v_last = v_prev.clone() + else: + v_last = torch.randn(1, dim) + return [(v_prev,), (v_last,)] + + +class TestStepCacheConfig: + def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self): + enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True)) + disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False)) + assert enabled.enabled is True + assert disabled.enabled is False + + +class TestStepCacheState: + def test_disabled_always_runs(self): + state = StepCacheState(StepCacheConfig(enabled=False)) + assert state.should_run_step(_make_prev_predictions(0.99)) is True + + def test_insufficient_history_always_runs(self): + state = StepCacheState(StepCacheConfig(min_history_steps=2)) + assert state.should_run_step([(torch.randn(1, 4),)]) is True + + def test_high_similarity_skips_with_countdown(self): + state = StepCacheState( + StepCacheConfig( + sim_thresholds=(0.95,), + skip_countdowns=(2,), + ) + ) + preds = _make_prev_predictions(0.99) + assert state.should_run_step(preds) is False + assert state.skip_countdown == 2 + assert state.should_run_step(preds) is False + assert state.should_run_step(preds) is True + + def test_low_similarity_runs(self): + state = StepCacheState(StepCacheConfig()) + preds = _make_prev_predictions(0.5) + assert state.should_run_step(preds) is True + + def test_trim_history(self): + state = StepCacheState(StepCacheConfig(max_history=2)) + history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)] + state.trim_history(history) + assert len(history) == 2 + + +class TestStepCacheBackend: + def test_enable_dreamzero_pipeline(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DreamZeroPipeline" + + backend = StepCacheBackend(DiffusionCacheConfig()) + backend.enable(mock_pipeline) + + assert backend.enabled is True + cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR) + assert isinstance(cfg, StepCacheConfig) + assert is_stepcache_active(mock_pipeline) + state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR) + assert isinstance(state, StepCacheState) + + def test_refresh_resets_state(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DreamZeroPipeline" + backend = StepCacheBackend(DiffusionCacheConfig()) + backend.enable(mock_pipeline) + state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR) + state.skip_countdown = 3 + backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False) + assert state.skip_countdown == 0 + + def test_enable_unsupported_pipeline_raises(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "QwenImagePipeline" + backend = StepCacheBackend(DiffusionCacheConfig()) + + with pytest.raises(ValueError, match="step_cache backend does not support"): + backend.enable(mock_pipeline) + + def test_selector_returns_backend(self): + backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]}) + assert isinstance(backend, StepCacheBackend) + assert backend.config.velocity_sim_thresholds == [0.9] + + def test_legacy_backend_aliases(self): + assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend) + assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend) diff --git a/tests/diffusion/diffusion_backend/test_diffusers_backend.py b/tests/diffusion/diffusion_backend/test_diffusers_backend.py index b64d6f738c2..babc23b4ad3 100644 --- a/tests/diffusion/diffusion_backend/test_diffusers_backend.py +++ b/tests/diffusion/diffusion_backend/test_diffusers_backend.py @@ -18,6 +18,7 @@ ) from vllm_omni.diffusion.models.diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.diffusion] @@ -48,7 +49,7 @@ def _make_request(**overrides) -> OmniDiffusionRequest: prompt_obj["negative_prompt"] = negative_prompt defaults = { - "prompts": [prompt_obj], + "prompt": prompt_obj, "sampling_params": OmniDiffusionSamplingParams( num_inference_steps=20, guidance_scale=7.5, @@ -66,6 +67,55 @@ def _make_request(**overrides) -> OmniDiffusionRequest: return OmniDiffusionRequest(**defaults) +def _patch_fake_diffusers_quantization_backends(mocker): + class FakePipelineQuantizationConfig: + def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None): + self.quant_mapping = quant_mapping + self.quant_backend = quant_backend + self.quant_kwargs = quant_kwargs + self.components_to_quantize = components_to_quantize + + class FakeTorchAoConfig: + def __init__(self, quant_type, modules_to_not_convert=None): + self.quant_type = quant_type + self.modules_to_not_convert = modules_to_not_convert + + @classmethod + def from_dict(cls, config): + return cls(quant_type=config["quant_type"]) + + class FakeInt8DynamicActivationInt8WeightConfig: + pass + + mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.PipelineQuantizationConfig", + FakePipelineQuantizationConfig, + ) + mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils.TorchAoConfig", + FakeTorchAoConfig, + ) + mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_torchao_quant_type_cls", + return_value=FakeInt8DynamicActivationInt8WeightConfig, + ) + mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.quantization_utils._get_diffusers_quantization_config_cls", + return_value=FakeTorchAoConfig, + ) + return FakePipelineQuantizationConfig, FakeTorchAoConfig, FakeInt8DynamicActivationInt8WeightConfig + + +class _TransformerBackedPipeline(DiffusionPipeline): + def __init__(self, transformer): + pass + + +def _make_batch(**overrides) -> DiffusionRequestBatch: + """Wrap a single request in a DiffusionRequestBatch, matching the forward contract.""" + return DiffusionRequestBatch(requests=[_make_request(**overrides)]) + + @pytest.mark.core_model @pytest.mark.cpu class TestPipelineArgumentsHandling: @@ -84,7 +134,7 @@ def test_adapter_forward_returns_output(self, mocker): "__call__", return_value=MockPipelineOutput(image=stub_image), ) - output = adapter.forward(request) + output = adapter.forward(DiffusionRequestBatch(requests=[request])) assert isinstance(output, DiffusionOutput) assert isinstance(output.output, MockPipelineOutput) @@ -92,7 +142,7 @@ def test_adapter_forward_returns_output(self, mocker): @pytest.mark.parametrize( "feature_id", - ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"], + ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "unsupported_quantization"], ) def test_adapter_guard_unsupported_feature(self, feature_id): if feature_id == "cfg_parallel": @@ -122,14 +172,188 @@ def test_adapter_guard_unsupported_feature(self, feature_id): ) elif feature_id == "enforce_eager": od_config = _make_od_config(enforce_eager=True) - elif feature_id == "quantization": - od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8")) + elif feature_id == "unsupported_quantization": + od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="gguf")) else: raise ValueError(f"Unknown feature ID: {feature_id}") with pytest.raises(NotImplementedError): DiffusersAdapterPipeline(od_config=od_config) + def test_adapter_load_weights_injects_supported_quantization_config(self, mocker): + ( + FakePipelineQuantizationConfig, + FakeTorchAoConfig, + FakeInt8DynamicActivationInt8WeightConfig, + ) = _patch_fake_diffusers_quantization_backends(mocker) + + class MockPipeline: + def __call__(self, prompt=None): + return None + + def to(self, device): + return self + + mock_from_pretrained = mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained", + return_value=MockPipeline(), + ) + mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config", + return_value={ + "_class_name": "TransformerBackedPipeline", + "scheduler": ["diffusers", "SchedulerMixin"], + "transformer": ["diffusers", "Transformer2DModel"], + "transformer_2": ["diffusers", "Transformer2DModel"], + "vae": ["diffusers", "AutoencoderKL"], + }, + ) + + adapter = DiffusersAdapterPipeline( + od_config=_make_od_config( + quantization_config=SimpleNamespace(quant_method="int8"), + diffusers_pipeline_cls=_TransformerBackedPipeline, + ) + ) + adapter.load_weights() + + pipeline_quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"] + torchao_config = pipeline_quant_config.quant_mapping["transformer"] + assert isinstance(pipeline_quant_config, FakePipelineQuantizationConfig) + assert isinstance(torchao_config, FakeTorchAoConfig) + assert isinstance(torchao_config.quant_type, FakeInt8DynamicActivationInt8WeightConfig) + assert isinstance( + pipeline_quant_config.quant_mapping["transformer_2"].quant_type, + FakeInt8DynamicActivationInt8WeightConfig, + ) + + @pytest.mark.parametrize( + ("pipeline_config", "diffusers_load_kwargs", "diffusers_pipeline_cls"), + [ + ( + { + "_class_name": "NoTransformerPipeline", + "scheduler": ["diffusers", "SchedulerMixin"], + "unet": ["diffusers", "UNet2DConditionModel"], + "vae": ["diffusers", "AutoencoderKL"], + }, + {}, + _TransformerBackedPipeline, + ), + ( + { + "_class_name": "TransformerPipeline", + "scheduler": ["diffusers", "SchedulerMixin"], + "transformer": ["diffusers", "Transformer2DModel"], + "vae": ["diffusers", "AutoencoderKL"], + }, + {"transformer": None}, + None, + ), + ], + ) + def test_adapter_load_weights_rejects_converted_quantization_without_transformer_component( + self, + pipeline_config, + diffusers_load_kwargs, + diffusers_pipeline_cls, + mocker, + ): + _patch_fake_diffusers_quantization_backends(mocker) + mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.load_config", + return_value=pipeline_config, + ) + mock_from_pretrained = mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained" + ) + + adapter = DiffusersAdapterPipeline( + od_config=_make_od_config( + quantization_config=SimpleNamespace(quant_method="int8"), + diffusers_load_kwargs=diffusers_load_kwargs, + diffusers_pipeline_cls=diffusers_pipeline_cls, + ) + ) + + with pytest.raises(NotImplementedError, match="transformer"): + adapter.load_weights() + mock_from_pretrained.assert_not_called() + + def test_adapter_rejects_dduf_file_with_converted_quantization(self): + with pytest.raises(NotImplementedError, match="dduf_file"): + DiffusersAdapterPipeline( + od_config=_make_od_config( + quantization_config=SimpleNamespace(quant_method="int8"), + diffusers_load_kwargs={"dduf_file": "model.dduf"}, + ) + ) + + def test_adapter_load_weights_preserves_diffusers_native_quantization_config(self, mocker): + class MockPipeline: + def __call__(self, prompt=None): + return None + + def to(self, device): + return self + + diffusers_native_quantization_config = object() + mock_from_pretrained = mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained", + return_value=MockPipeline(), + ) + + adapter = DiffusersAdapterPipeline( + od_config=_make_od_config( + quantization_config=SimpleNamespace(quant_method="gguf"), + diffusers_load_kwargs={ + "dduf_file": "model.dduf", + "quantization_config": diffusers_native_quantization_config, + }, + ) + ) + adapter.load_weights() + + kwargs = mock_from_pretrained.call_args.kwargs + assert kwargs["quantization_config"] is diffusers_native_quantization_config + assert kwargs["dduf_file"] == "model.dduf" + + def test_adapter_load_weights_builds_diffusers_native_quantization_config_from_dict(self, mocker): + FakePipelineQuantizationConfig, FakeTorchAoConfig, _ = _patch_fake_diffusers_quantization_backends(mocker) + + class MockPipeline: + def __call__(self, prompt=None): + return None + + def to(self, device): + return self + + mock_from_pretrained = mocker.patch( + "vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter.DiffusionPipeline.from_pretrained", + return_value=MockPipeline(), + ) + + adapter = DiffusersAdapterPipeline( + od_config=_make_od_config( + diffusers_load_kwargs={ + "quantization_config": { + "quant_mapping": { + "transformer": { + "quant_method": "torchao", + "quant_type": "fake-int8-config", + }, + }, + }, + }, + ) + ) + adapter.load_weights() + + quant_config = mock_from_pretrained.call_args.kwargs["quantization_config"] + assert isinstance(quant_config, FakePipelineQuantizationConfig) + assert isinstance(quant_config.quant_mapping["transformer"], FakeTorchAoConfig) + assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config" + def test_adapter_guard_unknown_output_type(self, mocker): """Test that the adapter wraps an unknown output type as-is. This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass.""" @@ -145,7 +369,7 @@ def test_adapter_guard_unknown_output_type(self, mocker): "__call__", return_value=raw_output, ) - output = adapter.forward(_make_request()) + output = adapter.forward(_make_batch()) assert isinstance(output, DiffusionOutput) assert output.output == raw_output @@ -205,7 +429,7 @@ def to(self, device): ), ) - kwargs = adapter._build_call_kwargs(req) + kwargs = adapter._build_call_kwargs(DiffusionRequestBatch(requests=[req])) assert kwargs["prompt"] == "a cat on mars" assert kwargs["negative_prompt"] == "low quality" @@ -403,7 +627,7 @@ def to(self, device): ), ) with pytest.raises(ValueError): - pipeline.forward(problematic_request) + pipeline.forward(DiffusionRequestBatch(requests=[problematic_request])) @pytest.mark.advanced_model diff --git a/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py b/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py new file mode 100644 index 00000000000..5faa42eaf9a --- /dev/null +++ b/tests/diffusion/diffusion_backend/test_diffusers_quantization_real.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Optional real-backend checks for Diffusers quantization config conversion.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _load_quantization_utils(): + module_path = ( + Path(__file__).resolve().parents[3] / "vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py" + ) + spec = importlib.util.spec_from_file_location("diffusers_quantization_utils_real_test", module_path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_convert_diffusers_quantization_config_from_real_torchao_dict(): + diffusers = pytest.importorskip("diffusers") + torchao_quantization = pytest.importorskip("torchao.quantization") + quant_type_cls = getattr(torchao_quantization, "Int8DynamicActivationInt8WeightConfig", None) + if quant_type_cls is None: + pytest.skip("torchao does not expose Int8DynamicActivationInt8WeightConfig") + + quantization_utils = _load_quantization_utils() + + real_torchao_config = diffusers.TorchAoConfig(quant_type=quant_type_cls()) + serialized_torchao_config = real_torchao_config.to_dict() + assert serialized_torchao_config["quant_type"].keys() == {"default"} + + load_kwargs = { + "quantization_config": { + "quant_mapping": { + "transformer": serialized_torchao_config, + }, + }, + } + + quantization_utils.convert_diffusers_quantization_config(load_kwargs) + + quant_config = load_kwargs["quantization_config"] + component_config = quant_config.quant_mapping["transformer"] + assert isinstance(quant_config, diffusers.PipelineQuantizationConfig) + assert isinstance(component_config, diffusers.TorchAoConfig) + assert isinstance( + component_config.quant_type, + quant_type_cls, + ) diff --git a/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py b/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py new file mode 100644 index 00000000000..6def7ebadbe --- /dev/null +++ b/tests/diffusion/diffusion_backend/test_diffusers_quantization_utils.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Focused tests for Diffusers backend quantization conversion helpers.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vllm_omni.diffusion.models.diffusers_adapter import quantization_utils + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _FakePipelineQuantizationConfig: + def __init__(self, *, quant_mapping=None, quant_backend=None, quant_kwargs=None, components_to_quantize=None): + self.quant_mapping = quant_mapping + self.quant_backend = quant_backend + self.quant_kwargs = quant_kwargs + self.components_to_quantize = components_to_quantize + + +class _FakeTorchAoConfig: + def __init__(self, quant_type): + self.quant_type = quant_type + + @classmethod + def from_dict(cls, config): + return cls(quant_type=config["quant_type"]) + + +class _FakeFloat8DynamicActivationFloat8WeightConfig: + pass + + +class _FakeInt8DynamicActivationInt8WeightConfig: + pass + + +@pytest.fixture(autouse=True) +def patch_quantization_backends(monkeypatch): + monkeypatch.setattr( + quantization_utils, + "PipelineQuantizationConfig", + _FakePipelineQuantizationConfig, + ) + monkeypatch.setattr(quantization_utils, "TorchAoConfig", _FakeTorchAoConfig) + + def fake_get_torchao_quant_type_cls(class_name: str): + return { + "Float8DynamicActivationFloat8WeightConfig": _FakeFloat8DynamicActivationFloat8WeightConfig, + "Int8DynamicActivationInt8WeightConfig": _FakeInt8DynamicActivationInt8WeightConfig, + }[class_name] + + monkeypatch.setattr(quantization_utils, "_get_torchao_quant_type_cls", fake_get_torchao_quant_type_cls) + monkeypatch.setattr(quantization_utils, "_get_diffusers_quantization_config_cls", lambda _: _FakeTorchAoConfig) + + +def _quant_config(method: str, **kwargs): + return SimpleNamespace(get_name=lambda: method, **kwargs) + + +@pytest.mark.parametrize( + ("method", "expected_type"), + [ + ("fp8", _FakeFloat8DynamicActivationFloat8WeightConfig), + ("int8", _FakeInt8DynamicActivationInt8WeightConfig), + ], +) +def test_apply_injects_converted_quantization_config(method, expected_type): + od_config = SimpleNamespace(quantization_config=_quant_config(method)) + load_kwargs = {} + + injected = quantization_utils.apply_diffusers_quantization_config( + od_config, + load_kwargs, + {"transformer": ["diffusers", "Transformer2DModel"]}, + ) + + assert injected is True + torchao_config = load_kwargs["quantization_config"].quant_mapping["transformer"] + assert isinstance(torchao_config.quant_type, expected_type) + + +def test_apply_injects_converted_quantization_config_for_transformer_2(): + od_config = SimpleNamespace(quantization_config=_quant_config("int8")) + load_kwargs = {} + + quantization_utils.apply_diffusers_quantization_config( + od_config, + load_kwargs, + { + "transformer": ["diffusers", "Transformer2DModel"], + "transformer_2": ["diffusers", "Transformer2DModel"], + }, + ) + + assert sorted(load_kwargs["quantization_config"].quant_mapping) == ["transformer", "transformer_2"] + + +@pytest.mark.parametrize("method", ["gguf", "modelopt", "mxfp4", "mxfp8", "inc"]) +def test_unsupported_methods_fail_explicitly(method): + with pytest.raises(NotImplementedError, match=method): + quantization_utils.ensure_supported_diffusers_quantization(_quant_config(method)) + + +@pytest.mark.parametrize( + "quant_config", + [ + _quant_config("fp8", activation_scheme="static"), + _quant_config("int8", is_checkpoint_int8_serialized=True), + _quant_config("fp8", weight_block_size=[128, 128]), + _quant_config("int8", ignored_layers=["transformer.proj_out"]), + ], +) +def test_ambiguous_mappings_fail_explicitly(quant_config): + with pytest.raises(NotImplementedError): + quantization_utils.ensure_supported_diffusers_quantization(quant_config) + + +def test_apply_preserves_diffusers_load_kwargs_quantization_config(mocker): + od_config = SimpleNamespace(quantization_config=_quant_config("fp8")) + existing = object() + load_kwargs = {"quantization_config": existing} + mock_warning = mocker.patch.object(quantization_utils.logger, "warning") + + injected = quantization_utils.apply_diffusers_quantization_config(od_config, load_kwargs, {}) + + assert injected is False + assert load_kwargs["quantization_config"] is existing + assert "Using the Diffusers-native quantization_config" in mock_warning.call_args.args[0] + + +def test_convert_diffusers_quantization_config_from_dict(): + load_kwargs = { + "quantization_config": { + "quant_mapping": { + "transformer": { + "quant_method": "torchao", + "quant_type": "fake-int8-config", + } + } + } + } + + quantization_utils.convert_diffusers_quantization_config(load_kwargs) + + quant_config = load_kwargs["quantization_config"] + assert isinstance(quant_config, _FakePipelineQuantizationConfig) + assert isinstance(quant_config.quant_mapping["transformer"], _FakeTorchAoConfig) + assert quant_config.quant_mapping["transformer"].quant_type == "fake-int8-config" diff --git a/tests/diffusion/distributed/test_cfg_parallel.py b/tests/diffusion/distributed/test_cfg_parallel.py index bf709618de2..af81f00efd4 100644 --- a/tests/diffusion/distributed/test_cfg_parallel.py +++ b/tests/diffusion/distributed/test_cfg_parallel.py @@ -208,6 +208,7 @@ def _test_cfg_parallel_worker( positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, cfg_normalize=test_config["cfg_normalize"], + kwargs=test_config["kwargs"], ) # CFG parallel returns the combined prediction on every rank. @@ -284,6 +285,7 @@ def _test_cfg_sequential_worker( positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs, cfg_normalize=test_config["cfg_normalize"], + kwargs=test_config["kwargs"], ) # Sequential CFG always returns output @@ -297,7 +299,10 @@ def _test_cfg_sequential_worker( @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("batch_size", [2]) @pytest.mark.parametrize("cfg_normalize", [False, True]) -def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool): +@pytest.mark.parametrize("kwargs", [None, {"step_i": 1}]) +def test_predict_noise_maybe_with_cfg( + cfg_parallel_size: int, dtype: torch.dtype, batch_size: int, cfg_normalize: bool, kwargs: dict | None +): """ Test that predict_noise_maybe_with_cfg produces identical results with and without CFG parallel. @@ -322,6 +327,7 @@ def test_predict_noise_maybe_with_cfg(cfg_parallel_size: int, dtype: torch.dtype "cfg_normalize": cfg_normalize, "model_seed": 42, # Fixed seed for model initialization "input_seed": 123, # Fixed seed for input generation + "kwargs": kwargs, # Additional kwargs to test passing through CFG parallel } mp_context = torch.multiprocessing.get_context("spawn") diff --git a/tests/diffusion/distributed/test_expert_parallel_layout.py b/tests/diffusion/distributed/test_expert_parallel_layout.py new file mode 100644 index 00000000000..c4d3c6ae71d --- /dev/null +++ b/tests/diffusion/distributed/test_expert_parallel_layout.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +import vllm_omni.diffusion.distributed.parallel_state as parallel_state + + +class _FakeGroup: + def __init__( + self, + group_ranks: list[list[int]], + local_rank: int, + parallel_mode: str, + **kwargs, + ) -> None: + self.group_ranks = group_ranks + self.parallel_mode = parallel_mode + self.device_group = object() + self.device_communicator = kwargs.get("device_communicator") + reduce_scatter = kwargs.get("reduce_scatter") + if reduce_scatter is not None: + self.reduce_scatter = reduce_scatter + self.ulysses_group = kwargs.get("ulysses_group") + self.ring_group = kwargs.get("ring_group") + self.local_group = next(group for group in group_ranks if local_rank in group) + self.world_size = len(self.local_group) + self.rank_in_group = self.local_group.index(local_rank) + + def destroy(self) -> None: + pass + + +@pytest.mark.cpu +@pytest.mark.core_model +def test_moe_ep_maps_diffusion_sp_cfg_dp_to_vllm_groups(monkeypatch): + """MoE+EP rank layout should map SP->PCP, CFG*DP->DP, and TP*SP*CFG*DP->EP.""" + local_rank = 0 + world_size = 32 + created_groups: list[_FakeGroup] = [] + + def fake_init_model_parallel_group( + group_ranks, + local_rank, + backend, + parallel_mode=None, + group_name=None, + **kwargs, + ): + del backend, group_name + group = _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + parallel_mode or "", + **kwargs, + ) + created_groups.append(group) + return group + + def fake_init_vllm_model_parallel_group( + group_ranks, + local_rank, + backend, + group_name, + ): + del backend + group = _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + f"vllm_{group_name}", + device_communicator=object(), + reduce_scatter=lambda tensor, **kwargs: tensor, + ) + created_groups.append(group) + return group + + fake_world_group = SimpleNamespace( + rank_in_group=local_rank, + local_rank=local_rank, + device_group=object(), + ) + fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=True)) + monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo") + monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_vllm_model_parallel_group", fake_init_vllm_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + monkeypatch.setattr(parallel_state, name, None) + for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + + parallel_state.initialize_model_parallel( + tensor_parallel_size=2, + sequence_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + pipeline_parallel_size=2, + cfg_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=True, + backend="gloo", + ) + + assert parallel_state.vllm_parallel_state._PCP is not parallel_state._SP + assert parallel_state.vllm_parallel_state._PCP.world_size == 2 + assert parallel_state._DP.world_size == 2 + assert parallel_state.vllm_parallel_state._DP is not parallel_state._DP + assert parallel_state.vllm_parallel_state._DP.world_size == 4 + assert parallel_state.vllm_parallel_state._EP.world_size == 16 + assert parallel_state.vllm_parallel_state._TP.world_size == 2 + assert parallel_state._PP.world_size == 2 + + assert parallel_state.vllm_parallel_state._PCP.device_communicator is not None + assert parallel_state.vllm_parallel_state._DP.device_communicator is not None + assert parallel_state.vllm_parallel_state._EP.device_communicator is not None + assert hasattr(parallel_state.vllm_parallel_state._PCP, "reduce_scatter") + assert hasattr(parallel_state.vllm_parallel_state._DP, "reduce_scatter") + assert hasattr(parallel_state.vllm_parallel_state._EP, "reduce_scatter") + + assert parallel_state.vllm_parallel_state._PCP.local_group == [0, 2] + assert parallel_state.vllm_parallel_state._DP.local_group == [0, 8, 16, 24] + assert parallel_state.vllm_parallel_state._EP.local_group == [ + 0, + 1, + 2, + 3, + 8, + 9, + 10, + 11, + 16, + 17, + 18, + 19, + 24, + 25, + 26, + 27, + ] + assert parallel_state.get_expert_parallel_group_ranks() == [ + [ + 0, + 1, + 2, + 3, + 8, + 9, + 10, + 11, + 16, + 17, + 18, + 19, + 24, + 25, + 26, + 27, + ], + [ + 4, + 5, + 6, + 7, + 12, + 13, + 14, + 15, + 20, + 21, + 22, + 23, + 28, + 29, + 30, + 31, + ], + ] + + vllm_group_names = [group.parallel_mode for group in created_groups if group.parallel_mode.startswith("vllm_")] + assert vllm_group_names == ["vllm_pcp", "vllm_tp", "vllm_dp", "vllm_ep"] + ep_groups = [group.local_group for group in created_groups if group.parallel_mode == "vllm_ep"] + assert ep_groups == [parallel_state.vllm_parallel_state._EP.local_group] + + +@pytest.mark.cpu +@pytest.mark.core_model +def test_cfg_parallel_keeps_diffusion_dp_without_ep(monkeypatch): + """vLLM DP should keep diffusion DP when expert parallelism is not enabled.""" + local_rank = 0 + world_size = 8 + + def fake_init_model_parallel_group( + group_ranks, + local_rank, + backend, + parallel_mode=None, + group_name=None, + **kwargs, + ): + del backend, group_name + return _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + parallel_mode or "", + **kwargs, + ) + + fake_world_group = SimpleNamespace( + rank_in_group=local_rank, + local_rank=local_rank, + device_group=object(), + ) + monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo") + monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + monkeypatch.setattr(parallel_state, name, None) + for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + + parallel_state.initialize_model_parallel( + tensor_parallel_size=2, + sequence_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + pipeline_parallel_size=1, + cfg_parallel_size=2, + data_parallel_size=2, + enable_expert_parallel=False, + backend="gloo", + ) + + assert parallel_state._DP.world_size == 2 + assert parallel_state.vllm_parallel_state._DP is parallel_state._DP + assert parallel_state.vllm_parallel_state._DP.world_size == 2 + assert parallel_state.vllm_parallel_state._DP.local_group == [0, 4] + assert parallel_state.vllm_parallel_state._PCP is None + assert parallel_state.vllm_parallel_state._EP is None + assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None + + +@pytest.mark.cpu +@pytest.mark.core_model +def test_non_moe_ep_fails_before_vllm_ep_remap(monkeypatch): + """Non-MoE diffusion configs should not create vLLM PCP/DP/EP remap state.""" + local_rank = 0 + world_size = 8 + + def fake_init_model_parallel_group( + group_ranks, + local_rank, + backend, + parallel_mode=None, + group_name=None, + **kwargs, + ): + del backend, group_name + return _FakeGroup( + [list(ranks) for ranks in group_ranks], + local_rank, + parallel_mode or "", + **kwargs, + ) + + fake_world_group = SimpleNamespace( + rank_in_group=local_rank, + local_rank=local_rank, + device_group=object(), + ) + fake_forward_context = SimpleNamespace(omni_diffusion_config=SimpleNamespace(is_moe=False)) + monkeypatch.setattr(parallel_state.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(parallel_state.torch.distributed, "get_world_size", lambda: world_size) + monkeypatch.setattr(parallel_state.torch.distributed, "get_backend", lambda *_args, **_kwargs: "gloo") + monkeypatch.setattr(parallel_state.torch.distributed, "new_group", lambda ranks: tuple(ranks)) + monkeypatch.setattr(parallel_state, "get_world_group", lambda: fake_world_group) + monkeypatch.setattr(parallel_state, "get_forward_context", lambda: fake_forward_context) + monkeypatch.setattr(parallel_state, "init_model_parallel_group", fake_init_model_parallel_group) + monkeypatch.setattr(parallel_state, "init_dit_group", lambda *_args, **_kwargs: None) + + for name in ("_DP", "_CFG", "_SP", "_PP", "_FS", "_EXPERT_PARALLEL_GROUP_RANKS"): + monkeypatch.setattr(parallel_state, name, None) + for name in ("_TP", "_PCP", "_DP", "_EP", "_PP"): + monkeypatch.setattr(parallel_state.vllm_parallel_state, name, None, raising=False) + + with pytest.raises(RuntimeError, match="Expert parallelism enabled for a non-MoE model"): + parallel_state.initialize_model_parallel( + tensor_parallel_size=2, + sequence_parallel_size=2, + ulysses_degree=2, + ring_degree=1, + pipeline_parallel_size=1, + cfg_parallel_size=2, + data_parallel_size=1, + enable_expert_parallel=True, + backend="gloo", + ) + + assert parallel_state.vllm_parallel_state._PCP is None + assert parallel_state.vllm_parallel_state._EP is None + assert parallel_state._EXPERT_PARALLEL_GROUP_RANKS is None diff --git a/tests/diffusion/distributed/test_hsdp.py b/tests/diffusion/distributed/test_hsdp.py index e6d782c0a23..91c6784fcd4 100644 --- a/tests/diffusion/distributed/test_hsdp.py +++ b/tests/diffusion/distributed/test_hsdp.py @@ -1,19 +1,97 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities. +"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities.""" -These tests verify HSDP configuration logic without requiring a distributed environment. -""" +import gc +import os +import socket import pytest +import torch +import torch.distributed as dist import torch.nn as nn +from torch.distributed.tensor import DeviceMesh, DTensor from vllm_omni.diffusion.data import DiffusionParallelConfig -from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig +from vllm_omni.diffusion.distributed.hsdp import ( + HSDPInferenceConfig, + _unshardable_parameters, + shard_model, +) pytestmark = [pytest.mark.diffusion, pytest.mark.parallel, pytest.mark.cpu, pytest.mark.core_model] +class _PackedBlock(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.ones(2, 2)) + self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False) + self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False) + + +class _PackedModel(nn.Module): + def __init__(self): + super().__init__() + self.block = _PackedBlock() + self.root_weight = nn.Parameter(torch.ones(2)) + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +@pytest.fixture(scope="module") +def cpu_process_group(): + if dist.is_initialized(): + yield + return + + master_port = _find_free_port() + os.environ.update( + { + "RANK": "0", + "LOCAL_RANK": "0", + "WORLD_SIZE": "1", + "MASTER_ADDR": "127.0.0.1", + "MASTER_PORT": str(master_port), + } + ) + dist.init_process_group("gloo", rank=0, world_size=1) + try: + yield + finally: + dist.destroy_process_group() + for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"): + os.environ.pop(key, None) + gc.collect() + + +def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group): + model = _PackedModel() + ignored_params = _unshardable_parameters(model) + expected_ignored = {model.block.packed_weight, model.block.input_global_scale} + assert ignored_params == expected_ignored + + packed_weight = model.block.packed_weight + input_global_scale = model.block.input_global_scale + shard_model( + model, + mesh=DeviceMesh("cpu", [0]), + hsdp_shard_conditions=[lambda name, _module: name == "block"], + ignored_params=ignored_params, + ) + + assert isinstance(model.block.weight, DTensor) + assert isinstance(model.root_weight, DTensor) + assert model.block.packed_weight is packed_weight + assert model.block.input_global_scale is input_global_scale + assert not isinstance(model.block.packed_weight, DTensor) + assert not isinstance(model.block.input_global_scale, DTensor) + + class TestHSDPInferenceConfig: """Tests for HSDPInferenceConfig dataclass.""" diff --git a/tests/diffusion/distributed/test_sequence_parallel.py b/tests/diffusion/distributed/test_sequence_parallel.py index 9f76b3b75c5..c02667b9351 100644 --- a/tests/diffusion/distributed/test_sequence_parallel.py +++ b/tests/diffusion/distributed/test_sequence_parallel.py @@ -23,6 +23,7 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.distributed.utils import build_local_sp_padding_mask from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -412,3 +413,85 @@ def test_sp_correctness_advanced(model_name: str): f"{baseline_str:<12} {r['sp_ms']:.0f}ms{'':<7} {speedup_str:<10} {status}" ) print("=" * 70) + + +@pytest.mark.skipif( + not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA and XPU", +) +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.core_model +def test_local_sp_padding_mask(monkeypatch: pytest.MonkeyPatch) -> None: + """A partially padded SP shard must receive a local-length mask.""" + mask = build_local_sp_padding_mask( + batch_size=2, + local_seq_len=4, + sp_original_seq_len=5, + sp_padding_size=3, + sequence_parallel_rank=1, + device=torch.device(current_omni_platform.device_type), + ) + + expected = torch.tensor( + [ + [True, False, False, False], + [True, False, False, False], + ], + dtype=torch.bool, + device=mask.device, + ) + assert mask is not None + assert mask.shape == (2, 4) + assert torch.equal(mask, expected) + + +@pytest.mark.skipif( + not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA and XPU", +) +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.core_model +def test_local_sp_padding_mask_no_padding(monkeypatch: pytest.MonkeyPatch) -> None: + """A rank whose local shard contains no padding should not get a mask.""" + mask = build_local_sp_padding_mask( + batch_size=2, + local_seq_len=4, + sp_original_seq_len=5, + sp_padding_size=3, + sequence_parallel_rank=0, + device=torch.device(current_omni_platform.device_type), + ) + + assert mask is None + + +@pytest.mark.skipif( + not (current_omni_platform.is_cuda() or current_omni_platform.is_xpu()), + reason="Only tested on CUDA and XPU", +) +@pytest.mark.diffusion +@pytest.mark.parallel +@pytest.mark.core_model +def test_wan_sp_plan() -> None: + """Wan2.2 must shard hidden states before CacheDiT-wrapped transformer blocks.""" + try: + from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan + from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel + except ImportError as exc: + pytest.skip(f"WanTransformer3DModel not available: {exc}") + + plan = getattr(WanTransformer3DModel, "_sp_plan", None) + + assert plan is not None + assert "_sp_shard_point" in plan + assert "blocks.0" not in plan + + shard_plan = plan["_sp_shard_point"] + assert 0 in shard_plan + assert shard_plan[0].split_dim == 1 + assert shard_plan[0].expected_dims == 3 + assert shard_plan[0].split_output is True + + validate_sp_plan(plan) diff --git a/tests/diffusion/distributed/test_wan_spatial_shard.py b/tests/diffusion/distributed/test_wan_spatial_shard.py new file mode 100644 index 00000000000..72a2a45f574 --- /dev/null +++ b/tests/diffusion/distributed/test_wan_spatial_shard.py @@ -0,0 +1,446 @@ +import os +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( + DistributedAutoencoderKLWan, +) +from vllm_omni.platforms import current_omni_platform + +# CPU unit tests are marked core_model + cpu. The multi-GPU correctness test at +# the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel. + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_for_parallel_decode_pads_uneven_height(): + x = torch.arange(1 * 1 * 1 * 5 * 2, dtype=torch.float32).reshape(1, 1, 1, 5, 2) + + local, expected_height = wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=2, + rank=2, + world_size=3, + ) + + assert expected_height == 20 + assert local.shape == (1, 1, 1, 2, 2) + assert torch.equal(local[..., 0, :], x[..., 4, :]) + assert torch.equal(local[..., 1, :], torch.zeros_like(local[..., 1, :])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_for_parallel_decode_pads_uneven_width(): + x = torch.arange(1 * 1 * 1 * 2 * 5, dtype=torch.float32).reshape(1, 1, 1, 2, 5) + + local, expected_width = wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=2, + split_dim="width", + rank=2, + world_size=3, + ) + + assert expected_width == 20 + assert local.shape == (1, 1, 1, 2, 2) + assert torch.equal(local[..., :, 0], x[..., :, 4]) + assert torch.equal(local[..., :, 1], torch.zeros_like(local[..., :, 1])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_for_parallel_decode_rejects_invalid_split_dim(): + x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32) + + with pytest.raises(ValueError, match="split_dim"): + wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=1, + split_dim="depth", + rank=0, + world_size=2, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_for_parallel_decode_rejects_zero_world_size(): + x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32) + + with pytest.raises(ValueError, match="world_size"): + wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=1, + rank=0, + world_size=0, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_for_parallel_decode_rejects_rank_out_of_range(): + x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32) + + with pytest.raises(ValueError, match="rank"): + wan_spatial_shard.split_for_parallel_decode( + x, + upsample_count=1, + rank=3, + world_size=3, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_gather_and_trim_height(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3)) + + def fake_all_gather(gathered, x, group=None): + for idx, output in enumerate(gathered): + output.copy_(x + idx) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.zeros((1, 1, 1, 2, 1), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object()) + + assert out.shape == (1, 1, 1, 5, 1) + assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_gather_and_trim_width(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3)) + + def fake_all_gather(gathered, x, group=None): + for idx, output in enumerate(gathered): + output.copy_(x + idx) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.zeros((1, 1, 1, 1, 2), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="width", group=object()) + + assert out.shape == (1, 1, 1, 1, 5) + assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_gather_and_trim_rank0_only_assembles_on_rank0(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3)) + + def fake_all_gather(gathered, x, group=None): + for idx, output in enumerate(gathered): + output.copy_(x + idx) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.zeros((1, 1, 1, 2, 1), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object(), dst=0) + + assert out.shape == (1, 1, 1, 5, 1) + assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_gather_and_trim_rank0_only_returns_empty_on_non_zero_rank(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) + + gathered_sizes = [] + + def fake_all_gather(gathered, x, group=None): + # Every rank must still take part in the collective even when it discards the result. + gathered_sizes.append(len(gathered)) + for output in gathered: + output.copy_(x) + + monkeypatch.setattr(wan_spatial_shard.dist, "all_gather", fake_all_gather) + + x = torch.ones((1, 1, 1, 2, 1), dtype=torch.float32) + out = wan_spatial_shard.gather_and_trim_extent(x, expected_extent=5, split_dim="height", group=object(), dst=0) + + assert gathered_sizes == [3] + assert out.numel() == 0 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_reshard_from_trimmed_height_pads_invalid_rows(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + + x = torch.arange(5, dtype=torch.float32).reshape(1, 1, 1, 5, 1) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=5, + local_input_extent=2, + split_dim="height", + rank=2, + world_size=3, + ) + ) + try: + out = wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=2, + split_dim="height", + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + assert out.shape == (1, 1, 1, 2, 1) + assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + + x = torch.arange(5, dtype=torch.float32).reshape(1, 1, 1, 1, 5) + token = wan_spatial_shard._SPATIAL_SHARD_CONTEXT.set( + wan_spatial_shard.SpatialShardContext( + input_extent=5, + local_input_extent=2, + split_dim="width", + rank=2, + world_size=3, + ) + ) + try: + out = wan_spatial_shard.reshard_from_trimmed_extent( + x, + local_extent=2, + split_dim="width", + group=object(), + ) + finally: + wan_spatial_shard._SPATIAL_SHARD_CONTEXT.reset(token) + + assert out.shape == (1, 1, 1, 1, 2) + assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1)) + + x = torch.randn((1, 1, 1, 4, 2)) + out, recv_top, recv_bottom = wan_spatial_shard.halo_exchange( + x, + group=object(), + halo_size=1, + ) + + assert out is x + assert recv_top is None + assert recv_bottom is None + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.MonkeyPatch): + x = torch.ones((1, 1, 2, 2)) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) + mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object()) + mid = mid_rank_pad(x) + assert mid.shape == (1, 1, 2, 3) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((0, 1, 1, 1), group=object()) + last = last_rank_pad(x) + assert last.shape == (1, 1, 3, 3) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.MonkeyPatch): + x = torch.ones((1, 1, 2, 2)) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3)) + mid_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width") + mid = mid_rank_pad(x) + assert mid.shape == (1, 1, 2, 2) + + monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3)) + last_rank_pad = wan_spatial_shard.WanDistZeroPad2d((1, 1, 0, 0), group=object(), split_dim="width") + last = last_rank_pad(x) + assert last.shape == (1, 1, 2, 3) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size", + lambda group=None: 4, + ) + + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.use_tiling = True + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="spatial_shard_height") + vae.is_distributed_enabled = lambda: True + + z = torch.zeros((1, 16, 1, 8, 8)) + + assert vae._spatial_shard_decode_split_dim() == "height" + assert vae._spatial_shard_decode_enabled(z) is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr( + "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size", + lambda group=None: 2, + ) + + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="spatial_shard_width") + vae.is_distributed_enabled = lambda: True + + z = torch.zeros((1, 16, 1, 8, 8)) + + assert vae._spatial_shard_decode_split_dim() == "width" + assert vae._spatial_shard_decode_enabled(z) is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_tile_mode_disables_spatial_shard_decode(): + vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan) + vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="tile") + vae.is_distributed_enabled = lambda: True + + z = torch.zeros((1, 16, 1, 8, 8)) + + assert vae._spatial_shard_decode_split_dim() is None + assert vae._spatial_shard_decode_enabled(z) is False + + +# ============================================================================= +# Multi-GPU numerical-correctness test (nightly Diffusion Test group, H100 x2) +# +# Spawns a small process group and verifies that spatial_shard_height/spatial_shard_width decode match +# a single-process (non-distributed) reference decode of the same latent within +# tolerance. Requires >= 2 accelerator devices and downloads a Wan VAE. +# ============================================================================= + +_SPATIAL_SHARD_MODEL = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" +_SPATIAL_SHARD_SUBFOLDER = "vae" +_SPATIAL_SHARD_WORLD_SIZE = 2 +_SPATIAL_SHARD_LATENT_FRAMES = 5 +_SPATIAL_SHARD_LATENT_HEIGHT = 60 +_SPATIAL_SHARD_LATENT_WIDTH = 104 +_SPATIAL_SHARD_TOLERANCE = 3e-2 + + +def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_port: str) -> None: + from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = master_port + device = current_omni_platform.get_torch_device(rank) + current_omni_platform.set_device(device) + dtype = torch.float32 + + backend = current_omni_platform.dist_backend + init_distributed_environment(world_size=_SPATIAL_SHARD_WORLD_SIZE, rank=rank, local_rank=rank, backend=backend) + initialize_model_parallel( + sequence_parallel_size=_SPATIAL_SHARD_WORLD_SIZE, ulysses_degree=_SPATIAL_SHARD_WORLD_SIZE, backend=backend + ) + + try: + vae = DistributedAutoencoderKLWan.from_pretrained( + _SPATIAL_SHARD_MODEL, subfolder=_SPATIAL_SHARD_SUBFOLDER, torch_dtype=dtype + ) + vae.to(device=device, dtype=dtype) + vae.eval() + + generator = torch.Generator(device=device).manual_seed(0) + latents = torch.randn( + ( + 1, + vae.config.z_dim, + _SPATIAL_SHARD_LATENT_FRAMES, + _SPATIAL_SHARD_LATENT_HEIGHT, + _SPATIAL_SHARD_LATENT_WIDTH, + ), + generator=generator, + device=device, + dtype=dtype, + ) + + with torch.inference_mode(): + # Ground-truth reference: standard non-parallel, untiled decode (computed identically on + # every rank). Tiling must be OFF so neither the tile-parallel nor the single-GPU tiled + # path is exercised; otherwise we would be comparing SP against tiled decode. + vae.use_tiling = False + vae.set_parallel_size(1, mode="tile") + reference = vae.decode(latents, return_dict=False)[0].float() + + # Spatially-sharded decode across the full group (requires tiling to enter tiled_decode). + vae.use_tiling = True + vae.set_parallel_size(_SPATIAL_SHARD_WORLD_SIZE, mode=f"spatial_shard_{split_dim}") + sharded = vae.decode(latents, return_dict=False)[0].float() + + # Only rank 0 assembles the full decoded sample (matching broadcast_result=False); + # non-zero ranks return an empty placeholder, so the comparison runs on rank 0 only. + if rank == 0: + diff = (sharded - reference).abs() + return_dict["max_abs_diff"] = diff.max().item() + return_dict["mean_abs_diff"] = diff.mean().item() + return_dict["shape"] = tuple(sharded.shape) + finally: + destroy_model_parallel() + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.parallel +@hardware_test(res={"cuda": "H100"}, num_cards=_SPATIAL_SHARD_WORLD_SIZE) +@pytest.mark.parametrize("split_dim", ["height", "width"]) +def test_spatial_shard_decode_matches_reference(split_dim: str): + manager = mp.get_context("spawn").Manager() + return_dict = manager.dict() + # Use a per-split-dim port to avoid collisions across parametrized runs. + master_port = str(29500 + (1 if split_dim == "width" else 0)) + + mp.spawn( + _spatial_shard_decode_worker, + args=(split_dim, return_dict, master_port), + nprocs=_SPATIAL_SHARD_WORLD_SIZE, + join=True, + ) + + assert "max_abs_diff" in return_dict, "rank 0 did not report a result" + max_abs_diff = return_dict["max_abs_diff"] + mean_abs_diff = return_dict["mean_abs_diff"] + print( + f"spatial_shard_{split_dim} vs reference: max_abs_diff={max_abs_diff:.6e} " + f"mean_abs_diff={mean_abs_diff:.6e} shape={return_dict.get('shape')}" + ) + assert max_abs_diff <= _SPATIAL_SHARD_TOLERANCE, ( + f"spatial_shard_{split_dim} max_abs_diff {max_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}" + ) + assert mean_abs_diff <= _SPATIAL_SHARD_TOLERANCE, ( + f"spatial_shard_{split_dim} mean_abs_diff {mean_abs_diff} exceeds {_SPATIAL_SHARD_TOLERANCE}" + ) diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py new file mode 100644 index 00000000000..babbc7abc74 --- /dev/null +++ b/tests/diffusion/kernels/mot/test_mot_linear.py @@ -0,0 +1,689 @@ +# ruff: noqa: N803, E741 +"""Layer-level correctness & performance test for MoT parallel linear layers. + +Compares two equivalent computation paths: + - Reference: 2x standard vLLM parallel linear layers + PyTorch index + scatter/gather (text_linear(x[text_idx]) + vae_linear(x[vae_idx])) + - Target: 1x MoT fused parallel linear layer + (mot_linear(x, text_indices, vae_indices)) + +The reference path uses cuBLAS GEMM (always auto-tuned by cuBLAS). +The MoT path uses a fused Triton kernel whose tile config is loaded +from a JSON file matched by ``device + dtype``. If no tuned +config is found for the current GPU, the kernel falls back to a conservative +default and a warning is printed — the correctness test still passes +but the performance comparison is NOT representative. + +Usage:: + pytest tests/diffusion/kernels/mot/test_mot_linear.py -v -s +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass + +import pytest +import torch +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, +) + +from vllm_omni.diffusion.distributed.parallel_state import ( + destroy_distributed_env, + init_distributed_environment, + initialize_model_parallel, + model_parallel_is_initialized, +) +from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import ( + MoTQKVParallelLinear, +) +from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import ( + MoTRowParallelLinear, +) +from vllm_omni.diffusion.layers.mot.ops.mot_gemm import ( + get_best_mot_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +# BAGEL-7B-MoT architecture parameters +_BAGEL_HEAD_SIZE = 128 +_BAGEL_TOTAL_NUM_HEADS = 28 +_BAGEL_TOTAL_NUM_KV_HEADS = 4 +_VAE_CHUNK_SIZE = 1024 # the token number of one image +_IMAGE_NUM = [1, 2, 4, 8] + + +@pytest.fixture(scope="module", autouse=True) +def _init_single_rank_tp_env(): + """Initialize single-rank distributed/TP env for vLLM linear params.""" + os.environ.setdefault("RANK", "0") + os.environ.setdefault("LOCAL_RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + os.environ.setdefault("MASTER_ADDR", "127.0.0.1") + os.environ.setdefault("MASTER_PORT", "29501") + + if not torch.distributed.is_initialized(): + init_distributed_environment(world_size=1, rank=0, local_rank=0) + + if not model_parallel_is_initialized(): + initialize_model_parallel( + data_parallel_size=1, + cfg_parallel_size=1, + sequence_parallel_size=1, + ulysses_degree=1, + ring_degree=1, + tensor_parallel_size=1, + pipeline_parallel_size=1, + ) + + yield + destroy_distributed_env() + + +# --------------------------------------------------------------------------- +# DType configuration — extensible for future quantized types +# --------------------------------------------------------------------------- + + +@dataclass +class DTypeConfig: + """Parsed dtype configuration for a test case.""" + + torch_dtype: torch.dtype + use_fp8_w8a8: bool = False + use_int8_w8a16: bool = False + use_int4_w4a16: bool = False + + +def _parse_dtype(dtype_str: str) -> DTypeConfig: + """Parse a dtype string into quantization flags and torch dtype. + + Supported now: + "w16a16_bf16" — BF16 weights & activations (no quantization) + "w16a16_fp16" — FP16 weights & activations (no quantization) + Reserved for future: + "fp8_w8a8" — FP8 W8A8 quantization + "int8_w8a16" — INT8 weight-only quantization + "int4_w4a16" — INT4 weight-only quantization + """ + supported: dict[str, DTypeConfig] = { + "w16a16_bf16": DTypeConfig(torch_dtype=torch.bfloat16), + "w16a16_fp16": DTypeConfig(torch_dtype=torch.float16), + } + if dtype_str in supported: + return supported[dtype_str] + pytest.skip(f"Quantized dtype '{dtype_str}' not yet implemented in layer test") + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _report_mot_config(K: int, N: int, M: int): + """Print which Triton tile config the MoT kernel will use.""" + loaded_m_key, config = get_best_mot_config(M, N, K, None) + if loaded_m_key == -1: + print( + " [config] WARNING: No tuned config found — " + "using conservative default. " + "Performance numbers are NOT representative. " + "Run mot_linear_benchmarks.py --tune to generate configs." + ) + else: + print(f" [config] Tuned config loaded (actual M={M}, loaded M={loaded_m_key}) config = {config})") + + +def _make_indices(image_num: int, vae_chunk_size: int, device: str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]: + """ + Simulate exact Bagel-MoT distributions for image generation: + Pattern per image like: [1 Text] + [4096 VAE] + [1 Text] + Returns text_indices, vae_indices, and the exact total M. + """ + text_idx_list = [] + vae_idx_list = [] + + current_idx = 0 + for _ in range(image_num): + text_idx_list.append(current_idx) + current_idx += 1 + + vae_idx_list.extend(range(current_idx, current_idx + vae_chunk_size)) + current_idx += vae_chunk_size + + text_idx_list.append(current_idx) + current_idx += 1 + + text_indices = torch.tensor(text_idx_list, dtype=torch.long, device=device) + vae_indices = torch.tensor(vae_idx_list, dtype=torch.long, device=device) + + exact_M = current_idx # exact_M = image_num * (vae_chunk_size + 2) + + return text_indices, vae_indices, exact_M + + +def _benchmark(fn, warmup: int = 20, iters: int = 100) -> float: + """Return mean latency in milliseconds.""" + cache_flusher = torch.empty(int(256 * 1024 * 1024 / 4), dtype=torch.int32, device="cuda") + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + + t_flush_start = time.perf_counter() + for _ in range(iters): + cache_flusher.zero_() + torch.accelerator.synchronize() + flush_time_total = time.perf_counter() - t_flush_start + + # 3. Measure the total time of "flush cache + operator execution" + torch.accelerator.synchronize() + t_total_start = time.perf_counter() + for _ in range(iters): + cache_flusher.zero_() + fn() + torch.accelerator.synchronize() + total_time = time.perf_counter() - t_total_start + + # 4. Asynchronous subtraction separation + # The E2E time of the pure operator = total time - flush time + pure_fn_time_total = total_time - flush_time_total + avg_ms = (pure_fn_time_total / iters) * 1000.0 + + del cache_flusher + return max(avg_ms, 0.001) + + +def _sync_weights(ref_text, ref_vae, mot_layer): + """Assign the same random weights to reference layers and MoT layer.""" + with torch.no_grad(): + W_text = torch.randn_like(ref_text.weight) * 0.02 + W_vae = torch.randn_like(ref_vae.weight) * 0.02 + ref_text.weight.copy_(W_text) + ref_vae.weight.copy_(W_vae) + mot_layer.weight.copy_(W_text) + mot_layer.gen_exp.weight.copy_(W_vae) + if ref_text.bias is not None and mot_layer.bias is not None: + b_text = torch.randn_like(ref_text.bias) * 0.02 + ref_text.bias.copy_(b_text) + mot_layer.bias.copy_(b_text) + if ref_vae.bias is not None and mot_layer.gen_exp.bias is not None: + b_vae = torch.randn_like(ref_vae.bias) * 0.02 + ref_vae.bias.copy_(b_vae) + mot_layer.gen_exp.bias.copy_(b_vae) + + +def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear): + """Reference path: index-gather → 2x standard linear → index-scatter.""" + M = x.size(0) + out_text = text_linear(x[text_indices]) + out_vae = vae_linear(x[vae_indices]) + if isinstance(out_text, tuple): + out_text = out_text[0] + if isinstance(out_vae, tuple): + out_vae = out_vae[0] + N = out_text.size(-1) + output = torch.empty(M, N, dtype=x.dtype, device=x.device) + output[text_indices] = out_text + output[vae_indices] = out_vae + return output + + +def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str): + """Compare outputs, print metrics, and assert correctness. + + Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16). + We upcast to fp32 solely for computing error metrics with higher + arithmetic precision — the actual layer outputs remain bf16. + """ + # Upcast for numerically stable error computation only + ref_hp = ref.float() + mot_hp = mot.float() + + abs_err = (ref_hp - mot_hp).abs() + max_abs = abs_err.max().item() + + # Mixed metric: relative error where |ref| >= 1, absolute error otherwise + denom = ref_hp.abs().clamp(min=1.0) + max_rel = (abs_err / denom).max().item() + + cos_sim = ( + torch.nn.functional.cosine_similarity( + ref_hp, + mot_hp, + dim=-1, + ) + .min() + .item() + ) + + print(f"\n [{tag}] max_abs={max_abs:.4e} max_rel={max_rel:.4e} min_cos_sim={cos_sim:.6f}") + + # Cosine similarity is the primary correctness gate: robust to scale + # and accumulation-order differences between cuBLAS and Triton. + # For bf16 GEMM with K up to ~20k, cos_sim > 0.99 is easily achieved. + assert cos_sim > 0.98, f"Cosine similarity too low: {cos_sim:.6f}" + # Supplementary per-element check (generous to avoid flaky failures + # on extreme K dimensions like 18944) + assert max_rel < 0.1, f"Max relative error too large: {max_rel:.4e}" + + +def _run_timing( + ref_fn, + mot_fn, + tag: str, + warmup: int = 20, + iters: int = 100, +): + """Benchmark both paths and print timing comparison.""" + ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters) + mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters) + speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf") + print(f" [{tag}] Ref(2x linear): {ref_ms:.3f} ms | MoT(fused): {mot_ms:.3f} ms | Speedup: {speedup:.2f}x") + + +# ========================================================================= +# Test: qkv proj +# ========================================================================= + + +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize( + "image_num, K, N, dtype", + [(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM] + [(num, 3584, 4608, "w16a16_fp16") for num in _IMAGE_NUM], + ids=[f"img{num}_K3584_N4608_bf16" for num in _IMAGE_NUM] + [f"img{num}_K3584_N4608_fp16" for num in _IMAGE_NUM], +) +def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool): + dcfg = _parse_dtype(dtype) + torch.manual_seed(42) + + with set_current_vllm_config(VllmConfig()): + text_linear = QKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=bias, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + vae_linear = QKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=bias, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + mot_linear = MoTQKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=bias, + vae_bias=bias, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + + assert text_linear.output_size_per_partition == N, ( + f"Expected output_size_per_partition={N}, " + f"got {text_linear.output_size_per_partition}. " + f"Check head parameters." + ) + + _sync_weights(text_linear, vae_linear, mot_linear) + + text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE) + x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") + + tag = f"QKVParallel M={M} K={K} N={N}" + _report_mot_config(K, N, M) + + with torch.no_grad(): + ref = _reference_forward( + x, + text_idx, + vae_idx, + text_linear, + vae_linear, + ) + mot_out, _ = mot_linear(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, tag) + + with torch.no_grad(): + _run_timing( + lambda: _reference_forward( + x, + text_idx, + vae_idx, + text_linear, + vae_linear, + ), + lambda: mot_linear(x, text_idx, vae_idx), + tag, + ) + + +# ========================================================================= +# Test: o proj +# ========================================================================= + + +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize( + "image_num, K, N, dtype", + [(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM] + [(num, 3584, 3584, "w16a16_fp16") for num in _IMAGE_NUM], + ids=[f"img{num}_K3584_N3584_bf16" for num in _IMAGE_NUM] + [f"img{num}_K3584_N3584_fp16" for num in _IMAGE_NUM], +) +def test_mot_o_proj( + image_num: int, + K: int, + N: int, + dtype: str, + bias: bool, +): + dcfg = _parse_dtype(dtype) + torch.manual_seed(42) + + with set_current_vllm_config(VllmConfig()): + text_linear = RowParallelLinear( + K, + N, + bias=bias, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + vae_linear = RowParallelLinear( + K, + N, + bias=bias, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + mot_linear = MoTRowParallelLinear( + K, + N, + bias=bias, + vae_bias=bias, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + + _sync_weights(text_linear, vae_linear, mot_linear) + + text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE) + x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") + + tag = f"O Proj M={M} K={K} N={N}" + _report_mot_config(K, N, M) + + # Correctness (also warms up Triton JIT compilation) + with torch.no_grad(): + ref = _reference_forward( + x, + text_idx, + vae_idx, + text_linear, + vae_linear, + ) + mot_out, _ = mot_linear(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, tag) + + # Performance + with torch.no_grad(): + _run_timing( + lambda: _reference_forward( + x, + text_idx, + vae_idx, + text_linear, + vae_linear, + ), + lambda: mot_linear(x, text_idx, vae_idx), + tag, + ) + + +# ========================================================================= +# Test: und-mode (text_indices=None) — falls back to standard forward +# ========================================================================= + + +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize( + "K, N, dtype", + [ + (3584, 4608, "w16a16_bf16"), + (3584, 4608, "w16a16_fp16"), + ], + ids=["QKV_K3584_N4608_bf16", "QKV_K3584_N4608_fp16"], +) +def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool): + """und-mode: text_indices=None should produce same output as standard QKVParallelLinear.""" + dcfg = _parse_dtype(dtype) + torch.manual_seed(42) + M = 1026 + + with set_current_vllm_config(VllmConfig()): + ref_linear = QKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=bias, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + mot_linear = MoTQKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=bias, + vae_bias=bias, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + + with torch.no_grad(): + mot_linear.weight.copy_(ref_linear.weight) + if bias and ref_linear.bias is not None: + mot_linear.bias.copy_(ref_linear.bias) + + x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") + + with torch.no_grad(): + ref_out, _ = ref_linear(x) + mot_out, _ = mot_linear(x, text_indices=None, vae_indices=None) + + _check_and_report(ref_out, mot_out, f"QKV und-mode M={M} K={K}") + + +@pytest.mark.parametrize("bias", [True, False]) +@pytest.mark.parametrize( + "K, N, dtype", + [ + (3584, 3584, "w16a16_bf16"), + (3584, 3584, "w16a16_fp16"), + ], + ids=["Row_K3584_N3584_bf16", "Row_K3584_N3584_fp16"], +) +def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool): + """und-mode: text_indices=None should produce same output as standard RowParallelLinear.""" + dcfg = _parse_dtype(dtype) + torch.manual_seed(42) + M = 1026 + + with set_current_vllm_config(VllmConfig()): + ref_linear = RowParallelLinear( + K, + N, + bias=bias, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + mot_linear = MoTRowParallelLinear( + K, + N, + bias=bias, + vae_bias=bias, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + + with torch.no_grad(): + mot_linear.weight.copy_(ref_linear.weight) + if bias and ref_linear.bias is not None: + mot_linear.bias.copy_(ref_linear.bias) + + x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") + + with torch.no_grad(): + ref_out, _ = ref_linear(x) + mot_out, _ = mot_linear(x, text_indices=None, vae_indices=None) + + _check_and_report(ref_out, mot_out, f"Row und-mode M={M} K={K}") + + +# ========================================================================= +# Test: boundary cases (all-text, all-VAE) +# ========================================================================= + + +@pytest.mark.parametrize( + "boundary_mode", + ["all_text", "all_vae"], + ids=["all_text", "all_vae"], +) +@pytest.mark.parametrize("dtype", ["w16a16_bf16", "w16a16_fp16"]) +def test_mot_qkv_boundary(boundary_mode: str, dtype: str): + """Boundary: all tokens routed to one expert.""" + dcfg = _parse_dtype(dtype) + torch.manual_seed(42) + M = 512 + K = 3584 + + with set_current_vllm_config(VllmConfig()): + text_linear = QKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=False, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + vae_linear = QKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=False, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + mot_linear = MoTQKVParallelLinear( + hidden_size=K, + head_size=_BAGEL_HEAD_SIZE, + total_num_heads=_BAGEL_TOTAL_NUM_HEADS, + total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS, + bias=False, + vae_bias=False, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + + _sync_weights(text_linear, vae_linear, mot_linear) + + all_indices = torch.arange(M, dtype=torch.long, device="cuda") + empty_indices = torch.empty(0, dtype=torch.long, device="cuda") + + if boundary_mode == "all_text": + text_idx, vae_idx = all_indices, empty_indices + else: + text_idx, vae_idx = empty_indices, all_indices + + x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") + + with torch.no_grad(): + ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear) + mot_out, _ = mot_linear(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, f"QKV boundary={boundary_mode} M={M}") + + +@pytest.mark.parametrize( + "boundary_mode", + ["all_text", "all_vae"], + ids=["all_text", "all_vae"], +) +@pytest.mark.parametrize("dtype", ["w16a16_bf16", "w16a16_fp16"]) +def test_mot_row_boundary(boundary_mode: str, dtype: str): + """Boundary: all tokens routed to one expert for RowParallel.""" + dcfg = _parse_dtype(dtype) + torch.manual_seed(42) + M = 512 + K = 3584 + N = 3584 + + with set_current_vllm_config(VllmConfig()): + text_linear = RowParallelLinear( + K, + N, + bias=False, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + vae_linear = RowParallelLinear( + K, + N, + bias=False, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + mot_linear = MoTRowParallelLinear( + K, + N, + bias=False, + vae_bias=False, + input_is_parallel=True, + params_dtype=dcfg.torch_dtype, + disable_tp=True, + ).cuda() + + _sync_weights(text_linear, vae_linear, mot_linear) + + all_indices = torch.arange(M, dtype=torch.long, device="cuda") + empty_indices = torch.empty(0, dtype=torch.long, device="cuda") + + if boundary_mode == "all_text": + text_idx, vae_idx = all_indices, empty_indices + else: + text_idx, vae_idx = empty_indices, all_indices + + x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda") + + with torch.no_grad(): + ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear) + mot_out, _ = mot_linear(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, f"Row boundary={boundary_mode} M={M}") diff --git a/tests/diffusion/kernels/mot/test_mot_norm.py b/tests/diffusion/kernels/mot/test_mot_norm.py new file mode 100644 index 00000000000..1153f61126b --- /dev/null +++ b/tests/diffusion/kernels/mot/test_mot_norm.py @@ -0,0 +1,357 @@ +# ruff: noqa: N803, E741 +"""Layer-level correctness & performance test for MoTRMSNorm. + +Compares two equivalent computation paths: + - Reference: 2x vLLM IR RMSNorm op + PyTorch index scatter/gather + (ir.ops.rms_norm(x[text_idx], text_w) + ir.ops.rms_norm(x[vae_idx], vae_w)) + - Target: 1x MoTRMSNorm fused Triton kernel + (mot_norm(x, text_indices, vae_indices)) + +Usage:: + pytest tests/diffusion/kernels/mot/test_mot_norm.py -v -s +""" + +from __future__ import annotations + +import time + +import pytest +import torch +from vllm import ir +from vllm.config import VllmConfig, set_current_vllm_config + +from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +_EPS = 1e-6 + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_indices( + M: int, + text_ratio: float, + device: str = "cuda", +) -> tuple[torch.Tensor, torch.Tensor]: + M_text = max(1, int(M * text_ratio)) + perm = torch.randperm(M, device=device) + text_indices = perm[:M_text].sort().values + vae_indices = perm[M_text:].sort().values + return text_indices, vae_indices + + +def _benchmark(fn, warmup: int = 20, iters: int = 200) -> float: + """Return mean latency in milliseconds.""" + for _ in range(warmup): + fn() + torch.accelerator.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn() + torch.accelerator.synchronize() + return (time.perf_counter() - t0) / iters * 1000.0 + + +def _reference_forward( + x: torch.Tensor, + text_indices: torch.Tensor, + vae_indices: torch.Tensor, + text_weight: torch.Tensor, + vae_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Reference: index-gather -> 2x vLLM IR RMSNorm op -> scatter.""" + output = torch.empty_like(x) + output[text_indices] = ir.ops.rms_norm(x[text_indices], text_weight, eps) + output[vae_indices] = ir.ops.rms_norm(x[vae_indices], vae_weight, eps) + return output + + +def _reference_forward_head_norm( + x: torch.Tensor, + text_indices: torch.Tensor, + vae_indices: torch.Tensor, + text_weight: torch.Tensor, + vae_weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + """Reference for head_norm=True path. + + The MoT routing happens on token dimension while RMSNorm is applied + independently on each head's last dimension. + """ + output = torch.empty_like(x) + hidden_size = x.shape[-1] + output[text_indices] = ir.ops.rms_norm( + x[text_indices].reshape(-1, hidden_size), + text_weight, + eps, + ).reshape_as(x[text_indices]) + output[vae_indices] = ir.ops.rms_norm( + x[vae_indices].reshape(-1, hidden_size), + vae_weight, + eps, + ).reshape_as(x[vae_indices]) + return output + + +def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str): + """Compare outputs, print metrics, and assert correctness. + + Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16). + We upcast to fp32 solely for computing error metrics with higher + arithmetic precision — the actual layer outputs remain bf16. + """ + ref_hp = ref.float() + mot_hp = mot.float() + + abs_err = (ref_hp - mot_hp).abs() + max_abs = abs_err.max().item() + + denom = ref_hp.abs().clamp(min=1.0) + max_rel = (abs_err / denom).max().item() + + cos_sim = ( + torch.nn.functional.cosine_similarity( + ref_hp, + mot_hp, + dim=-1, + ) + .min() + .item() + ) + + print(f"\n [{tag}] max_abs={max_abs:.4e} max_rel={max_rel:.4e} min_cos_sim={cos_sim:.6f}") + + # RMSNorm is element-wise (no cross-element accumulation like GEMM), + # so the error between two fp32-accumulating implementations should + # be very small — well within bf16 rounding. + assert cos_sim > 0.99, f"Cosine similarity too low: {cos_sim:.6f}" + assert max_rel < 0.05, f"Max relative error too large: {max_rel:.4e}" + + +def _run_timing( + ref_fn, + mot_fn, + tag: str, + warmup: int = 20, + iters: int = 200, +): + """Benchmark both paths and print timing comparison.""" + ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters) + mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters) + speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf") + print(f" [{tag}] Ref(2x rms_norm): {ref_ms:.3f} ms | MoT(fused): {mot_ms:.3f} ms | Speedup: {speedup:.2f}x") + + +# ========================================================================= +# Test: MoTRMSNorm vs 2x vLLM rms_norm +# ========================================================================= + + +@pytest.mark.parametrize( + "M, hidden_size, text_ratio, dtype", + [ + (2048, 3584, 0.01, torch.bfloat16), + (8192, 3584, 0.01, torch.bfloat16), + (2048, 128, 0.01, torch.bfloat16), + (8192, 128, 0.01, torch.bfloat16), + (2048, 3584, 0.01, torch.float16), + (8192, 3584, 0.01, torch.float16), + ], + ids=[ + "M2048_H3584_layernorm_bf16", + "M8192_H3584_layernorm_bf16", + "M2048_H128_qknorm_bf16", + "M8192_H128_qknorm_bf16", + "M2048_H3584_layernorm_fp16", + "M8192_H3584_layernorm_fp16", + ], +) +def test_mot_rms_norm(M: int, hidden_size: int, text_ratio: float, dtype: torch.dtype): + torch.manual_seed(42) + + # --- Build MoT layer --- + mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda() + with torch.no_grad(): + W_text = torch.randn(hidden_size, dtype=dtype, device="cuda") + W_vae = torch.randn(hidden_size, dtype=dtype, device="cuda") + mot_norm.weight.data.copy_(W_text) + mot_norm.gen_weight.data.copy_(W_vae) + + # --- Build reference weights (same data) --- + ref_text_weight = W_text.clone() + ref_vae_weight = W_vae.clone() + + # --- Inputs --- + x = torch.randn(M, hidden_size, dtype=dtype, device="cuda") + text_idx, vae_idx = _make_indices(M, text_ratio) + + tag = f"RMSNorm M={M} H={hidden_size}" + + # vLLM's IR RMSNorm op may inspect global config; wrap in VllmConfig + # context for safety. + with set_current_vllm_config(VllmConfig()): + # --- Correctness (also warms up Triton JIT) --- + with torch.no_grad(): + ref = _reference_forward( + x, + text_idx, + vae_idx, + ref_text_weight, + ref_vae_weight, + _EPS, + ) + mot_out = mot_norm(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, tag) + + # --- Performance --- + with torch.no_grad(): + _run_timing( + lambda: _reference_forward( + x, + text_idx, + vae_idx, + ref_text_weight, + ref_vae_weight, + _EPS, + ), + lambda: mot_norm(x, text_idx, vae_idx), + tag, + ) + + +@pytest.mark.parametrize( + "M, num_heads, head_dim, text_ratio", + [ + (2048, 28, 128, 0.01), + ], + ids=[ + "M2048_NH28_HD128_qknorm_head_norm", + ], +) +def test_mot_rms_norm_head_norm( + M: int, + num_heads: int, + head_dim: int, + text_ratio: float, +): + torch.manual_seed(42) + dtype = torch.bfloat16 + + # --- Build MoT layer (head_norm path) --- + mot_norm = MoTRMSNorm(head_dim, head_norm=True, eps=_EPS).cuda() + with torch.no_grad(): + W_text = torch.randn(head_dim, dtype=dtype, device="cuda") + W_vae = torch.randn(head_dim, dtype=dtype, device="cuda") + mot_norm.weight.data.copy_(W_text) + mot_norm.gen_weight.data.copy_(W_vae) + + # --- Build reference weights (same data) --- + ref_text_weight = W_text.clone() + ref_vae_weight = W_vae.clone() + + # --- Inputs --- + x = torch.randn(M, num_heads, head_dim, dtype=dtype, device="cuda") + text_idx, vae_idx = _make_indices(M, text_ratio) + + tag = f"RMSNorm(head_norm=True) M={M} NH={num_heads} HD={head_dim}" + + with set_current_vllm_config(VllmConfig()): + # --- Correctness (also warms up Triton JIT) --- + with torch.no_grad(): + ref = _reference_forward_head_norm( + x, + text_idx, + vae_idx, + ref_text_weight, + ref_vae_weight, + _EPS, + ) + mot_out = mot_norm(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, tag) + + +# ========================================================================= +# Test: und-mode (text_indices=None) — degrades to standard RMSNorm +# ========================================================================= + + +@pytest.mark.parametrize( + "M, hidden_size, dtype", + [ + (2048, 3584, torch.bfloat16), + (2048, 3584, torch.float16), + (2048, 128, torch.bfloat16), + ], + ids=["M2048_H3584_bf16", "M2048_H3584_fp16", "M2048_H128_bf16"], +) +def test_mot_rms_norm_und_mode(M: int, hidden_size: int, dtype: torch.dtype): + """und-mode: text_indices=None should apply self.weight to all tokens (standard RMSNorm).""" + torch.manual_seed(42) + + mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda() + with torch.no_grad(): + W_text = torch.randn(hidden_size, dtype=dtype, device="cuda") + mot_norm.weight.data.copy_(W_text) + + x = torch.randn(M, hidden_size, dtype=dtype, device="cuda") + + with set_current_vllm_config(VllmConfig()): + with torch.no_grad(): + ref = ir.ops.rms_norm(x, W_text, _EPS) + mot_out = mot_norm(x, text_indices=None, vae_indices=None) + + _check_and_report(ref, mot_out, f"RMSNorm und-mode M={M} H={hidden_size}") + + +# ========================================================================= +# Test: boundary cases (all-text, all-VAE) +# ========================================================================= + + +@pytest.mark.parametrize( + "boundary_mode", + ["all_text", "all_vae"], + ids=["all_text", "all_vae"], +) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +def test_mot_rms_norm_boundary(boundary_mode: str, dtype: torch.dtype): + """Boundary: all tokens routed to one expert.""" + torch.manual_seed(42) + M = 1024 + hidden_size = 3584 + + mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda() + with torch.no_grad(): + W_text = torch.randn(hidden_size, dtype=dtype, device="cuda") + W_vae = torch.randn(hidden_size, dtype=dtype, device="cuda") + mot_norm.weight.data.copy_(W_text) + mot_norm.gen_weight.data.copy_(W_vae) + + x = torch.randn(M, hidden_size, dtype=dtype, device="cuda") + all_indices = torch.arange(M, dtype=torch.long, device="cuda") + empty_indices = torch.empty(0, dtype=torch.long, device="cuda") + + if boundary_mode == "all_text": + text_idx, vae_idx = all_indices, empty_indices + expected_weight = W_text + else: + text_idx, vae_idx = empty_indices, all_indices + expected_weight = W_vae + + with set_current_vllm_config(VllmConfig()): + with torch.no_grad(): + ref = ir.ops.rms_norm(x, expected_weight, _EPS) + mot_out = mot_norm(x, text_idx, vae_idx) + + _check_and_report(ref, mot_out, f"RMSNorm boundary={boundary_mode} M={M}") diff --git a/tests/diffusion/layers/test_rope_broadcast.py b/tests/diffusion/layers/test_rope_broadcast.py new file mode 100644 index 00000000000..29fcf180428 --- /dev/null +++ b/tests/diffusion/layers/test_rope_broadcast.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shape broadcasting regression tests for RotaryEmbedding.forward_native. + +Covers the batch-dimension shape mismatch fix: forward_native now strips the +batch dim from 3D cos/sin before applying RoPE, consistent with forward_cuda, +forward_hip, and apply_rotary_emb_mindiesd paths. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_inputs( + batch: int, + seq_len: int, + num_heads: int, + head_dim: int, + cos_batch: int | None = None, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Generate x, cos, sin tensors for RoPE testing. + + Args: + batch: Batch dimension of x. + seq_len: Sequence length. + num_heads: Number of attention heads. + head_dim: Head dimension (must be even). + cos_batch: Batch dimension of cos/sin. Defaults to batch if None. + dtype: Data type for the tensors. + """ + if cos_batch is None: + cos_batch = batch + + torch.manual_seed(42) + x = torch.randn(batch, seq_len, num_heads, head_dim, dtype=dtype) + cos = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype) + sin = torch.randn(cos_batch, seq_len, head_dim // 2, dtype=dtype) + return x, cos, sin + + +class TestRotaryEmbeddingNativeShapeRegression: + """Regression tests for forward_native shape handling (#8e297d81).""" + + @pytest.fixture + def rope_neox(self) -> RotaryEmbedding: + """NeoX-style (non-interleaved) RoPE.""" + return RotaryEmbedding(is_neox_style=True) + + @pytest.fixture + def rope_gptj(self) -> RotaryEmbedding: + """GPT-J style (interleaved) RoPE.""" + return RotaryEmbedding(is_neox_style=False) + + # ── 2D cos/sin (no batch dim) ────────────────────────────────────── + + def test_2d_cos_sin_unchanged(self, rope_neox: RotaryEmbedding) -> None: + """2D cos/sin [S, D/2] pass through unchanged (no-op path).""" + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=2) + # Manually drop batch dim to create 2D cos/sin + cos_2d = cos[0] + sin_2d = sin[0] + assert cos_2d.dim() == 2 + output = rope_neox.forward_native(x, cos_2d, sin_2d) + assert output.shape == x.shape + assert not torch.isnan(output).any() + assert not torch.isinf(output).any() + + # ── 3D cos/sin — batch-dim stripping ────────────────────────────── + + @pytest.mark.parametrize("cos_batch", [1, 2, 4]) + def test_3d_cos_sin_no_error(self, rope_neox: RotaryEmbedding, cos_batch: int) -> None: + """3D cos/sin with batch dim should not raise errors.""" + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch) + output = rope_neox.forward_native(x, cos, sin) + assert output.shape == x.shape + + @pytest.mark.parametrize("interleaved", [False, True]) + def test_3d_cos_identical_to_2d(self, interleaved: bool) -> None: + """3D cos/sin with identical batch slices == 2D cos/sin result.""" + rope = RotaryEmbedding(is_neox_style=not interleaved) + B, S, H, D = 2, 16, 8, 64 + x = torch.randn(B, S, H, D) + cos_2d = torch.randn(S, D // 2) + sin_2d = torch.randn(S, D // 2) + + # 3D variant: expand the same 2D data across batch + cos_3d = cos_2d.unsqueeze(0).expand(B, -1, -1) + sin_3d = sin_2d.unsqueeze(0).expand(B, -1, -1) + + out_2d = rope.forward_native(x, cos_2d, sin_2d) + out_3d = rope.forward_native(x, cos_3d, sin_3d) + + torch.testing.assert_close(out_3d, out_2d, atol=0, rtol=0) + + # ── Batch size mismatch — the regression scenario ───────────────── + + @pytest.mark.parametrize( + "x_batch,cos_batch", + [ + (2, 1), # cos has fewer batch elements + (1, 2), # cos has more batch elements (takes [0]) + (4, 1), # large x batch, single cos + ], + ) + def test_batch_size_mismatch_regression( + self, + rope_neox: RotaryEmbedding, + x_batch: int, + cos_batch: int, + ) -> None: + """forward_native works when cos batch dim differs from x batch dim. + + This is the exact regression scenario — before the fix, a 3D cos/sin + with a different batch size than x would cause a shape mismatch error + inside apply_rotary_emb_torch during broadcast. + """ + x, cos, sin = _make_inputs(batch=x_batch, seq_len=16, num_heads=8, head_dim=64, cos_batch=cos_batch) + output = rope_neox.forward_native(x, cos, sin) + assert output.shape == x.shape + assert output.dtype == x.dtype + + # ── Batch consistency — all elements get same rotation ──────────── + + def test_2d_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None: + """With 2D cos, identical x slices → identical output slices.""" + B, S, H, D = 4, 16, 8, 64 + x_template = torch.randn(1, S, H, D) + x = x_template.expand(B, -1, -1, -1) # all batch elements identical + cos = torch.randn(S, D // 2) + sin = torch.randn(S, D // 2) + + output = rope_neox.forward_native(x, cos, sin) + + # All output batch slices must be identical + for i in range(1, B): + torch.testing.assert_close(output[i], output[0], atol=0, rtol=0) + + def test_3d_identical_cos_broadcasts_identically_across_batch(self, rope_neox: RotaryEmbedding) -> None: + """With 3D cos (identical slices), identical x → identical output.""" + B, S, H, D = 4, 16, 8, 64 + x_template = torch.randn(1, S, H, D) + x = x_template.expand(B, -1, -1, -1) + cos = torch.randn(1, S, D // 2).expand(B, -1, -1) + sin = torch.randn(1, S, D // 2).expand(B, -1, -1) + + output = rope_neox.forward_native(x, cos, sin) + + for i in range(1, B): + torch.testing.assert_close(output[i], output[0], atol=0, rtol=0) + + # ── Output shape / dtype invariants ─────────────────────────────── + + def test_output_shape_preserved(self, rope_neox: RotaryEmbedding) -> None: + """Output shape must match input shape for all cos/sin dimensionalities.""" + B, S, H, D = 2, 32, 4, 128 + x = torch.randn(B, S, H, D) + + # 2D cos/sin + cos_2d = torch.randn(S, D // 2) + sin_2d = torch.randn(S, D // 2) + assert rope_neox.forward_native(x, cos_2d, sin_2d).shape == (B, S, H, D) + + # 3D cos/sin, matching batch + cos_3d = torch.randn(B, S, D // 2) + sin_3d = torch.randn(B, S, D // 2) + assert rope_neox.forward_native(x, cos_3d, sin_3d).shape == (B, S, H, D) + + # 3D cos/sin, different batch + cos_3d_1 = torch.randn(1, S, D // 2) + sin_3d_1 = torch.randn(1, S, D // 2) + assert rope_neox.forward_native(x, cos_3d_1, sin_3d_1).shape == (B, S, H, D) + + def test_output_dtype_preserved(self, rope_neox: RotaryEmbedding) -> None: + """Output dtype must match input dtype.""" + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64, dtype=dtype) + output = rope_neox.forward_native(x, cos, sin) + assert output.dtype == dtype, f"dtype mismatch: {output.dtype} != {dtype}" + + def test_output_no_nan_inf( + self, + rope_neox: RotaryEmbedding, + rope_gptj: RotaryEmbedding, + ) -> None: + """Output must be finite for both RoPE styles.""" + for rope in (rope_neox, rope_gptj): + x, cos, sin = _make_inputs(batch=2, seq_len=32, num_heads=4, head_dim=64) + output = rope.forward_native(x, cos, sin) + assert not torch.isnan(output).any(), f"NaN in {rope}" + assert not torch.isinf(output).any(), f"Inf in {rope}" + + def test_output_contiguous(self, rope_neox: RotaryEmbedding) -> None: + """Output should be contiguous.""" + x, cos, sin = _make_inputs(batch=2, seq_len=16, num_heads=8, head_dim=64) + output = rope_neox.forward_native(x, cos, sin) + assert output.is_contiguous() + + # ── CFG-style doubled batch ─────────────────────────────────────── + + def test_cfg_style_doubled_batch(self, rope_neox: RotaryEmbedding) -> None: + """CFG (classifier-free guidance) doubles batch — cos may have + original batch size while x has 2× batch.""" + B, S, H, D = 2, 16, 8, 64 + # Simulate CFG: x has 2× batch (cond + uncond concatenated) + x = torch.randn(B * 2, S, H, D) + # cos/sin only have the original batch + cos = torch.randn(B, S, D // 2) + sin = torch.randn(B, S, D // 2) + + output = rope_neox.forward_native(x, cos, sin) + assert output.shape == (B * 2, S, H, D) + + # ── RotaryEmbeddingWan forward_native is unchanged ──────────────── + + def test_wan_forward_native_unchanged(self) -> None: + """RotaryEmbeddingWan.forward_native still uses its own implementation + (unflatten + stack + flatten) and should not crash with 3D cos/sin.""" + from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan + + rope = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + B, S, H, D = 2, 16, 8, 64 + x = torch.randn(B, S, H, D) + cos = torch.randn(B, S, 1, D // 2) + sin = torch.randn(B, S, 1, D // 2) + + output = rope.forward_native(x, cos, sin) + assert output.shape == x.shape + assert not torch.isnan(output).any() diff --git a/tests/diffusion/lora/test_lora_manager.py b/tests/diffusion/lora/test_lora_manager.py index 785f5d84217..723e4b13a93 100644 --- a/tests/diffusion/lora/test_lora_manager.py +++ b/tests/diffusion/lora/test_lora_manager.py @@ -575,3 +575,45 @@ def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs): assert "bagel.language_model.qkv_proj" in manager._lora_modules # Verify the module was actually replaced in the tree (not just recorded) assert isinstance(pipeline.bagel.language_model.qkv_proj, _DummyBaseLayerWithLoRA) + + +def test_lora_manager_discovers_unet_component(monkeypatch): + """Verify that _replace_layers_with_lora finds layers under 'unet'.""" + import vllm_omni.diffusion.lora.manager as manager_mod + + monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA) + + def _fake_from_layer_diffusion(*, layer: torch.nn.Module, **_kwargs): + if isinstance(layer, _FakeLinearBase): + return _DummyBaseLayerWithLoRA(layer) + return layer + + replace_calls: list[str] = [] + + monkeypatch.setattr(manager_mod, "from_layer_diffusion", _fake_from_layer_diffusion) + monkeypatch.setattr( + manager_mod, + "replace_submodule", + lambda root, name, sub: fake_replace_submodule(root, name, sub, replace_calls), + ) + + # Pipeline with a 'unet' component (no 'transformer') + pipeline = torch.nn.Module() + pipeline.unet = torch.nn.Module() + pipeline.unet.down_block = torch.nn.Module() + pipeline.unet.down_block.proj = _FakeLinearBase() + + manager = DiffusionLoRAManager( + pipeline=pipeline, + device=torch.device("cpu"), + dtype=torch.bfloat16, + max_cached_adapters=1, + ) + + peft_helper = type("_PH", (), {"r": 1})() + manager._replace_layers_with_lora(peft_helper) + + assert "down_block.proj" in replace_calls + assert "unet.down_block.proj" in manager._lora_modules + # Verify the module was actually replaced in the tree (not just recorded) + assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA) diff --git a/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py deleted file mode 100644 index 0c81dc80141..00000000000 --- a/tests/diffusion/model_loader/gguf_adapters/test_flux2_klein_gguf_adapter.py +++ /dev/null @@ -1,112 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the Flux2-Klein GGUF adapter.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter -from vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein import ( - Flux2KleinGGUFAdapter, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - -def _make_od_config( - model_class_name: str = "Flux2Pipeline", - model_type: str = "flux", -): - return SimpleNamespace( - model_class_name=model_class_name, - tf_model_config={"model_type": model_type}, - ) - - -def _make_source(prefix: str = "", subfolder: str = "transformer"): - return SimpleNamespace(prefix=prefix, subfolder=subfolder) - - -def test_flux2_klein_adapter_selected_for_flux_family(): - adapter = get_gguf_adapter( - "dummy.gguf", - object(), - _make_source(), - _make_od_config(), - ) - - assert isinstance(adapter, Flux2KleinGGUFAdapter) - - -def test_flux2_klein_adapter_matches_flux_model_type(): - assert Flux2KleinGGUFAdapter.is_compatible( - _make_od_config(model_class_name="OtherPipeline", model_type="flux-dev"), - object(), - _make_source(), - ) - - -def test_flux2_klein_adapter_renames_core_projection_paths(monkeypatch: pytest.MonkeyPatch): - import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module - - monkeypatch.setattr( - flux_module, - "gguf_quant_weights_iterator", - lambda _path: iter( - [ - ("img_in.qweight", torch.ones((1, 1))), - ("time_in.in_layer.qweight_type", torch.tensor(1)), - ("double_blocks.0.img_attn.norm.query_norm.weight", torch.full((1, 1), 2.0)), - ("final_layer.linear.qweight", torch.full((1, 1), 3.0)), - ] - ), - ) - - adapter = Flux2KleinGGUFAdapter( - "dummy.gguf", - object(), - _make_source(), - _make_od_config(), - ) - - weights = list(adapter.weights_iterator()) - names = [name for name, _ in weights] - - assert "x_embedder.qweight" in names - assert "time_guidance_embed.timestep_embedder.linear_1.qweight_type" in names - assert "transformer_blocks.0.attn.norm_q.weight" in names - assert "proj_out.qweight" in names - - -def test_flux2_klein_adapter_swaps_final_adaln_shift_and_scale(monkeypatch: pytest.MonkeyPatch): - import vllm_omni.diffusion.model_loader.gguf_adapters.flux2_klein as flux_module - - monkeypatch.setattr( - flux_module, - "gguf_quant_weights_iterator", - lambda _path: iter( - [ - ( - "final_layer.adaLN_modulation.1.scale", - torch.tensor([1.0, 2.0, 3.0, 4.0]), - ) - ] - ), - ) - - adapter = Flux2KleinGGUFAdapter( - "dummy.gguf", - object(), - _make_source(), - _make_od_config(), - ) - - weights = list(adapter.weights_iterator()) - - assert len(weights) == 1 - assert weights[0][0] == "norm_out.linear.weight" - assert torch.equal(weights[0][1], torch.tensor([3.0, 4.0, 1.0, 2.0])) diff --git a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py deleted file mode 100644 index 6b7392c80e2..00000000000 --- a/tests/diffusion/model_loader/gguf_adapters/test_qwen_image_gguf_adapter.py +++ /dev/null @@ -1,184 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from types import SimpleNamespace - -import pytest - -from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter -from vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image import ( - QwenImageGGUFAdapter, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - -def _make_od_config( - model_class_name: str = "QwenImagePipeline", - model_type: str = "qwen_image", -): - return SimpleNamespace( - model_class_name=model_class_name, - tf_model_config={"model_type": model_type}, - ) - - -def _make_source(prefix: str = "", subfolder: str = "transformer"): - return SimpleNamespace(prefix=prefix, subfolder=subfolder) - - -class _FakeTransformer: - def __init__(self) -> None: - self._params = { - "img_in.qweight": 1, - "img_in.qweight_type": 1, - "txt_in.qweight": 1, - "txt_in.qweight_type": 1, - "proj_out.qweight": 1, - "proj_out.qweight_type": 1, - "time_text_embed.timestep_embedder.linear_1.qweight": 1, - "time_text_embed.timestep_embedder.linear_1.qweight_type": 1, - "time_text_embed.timestep_embedder.linear_2.qweight": 1, - "time_text_embed.timestep_embedder.linear_2.qweight_type": 1, - "transformer_blocks.0.img_mod.1.qweight": 1, - "transformer_blocks.0.img_mod.1.qweight_type": 1, - "transformer_blocks.0.txt_mod.1.qweight": 1, - "transformer_blocks.0.txt_mod.1.qweight_type": 1, - "transformer_blocks.0.attn.to_qkv.qweight": 1, - "transformer_blocks.0.attn.to_qkv.qweight_type": 1, - "transformer_blocks.0.attn.add_kv_proj.qweight": 1, - "transformer_blocks.0.attn.add_kv_proj.qweight_type": 1, - "transformer_blocks.0.attn.to_out.weight": 1, - "transformer_blocks.0.img_mlp.net.0.proj.qweight": 1, - "transformer_blocks.0.img_mlp.net.0.proj.qweight_type": 1, - } - - def named_parameters(self): - return list(self._params.items()) - - def named_buffers(self): - return [] - - def get_submodule(self, _prefix: str): - return self - - -def test_qwen_adapter_selected_for_qwen_image_family(): - adapter = get_gguf_adapter( - "dummy.gguf", - _FakeTransformer(), - _make_source(), - _make_od_config(), - ) - assert isinstance(adapter, QwenImageGGUFAdapter) - - -def test_qwen_adapter_matches_multiple_pipeline_variants(): - for model_class_name in ( - "QwenImagePipeline", - "QwenImageEditPipeline", - "QwenImageEditPlusPipeline", - "QwenImageLayeredPipeline", - ): - assert QwenImageGGUFAdapter.is_compatible( - _make_od_config(model_class_name=model_class_name), - _FakeTransformer(), - _make_source(), - ) - - -def test_qwen_adapter_preserves_split_projection_names(monkeypatch: pytest.MonkeyPatch): - import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module - - monkeypatch.setattr( - qwen_image_module, - "gguf_quant_weights_iterator", - lambda _path: iter( - [ - ("transformer_blocks.0.attn.to_q.qweight_type", 1), - ("transformer_blocks.0.attn.to_q.qweight", 2), - ("transformer_blocks.0.attn.to_k.qweight_type", 3), - ("transformer_blocks.0.attn.to_k.qweight", 4), - ("transformer_blocks.0.attn.to_out.0.qweight_type", 5), - ("transformer_blocks.0.attn.to_out.0.qweight", 6), - ] - ), - ) - - adapter = QwenImageGGUFAdapter( - "dummy.gguf", - _FakeTransformer(), - _make_source(), - _make_od_config(), - ) - - weights = list(adapter.weights_iterator()) - - assert ("transformer_blocks.0.attn.to_q.qweight_type", 1) in weights - assert ("transformer_blocks.0.attn.to_q.qweight", 2) in weights - assert ("transformer_blocks.0.attn.to_k.qweight_type", 3) in weights - assert ("transformer_blocks.0.attn.to_k.qweight", 4) in weights - assert ("transformer_blocks.0.attn.to_out.0.qweight_type", 5) in weights - assert ("transformer_blocks.0.attn.to_out.0.qweight", 6) in weights - - -def test_qwen_adapter_keeps_top_level_quantized_weights(monkeypatch: pytest.MonkeyPatch): - import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module - - monkeypatch.setattr( - qwen_image_module, - "gguf_quant_weights_iterator", - lambda _path: iter( - [ - ("img_in.qweight_type", 1), - ("img_in.qweight", 2), - ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3), - ("time_text_embed.timestep_embedder.linear_1.qweight", 4), - ] - ), - ) - - adapter = QwenImageGGUFAdapter( - "dummy.gguf", - _FakeTransformer(), - _make_source(), - _make_od_config(), - ) - - weights = list(adapter.weights_iterator()) - - assert ("img_in.qweight_type", 1) in weights - assert ("img_in.qweight", 2) in weights - assert ("time_text_embed.timestep_embedder.linear_1.qweight_type", 3) in weights - assert ("time_text_embed.timestep_embedder.linear_1.qweight", 4) in weights - - -def test_qwen_adapter_keeps_modulation_quantized_weights(monkeypatch: pytest.MonkeyPatch): - import vllm_omni.diffusion.model_loader.gguf_adapters.qwen_image as qwen_image_module - - monkeypatch.setattr( - qwen_image_module, - "gguf_quant_weights_iterator", - lambda _path: iter( - [ - ("transformer_blocks.0.img_mod.1.qweight_type", 1), - ("transformer_blocks.0.img_mod.1.qweight", 2), - ("transformer_blocks.0.txt_mod.1.qweight_type", 3), - ("transformer_blocks.0.txt_mod.1.qweight", 4), - ] - ), - ) - - adapter = QwenImageGGUFAdapter( - "dummy.gguf", - _FakeTransformer(), - _make_source(), - _make_od_config(), - ) - - weights = list(adapter.weights_iterator()) - - assert ("transformer_blocks.0.img_mod.1.qweight_type", 1) in weights - assert ("transformer_blocks.0.img_mod.1.qweight", 2) in weights - assert ("transformer_blocks.0.txt_mod.1.qweight_type", 3) in weights - assert ("transformer_blocks.0.txt_mod.1.qweight", 4) in weights diff --git a/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py b/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py deleted file mode 100644 index cccea56fb30..00000000000 --- a/tests/diffusion/model_loader/gguf_adapters/test_z_image_gguf_adapter.py +++ /dev/null @@ -1,83 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for the Z-Image GGUF adapter.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter -from vllm_omni.diffusion.model_loader.gguf_adapters.z_image import ( - ZImageGGUFAdapter, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - -def _make_od_config( - model_class_name: str = "ZImagePipeline", - model_type: str = "z_image", -): - return SimpleNamespace( - model_class_name=model_class_name, - tf_model_config={"model_type": model_type}, - ) - - -def _make_source(prefix: str = "", subfolder: str = "transformer"): - return SimpleNamespace(prefix=prefix, subfolder=subfolder) - - -def test_z_image_adapter_selected_for_z_image_family(): - adapter = get_gguf_adapter( - "dummy.gguf", - object(), - _make_source(), - _make_od_config(), - ) - - assert isinstance(adapter, ZImageGGUFAdapter) - - -def test_z_image_adapter_matches_model_type_variants(): - for model_type in ("z_image", "zimage", "z-image"): - assert ZImageGGUFAdapter.is_compatible( - _make_od_config(model_class_name="OtherPipeline", model_type=model_type), - object(), - _make_source(), - ) - - -def test_z_image_adapter_renames_known_gguf_tensor_paths(monkeypatch: pytest.MonkeyPatch): - import vllm_omni.diffusion.model_loader.gguf_adapters.z_image as z_image_module - - monkeypatch.setattr( - z_image_module, - "gguf_quant_weights_iterator", - lambda _path: iter( - [ - ("model.diffusion_model.final_layer.qweight", torch.ones((1, 1))), - ("model.diffusion_model.x_embedder.qweight_type", torch.tensor(1)), - ("transformer_blocks.0.attention.out.weight", torch.full((1, 1), 2.0)), - ("transformer_blocks.0.attention.qkv.qweight", torch.full((1, 1), 3.0)), - ] - ), - ) - - adapter = ZImageGGUFAdapter( - "dummy.gguf", - object(), - _make_source(), - _make_od_config(), - ) - - weights = list(adapter.weights_iterator()) - names = [name for name, _ in weights] - - assert "all_final_layer.2-1.qweight" in names - assert "all_x_embedder.2-1.qweight_type" in names - assert "transformer_blocks.0.attention.to_out.0.weight" in names - assert "transformer_blocks.0.attention.to_qkv.qweight" in names diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index 376dbe685b9..f370d5b497b 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -16,7 +16,6 @@ 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 @@ -111,27 +110,6 @@ def test_empty_source_prefix_keeps_full_model_strict_check(): loader.load_weights(model) -def test_qwen_model_class_selects_qwen_gguf_adapter(): - od_config = type( - "Config", - (), - { - "model_class_name": "QwenImagePipeline", - "tf_model_config": {"model_type": "qwen_image"}, - }, - )() - source = DiffusersPipelineLoader.ComponentSource( - model_or_path="dummy", - subfolder="transformer", - revision=None, - prefix="transformer.", - ) - - adapter = get_gguf_adapter("dummy.gguf", object(), source, od_config) - - assert adapter.__class__.__name__ == "QwenImageGGUFAdapter" - - class _ConfigAwareModel(nn.Module): def __init__(self, *, od_config): super().__init__() @@ -186,7 +164,6 @@ def __exit__(self, exc_type, exc, tb): 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: 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")) @@ -202,6 +179,59 @@ def __exit__(self, exc_type, exc, tb): assert get_current_diffusion_config_or_none() is None +def test_hsdp_processes_quantized_weights_before_sharding(mocker): + import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod + from vllm_omni.diffusion.offloader.module_collector import PipelineModules + + od_config = SimpleNamespace( + dtype=torch.float32, + parallel_config=SimpleNamespace( + use_hsdp=True, + hsdp_replicate_size=1, + hsdp_shard_size=2, + ), + quantization_config=None, + ) + loader = DiffusersPipelineLoader(LoadConfig(), od_config) + loader.quant_config = object() + + model = nn.Module() + model.transformer = nn.Linear(2, 2, bias=False) + events: list[str] = [] + + loader._init_from_load_format = mocker.Mock(return_value=model) # type: ignore[method-assign] + loader.load_weights = mocker.Mock(side_effect=lambda _model: events.append("load")) # type: ignore[method-assign] + loader._process_weights_after_loading = mocker.Mock( # type: ignore[method-assign] + side_effect=lambda _model, _device: events.append("process") + ) + mocker.patch.object( + loader_mod.ModuleDiscovery, + "discover", + return_value=PipelineModules( + dits=[model.transformer], + dit_names=["transformer"], + vaes=[], + encoders=[], + encoder_names=[], + resident_modules=[], + resident_names=[], + ), + ) + mocker.patch( + "vllm_omni.diffusion.quantization.hsdp_fp8.prepare_fp8_layers_for_fsdp", + side_effect=lambda _model: events.append("prepare"), + ) + mocker.patch.object( + loader_mod, + "apply_hsdp_to_model", + side_effect=lambda *_args, **_kwargs: events.append("shard"), + ) + + loader._load_model_with_hsdp(torch.device("cpu")) + + assert events == ["load", "process", "prepare", "shard"] + + 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( diff --git a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py b/tests/diffusion/model_loader/test_diffusers_loader_gguf.py deleted file mode 100644 index f747dca818b..00000000000 --- a/tests/diffusion/model_loader/test_diffusers_loader_gguf.py +++ /dev/null @@ -1,259 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for GGUF-specific DiffusersPipelineLoader behavior.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch -import torch.nn as nn - -import vllm_omni.diffusion.model_loader.diffusers_loader as loader_module -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader - -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - -class _DummyModel(nn.Module): - def __init__(self) -> None: - super().__init__() - self.transformer = nn.Linear(2, 2, bias=True) - self.vae = nn.Linear(2, 2, bias=False) - self.register_buffer("transformer_buffer", torch.ones(1)) - self.calls: list[list[str]] = [] - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path="dummy", - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - DiffusersPipelineLoader.ComponentSource( - model_or_path="dummy", - subfolder="vae", - revision=None, - prefix="vae.", - fall_back_to_pt=True, - ), - ] - - def load_weights(self, weights): - loadable = dict(self.named_parameters()) - loadable.update(dict(self.named_buffers())) - seen: list[str] = [] - loaded: set[str] = set() - for name, tensor in weights: - seen.append(name) - if name in loadable: - target = loadable[name] - target.data.copy_(tensor.to(dtype=target.dtype)) - loaded.add(name) - self.calls.append(seen) - return loaded - - -def _make_loader() -> DiffusersPipelineLoader: - loader = object.__new__(DiffusersPipelineLoader) - loader.load_config = SimpleNamespace( - download_dir="cache-dir", - 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 - - -def test_is_gguf_quantization_accepts_dict_config(): - loader = _make_loader() - loader.quant_config = { - "method": "gguf", - "gguf_model": "weights.gguf", - } - - assert loader._is_gguf_quantization() is True - - -def test_is_gguf_quantization_rejects_non_gguf_dict_config(): - loader = _make_loader() - loader.quant_config = { - "method": "fp8", - "gguf_model": "weights.gguf", - } - - assert loader._is_gguf_quantization() is False - - -def test_is_gguf_quantization_requires_gguf_model_for_dict_config(): - loader = _make_loader() - loader.quant_config = {"method": "gguf"} - - with pytest.raises(ValueError, match="gguf_model"): - loader._is_gguf_quantization() - - -def test_is_gguf_quantization_accepts_object_config(): - loader = _make_loader() - loader.quant_config = SimpleNamespace( - get_name=lambda: "gguf", - gguf_model="weights.gguf", - ) - - assert loader._is_gguf_quantization() is True - - -def test_is_gguf_quantization_uses_fallback_object_without_get_name(): - loader = _make_loader() - loader.quant_config = SimpleNamespace(gguf_model="weights.gguf") - - assert loader._is_gguf_quantization() is True - - -def test_get_model_loadable_names_collects_parameters_and_buffers(): - loader = _make_loader() - model = _DummyModel() - - names = loader._get_model_loadable_names(model) - - assert "transformer.weight" in names - assert "transformer.bias" in names - assert "transformer_buffer" in names - - -def test_resolve_gguf_model_path_returns_local_file(tmp_path): - loader = _make_loader() - gguf_file = tmp_path / "model.gguf" - gguf_file.write_bytes(b"gguf") - - resolved = loader._resolve_gguf_model_path(str(gguf_file), revision=None) - - assert resolved == str(gguf_file) - - -def test_resolve_gguf_model_path_downloads_explicit_gguf_filename(monkeypatch: pytest.MonkeyPatch): - loader = _make_loader() - - monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False) - monkeypatch.setattr( - loader_module, - "hf_hub_download", - lambda repo_id, filename, revision, cache_dir: f"{cache_dir}/{repo_id}/{filename}", - ) - - resolved = loader._resolve_gguf_model_path("owner/repo/model-Q4_0.gguf", revision="main") - - assert resolved == "cache-dir/owner/repo/model-Q4_0.gguf" - - -def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.MonkeyPatch): - loader = _make_loader() - - monkeypatch.setattr(loader_module.os.path, "isfile", lambda _path: False) - monkeypatch.setattr( - loader_module, - "download_gguf", - lambda repo_id, quant_type, cache_dir, revision, ignore_patterns: f"{cache_dir}/{repo_id}/{quant_type}.gguf", - ) - - resolved = loader._resolve_gguf_model_path("owner/repo:Q4_0", revision="main") - - assert resolved == "cache-dir/owner/repo/Q4_0.gguf" - - -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.", - ) - - class _Adapter: - def weights_iterator(self): - yield "weight", torch.ones((2, 2)) - yield "bias", torch.zeros(2) - - 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())) - - assert weights[0][0] == "transformer.weight" - assert torch.equal(weights[0][1], torch.ones((2, 2))) - assert weights[1][0] == "transformer.bias" - assert torch.equal(weights[1][1], torch.zeros(2)) - - -def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch): - loader = _make_loader() - model = _DummyModel() - - monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources)) - monkeypatch.setattr( - loader, - "_get_gguf_weights_iterator", - lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]), - ) - - def _hf_iter(source): - if source.subfolder == "transformer": - return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))]) - return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))]) - - monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter) - monkeypatch.setattr( - loader, - "_get_expected_parameter_names", - lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"}, - ) - - loaded = loader._load_weights_with_gguf(model) - - assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"} - assert model.calls[0] == ["transformer.weight"] - assert model.calls[1] == ["transformer.bias"] - assert model.calls[2] == ["vae.weight"] - - -def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch): - loader = _make_loader() - model = _DummyModel() - - monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources)) - monkeypatch.setattr( - loader, - "_get_gguf_weights_iterator", - lambda source, model: iter( - [ - ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)), - ("transformer.bias", torch.zeros(2, dtype=torch.float32)), - ] - ), - ) - - hf_calls: list[str] = [] - - def _hf_iter(source): - hf_calls.append(source.subfolder or "") - if source.subfolder == "transformer": - return iter([("transformer.bias", torch.zeros(2, dtype=torch.float32))]) - return iter([("vae.weight", torch.full((2, 2), 2.0, dtype=torch.float32))]) - - monkeypatch.setattr(loader, "_get_weights_iterator", _hf_iter) - monkeypatch.setattr( - loader, - "_get_expected_parameter_names", - lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"}, - ) - - loaded = loader._load_weights_with_gguf(model) - - assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"} - assert hf_calls == ["vae"] diff --git a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py index 513e80d2887..8c429a00f60 100644 --- a/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py +++ b/tests/diffusion/model_loader/test_modelopt_fp8_adapter.py @@ -9,6 +9,7 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( ModelOptFp8CheckpointAdapter, + ModelOptNvFp4CheckpointAdapter, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -42,6 +43,70 @@ def __init__(self) -> None: ) +class _RemappedModelOptModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.runtime = nn.Module() + self.runtime.proj = nn.Linear(2, 2, bias=False) + + @staticmethod + def remap_checkpoint_key(name: str) -> str: + return { + "transformer.orig.proj.weight": "runtime.proj.weight", + }.get(name, name) + + +class _QuantizedRemappedModelOptModel(nn.Module): + """FP8 target whose params live under a namespace the raw checkpoint keys + never use, reachable only via ``remap_checkpoint_key``. + + Mirrors the Cosmos3 case that regressed: the generic weights-mapper cannot + resolve ``transformer.orig.*`` to ``runtime.*``, so without the remap hook + the adapter treats every tensor as unresolved and silently drops the scales. + """ + + def __init__(self) -> None: + super().__init__() + self.runtime = nn.Module() + self.runtime.proj = nn.Module() + self.runtime.proj.register_parameter( + "weight", + nn.Parameter(torch.empty(2, 2, dtype=torch.float8_e4m3fn), requires_grad=False), + ) + self.runtime.proj.register_parameter( + "weight_scale", + nn.Parameter(torch.empty(1), requires_grad=False), + ) + self.runtime.proj.register_parameter( + "input_scale", + nn.Parameter(torch.empty(1), requires_grad=False), + ) + + @staticmethod + def remap_checkpoint_key(name: str) -> str: + prefix = "transformer.orig.proj." + if name.startswith(prefix): + return "runtime.proj." + name[len(prefix) :] + return name + + +class _RemappedNvFp4Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.runtime = nn.Module() + self.runtime.proj = nn.Module() + self.runtime.proj.register_parameter( + "weight", + nn.Parameter(torch.empty(1, 1, dtype=torch.uint8), requires_grad=False), + ) + for name in ("input_scale", "weight_scale", "weight_scale_2"): + self.runtime.proj.register_parameter(name, nn.Parameter(torch.empty(1), requires_grad=False)) + + @staticmethod + def remap_checkpoint_key(name: str) -> str: + return name.replace("transformer.orig.proj.", "runtime.proj.") + + def _make_source() -> SimpleNamespace: return SimpleNamespace( subfolder="transformer", @@ -92,3 +157,85 @@ def test_modelopt_adapter_keeps_scale_tensors_for_quantized_target(): "transformer.block.to_q.weight_scale", "transformer.block.to_q.input_scale", ] + + +def test_modelopt_adapter_uses_checkpoint_key_remap_for_target_dtype(): + model = _RemappedModelOptModel() + adapter = ModelOptFp8CheckpointAdapter(model, _make_source()) + fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn) + scale = torch.tensor([0.5], dtype=torch.float32) + + adapted = list( + adapter.adapt( + iter( + [ + ("transformer.orig.proj.weight", fp8_weight), + ("transformer.orig.proj.weight_scale", scale), + ] + ) + ) + ) + + assert [name for name, _ in adapted] == ["runtime.proj.weight"] + assert adapted[0][1].dtype == model.runtime.proj.weight.dtype + assert torch.allclose(adapted[0][1], fp8_weight.to(torch.float32) * scale) + + +def test_modelopt_adapter_remaps_and_keeps_scales_for_quantized_target(): + """Regression: scales were dropped for a quantized (FP8) target that only the + model's ``remap_checkpoint_key`` hook can resolve. Assert both scales are + emitted and the FP8 weight passes through unchanged (no dequantization).""" + model = _QuantizedRemappedModelOptModel() + adapter = ModelOptFp8CheckpointAdapter(model, _make_source()) + fp8_weight = torch.tensor([[2.0, -4.0], [1.0, 3.0]], dtype=torch.float32).to(torch.float8_e4m3fn) + weight_scale = torch.tensor([0.5], dtype=torch.float32) + input_scale = torch.tensor([1.0], dtype=torch.float32) + + adapted = list( + adapter.adapt( + iter( + [ + ("transformer.orig.proj.weight_scale", weight_scale), + ("transformer.orig.proj.input_scale", input_scale), + ("transformer.orig.proj.weight", fp8_weight), + ] + ) + ) + ) + + emitted = dict(adapted) + assert len(adapted) == len(emitted) == 3 + + # Both scales are retained and emitted under the remapped runtime names. + assert torch.equal(emitted["runtime.proj.weight_scale"], weight_scale) + assert torch.equal(emitted["runtime.proj.input_scale"], input_scale) + + # The FP8 weight is passed through unchanged under the remapped target name + # (the quantized param is FP8, so it must NOT be dequantized). Compare via + # float32 since torch.equal has no FP8 CPU kernel. + assert emitted["runtime.proj.weight"].dtype == torch.float8_e4m3fn + assert torch.equal(emitted["runtime.proj.weight"].to(torch.float32), fp8_weight.to(torch.float32)) + + +def test_modelopt_nvfp4_adapter_remaps_quantized_weights_and_scales(): + adapter = ModelOptNvFp4CheckpointAdapter(_RemappedNvFp4Model(), _make_source()) + prefix = "transformer.orig.proj" + checkpoint_tensors = [ + (f"{prefix}.input_scale", torch.tensor([1.0])), + (f"{prefix}.weight_scale", torch.tensor([0.5])), + (f"{prefix}.weight_scale_2", torch.tensor([0.25])), + (f"{prefix}.weight", torch.tensor([[1]], dtype=torch.uint8)), + ] + + adapted = list(adapter.adapt(iter(checkpoint_tensors))) + + assert [name for name, _ in adapted] == [ + name.replace(f"{prefix}.", "runtime.proj.") for name, _ in checkpoint_tensors + ] + + +def test_modelopt_nvfp4_adapter_rejects_unconsumed_pre_quant_scale(): + adapter = ModelOptNvFp4CheckpointAdapter(nn.Module(), _make_source()) + + with pytest.raises(ValueError, match="does not consume pre_quant_scale"): + list(adapter.adapt(iter([("transformer.proj.pre_quant_scale", torch.tensor([1.0]))]))) diff --git a/tests/diffusion/models/bagel/test_naive_cache.py b/tests/diffusion/models/bagel/test_naive_cache.py new file mode 100644 index 00000000000..26ad0daeede --- /dev/null +++ b/tests/diffusion/models/bagel/test_naive_cache.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for NaiveCache merge/split logic used in batched CFG.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.bagel.bagel_transformer import NaiveCache + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +NUM_LAYERS = 2 +NUM_KV_HEADS = 4 +HEAD_DIM = 8 + + +def _make_cache(num_layers, seq_len, num_kv_heads=NUM_KV_HEADS, head_dim=HEAD_DIM, seed=0): + """Create a NaiveCache with deterministic random data. seq_len=0 returns an empty cache.""" + gen = torch.Generator().manual_seed(seed) + cache = NaiveCache(num_layers) + if seq_len == 0: + return cache + for layer in range(num_layers): + cache.key_cache[layer] = torch.randn(seq_len, num_kv_heads, head_dim, generator=gen) + cache.value_cache[layer] = torch.randn(seq_len, num_kv_heads, head_dim, generator=gen) + return cache + + +def test_init_creates_none_entries(): + """Ensure the NaiveCache is initialized with None values per layer.""" + cache = NaiveCache(NUM_LAYERS) + assert cache.num_layers == NUM_LAYERS + for layer in range(NUM_LAYERS): + assert cache.key_cache[layer] is None + assert cache.value_cache[layer] is None + + +@pytest.mark.parametrize("seq_len", [0, 10]) +def test_seq_lens_empty(seq_len): + """Ensure that by default, we have 0 seq lens.""" + cache = _make_cache(NUM_LAYERS, seq_len=seq_len) + assert cache.seq_lens == seq_len + assert cache.num_layers == NUM_LAYERS + + +### Merge tests +def test_merge_two_equal_length(): + """Ensure that we can merge two NaiveCaches that are identically shaped.""" + c0 = _make_cache(NUM_LAYERS, seq_len=5, seed=0) + c1 = _make_cache(NUM_LAYERS, seq_len=5, seed=1) + merged = NaiveCache.merge([c0, c1]) + + assert merged.key_values_lens == [5, 5] + for layer in range(NUM_LAYERS): + assert merged.key_cache[layer].shape[0] == 10 + # the merged cache will just have keys and values per layer concatenated + assert torch.equal(merged.key_cache[layer][:5], c0.key_cache[layer]) + assert torch.equal(merged.key_cache[layer][5:], c1.key_cache[layer]) + assert torch.equal(merged.value_cache[layer][:5], c0.value_cache[layer]) + assert torch.equal(merged.value_cache[layer][5:], c1.value_cache[layer]) + + +def test_merge_three_zero_len_cache(): + """Ensure we handle zero length cache correctly in merge.""" + # NOTE: This is relevant for text_cfg in Bagel, which has a len of 0 by default + gen_cache = _make_cache(NUM_LAYERS, seq_len=10, seed=0) + text_cfg_cache = _make_cache(NUM_LAYERS, seq_len=0) + img_cfg_cache = _make_cache(NUM_LAYERS, seq_len=7, seed=2) + merged = NaiveCache.merge([gen_cache, text_cfg_cache, img_cfg_cache]) + + assert merged.key_values_lens == [10, 0, 7] + for layer in range(NUM_LAYERS): + assert merged.key_cache[layer].shape[0] == 17 + assert torch.equal(merged.key_cache[layer][:10], gen_cache.key_cache[layer]) + assert torch.equal(merged.key_cache[layer][10:], img_cfg_cache.key_cache[layer]) + + +def test_merge_all_empty(): + """Ensure that merging empty caches is well defined.""" + caches = [_make_cache(NUM_LAYERS, seq_len=0) for _ in range(3)] + merged = NaiveCache.merge(caches) + + assert merged.key_values_lens == [0, 0, 0] + for layer in range(NUM_LAYERS): + assert merged.key_cache[layer] is None + assert merged.value_cache[layer] is None + + +def test_merge_single_cache(): + """Ensure merging one cache returns an identical cache.""" + c = _make_cache(NUM_LAYERS, seq_len=8, seed=42) + merged = NaiveCache.merge([c]) + + assert merged.key_values_lens == [8] + for layer in range(NUM_LAYERS): + assert torch.equal(merged.key_cache[layer], c.key_cache[layer]) + + +def test_merge_preserves_dtype(): + """Ensure merging doesn't modify dtypes.""" + cache = NaiveCache(NUM_LAYERS) + for layer in range(NUM_LAYERS): + cache.key_cache[layer] = torch.randn(5, NUM_KV_HEADS, HEAD_DIM, dtype=torch.bfloat16) + cache.value_cache[layer] = torch.randn(5, NUM_KV_HEADS, HEAD_DIM, dtype=torch.bfloat16) + + merged = NaiveCache.merge([cache]) + assert merged.key_cache[0].dtype == torch.bfloat16 + assert merged.value_cache[0].dtype == torch.bfloat16 + + +### Split tests +def test_split_all_nonzero(): + """Ensure NaiveCache splits in the simple case (all lens nonzero).""" + t = torch.randn(15, NUM_KV_HEADS, HEAD_DIM) + parts = NaiveCache.split_with_zeros(t, [5, 4, 6]) + + assert len(parts) == 3 + assert all(p is not None for p in parts) + assert parts[0].shape[0] == 5 + assert parts[1].shape[0] == 4 + assert parts[2].shape[0] == 6 + assert torch.equal(torch.cat(parts), t) + + +def test_split_with_zero(): + """Ensure NaiveCache split handles zero length correctly (used in Bagel).""" + t = torch.randn(17, NUM_KV_HEADS, HEAD_DIM) + parts = NaiveCache.split_with_zeros(t, [10, 0, 7]) + + assert parts[0].shape[0] == 10 + assert parts[1] is None + assert parts[2].shape[0] == 7 + assert torch.equal(torch.cat([parts[0], parts[2]]), t) + + +def test_split_wrong_sum_raises(): + """Ensure NaiveCache raises if splits don't match the sum of dims on axis 0.""" + t = torch.randn(10, NUM_KV_HEADS, HEAD_DIM) + with pytest.raises(ValueError, match="dim 0"): + NaiveCache.split_with_zeros(t, [5, 3]) + + +def test_split_negative_length_raises(): + """Ensure NaiveCache raises if splits have any negative values.""" + t = torch.randn(10, NUM_KV_HEADS, HEAD_DIM) + with pytest.raises(ValueError, match="greater than or equal to zero"): + NaiveCache.split_with_zeros(t, [5, -1, 6]) + + +def test_split_preserves_dtype(): + """Ensure NaiveCache split preserves dtype.""" + t = torch.randn(10, NUM_KV_HEADS, HEAD_DIM, dtype=torch.bfloat16) + parts = NaiveCache.split_with_zeros(t, [4, 6]) + assert parts[0].dtype == torch.bfloat16 + assert parts[1].dtype == torch.bfloat16 + + +### from_object tests (for kv cache transfer) +def test_from_object_passthrough(): + """Ensure a NaiveCache input is returned as is.""" + cache = _make_cache(NUM_LAYERS, seq_len=5) + assert NaiveCache.from_object(cache) is cache + + +def test_from_object_converts_simple_namespace(): + """Ensure SimpleNamespace with list-based caches converts to NaiveCache.""" + keys = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(NUM_LAYERS)] + values = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(NUM_LAYERS)] + ns = SimpleNamespace(key_cache=keys, value_cache=values) + + cache = NaiveCache.from_object(ns) + + assert isinstance(cache, NaiveCache) + assert cache.num_layers == NUM_LAYERS + for i in range(NUM_LAYERS): + assert torch.equal(cache.key_cache[i], keys[i]) + assert torch.equal(cache.value_cache[i], values[i]) + + +def test_from_object_mismatched_lengths_raises(): + """Ensure mismatched key/value cache lengths raise due to strict=True in zip.""" + keys = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(2)] + values = [torch.randn(5, NUM_KV_HEADS, HEAD_DIM) for _ in range(3)] + ns = SimpleNamespace(key_cache=keys, value_cache=values) + + with pytest.raises(ValueError): + NaiveCache.from_object(ns) + + +### End to end test for split / merge +def test_round_trip_two_populated(): + """Roundtrip test for merging and resplitting two simple caches.""" + c0 = _make_cache(NUM_LAYERS, seq_len=5, seed=0) + c1 = _make_cache(NUM_LAYERS, seq_len=8, seed=1) + merged = NaiveCache.merge([c0, c1]) + + for layer in range(NUM_LAYERS): + k_parts = NaiveCache.split_with_zeros(merged.key_cache[layer], merged.key_values_lens) + v_parts = NaiveCache.split_with_zeros(merged.value_cache[layer], merged.key_values_lens) + assert torch.equal(k_parts[0], c0.key_cache[layer]) + assert torch.equal(k_parts[1], c1.key_cache[layer]) + assert torch.equal(v_parts[0], c0.value_cache[layer]) + assert torch.equal(v_parts[1], c1.value_cache[layer]) + + +def test_round_trip_three_branches_with_zero_cfg(): + """Roundtrip test with a zero entry (i.e., same as Bagel's gen/text_cfg/img_cfg case).""" + gen_cache = _make_cache(NUM_LAYERS, seq_len=10, seed=0) + text_cfg_cache = _make_cache(NUM_LAYERS, seq_len=0) + img_cfg_cache = _make_cache(NUM_LAYERS, seq_len=7, seed=2) + merged = NaiveCache.merge([gen_cache, text_cfg_cache, img_cfg_cache]) + + for layer in range(NUM_LAYERS): + k_parts = NaiveCache.split_with_zeros(merged.key_cache[layer], merged.key_values_lens) + v_parts = NaiveCache.split_with_zeros(merged.value_cache[layer], merged.key_values_lens) + assert torch.equal(k_parts[0], gen_cache.key_cache[layer]) + assert k_parts[1] is None + assert torch.equal(k_parts[2], img_cfg_cache.key_cache[layer]) + assert torch.equal(v_parts[0], gen_cache.value_cache[layer]) + assert v_parts[1] is None + assert torch.equal(v_parts[2], img_cfg_cache.value_cache[layer]) diff --git a/tests/diffusion/models/bagel/test_trajectory_recording.py b/tests/diffusion/models/bagel/test_trajectory_recording.py index aa86f033a72..f3dbc0df2e8 100644 --- a/tests/diffusion/models/bagel/test_trajectory_recording.py +++ b/tests/diffusion/models/bagel/test_trajectory_recording.py @@ -19,15 +19,18 @@ NUM_TOKENS = 8 HIDDEN_DIM = 16 NUM_TIMESTEPS = 5 -# generate_image samples (num_timesteps + 1) points then drops the last, matching -# upstream lance: see Bagel.generate_image in bagel_transformer.py. -EXPECTED_STEPS = NUM_TIMESTEPS +# Official BAGEL samples num_timesteps points then drops the terminal t=0, +# yielding num_timesteps - 1 denoise steps. +EXPECTED_STEPS = NUM_TIMESTEPS - 1 def _make_mock_bagel(mocker: MockerFixture): """Create a mock Bagel with forward returning constant velocity.""" mock = mocker.MagicMock(spec=Bagel) mock._sp_size = 1 + # MagicMock would otherwise auto-return a truthy stub for this flag; pin it + # to the official BAGEL schedule convention (num_timesteps points). + mock._denoise_schedule_extra_step = False # forward returns a small constant velocity so x_t changes each step def fake_forward(self, x_t, **kwargs): @@ -179,7 +182,7 @@ def test_trajectory_timesteps_match_expected_schedule(self, bagel_and_args): bagel, args = bagel_and_args # Recompute the expected timestep schedule (mirrors generate_image logic) - ts = torch.linspace(1, 0, args["num_timesteps"] + 1) + ts = torch.linspace(1, 0, args["num_timesteps"]) shift = args.get("timestep_shift", 1.0) ts = shift * ts / (1 + (shift - 1) * ts) expected_timesteps = ts[:-1] # last element is dropped @@ -197,7 +200,7 @@ def test_latent_count_equals_timesteps_plus_one(self, bagel_and_args): bagel, args = bagel_and_args # Compute the number of denoising steps the same way generate_image does - ts = torch.linspace(1, 0, args["num_timesteps"] + 1) + ts = torch.linspace(1, 0, args["num_timesteps"]) shift = args.get("timestep_shift", 1.0) ts = shift * ts / (1 + (shift - 1) * ts) num_steps = len(ts) - 1 # timesteps = ts[:-1] diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index f77f63d5077..404979a78c9 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -9,24 +9,63 @@ from types import SimpleNamespace from typing import Any +import numpy as np import pytest import torch from PIL import Image from torch import nn +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_pipeline_declares_layerwise_offload_components() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + assert Cosmos3OmniDiffusersPipeline._dit_modules == ["transformer.language_model", "transformer"] + assert Cosmos3OmniDiffusersPipeline._encoder_modules == [] + assert Cosmos3OmniDiffusersPipeline._vae_modules == ["vae"] + assert Cosmos3OmniDiffusersPipeline._resident_modules == [] + assert hasattr(Cosmos3OmniDiffusersPipeline, "enable_omni_model_cpu_offload") + + class StubScheduler: - def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None: + def __init__( + self, + timesteps: list[int] | None = None, + *, + flow_shift: float = 1.0, + ) -> None: self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64) - self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift) - self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + self.config = SimpleNamespace( + num_train_timesteps=1000, + flow_shift=flow_shift, + ) + self.set_timesteps_calls: list[dict[str, Any]] = [] self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] - def set_timesteps(self, num_steps: int, device: torch.device) -> None: - self.set_timesteps_calls.append((num_steps, device)) - self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device) + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + *, + shift: float | None = None, + sigmas: list[float] | None = None, + ) -> None: + self.set_timesteps_calls.append( + { + "num_inference_steps": num_inference_steps, + "device": device, + "shift": shift, + "sigmas": sigmas, + } + ) + if sigmas is not None: + self.timesteps = torch.tensor(sigmas, device=device) + else: + assert num_inference_steps is not None + self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device) def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs): del kwargs @@ -53,8 +92,10 @@ def __init__(self, z_dim: int = 2, *, temporal: int = 4, spatial: int = 8) -> No latents_mean=[0.0] * z_dim, latents_std=[1.0] * z_dim, ) + self.encode_input_shapes: list[tuple[int, ...]] = [] def encode(self, video: torch.Tensor): + self.encode_input_shapes.append(tuple(video.shape)) latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1 latent_height = video.shape[-2] // self.config.scale_factor_spatial latent_width = video.shape[-1] // self.config.scale_factor_spatial @@ -131,9 +172,12 @@ def forward( ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0 sound_latents = kwargs.get("sound_latents") + control_latents = kwargs.get("control_latents") + control_bonus = 100 if control_latents is not None else 0 self.calls.append( { "token": token, + "has_control": control_latents is not None, "timestep": timestep.clone(), "text_mask": text_mask.clone(), "cache_before": self.cached_kv, @@ -145,7 +189,7 @@ def forward( self.cached_kv = [(marker, marker + 100)] self.cached_freqs_gen = (marker + 200, marker + 300) action_latents = kwargs.get("action_latents") - outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))] + outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token + control_bonus))] if action_latents is not None: outputs.append(torch.full_like(action_latents, float(token + 20))) if sound_latents is not None: @@ -172,6 +216,7 @@ def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch): def make_cosmos3_pipeline(): def _make(): from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT, Cosmos3OmniDiffusersPipeline, ) @@ -185,11 +230,13 @@ def _make(): pipeline.vae_scale_factor_temporal = 4 pipeline.vae_scale_factor_spatial = 8 pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0) - pipeline._base_scheduler_config = pipeline.scheduler.config - pipeline._engine_init_flow_shift = 1.0 - pipeline._current_flow_shift = 1.0 + pipeline._engine_init_flow_shift = COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT + pipeline._current_flow_shift = COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT + pipeline.is_distilled_model = False + pipeline.is_edge_model = False pipeline._guidance_scale = None pipeline._num_timesteps = None + pipeline._cosmos3_branch_caches = None pipeline._cache_dit_requires_paired_cfg = False pipeline._sound_tokenizer = None pipeline.progress_bar = passthrough_progress_bar @@ -198,6 +245,13 @@ def _make(): return _make +@pytest.fixture +def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.distributed import cfg_parallel + + monkeypatch.setattr(cfg_parallel, "get_classifier_free_guidance_world_size", lambda: 1) + + def make_sampling_params(**overrides: Any) -> SimpleNamespace: values = { "height": None, @@ -205,6 +259,7 @@ def make_sampling_params(**overrides: Any) -> SimpleNamespace: "num_frames": None, "num_inference_steps": None, "guidance_scale": None, + "guidance_scale_provided": False, "generator": None, "seed": 123, "num_outputs_per_prompt": 1, @@ -217,6 +272,31 @@ def make_sampling_params(**overrides: Any) -> SimpleNamespace: return SimpleNamespace(**values) +def make_request_batch(prompt: Any, sampling_params: SimpleNamespace) -> DiffusionRequestBatch: + if isinstance(prompt, list): + return DiffusionRequestBatch( + requests=[ + SimpleNamespace( + prompt=item, + request_id=f"cosmos3-test-{idx}", + sampling_params=sampling_params, + kv_sender_info=None, + ) + for idx, item in enumerate(prompt) + ] + ) + return DiffusionRequestBatch( + requests=[ + SimpleNamespace( + prompt=prompt, + request_id="cosmos3-test", + sampling_params=sampling_params, + kv_sender_info=None, + ) + ] + ) + + def _ids(value: int) -> torch.Tensor: return torch.tensor([[value]], dtype=torch.long) @@ -225,13 +305,171 @@ def _mask() -> torch.Tensor: return torch.ones(1, 1, dtype=torch.long) +def _capture_tokenize_calls(pipeline: Any) -> list[dict[str, Any]]: + calls: list[dict[str, Any]] = [] + + def _tokenize( + text: str, + max_sequence_length: int, + use_system_prompt: bool = False, + system_prompt: str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + index = len(calls) + 1 + calls.append( + { + "text": text, + "max_sequence_length": max_sequence_length, + "use_system_prompt": use_system_prompt, + "system_prompt": system_prompt, + } + ) + return _ids(index), torch.full((1, 1), index, dtype=torch.long) + + pipeline._tokenize_prompt = _tokenize + return calls + + +@pytest.mark.parametrize( + ("provided", "value", "default", "is_distilled", "expected"), + [ + (False, 1.0, 7.0, False, 7.0), + (True, 1.0, 7.0, False, 1.0), + (True, 4.5, 7.0, False, 4.5), + (False, 1.0, 7.0, True, 1.0), + (True, 4.5, 7.0, True, 1.0), + ], +) +def test_resolve_guidance_scale( + make_cosmos3_pipeline, + provided: bool, + value: float, + default: float, + is_distilled: bool, + expected: float, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = is_distilled + sp = make_sampling_params( + guidance_scale=value, + guidance_scale_provided=provided, + ) + + assert pipeline._resolve_guidance_scale(sp, default) == expected + + +def test_distilled_generation_accepts_t2i_and_i2v(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + common = { + "action_enabled": False, + "transfer_config": None, + "is_v2v": False, + "sound_enabled": False, + } + + pipeline._validate_distilled_generation_mode( + is_t2i=True, + image_tensor=None, + **common, + ) + pipeline._validate_distilled_generation_mode( + is_t2i=False, + image_tensor=torch.zeros(1), + **common, + ) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"action_enabled": True}, "action requests are unsupported"), + ({"transfer_config": SimpleNamespace()}, "transfer requests are unsupported"), + ({"is_v2v": True}, "video-to-video requests are unsupported"), + ({"sound_enabled": True}, "sound generation is unsupported"), + ( + {"is_t2i": False, "image_tensor": None}, + "text-to-video requests are unsupported", + ), + ( + {"is_t2i": True, "image_tensor": torch.zeros(1)}, + "image-conditioned image generation is unsupported", + ), + ], +) +def test_distilled_generation_rejects_unsupported_modes( + make_cosmos3_pipeline, + overrides: dict[str, Any], + message: str, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + kwargs = { + "is_t2i": True, + "image_tensor": None, + "action_enabled": False, + "transfer_config": None, + "is_v2v": False, + "sound_enabled": False, + } + kwargs.update(overrides) + + with pytest.raises(ValueError, match=message): + pipeline._validate_distilled_generation_mode(**kwargs) + + +def test_distilled_generation_rejects_robolab_policy(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + + with pytest.raises(ValueError, match="RoboLab/action policy requests are unsupported"): + pipeline._forward_robolab_policy(make_sampling_params(), None, 0.0) + + +@pytest.mark.parametrize( + ("prompt", "sampling_params", "message"), + [ + ( + {"prompt": "x", "modalities": ["video"], "generate_sound": True}, + make_sampling_params(), + "do not support sound generation", + ), + ( + {"prompt": "x", "modalities": ["video"]}, + make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}), + "do not support transfer inference", + ), + ( + { + "prompt": "x", + "modalities": ["video"], + "additional_information": { + "preprocessed_video": torch.zeros(1, 3, 5, 16, 16), + }, + }, + make_sampling_params(height=16, width=16, num_frames=5), + "do not support video-to-video generation", + ), + ], +) +def test_edge_forward_rejects_unsupported_modes( + make_cosmos3_pipeline, + prompt: dict[str, Any], + sampling_params: SimpleNamespace, + message: str, +) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_edge_model = True + + with pytest.raises(ValueError, match=message): + pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params)) + + def test_pipeline_registered_and_exported() -> None: from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS from vllm_omni.diffusion.models import cosmos3 from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.registry import ( - _DIFFUSION_ACTION_POST_PROCESS_FUNCS, _DIFFUSION_IR_OP_PRIORITY_FUNCS, _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS, @@ -241,20 +479,45 @@ def test_pipeline_registered_and_exported() -> None: assert issubclass(Cosmos3OmniDiffusersPipeline, nn.Module) assert issubclass(Cosmos3OmniDiffusersPipeline, ProgressBarMixin) assert Cosmos3OmniDiffusersPipeline.support_image_input is True - assert _DIFFUSION_MODELS["Cosmos3OmniDiffusersPipeline"] == ( - "cosmos3", - "pipeline_cosmos3", - "Cosmos3OmniDiffusersPipeline", - ) - assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func" - assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func" - assert ( - _DIFFUSION_ACTION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_action_post_process_func" - ) - assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_ir_op_priority_func" - assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__ + for pipeline_name in ("Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"): + assert _DIFFUSION_MODELS[pipeline_name] == ( + "cosmos3", + "pipeline_cosmos3", + "Cosmos3OmniDiffusersPipeline", + ) + assert _DIFFUSION_PRE_PROCESS_FUNCS[pipeline_name] == "get_cosmos3_pre_process_func" + assert _DIFFUSION_POST_PROCESS_FUNCS[pipeline_name] == "get_cosmos3_post_process_func" + assert _DIFFUSION_IR_OP_PRIORITY_FUNCS[pipeline_name] == "get_cosmos3_ir_op_priority_func" + assert pipeline_name in CUSTOM_DIT_ENABLERS + + +@pytest.mark.parametrize( + "pipeline_name", + ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"], +) +def test_cosmos3_model_index_resolves_pipeline( + tmp_path, + pipeline_name: str, +) -> None: + import json + + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + from vllm_omni.diffusion.registry import DiffusionModelRegistry + + (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": pipeline_name})) + + config = OmniDiffusionConfig(model=str(tmp_path)) + config.enrich_config() + + assert config.model_class_name == pipeline_name + resolved_pipeline_cls = DiffusionModelRegistry._try_load_model_cls(config.model_class_name) + assert resolved_pipeline_cls is Cosmos3OmniDiffusersPipeline + @pytest.fixture def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch): @@ -276,11 +539,24 @@ def to(self, _device): return self class _StubDiffusersScheduler: - config = SimpleNamespace(flow_shift=1.0) + load_config_calls: list[dict[str, Any]] = [] + from_config_calls: list[dict[str, Any]] = [] + + def __init__(self, *, flow_shift: float = 1.0) -> None: + self.config = SimpleNamespace(flow_shift=flow_shift) @classmethod - def from_pretrained(cls, *args, **kwargs): - return cls() + def load_config(cls, *args, **kwargs): + cls.load_config_calls.append({"args": args, "kwargs": dict(kwargs)}) + return { + "_class_name": "UniPCMultistepScheduler", + "flow_shift": 1.0, + } + + @classmethod + def from_config(cls, config, **kwargs): + cls.from_config_calls.append({"config": config, "kwargs": dict(kwargs)}) + return cls(flow_shift=kwargs.get("shift", config.get("flow_shift", 1.0))) class _StubVideoProcessor: def __init__(self, *args, **kwargs) -> None: @@ -288,12 +564,18 @@ def __init__(self, *args, **kwargs) -> None: monkeypatch.setattr(pipeline_cosmos3, "AutoTokenizer", _StubAutoTokenizer) monkeypatch.setattr(pipeline_cosmos3, "DistributedAutoencoderKLWan", _StubDiffusersVAE) - monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler) + monkeypatch.setattr(pipeline_cosmos3, "FlowUniPCMultistepScheduler", _StubDiffusersScheduler) monkeypatch.setattr(pipeline_cosmos3, "VideoProcessor", _StubVideoProcessor) monkeypatch.setattr(pipeline_cosmos3, "get_local_device", lambda: torch.device("cpu")) + return _StubDiffusersScheduler -def _make_od_config(*, sound_gen: bool) -> SimpleNamespace: +def _make_od_config( + *, + sound_gen: bool, + tf_model_config_overrides: dict[str, Any] | None = None, + model_config: dict[str, Any] | None = None, +) -> SimpleNamespace: tf_model_config = { "hidden_size": 8, "num_hidden_layers": 0, @@ -308,6 +590,8 @@ def _make_od_config(*, sound_gen: bool) -> SimpleNamespace: } if sound_gen: tf_model_config["sound_gen"] = True + if tf_model_config_overrides: + tf_model_config.update(tf_model_config_overrides) return SimpleNamespace( enable_cpu_offload=False, enable_diffusion_pipeline_profiler=False, @@ -316,7 +600,7 @@ def _make_od_config(*, sound_gen: bool) -> SimpleNamespace: flow_shift=None, quantization_config=None, custom_pipeline_args={}, - model_config={}, + model_config=model_config or {}, tf_model_config=tf_model_config, ) @@ -332,6 +616,311 @@ def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_in assert not hasattr(pipeline.transformer, "audio_proj_out") +def test_pipeline_init_uses_flow_unipc_with_cosmos3_defaults(stub_real_pipeline_init) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + + od_config = _make_od_config(sound_gen=False) + od_config.flow_shift = 2.5 + + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) + + assert len(stub_real_pipeline_init.load_config_calls) == 1 + assert len(stub_real_pipeline_init.from_config_calls) == 1 + call = stub_real_pipeline_init.from_config_calls[0] + assert call["config"]["_class_name"] == "UniPCMultistepScheduler" + assert call["kwargs"] == { + "shift": 1.0, + "use_dynamic_shifting": False, + "prediction_type": "flow_prediction", + } + assert pipeline.is_distilled_model is False + assert pipeline._engine_init_flow_shift == 2.5 + + +@pytest.mark.parametrize( + ("scheduler_class_name", "expected_distilled"), + [ + ("FlowMatchEulerDiscreteScheduler", True), + ("UniPCMultistepScheduler", False), + ], +) +def test_pipeline_resolves_scheduler_class_from_checkpoint_file( + tmp_path, + stub_real_pipeline_init, + monkeypatch: pytest.MonkeyPatch, + scheduler_class_name: str, + expected_distilled: bool, +) -> None: + import json + + from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3 + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler as RealFlowUniPCMultistepScheduler, + ) + + t_list = [1.0, 0.75, 0.5, 0.25] + scheduler_dir = tmp_path / "scheduler" + scheduler_dir.mkdir() + scheduler_config = {"_class_name": scheduler_class_name} + if expected_distilled: + scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list} + (scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config)) + + class StubFlowUniPCScheduler(StubScheduler): + from_config_calls: list[tuple[Any, dict[str, Any]]] = [] + + @classmethod + def load_config(cls, *args, **kwargs): + return RealFlowUniPCMultistepScheduler.load_config(*args, **kwargs) + + @classmethod + def from_config(cls, config, **kwargs): + cls.from_config_calls.append((config, dict(kwargs))) + return cls() + + class StubFlowMatchScheduler(StubScheduler): + from_config_calls: list[tuple[Any, dict[str, Any]]] = [] + + @classmethod + def from_config(cls, config, **kwargs): + cls.from_config_calls.append((config, dict(kwargs))) + scheduler = cls() + scheduler.config.fixed_step_sampler_config = config["fixed_step_sampler_config"] + return scheduler + + monkeypatch.setattr(pipeline_cosmos3, "FlowUniPCMultistepScheduler", StubFlowUniPCScheduler) + monkeypatch.setattr( + pipeline_cosmos3, + "FlowMatchEulerDiscreteScheduler", + StubFlowMatchScheduler, + ) + + od_config = _make_od_config(sound_gen=False) + od_config.model = str(tmp_path) + pipeline = Cosmos3OmniDiffusersPipeline(od_config=od_config) + + assert pipeline.is_distilled_model is expected_distilled + if expected_distilled: + assert len(StubFlowMatchScheduler.from_config_calls) == 1 + assert StubFlowMatchScheduler.from_config_calls[0][1] == {"stochastic_sampling": True} + assert StubFlowUniPCScheduler.from_config_calls == [] + assert pipeline._scheduler_init_t_list == t_list + else: + assert len(StubFlowUniPCScheduler.from_config_calls) == 1 + assert StubFlowMatchScheduler.from_config_calls == [] + assert not hasattr(pipeline, "_scheduler_init_t_list") + + +def test_distilled_pipeline_initializes_sde_scheduler( + stub_real_pipeline_init, + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + FlowMatchEulerDiscreteScheduler, + ) + + t_list = [1.0, 0.9375, 0.8333333333333334, 0.625] + scheduler_config = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "shift": 1.0, + "stochastic_sampling": False, + "fixed_step_sampler_config": { + "sample_type": "sde", + "t_list": t_list, + }, + } + monkeypatch.setattr( + stub_real_pipeline_init, + "load_config", + classmethod(lambda cls, *args, **kwargs: scheduler_config), + ) + + pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False)) + assert pipeline.is_distilled_model is True + assert isinstance(pipeline.scheduler, FlowMatchEulerDiscreteScheduler) + assert pipeline.scheduler.config.stochastic_sampling is True + assert pipeline._scheduler_init_t_list == t_list + + +@pytest.mark.parametrize( + ("fixed_step_config", "error_pattern"), + [ + pytest.param( + {"t_list": [1.0, 0.5]}, + r"fixed_step_sampler_config\.sample_type=sde", + id="missing-sample-type", + ), + pytest.param( + {"sample_type": "ode", "t_list": [1.0, 0.5]}, + r"fixed_step_sampler_config\.sample_type=sde", + id="unsupported-sample-type", + ), + pytest.param( + {"sample_type": "sde"}, + r"non-empty fixed_step_sampler_config\.t_list", + id="missing-t-list", + ), + pytest.param( + {"sample_type": "sde", "t_list": []}, + r"non-empty fixed_step_sampler_config\.t_list", + id="empty-t-list", + ), + ], +) +def test_distilled_scheduler_validates_fixed_step_config( + stub_real_pipeline_init, + monkeypatch: pytest.MonkeyPatch, + fixed_step_config: dict[str, Any], + error_pattern: str, +) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + ) + + scheduler_config = { + "_class_name": "FlowMatchEulerDiscreteScheduler", + "fixed_step_sampler_config": fixed_step_config, + } + monkeypatch.setattr( + stub_real_pipeline_init, + "load_config", + classmethod(lambda cls, *args, **kwargs: scheduler_config), + ) + + with pytest.raises(ValueError, match=error_pattern): + Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False)) + + +def test_pipeline_init_selects_edge_transformer_from_backbone_type(stub_real_pipeline_init) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT, + Cosmos3OmniDiffusersPipeline, + ) + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + COSMOS3_EDGE_BACKBONE_TYPE, + Cosmos3EdgeVFMTransformer, + ) + + pipeline = Cosmos3OmniDiffusersPipeline( + od_config=_make_od_config( + sound_gen=False, + tf_model_config_overrides={ + "backbone_type": COSMOS3_EDGE_BACKBONE_TYPE, + "qk_norm_for_text": False, + "latent_channel": 48, + "latent_patch_size": 2, + "temporal_compression_factor": 4, + "layer_norm_epsilon": 1e-5, + }, + ) + ) + + assert isinstance(pipeline.transformer, Cosmos3EdgeVFMTransformer) + assert pipeline.is_edge_model is True + assert pipeline._engine_init_flow_shift == COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT + + +def test_pipeline_init_does_not_select_edge_from_model_index_class_name(stub_real_pipeline_init) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer + + pipeline = Cosmos3OmniDiffusersPipeline( + od_config=_make_od_config( + sound_gen=False, + model_config={"_class_name": "Cosmos3EdgeOmniDiffusersPipeline"}, + ) + ) + + assert isinstance(pipeline.transformer, Cosmos3VFMTransformer) + assert not isinstance(pipeline.transformer, Cosmos3EdgeVFMTransformer) + + +def test_flow_unipc_reuses_scheduler_and_forwards_each_request_shift(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + scheduler = pipeline.scheduler + + assert pipeline._engine_init_flow_shift == 10.0 + assert pipeline._current_flow_shift == 10.0 + assert scheduler.config.flow_shift == 1.0 + + for shift in (3.0, 10.0): + pipeline._set_flow_shift(shift) + pipeline._set_timesteps(4, torch.device("cpu"), shift=pipeline._current_flow_shift) + + assert pipeline.scheduler is scheduler + assert pipeline._current_flow_shift == 10.0 + assert [call["shift"] for call in scheduler.set_timesteps_calls] == [3.0, 10.0] + assert all(call["num_inference_steps"] == 4 for call in scheduler.set_timesteps_calls) + assert all(call["sigmas"] is None for call in scheduler.set_timesteps_calls) + + +def test_flow_unipc_reproducible_with_same_seed(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, + ) + + pipeline = make_cosmos3_pipeline() + pipeline.scheduler = FlowUniPCMultistepScheduler( + shift=1.0, + use_dynamic_shifting=False, + prediction_type="flow_prediction", + ) + pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: ( + _ids(1), + _mask(), + _ids(0), + _mask(), + ) + + def run(seed: int) -> torch.Tensor: + request = make_request_batch( + {"prompt": "A test video.", "modalities": ["video"]}, + make_sampling_params( + seed=seed, + guidance_scale=1.0, + guidance_scale_provided=True, + num_inference_steps=4, + num_frames=5, + height=16, + width=16, + extra_args={"flow_shift": 3.0}, + ), + ) + output = pipeline.forward(request) + return output.output["video"] + + first = run(123) + second = run(123) + different_seed = run(456) + + torch.testing.assert_close(first, second) + assert not torch.equal(first, different_seed) + + +def test_distilled_set_timesteps_uses_fixed_sigma_schedule(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + pipeline.is_distilled_model = True + pipeline._scheduler_init_t_list = [1.0, 0.75, 0.5, 0.25] + + pipeline._set_timesteps( + num_inference_steps=99, + device=torch.device("cpu"), + shift=7.0, + ) + + assert pipeline.scheduler.set_timesteps_calls == [ + { + "num_inference_steps": None, + "device": torch.device("cpu"), + "shift": None, + "sigmas": pipeline._scheduler_init_t_list, + } + ] + + def test_pipeline_init_passes_tokenizer_attrs_into_transformer( stub_real_pipeline_init, monkeypatch: pytest.MonkeyPatch, @@ -363,38 +952,193 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None: preprocess = get_cosmos3_pre_process_func(SimpleNamespace()) i2v = SimpleNamespace( - prompts=[{"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}}], - sampling_params=SimpleNamespace(height=None, width=None, extra_args={}), + prompt={"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}}, + sampling_params=make_sampling_params(height=None, width=None, extra_args={}), ) result = preprocess(i2v) assert (result.sampling_params.height, result.sampling_params.width) == (672, 1344) - assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344) + assert tuple(result.prompt["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344) frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue")] action = SimpleNamespace( - prompts=[{"prompt": "Move.", "multi_modal_data": {"video": frames}}], - sampling_params=SimpleNamespace(height=16, width=32, extra_args={"action_mode": "forward_dynamics"}), + prompt={"prompt": "Move.", "multi_modal_data": {"video": frames}}, + sampling_params=make_sampling_params(height=16, width=32, extra_args={"action_mode": "forward_dynamics"}), ) - additional = preprocess(action).prompts[0]["additional_information"] + additional = preprocess(action).prompt["additional_information"] 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( + prompt={"prompt": "Continue.", "multi_modal_data": {"video": frames}}, + sampling_params=make_sampling_params( height=16, width=32, extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"}, ), ) - additional = preprocess(v2v).prompts[0]["additional_information"] + additional = preprocess(v2v).prompt["additional_information"] assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32) assert additional["condition_frame_indexes_vision"] == [0, 1] +def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + Cosmos3OmniDiffusersPipeline, + get_cosmos3_pre_process_func, + ) + + cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True})) + assert cfg is not None + assert list(cfg.hints) == ["edge"] + assert cfg.guidance_scale == 3.0 + assert cfg.control_guidance == 1.5 + assert cfg.flow_shift == 10.0 + assert cfg.num_video_frames_per_chunk == 93 + assert cfg.share_vision_temporal_positions is True + # fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies. + defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True})) + assert defaulted_fps_cfg is not None + assert defaulted_fps_cfg.fps == 10 + # fps provided (frame_rate set) -> the user value wins over the preset default. + explicit_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(frame_rate=24.0, extra_args={"wsm": True})) + assert explicit_fps_cfg is not None + assert explicit_fps_cfg.fps == 24.0 + assert ( + Cosmos3OmniDiffusersPipeline.reference_video_decode_spec(extra_args={"edge": True, "max_frames": 4}).max_frames + == 4 + ) + frames_for_pad = torch.arange(3 * 3, dtype=torch.uint8).reshape(1, 3, 1, 3) + assert transfer.pad_temporal_frames(frames_for_pad, 5)[0, :, 0, 0].tolist() == [0, 3, 6, 6, 3] + + real_import_module = transfer.importlib.import_module + + def raise_missing_cv2(name: str, *args: Any, **kwargs: Any): + if name == "cv2": + raise ImportError("missing cv2") + return real_import_module(name, *args, **kwargs) + + monkeypatch.setattr(transfer.importlib, "import_module", raise_missing_cv2) + with pytest.raises(ImportError, match="opencv-python"): + transfer.load_or_compute_control_frames( + cfg.hints["edge"], + height=8, + width=8, + max_frames=1, + input_frames=torch.zeros(3, 1, 8, 8, dtype=torch.uint8), + ) + + precomputed = torch.zeros(3, 2, 8, 8, dtype=torch.uint8) + precomputed_cfg = transfer.resolve_transfer_config( + make_sampling_params(extra_args={"edge": {"control": precomputed}}) + ) + assert precomputed_cfg is not None + loaded = transfer.load_or_compute_control_frames( + precomputed_cfg.hints["edge"], + height=8, + width=8, + max_frames=2, + input_frames=None, + ) + assert tuple(loaded.shape) == (3, 2, 8, 8) + + preprocess = get_cosmos3_pre_process_func(SimpleNamespace()) + + class FramesWithFps(list): + fps = 12.5 + + frames = FramesWithFps(Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "black")) + prompt = {"prompt": "transfer", "multi_modal_data": {"video": frames}} + request = SimpleNamespace( + prompt=prompt, + sampling_params=SimpleNamespace( + height=16, + width=32, + extra_args={"edge": True, "max_frames": 4, "resolution": "256"}, + ), + ) + additional = preprocess(request).prompt["additional_information"] + assert (request.sampling_params.height, request.sampling_params.width) == (192, 320) + assert tuple(additional["preprocessed_transfer_video"].shape) == (1, 3, 4, 192, 320) + assert additional["transfer_input_fps"] == 12.5 + assert "preprocessed_video" not in additional + + +def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None: + """When fps and frame_rate differ, transfer must pick frame_rate -- the same + precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real + sampling-params dataclass so the resolved_frame_rate property is exercised.""" + from vllm_omni.diffusion.models.cosmos3 import transfer + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sp = OmniDiffusionSamplingParams(fps=24, frame_rate=12.0, extra_args={"edge": True}) + assert sp.resolved_frame_rate == 12.0 + cfg = transfer.resolve_transfer_config(sp) + assert cfg is not None + # edge has no preset fps default, so cfg.fps comes straight from fps resolution. + assert cfg.fps == sp.resolved_frame_rate == 12.0 + + +def test_transfer_edge_uses_rgb_canny(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + + class FakeCv2: + def __init__(self) -> None: + self.canny_inputs: list[np.ndarray] = [] + + def Canny(self, image, lower, upper): + assert (lower, upper) == (100, 200) + self.canny_inputs.append(image.copy()) + return np.zeros(image.shape[:2], dtype=np.uint8) + + fake_cv2 = FakeCv2() + monkeypatch.setattr(transfer, "_import_cv2", lambda _hint_key: fake_cv2) + + frames = torch.zeros(3, 1, 4, 5, dtype=torch.uint8) + frames[0] = 255 + edge = transfer.make_edge_control(frames, "medium") + + assert tuple(edge.shape) == (3, 1, 4, 5) + assert len(fake_cv2.canny_inputs) == 1 + assert fake_cv2.canny_inputs[0].shape == (4, 5, 3) + + +def test_transfer_blur_uses_scaled_bilateral(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + + class FakeCv2: + INTER_AREA = 1 + INTER_LINEAR = 2 + INTER_CUBIC = 3 + + def __init__(self) -> None: + self.bilateral_calls: list[tuple[tuple[int, int, int], int, float, float]] = [] + + def resize(self, image, size, interpolation): + del interpolation + width, height = size + return np.zeros((height, width, image.shape[2]), dtype=image.dtype) + + def bilateralFilter(self, image, diameter, sigma_color, sigma_space): + self.bilateral_calls.append((image.shape, diameter, sigma_color, sigma_space)) + return image + + def GaussianBlur(self, *args, **kwargs): + raise AssertionError("Cosmos3 transfer blur should use bilateralFilter, not GaussianBlur.") + + fake_cv2 = FakeCv2() + monkeypatch.setattr(transfer, "_import_cv2", lambda _hint_key: fake_cv2) + + frames = torch.zeros(3, 1, 72, 72, dtype=torch.uint8) + blurred = transfer.make_blur_control(frames, "high") + + assert tuple(blurred.shape) == (3, 1, 72, 72) + assert fake_cv2.bilateral_calls == [((72, 72, 3), 3, 15.0, 10.0)] + + def test_postprocess_handles_image_video_audio_and_validation() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func @@ -423,11 +1167,11 @@ def test_postprocess_handles_image_video_audio_and_validation() -> None: def test_action_postprocess_handles_robolab_policy_outputs() -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( RoboLabPolicyInputs, - get_cosmos3_action_post_process_func, + get_cosmos3_post_process_func, make_robolab_action_postprocess_inputs, ) - func = get_cosmos3_action_post_process_func(SimpleNamespace()) + func = get_cosmos3_post_process_func(SimpleNamespace()) inputs = RoboLabPolicyInputs( prompt="Pick the cube.", video_tensor=torch.zeros(1, 3, 3, 16, 16), @@ -451,13 +1195,41 @@ def test_action_postprocess_handles_robolab_policy_outputs() -> None: ) action = torch.tensor([[[0.0, 0.25], [1.0, 0.75]]]) - custom_output = {"robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs)} - processed = func(action, custom_output=custom_output) + processed = func( + { + "payload": { + "actions": action, + }, + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + "internal": { + "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs), + }, + }, + } + ) - assert processed.shape == (1, 2) - assert processed.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype - torch.testing.assert_close(torch.from_numpy(processed), torch.tensor([[1.0, 0.25]])) - assert "robolab_action_postprocess" not in custom_output + processed_action = processed["payload"]["actions"] + assert processed_action.shape == (1, 2) + assert processed_action.dtype == torch.zeros((), dtype=torch.float32).numpy().dtype + torch.testing.assert_close(torch.from_numpy(processed_action), torch.tensor([[1.0, 0.25]])) + assert processed["metadata"] == { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + } def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.MonkeyPatch) -> None: @@ -487,12 +1259,35 @@ class FakeIrOpPriorityConfig: assert merged.custom_op == ["platform_kernel", "native"] -def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None: - from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline +def test_format_and_tokenize_prompts_leaves_plain_prompts_unchanged_by_default(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_SYSTEM_PROMPT pipeline = make_cosmos3_pipeline() - captured: list[str] = [] - pipeline._tokenize_prompt = lambda text, *args, **kwargs: (captured.append(text) or _ids(len(captured)), _mask()) + calls = _capture_tokenize_calls(pipeline) + + result = pipeline._format_and_tokenize_prompts( + " A robot. ", + " bad. ", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + use_system_prompt=False, + is_t2i=False, + ) + + assert [call["text"] for call in calls] == ["A robot.", "bad."] + assert all(call["max_sequence_length"] == 32 for call in calls) + assert all(call["use_system_prompt"] is False for call in calls) + assert all(call["system_prompt"] == COSMOS3_SYSTEM_PROMPT for call in calls) + assert [tensor.item() for tensor in result] == [1, 1, 2, 2] + + +def test_format_and_tokenize_prompts_applies_video_templates_and_system_override(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) pipeline._format_and_tokenize_prompts( "A robot", @@ -502,21 +1297,167 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No height=720, width=1280, max_sequence_length=32, + sp=SimpleNamespace( + system_prompt="direct system prompt", + extra_args={ + "use_duration_template": True, + "use_resolution_template": True, + "system_prompt": "API system prompt", + }, + ), + use_system_prompt=True, + is_t2i=False, + ) + + assert calls == [ + { + "text": ("A robot. The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution."), + "max_sequence_length": 32, + "use_system_prompt": True, + "system_prompt": "API system prompt", + }, + { + "text": ( + "bad. The video is not 2.0 seconds long and is not of 24 FPS. This video is not of 720x1280 resolution." + ), + "max_sequence_length": 32, + "use_system_prompt": True, + "system_prompt": "API system prompt", + }, + ] + + +def test_format_and_tokenize_prompts_uses_image_templates_for_t2i(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_T2I_SYSTEM_PROMPT + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + "A robot", + "bad", + num_frames=1, + frame_rate=24, + height=1024, + width=768, + max_sequence_length=64, sp=SimpleNamespace( extra_args={ - "negative_metadata_mode": "inverse", - # Duration/resolution metadata templates are off by default - # (commit "Change resolution and duration templates to off by - # default"); enable them explicitly to exercise the formatting. "use_duration_template": True, "use_resolution_template": True, } ), use_system_prompt=True, + is_t2i=True, + ) + + assert [call["text"] for call in calls] == [ + "A robot. This image is of 1024x768 resolution.", + "bad. This image is not of 1024x768 resolution.", + ] + assert all("seconds" not in call["text"] and "FPS" not in call["text"] for call in calls) + assert all(call["system_prompt"] == COSMOS3_T2I_SYSTEM_PROMPT for call in calls) + + +def test_format_and_tokenize_prompts_rewrites_json_object_metadata(make_cosmos3_pipeline) -> None: + import json + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + ('{"caption": "A robot", "duration": "old", "fps": 1, "resolution": {"H": 1, "W": 2}, "aspect_ratio": "1:1"}'), + "bad", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace( + extra_args={ + "aspect_ratio": "16:9", + "use_duration_template": True, + "use_resolution_template": True, + } + ), + is_t2i=False, + ) + + assert json.loads(calls[0]["text"]) == { + "caption": "A robot", + "duration": "2s", + "fps": 24.0, + "resolution": {"H": 720, "W": 1280}, + "aspect_ratio": "16:9", + } + assert "The video is 2.0 seconds long" not in calls[0]["text"] + assert calls[1]["text"] == ( + "bad. The video is not 2.0 seconds long and is not of 24 FPS. This video is not of 720x1280 resolution." + ) + + +def test_format_and_tokenize_prompts_removes_video_metadata_from_t2i_json(make_cosmos3_pipeline) -> None: + import json + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + '{"caption": "A robot", "duration": "2s", "fps": 24}', + "", + num_frames=1, + frame_rate=24, + height=1024, + width=768, + max_sequence_length=32, + sp=SimpleNamespace(extra_args={}), + is_t2i=True, + ) + + assert json.loads(calls[0]["text"]) == { + "caption": "A robot", + "resolution": {"H": 1024, "W": 768}, + } + + +@pytest.mark.parametrize( + "prompt", + [ + "{malformed", + '["A robot"]', + ], +) +def test_format_and_tokenize_prompts_falls_back_for_non_object_json( + make_cosmos3_pipeline, + prompt: str, +) -> None: + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + + pipeline._format_and_tokenize_prompts( + prompt, + "", + num_frames=48, + frame_rate=24, + height=720, + width=1280, + max_sequence_length=32, + sp=SimpleNamespace( + extra_args={ + "use_duration_template": True, + "use_resolution_template": True, + } + ), is_t2i=False, ) - assert "The video is 2.0 seconds long" in captured[0] - assert "The video is not 2.0 seconds long" in captured[1] + + assert calls[0]["text"] == ( + f"{prompt}. The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution." + ) + + +def test_checkpoint_key_remap() -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline remaps = { "embed_tokens.weight": "transformer.language_model.embed_tokens.weight", @@ -528,9 +1469,14 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No "layers.3.self_attn.to_q.weight": "transformer.language_model.layers.3.self_attn.to_q.weight", "layers.3.self_attn.to_out.weight": "transformer.language_model.layers.3.self_attn.to_out.weight", "layers.3.self_attn.norm_q.weight": "transformer.language_model.layers.3.self_attn.norm_q.weight", + "layers.3.self_attn.k_norm_und_for_gen.weight": ( + "transformer.language_model.layers.3.self_attn.k_norm_und_for_gen.weight" + ), "layers.3.self_attn.add_q_proj.weight": "transformer.gen_layers.3.cross_attention.to_q.weight", "layers.3.self_attn.to_add_out.weight": "transformer.gen_layers.3.cross_attention.to_out.weight", "layers.3.self_attn.norm_added_q.weight": "transformer.gen_layers.3.cross_attention.norm_q.weight", + "layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight", + "layers.3.mlp_moe_gen.down_proj.weight": "transformer.gen_layers.3.mlp.down_proj.weight", "transformer.model.layers.3.self_attn.add_k_proj.weight": ( "transformer.gen_layers.3.cross_attention.to_k.weight" ), @@ -543,7 +1489,7 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline) latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0)) assert latents.shape == (1, 2, 2, 2, 3) - pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full((1, 2, 2, 2, 3), 5.0) + pipeline._encode_conditioning_image_latent = lambda *args, **kwargs: torch.full((1, 2, 1, 2, 3), 5.0) i2v_latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v( torch.zeros(1, 3, 16, 24), 16, 24, 5, torch.Generator(device="cpu").manual_seed(0) ) @@ -593,6 +1539,87 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline) torch.testing.assert_close(action, clean) +def test_prepare_latents_i2v_encodes_only_conditioning_frame(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + calls: list[tuple[str, tuple[int, ...]]] = [] + + def record_to_vae_device(tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor: + assert pin_cpu is True + calls.append(("to_vae_device", tuple(tensor.shape))) + return tensor + + class RecordingVAE(StubCosmos3VAE): + def encode(self, video: torch.Tensor): + calls.append(("encode", tuple(video.shape))) + return super().encode(video) + + pipeline._to_vae_device = record_to_vae_device + pipeline.vae = RecordingVAE(z_dim=2) + generator = torch.Generator(device="cpu").manual_seed(0) + + latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v( + torch.zeros(1, 3, 16, 24), + 16, + 24, + 9, + generator, + ) + + assert calls == [ + ("to_vae_device", (1, 3, 16, 24)), + ("encode", (1, 3, 1, 16, 24)), + ] + assert pipeline.vae.encode_input_shapes[-1] == (1, 3, 1, 16, 24) + assert latents.shape == (1, 2, 3, 2, 3) + assert image_latent.shape == (1, 2, 1, 2, 3) + torch.testing.assert_close(latents[:, :, 0:1], image_latent) + assert not torch.allclose(latents[:, :, 1:], image_latent.expand(-1, -1, 2, -1, -1)) + assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]] + + +@pytest.mark.parametrize("mode", ["policy", "forward_dynamics"]) +def test_prepare_action_video_latents_encodes_only_conditioning_frame(make_cosmos3_pipeline, mode: str) -> None: + pipeline = make_cosmos3_pipeline() + video = torch.zeros(1, 3, 9, 24, 32) + + latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video( + video, + mode, + 24, + 32, + 9, + torch.Generator(device="cpu").manual_seed(0), + image_size=torch.tensor([1, 3, 16, 24]), + ) + + assert pipeline.vae.encode_input_shapes == [(1, 3, 1, 24, 32)] + assert latents.shape == (1, 2, 3, 2, 3) + assert condition_latents.shape == latents.shape + torch.testing.assert_close(condition_latents[:, :, 0], torch.ones(1, 2, 2, 3)) + assert torch.count_nonzero(condition_latents[:, :, 1:]) == 0 + torch.testing.assert_close(latents[:, :, 0:1], condition_latents[:, :, 0:1]) + assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]], [[1.0]]]]] + + +def test_prepare_inverse_dynamics_latents_encodes_full_video(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + video = torch.zeros(1, 3, 9, 16, 24) + + latents, velocity_mask, condition_latents = pipeline._prepare_latents_action_video( + video, + "inverse_dynamics", + 16, + 24, + 9, + torch.Generator(device="cpu").manual_seed(0), + ) + + assert pipeline.vae.encode_input_shapes == [(1, 3, 9, 16, 24)] + assert condition_latents.shape == (1, 2, 3, 2, 3) + torch.testing.assert_close(latents, condition_latents) + assert torch.count_nonzero(velocity_mask) == 0 + + def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None: pipeline = make_cosmos3_pipeline() latents = torch.zeros(1, 2, 1, 1, 1) @@ -624,6 +1651,9 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N image_latent=torch.full((1, 2, 1, 1, 1), 7.0), ) torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0)) + i2v_noise = pipeline.scheduler.step_calls[-1][0] + torch.testing.assert_close(i2v_noise[:, :, 0:1], torch.zeros(1, 2, 1, 1, 1)) + torch.testing.assert_close(i2v_noise[:, :, 1:2], torch.full((1, 2, 1, 1, 1), 2.0)) pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4) video_result, action_result = pipeline.diffuse( @@ -643,6 +1673,274 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result)) +def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + pipeline = make_cosmos3_pipeline() + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + result = pipeline.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.5, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + + assert [(call["token"], call["has_control"]) for call in pipeline.transformer.calls] == [ + (2, True), + (2, False), + (1, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 254.0)) + + +def test_diffuse_transfer_skips_idle_cfg_branches(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + control_only = make_cosmos3_pipeline() + control_result = control_only.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=1.0, + control_guidance=1.5, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + assert [(call["token"], call["has_control"]) for call in control_only.transformer.calls] == [ + (2, True), + (2, False), + ] + torch.testing.assert_close(control_result, torch.full_like(latents, 152.0)) + + text_only = make_cosmos3_pipeline() + text_result = text_only.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.0, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + ) + assert [(call["token"], call["has_control"]) for call in text_only.transformer.calls] == [ + (2, True), + (1, True), + ] + torch.testing.assert_close(text_result, torch.full_like(latents, 104.0)) + + +def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None: + pipeline = make_cosmos3_pipeline() + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + + result = pipeline.diffuse_transfer( + latents=latents, + timesteps=torch.tensor([900, 500, 100]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + control_guidance=1.5, + control_guidance_interval=(400.0, 1000.0), + control_latents=[torch.zeros_like(latents)], + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + guidance_interval=(800.0, 1000.0), + ) + + assert [(call["token"], call["has_control"]) for call in pipeline.transformer.calls] == [ + (2, True), + (2, False), + (1, True), + (2, True), + (2, False), + (2, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 508.0)) + + +@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)]) +def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None: + pipeline = make_cosmos3_pipeline() + captured: dict[str, Any] = {} + + def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): + del prompt, negative_prompt, num_frames, height, width, args, kwargs + captured["format_frame_rate"] = frame_rate + return _ids(2), _mask(), _ids(1), _mask() + + def fake_encode(video: torch.Tensor) -> torch.Tensor: + latent_frames = (video.shape[2] - 1) // pipeline.vae_scale_factor_temporal + 1 + return torch.ones( + 1, + 2, + latent_frames, + max(1, video.shape[-2] // pipeline.vae_scale_factor_spatial), + max(1, video.shape[-1] // pipeline.vae_scale_factor_spatial), + ) + + def fake_prepare(target_norm, current_conditional_frames, generator): + del current_conditional_frames, generator + latent = fake_encode(target_norm) + velocity_mask = torch.ones(1, 1, latent.shape[2], 1, 1) + return torch.zeros_like(latent), velocity_mask, torch.zeros_like(latent) + + def fake_diffuse_transfer(**kwargs): + captured["shared_kwargs"] = kwargs["shared_kwargs"] + return kwargs["latents"] + + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._format_and_tokenize_prompts = fake_format + pipeline._encode_video_tensor = fake_encode + pipeline._prepare_transfer_latents = fake_prepare + pipeline.diffuse_transfer = fake_diffuse_transfer + + def fake_set_flow_shift(target): + captured.setdefault("flow_shifts", []).append(target) + pipeline._current_flow_shift = float(target) + + pipeline._set_flow_shift = fake_set_flow_shift + pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta") + + control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8) + request = SimpleNamespace( + prompts=[ + { + "prompt": "transfer", + "modalities": ["video"], + "additional_information": { + "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16), + "transfer_input_fps": 8.0, + }, + } + ], + sampling_params=make_sampling_params( + height=16, + width=16, + # fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10). + extra_args={ + hint_key: {"control": control}, + "max_frames": 5, + "num_video_frames_per_chunk": 5, + "show_control_condition": True, + }, + ), + ) + + output = pipeline.forward(request) + + assert captured["format_frame_rate"] == expected_fps + assert captured["shared_kwargs"]["fps"] == expected_fps + assert captured["flow_shifts"] == [10.0] + # Transfer applies the V2V flow shift when building its timestep schedule. + assert [call["shift"] for call in pipeline.scheduler.set_timesteps_calls] == [10.0] + assert output.output["metadata"]["video"]["fps"] == expected_fps + assert output.output["payload"]["video"].device.type == "meta" + + +def test_forward_transfer_runs_multichunk_overlap_path( + make_cosmos3_pipeline, + sequential_cfg_parallel, +) -> None: + pipeline = make_cosmos3_pipeline() + captured: dict[str, Any] = {"targets": [], "conditional_frames": []} + + pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16) + pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask()) + pipeline._set_flow_shift = lambda target, **_kwargs: captured.setdefault("flow_shifts", []).append(target) + + original_prepare = pipeline._prepare_transfer_latents + + def recording_prepare(target_norm, current_conditional_frames, generator): + captured["targets"].append(target_norm.detach().clone()) + captured["conditional_frames"].append(current_conditional_frames) + return original_prepare(target_norm, current_conditional_frames, generator) + + pipeline._prepare_transfer_latents = recording_prepare + + decoded_chunks = [ + torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2], dtype=torch.float32), + torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5], dtype=torch.float32), + ] + + def fake_decode(latents): + chunk_values = decoded_chunks[len(captured.setdefault("decode_calls", []))] + captured["decode_calls"].append(latents.detach().clone()) + return chunk_values.view(1, 1, 5, 1, 1).expand(1, 3, 5, 16, 16).clone() + + pipeline._decode_latents = fake_decode + + transfer_video = torch.zeros(1, 3, 8, 16, 16) + transfer_video[:, :, 1] = 1.0 + control = torch.zeros(3, 8, 16, 16, dtype=torch.uint8) + request = SimpleNamespace( + prompts=[ + { + "prompt": "transfer", + "modalities": ["video"], + "additional_information": { + "preprocessed_transfer_video": transfer_video, + "transfer_input_fps": 8.0, + }, + } + ], + sampling_params=make_sampling_params( + height=16, + width=16, + num_inference_steps=1, + guidance_scale=1.0, + extra_args={ + "edge": {"control": control}, + "control_guidance": 1.0, + "max_frames": 8, + "num_video_frames_per_chunk": 5, + "num_conditional_frames": 1, + "num_first_chunk_conditional_frames": 2, + }, + ), + ) + + output = pipeline.forward(request) + + assert captured["conditional_frames"] == [2, 1] + assert len(captured["decode_calls"]) == 2 + assert output.output["payload"]["video"].shape == (1, 3, 8, 16, 16) + torch.testing.assert_close( + output.output["payload"]["video"][0, 0, :, 0, 0], + torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]), + ) + assert output.output["metadata"]["transfer"]["controls"]["edge"].shape == (1, 3, 8, 16, 16) + torch.testing.assert_close(captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0)) + torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0)) + torch.testing.assert_close(captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0)) + torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2)) + + def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None: """With cache-dit active the uncond pass is kept even outside the guidance interval (so cache-dit's has_separate_cfg parity stays in phase), and the @@ -680,7 +1978,18 @@ class TestForwardRouting: def _install_forward_stubs(self, pipeline): captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []} - def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): + def fake_format( + prompt, + negative_prompt, + num_frames, + frame_rate, + height, + width, + max_sequence_length, + sp, + use_system_prompt=False, + is_t2i=False, + ): captured["format"] = { "prompt": prompt, "negative_prompt": negative_prompt, @@ -688,7 +1997,8 @@ def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, "frame_rate": frame_rate, "height": height, "width": width, - "is_t2i": kwargs["is_t2i"], + "use_system_prompt": use_system_prompt, + "is_t2i": is_t2i, } return _ids(2), _mask(), _ids(1), _mask() @@ -707,7 +2017,12 @@ def fake_diffuse(**kwargs): pipeline._format_and_tokenize_prompts = fake_format pipeline._prepare_latents = fake_prepare - pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target) + + def fake_set_flow_shift(target): + captured.setdefault("flow_shifts", []).append(target) + pipeline._current_flow_shift = float(target) + + pipeline._set_flow_shift = fake_set_flow_shift pipeline.diffuse = fake_diffuse pipeline._decode_latents = lambda latents: latents return captured @@ -718,12 +2033,24 @@ def fake_diffuse(**kwargs): ( {"prompt": "A painted robot", "modalities": ["image"]}, make_sampling_params(num_outputs_per_prompt=2), - {"key": "image", "is_t2i": True, "flow": [3.0], "steps": [50, 50], "frames": 1}, + { + "key": "image", + "is_t2i": True, + "flow": [3.0], + "steps": [50, 50], + "frames": 1, + }, ), ( "A warehouse robot", make_sampling_params(), - {"key": "video", "is_t2i": False, "flow": [1.0], "steps": [35], "frames": 189}, + { + "key": "video", + "is_t2i": False, + "flow": [10.0], + "steps": [35], + "frames": 189, + }, ), ], ) @@ -737,13 +2064,14 @@ def test_forward_defaults_and_mode_selection( pipeline = make_cosmos3_pipeline() captured = self._install_forward_stubs(pipeline) - output = pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params)) + output = pipeline.forward(make_request_batch(prompt, sampling_params)) assert expected["key"] in output.output assert captured["format"]["is_t2i"] is expected["is_t2i"] assert captured["format"]["num_frames"] == expected["frames"] assert captured["flow_shifts"] == expected["flow"] - assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"] + assert [call["num_inference_steps"] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"] + assert all(call["shift"] == expected["flow"][0] for call in pipeline.scheduler.set_timesteps_calls) def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None: pipeline = make_cosmos3_pipeline() @@ -757,15 +2085,13 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non torch.zeros(1, 2, 1, 1, 1), ) pipeline.forward( - SimpleNamespace( - prompts=[ - { - "prompt": "move", - "modalities": ["video"], - "additional_information": {"preprocessed_image": image_tensor}, - } - ], - sampling_params=make_sampling_params(height=16, width=16, num_frames=5), + make_request_batch( + { + "prompt": "move", + "modalities": ["video"], + "additional_information": {"preprocessed_image": image_tensor}, + }, + make_sampling_params(height=16, width=16, num_frames=5), ) ) assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask @@ -779,18 +2105,16 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non 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), + make_request_batch( + { + "prompt": "continue", + "modalities": ["video"], + "additional_information": { + "preprocessed_video": video_tensor, + "condition_frame_indexes_vision": [0], + }, + }, + make_sampling_params(height=16, width=16, num_frames=5), ) ) assert captured["flow_shifts"][-1] == 10.0 @@ -801,12 +2125,12 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non 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) - pipeline._prepare_sound_latents = lambda *args: (sound_latents, 4) + pipeline._prepare_sound_latents = lambda *args, **kwargs: (sound_latents, 4) pipeline._decode_sound_latents = lambda *args: torch.ones(1, 2, 20) output = pipeline.forward( - SimpleNamespace( - prompts=[{"prompt": "A robot", "modalities": ["video"], "generate_sound": True}], - sampling_params=make_sampling_params(num_frames=9, frame_rate=3.0), + make_request_batch( + {"prompt": "A robot", "modalities": ["video"], "generate_sound": True}, + make_sampling_params(num_frames=9, frame_rate=3.0), ) ) assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents @@ -814,15 +2138,13 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4) output = pipeline.forward( - SimpleNamespace( - prompts=[ - { - "prompt": "Pick the block.", - "modalities": ["video"], - "additional_information": {"preprocessed_image": image_tensor}, - } - ], - sampling_params=make_sampling_params( + make_request_batch( + { + "prompt": "Pick the block.", + "modalities": ["video"], + "additional_information": {"preprocessed_image": image_tensor}, + }, + make_sampling_params( height=16, width=16, extra_args={ @@ -835,8 +2157,13 @@ def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> Non ) ) assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7] - assert output.custom_output["action"].shape == (1, 2, 2) - assert "action_only_output" not in output.custom_output + assert output.output["payload"]["actions"].shape == (1, 2, 2) + assert output.output["metadata"]["actions"] == { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + } + assert "common" not in output.output["metadata"] def test_forward_dispatches_robolab_policy_flow( self, @@ -901,7 +2228,7 @@ def fake_prepare_action_video(*args, **kwargs): AssertionError("RoboLab should not decode video") ) - output = pipeline.forward(SimpleNamespace(prompts=["ignored"], sampling_params=make_sampling_params())) + output = pipeline.forward(make_request_batch("ignored", make_sampling_params())) assert captured["format"] == { "prompt": "Pick the cube.", @@ -910,6 +2237,7 @@ def fake_prepare_action_video(*args, **kwargs): "frame_rate": 15.0, "height": 16, "width": 16, + "use_system_prompt": False, "is_t2i": False, } assert "flow_shifts" not in captured @@ -919,23 +2247,46 @@ def fake_prepare_action_video(*args, **kwargs): assert captured["prepare_action_video"]["kwargs"] == {"image_size": None} assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7] assert captured["diffuse_calls"][-1]["timesteps"].tolist() == [4, 3, 2, 1] - assert output.output == {} - assert output.custom_output["action_only_output"] is True - assert output.custom_output["action"].shape == (1, 2, 2) - assert "actions" not in output.custom_output - assert "robolab_action_postprocess" in output.custom_output - assert "robolab_policy_inputs" not in output.custom_output + assert output.output["payload"]["actions"].shape == (1, 2, 2) + assert output.output["metadata"]["actions"] == { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + } + assert output.output["metadata"]["common"]["action_only_output"] is True + assert "robolab_action_postprocess" in output.output["metadata"]["internal"] + assert "robolab_policy_inputs" not in output.output["metadata"] @pytest.mark.parametrize( ("prompt", "sampling_params", "message"), [ (["one", "two"], make_sampling_params(), "single prompt"), - ([{"prompt": "one", "modalities": ["image", "video"]}], make_sampling_params(), "both image and video"), + ({"prompt": "one", "modalities": ["image", "video"]}, make_sampling_params(), "both image and video"), ( - [{"prompt": "x", "modalities": ["image"], "generate_sound": True}], + {"prompt": "x", "modalities": ["image"], "generate_sound": True}, make_sampling_params(), "only for video", ), + ( + {"prompt": "x", "modalities": ["image"]}, + make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}), + "transfer inference is supported only for video outputs", + ), + ( + {"prompt": "x", "modalities": ["video"], "generate_sound": True}, + make_sampling_params(extra_args={"edge": {"control_path": "/tmp/control.mp4"}}), + "cannot be combined with sound generation", + ), + ( + {"prompt": "x", "modalities": ["video"]}, + make_sampling_params( + extra_args={ + "edge": {"control_path": "/tmp/control.mp4"}, + "action_mode": "policy", + } + ), + "cannot be combined with action generation", + ), ], ) def test_forward_rejects_invalid_public_requests( @@ -949,4 +2300,4 @@ def test_forward_rejects_invalid_public_requests( pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3) with pytest.raises(ValueError, match=message): - pipeline.forward(SimpleNamespace(prompts=prompt, sampling_params=sampling_params)) + pipeline.forward(make_request_batch(prompt, sampling_params)) diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 62411fc67ff..cb59dffe5a3 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -10,9 +10,38 @@ import torch from torch import nn +from vllm_omni.platforms import current_omni_platform + pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +@pytest.fixture(autouse=True) +def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + """Provide the TP metadata required by vLLM parallel linear layers.""" + from vllm.model_executor import parameter + from vllm.model_executor.layers import linear + + from vllm_omni.diffusion.models.cosmos3 import ( + transformer_cosmos3, + transformer_cosmos3_edge, + ) + + monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1) + monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1) + + +@pytest.fixture +def accelerator_device() -> torch.device: + """Provide an accelerator device, skipping when none is available.""" + if current_omni_platform.get_device_count() == 0: + pytest.skip("Accelerator required for this test") + return current_omni_platform.get_torch_device(0) + + def _tiny_cosmos3_config(**overrides): config = { "hidden_size": 8, @@ -30,6 +59,22 @@ def _tiny_cosmos3_config(**overrides): return config +def _tiny_cosmos3_edge_config(**overrides): + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + COSMOS3_EDGE_BACKBONE_TYPE, + ) + + config = _tiny_cosmos3_config( + latent_patch_size=2, + latent_channel=48, + temporal_compression_factor=4, + backbone_type=COSMOS3_EDGE_BACKBONE_TYPE, + qk_norm_for_text=False, + ) + config.update(overrides) + return config + + def test_mrope_position_ids_cover_text_video_sound_and_action() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( compute_mrope_position_ids_action, @@ -68,6 +113,16 @@ def test_mrope_position_ids_cover_text_video_sound_and_action() -> None: assert action_offset == 14 +def test_timestep_embedder_stores_frequencies_in_fp32() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( + TimestepEmbedder, + ) + + embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16) + + assert embedder.freqs.dtype == torch.float32 + + @pytest.mark.parametrize( ("key", "value"), [ @@ -87,29 +142,245 @@ def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> Cosmos3VFMTransformer._validate_supported_config(None) -def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: - from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer +def test_edge_config_resolves_nemotron_defaults() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) - model = object.__new__(Cosmos3VFMTransformer) + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + layer_norm_epsilon=1e-5, + rope_scaling={}, + rope_parameters={"mrope_section": [24, 20, 20]}, + ), + dtype=torch.float32, + ) + ) + + assert model.rms_norm_eps == 1e-5 + assert model.rope_theta == 100_000_000 + assert model.mrope_section == [24, 20, 20] + assert model.latent_channel_size == 48 + assert model.latent_patch_size == 2 + assert model.temporal_compression_factor == 4 + + +def test_edge_config_requires_backbone_type() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + config = _tiny_cosmos3_edge_config() + config.pop("backbone_type") + + with pytest.raises(ValueError, match="must declare backbone_type"): + Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=config, + dtype=torch.float32, + ) + ) + + +def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + Cosmos3Relu2MLP, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + num_hidden_layers=1, + use_k_norm_und_for_gen=False, + ), + dtype=torch.float32, + ) + ) + + layer = model.language_model.layers[0] + gen_layer = model.gen_layers[0] + assert isinstance(layer.mlp, Cosmos3Relu2MLP) + assert not hasattr(layer.mlp, "gate_proj") + assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP) + assert not hasattr(gen_layer.mlp, "gate_proj") + assert not hasattr(layer.self_attn, "norm_q") + assert not hasattr(layer.self_attn, "norm_k") + assert layer.self_attn.k_norm_und_for_gen is None + + +def test_edge_creates_gen_k_norm_when_enabled() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + num_hidden_layers=1, + qk_norm_for_diffusion=True, + use_k_norm_und_for_gen=True, + ), + dtype=torch.float32, + ) + ) + + k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen + assert model.use_k_norm_und_for_gen is True + assert isinstance(k_norm, RMSNorm) + assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm) + assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm) + + +def test_edge_supports_no_gen_qk_norm_variant() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = Cosmos3EdgeVFMTransformer( + SimpleNamespace( + tf_model_config=_tiny_cosmos3_edge_config( + num_hidden_layers=1, + qk_norm_for_diffusion=False, + use_k_norm_und_for_gen=False, + ), + dtype=torch.float32, + ) + ) + + assert model.qk_norm_for_diffusion is False + assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None + assert model.gen_layers[0].cross_attention.qk_norm is False + assert not hasattr(model.gen_layers[0].cross_attention, "norm_q") + assert not hasattr(model.gen_layers[0].cross_attention, "norm_k") + + +def test_edge_validates_required_relu2_weights() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeVFMTransformer, + ) + + model = object.__new__(Cosmos3EdgeVFMTransformer) nn.Module.__init__(model) - model.language_model = nn.Module() - model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)]) - model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) - model.norm_moe_gen = nn.LayerNorm(2) + model.num_hidden_layers = 1 + model.qk_norm_for_diffusion = True + model.use_k_norm_und_for_gen = True + + complete = { + "transformer.language_model.layers.0.mlp.up_proj.weight", + "transformer.language_model.layers.0.mlp.down_proj.weight", + "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight", + "transformer.gen_layers.0.mlp.up_proj.weight", + "transformer.gen_layers.0.mlp.down_proj.weight", + } + model.validate_loaded_weights(complete) + + mlp_weights = complete - { + "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight", + } + for missing_weight in mlp_weights: + with pytest.raises(ValueError, match="missing required weights"): + model.validate_loaded_weights(complete - {missing_weight}) - matched = [ - name - for name, module in model.named_modules() - if any(condition(name, module) for condition in model._hsdp_shard_conditions) - ] - assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] - assert Cosmos3VFMTransformer._layerwise_offload_blocks_attrs == ["gen_layers"] - assert Cosmos3VFMTransformer._repeated_blocks == ["Cosmos3GenDecoderLayer"] + missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"} + with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"): + model.validate_loaded_weights(missing_k_norm) + + model.use_k_norm_und_for_gen = False + model.validate_loaded_weights(missing_k_norm) - model.latent_patch_size = 2 - model.latent_channel_size = 3 - latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5) - torch.testing.assert_close(model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents) + +def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeCausalAttention, + ) + + class ScaleProjection(nn.Module): + def __init__(self, scale: float) -> None: + super().__init__() + self.scale = scale + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * self.scale + + class IdentityOutput(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + class RawKAttention(nn.Module): + def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + del v + return q + k + + attn = object.__new__(Cosmos3EdgeCausalAttention) + nn.Module.__init__(attn) + attn.hidden_size = 4 + attn.num_heads = 2 + attn.num_kv_heads = 2 + attn.head_dim = 2 + attn.num_heads_local = 2 + attn.num_kv_heads_local = 2 + attn.to_q = ScaleProjection(1.0) + attn.to_k = ScaleProjection(1.0) + attn.to_v = ScaleProjection(1.0) + attn.to_out = IdentityOutput() + attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5) + attn.attn = RawKAttention() + + hidden_states = torch.tensor( + [[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]], + dtype=torch.float32, + ) + freqs_cos = torch.ones(1, 2, 1, 2) + freqs_sin = torch.zeros(1, 2, 1, 2) + + reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin) + attn.to_k.scale = 3.0 + scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin) + + torch.testing.assert_close(gen_k, scaled_gen_k) + torch.testing.assert_close(gen_v, scaled_gen_v) + assert not torch.allclose(reasoner_out, scaled_reasoner_out) + + +def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( + Cosmos3EdgeLanguageModel, + Cosmos3EdgeVFMTransformer, + ) + + for transformer_cls, language_model_cls in ( + (Cosmos3VFMTransformer, Cosmos3LanguageModel), + (Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel), + ): + model = object.__new__(transformer_cls) + nn.Module.__init__(model) + model.language_model = nn.Module() + model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)]) + model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) + model.norm_moe_gen = nn.LayerNorm(2) + + matched = [ + name + for name, module in model.named_modules() + if any(condition(name, module) for condition in model._hsdp_shard_conditions) + ] + assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] + assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"] + assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"] + assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"] + + model.latent_patch_size = 2 + model.latent_channel_size = 3 + latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5) + torch.testing.assert_close( + model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), + latents, + ) def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None: @@ -131,6 +402,140 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No assert tuple(output.shape) == (1, 2, 1, 2, 2) +def test_qwen_gen_mlp_remains_gated() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer + + model = Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32) + ) + + assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP) + assert hasattr(model.gen_layers[0].mlp, "gate_proj") + + +def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) + + class ToyReasonerLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + def forward(self, hidden: torch.Tensor, freqs): + del freqs + kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1) + return hidden + self.weight.to(hidden.dtype), kv, kv + + class ToyGeneratorLayer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + def forward(self, hidden: torch.Tensor, **kwargs): + del kwargs + return hidden + self.weight.to(hidden.dtype) + + model = transformer_cosmos3.Cosmos3VFMTransformer( + SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) + ) + model.language_model.layers = nn.ModuleList([ToyReasonerLayer()]) + model.gen_layers = nn.ModuleList([ToyGeneratorLayer()]) + model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False) + + assert model._model_cpu_offload_enabled + assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"} + assert model.device == torch.device("cpu") + + 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 tuple(output.shape) == (1, 2, 1, 2, 2) + # The reasoner runs once then the generator component stays resident for GEN. + assert model._active_model_cpu_offload_component == "generator" + + model.disable_model_cpu_offload() + assert not model._model_cpu_offload_enabled + + +def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device( + monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device +) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + 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) + ) + model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)]) + model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) + model.to(accelerator_device) + + reasoner_param = model.language_model.layers[0].weight + generator_param = model.gen_layers[0].weight + + model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False) + + assert model._model_cpu_offload_enabled + # On enable, every group is parked on CPU until a phase activates it. + assert reasoner_param.device.type == "cpu" + assert generator_param.device.type == "cpu" + + model._activate_model_cpu_offload_component("reasoner") + assert reasoner_param.device == accelerator_device + assert generator_param.device.type == "cpu" + + model._activate_model_cpu_offload_component("generator") + assert reasoner_param.device.type == "cpu" + assert generator_param.device == accelerator_device + + model.disable_model_cpu_offload() + assert not model._model_cpu_offload_enabled + # Disable restores both components onto the device. + assert reasoner_param.device == accelerator_device + assert generator_param.device == accelerator_device + + +def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + 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) + ) + hidden_states = torch.zeros(1, 2, 1, 2, 2) + output = model( + hidden_states=hidden_states, + 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, + control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], + ) + + assert tuple(output.shape) == tuple(hidden_states.shape) + with pytest.raises(ValueError, match="control latent shape"): + model( + hidden_states=hidden_states, + 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, + control_latents=[torch.zeros(1, 2, 2, 2, 2)], + ) + + def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 @@ -298,6 +703,27 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch ) +def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.distributed import parallel_state + from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 + + model = object.__new__(transformer_cosmos3.Cosmos3VFMTransformer) + model.latent_patch_size = 2 + vs = (3, 16, 16) + + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 1) + assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 97 + + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 2) + assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 98 + assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=98) == 98 + + # ulysses=4, with a transfer control folded into the vision base. + monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 4) + padded = model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97, num_vision_items=2) + assert (2 * 192 + padded) % 4 == 0 + + def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer @@ -354,3 +780,33 @@ def __call__(self, x, position_ids): _, gen_pos = rotary.position_ids assert gen_pos.shape == (3, 1, 5) assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102] + + rotary.position_ids.clear() + model._compute_rope_freqs( + text_mask=torch.tensor([[1, 1]], dtype=torch.long), + t=2, + hp=1, + wp=1, + fps=24.0, + device=torch.device("cpu"), + dtype=torch.float32, + num_vision_items=3, + share_vision_temporal_positions=True, + ) + _, shared_gen_pos = rotary.position_ids + assert shared_gen_pos[0, 0].tolist() == [102, 103, 102, 103, 102, 103] + + rotary.position_ids.clear() + model._compute_rope_freqs( + text_mask=torch.tensor([[1, 1]], dtype=torch.long), + t=2, + hp=1, + wp=1, + fps=24.0, + device=torch.device("cpu"), + dtype=torch.float32, + num_vision_items=3, + share_vision_temporal_positions=False, + ) + _, offset_gen_pos = rotary.position_ids + assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107] diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py index 1365a3477c9..5f8de588ede 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py +++ b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py @@ -6,12 +6,17 @@ import torch from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2I2VDMD2Pipeline, + LTX2ImageToVideoPipeline, + LTX2Pipeline, + LTX2T2VDMD2Pipeline, +) from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline from vllm_omni.diffusion.request import OmniDiffusionRequest, OmniDiffusionSamplingParams +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -43,10 +48,10 @@ def _mock_base_init(self, *a, **kw): return pipeline -def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest: +def _make_request(prompt=None, **sp_kwargs) -> OmniDiffusionRequest: sp = OmniDiffusionSamplingParams(**sp_kwargs) return OmniDiffusionRequest( - prompts=prompts or [{"prompt": "a cat dancing"}], + prompt=prompt if prompt is not None else {"prompt": "a cat dancing"}, sampling_params=sp, request_id="dmd2-sanitize", ) @@ -133,35 +138,45 @@ def test_is_cfg_negative_forced_false(pipeline): def test_negative_prompt_stripped_from_prompt_dict(pipeline): - req = _make_request(prompts=[{"prompt": "a cat", "negative_prompt": "blurry"}]) + req = _make_request(prompt={"prompt": "a cat", "negative_prompt": "blurry"}) pipeline._sanitize_dmd2_request(req) - assert "negative_prompt" not in req.prompts[0] - assert req.prompts[0]["prompt"] == "a cat" + assert "negative_prompt" not in req.prompt + assert req.prompt["prompt"] == "a cat" def test_no_negative_prompt_unchanged(pipeline): - req = _make_request(prompts=[{"prompt": "a cat"}]) + req = _make_request(prompt={"prompt": "a cat"}) pipeline._sanitize_dmd2_request(req) - assert req.prompts[0] == {"prompt": "a cat"} + assert req.prompt == {"prompt": "a cat"} def test_string_prompt_not_mutated(pipeline): """String prompts (not dicts) must pass through unchanged.""" - req = _make_request(prompts=["a cat dancing"]) - pipeline._sanitize_dmd2_request(req) - assert req.prompts == ["a cat dancing"] - - -def test_multiple_prompts_all_sanitized(pipeline): - req = _make_request( - prompts=[ - {"prompt": "a cat", "negative_prompt": "blurry"}, - {"prompt": "a dog", "negative_prompt": "ugly"}, - ] - ) - pipeline._sanitize_dmd2_request(req) - for p in req.prompts: - assert "negative_prompt" not in p + req = _make_request(prompt="a cat dancing") + pipeline._sanitize_dmd2_request(req) + assert req.prompt == "a cat dancing" + + +def test_request_batch_sanitizes_each_request(pipeline): + requests = [ + _make_request( + prompt={"prompt": f"prompt-{index}", "negative_prompt": "blurry"}, + num_inference_steps=40 + index, + guidance_scale=5.0 + index, + guidance_scale_provided=True, + ) + for index in range(2) + ] + + pipeline._sanitize_dmd2_request(DiffusionRequestBatch(requests)) + + for request in requests: + assert request.sampling_params.num_inference_steps == pipeline.dmd2_config.num_inference_steps + assert request.sampling_params.guidance_scale == pipeline.dmd2_config.guidance_scale + assert request.sampling_params.guidance_scale_provided is False + assert request.sampling_params.do_classifier_free_guidance is False + assert request.sampling_params.is_cfg_negative is False + assert "negative_prompt" not in request.prompt # --------------------------------------------------------------------------- diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py index e5a1b6d8361..a395af95981 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py +++ b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py @@ -6,8 +6,12 @@ import torch from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX2T2VDMD2Pipeline -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2I2VDMD2Pipeline, + LTX2ImageToVideoPipeline, + LTX2Pipeline, + LTX2T2VDMD2Pipeline, +) from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImageDMD2Pipeline, QwenImagePipeline from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, WanT2VDMD2Pipeline from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import Wan22I2VPipeline, WanI2VDMD2Pipeline @@ -48,7 +52,7 @@ def _mock_base_init(self, *a, **kw): def _make_request(**sp_kwargs) -> OmniDiffusionRequest: sp = OmniDiffusionSamplingParams(**sp_kwargs) - return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler") + return OmniDiffusionRequest(prompt={"prompt": "a cat"}, sampling_params=sp, request_id="dmd2-scheduler") @pytest.fixture( diff --git a/tests/diffusion/models/flux/test_flux_pipeline.py b/tests/diffusion/models/flux/test_flux_pipeline.py new file mode 100644 index 00000000000..5a06133f4e5 --- /dev/null +++ b/tests/diffusion/models/flux/test_flux_pipeline.py @@ -0,0 +1,184 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.flux.pipeline_flux import FluxPipeline +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_flux_sampling(**overrides): + values = { + "height": 32, + "width": 32, + "num_inference_steps": 2, + "sigmas": None, + "guidance_scale": 3.5, + "generator": None, + "true_cfg_scale": 4.0, + "num_outputs_per_prompt": 0, + "latents": None, + "output_type": None, + "max_sequence_length": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _make_flux_pipeline(): + pipeline = object.__new__(FluxPipeline) + nn.Module.__init__(pipeline) + pipeline.default_sample_size = 128 + pipeline.vae_scale_factor = 8 + pipeline.device = torch.device("cpu") + pipeline.text_encoder = None + pipeline.transformer = SimpleNamespace( + in_channels=4, + guidance_embeds=False, + dtype=torch.float32, + ) + return pipeline + + +def test_forward_collates_request_prompt_tensors_for_flux(monkeypatch): + monkeypatch.setattr( + "vllm_omni.diffusion.models.flux.pipeline_flux.get_classifier_free_guidance_world_size", + lambda: 1, + ) + pipeline = _make_flux_pipeline() + encode_calls = [] + prepare_latents_call = {} + diffuse_call = {} + + def _fake_encode_prompt(**kwargs): + encode_calls.append(kwargs) + return ( + kwargs["prompt_embeds"], + kwargs["pooled_prompt_embeds"], + torch.zeros(kwargs["prompt_embeds"].shape[1], 3), + ) + + def _fake_prepare_latents( + batch_size, + num_channels_latents, + height, + width, + dtype, + device, + generator, + latents, + ): + prepare_latents_call.update( + { + "batch_size": batch_size, + "num_channels_latents": num_channels_latents, + "height": height, + "width": width, + "dtype": dtype, + "device": device, + "generator": generator, + "latents": latents, + } + ) + return torch.zeros(batch_size, 1, 1), torch.zeros(1, 3) + + def _fake_diffuse( + prompt_embeds, + pooled_prompt_embeds, + negative_prompt_embeds, + negative_pooled_prompt_embeds, + *args, + **kwargs, + ): + diffuse_call.update( + { + "prompt_embeds": prompt_embeds, + "pooled_prompt_embeds": pooled_prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds, + } + ) + return torch.arange(2, dtype=torch.float32).view(2, 1) + + pipeline.encode_prompt = _fake_encode_prompt + pipeline.prepare_latents = _fake_prepare_latents + pipeline.prepare_timesteps = lambda *args, **kwargs: (torch.tensor([1.0]), 1) + pipeline.diffuse = _fake_diffuse + + prompt_embeds_a = torch.zeros(2, 3) + prompt_embeds_b = torch.ones(2, 3) + pooled_prompt_embeds_a = torch.full((4,), 2.0) + pooled_prompt_embeds_b = torch.full((4,), 3.0) + negative_prompt_embeds_a = torch.full((2, 3), 4.0) + negative_prompt_embeds_b = torch.full((2, 3), 5.0) + negative_pooled_prompt_embeds_a = torch.full((4,), 6.0) + negative_pooled_prompt_embeds_b = torch.full((4,), 7.0) + latents_a = torch.zeros(1, 1, 1) + latents_b = torch.ones(1, 1, 1) + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="flux-prompt-a", + prompt={ + "prompt": "prompt-a", + "negative_prompt": "negative-a", + "prompt_embeds": prompt_embeds_a, + "pooled_prompt_embeds": pooled_prompt_embeds_a, + "negative_prompt_embeds": negative_prompt_embeds_a, + "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a, + }, + sampling_params=_make_flux_sampling(generator=gen_a, latents=latents_a, output_type="latent"), + ), + SimpleNamespace( + request_id="flux-prompt-b", + prompt={ + "prompt": "prompt-b", + "negative_prompt": "negative-b", + "additional_information": { + "prompt_embeds": [prompt_embeds_b], + "pooled_prompt_embeds": [pooled_prompt_embeds_b], + "negative_prompt_embeds": [negative_prompt_embeds_b], + "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b], + }, + }, + sampling_params=_make_flux_sampling(generator=gen_b, latents=latents_b, output_type="latent"), + ), + ] + ) + + outputs = pipeline.forward(batch) + + assert encode_calls[0]["prompt"] is None + assert encode_calls[0]["prompt_2"] is None + torch.testing.assert_close( + encode_calls[0]["prompt_embeds"], + torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + encode_calls[0]["pooled_prompt_embeds"], + torch.stack([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=0), + ) + assert encode_calls[1]["prompt"] is None + assert encode_calls[1]["prompt_2"] is None + torch.testing.assert_close( + encode_calls[1]["prompt_embeds"], + torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + encode_calls[1]["pooled_prompt_embeds"], + torch.stack([negative_pooled_prompt_embeds_a, negative_pooled_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + diffuse_call["negative_pooled_prompt_embeds"], + torch.stack([negative_pooled_prompt_embeds_a, negative_pooled_prompt_embeds_b], dim=0), + ) + assert prepare_latents_call["generator"] == [gen_a, gen_b] + torch.testing.assert_close(prepare_latents_call["latents"], torch.cat([latents_a, latents_b], dim=0)) + torch.testing.assert_close(outputs[0].output, torch.tensor([[0.0]])) + torch.testing.assert_close(outputs[1].output, torch.tensor([[1.0]])) diff --git a/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py b/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py index 3c6b5bada06..5afcc4e7830 100644 --- a/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py +++ b/tests/diffusion/models/flux2/test_flux2_klein_num_inference_steps.py @@ -39,7 +39,7 @@ def _make_minimal_request( ) req = MagicMock() req.sampling_params = params - req.prompts = [prompt] + req.prompt = prompt req.multi_modal_data = {} return req diff --git a/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py new file mode 100644 index 00000000000..3900c0c94e9 --- /dev/null +++ b/tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py @@ -0,0 +1,50 @@ +from unittest.mock import Mock, patch + +import pytest +import torch + +from vllm_omni.diffusion.models.flux2.pipeline_flux2 import ( + Flux2Pipeline, + _resolve_component_quant_config, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_resolve_component_quant_config_routes_components(): + text_encoder_config = object() + component_config = Mock() + component_config.resolve.side_effect = { + "text_encoder": text_encoder_config, + "transformer": None, + }.get + + assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config + assert _resolve_component_quant_config(component_config, "transformer") is None + + +def test_resolve_component_quant_config_preserves_global_config(): + global_config = object() + assert _resolve_component_quant_config(global_config, "transformer") is global_config + assert _resolve_component_quant_config(None, "text_encoder") is None + + +def test_offload_transformer_before_text_encoder_weights(): + pipeline = object.__new__(Flux2Pipeline) + pipeline.transformer = Mock() + weights = [ + ("transformer.layer.weight", torch.empty(1)), + ("text_encoder.layer.weight", torch.empty(1)), + ("text_encoder.layer.bias", torch.empty(1)), + ] + + iterator = iter(pipeline._offload_transformer_before_text_encoder(weights)) + assert next(iterator)[0] == "transformer.layer.weight" + pipeline.transformer.to.assert_not_called() + + with patch("vllm_omni.diffusion.models.flux2.pipeline_flux2.current_omni_platform.empty_cache") as empty_cache: + assert next(iterator)[0] == "text_encoder.layer.weight" + assert next(iterator)[0] == "text_encoder.layer.bias" + + pipeline.transformer.to.assert_called_once_with("cpu") + empty_cache.assert_called_once_with() diff --git a/vllm_omni/model_executor/stage_configs/__init__.py b/tests/diffusion/models/gr00t/__init__.py similarity index 100% rename from vllm_omni/model_executor/stage_configs/__init__.py rename to tests/diffusion/models/gr00t/__init__.py diff --git a/tests/diffusion/models/gr00t/test_pipeline.py b/tests/diffusion/models/gr00t/test_pipeline.py new file mode 100644 index 00000000000..3221121074f --- /dev/null +++ b/tests/diffusion/models/gr00t/test_pipeline.py @@ -0,0 +1,141 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +from vllm_omni.diffusion.models.gr00t import pipeline_gr00t +from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeGr00tPolicy: + instances = [] + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.reset_calls = 0 + self.seen_obs = None + self.embodiment_tag = SimpleNamespace(value="fake_embodiment") + self.language_key = "annotation.language.language_instruction" + self.modality_configs = { + "action": SimpleNamespace( + delta_indices=[0, 1], + modality_keys=["arm", "gripper"], + ) + } + self.processor = SimpleNamespace( + state_action_processor=SimpleNamespace( + norm_params={ + "fake_embodiment": { + "action": { + "arm": {"dim": np.array(2)}, + "gripper": {"dim": np.array(1)}, + } + } + } + ) + ) + FakeGr00tPolicy.instances.append(self) + + def get_action(self, obs): + self.seen_obs = obs + return { + "arm": np.array([[[1.0, 2.0]]], dtype=np.float64), + "gripper": [[[3.0]]], + }, {"latency_ms": 1.0} + + def reset(self): + self.reset_calls += 1 + return {"reset": True} + + +@pytest.fixture(autouse=True) +def fake_gr00t_policy(monkeypatch): + FakeGr00tPolicy.instances.clear() + monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", FakeGr00tPolicy) + + +def _pipeline(): + od_config = SimpleNamespace( + model="nvidia/GR00T-N1.7-3B", + model_config={ + "embodiment_tag": "LIBERO_PANDA", + "strict": False, + }, + custom_pipeline_args={}, + ) + return Gr00tN1d7Pipeline(od_config=od_config) + + +def test_pipeline_initializes_local_policy(): + pipeline = _pipeline() + + policy = FakeGr00tPolicy.instances[0] + assert policy.kwargs["model_path"] == "nvidia/GR00T-N1.7-3B" + assert policy.kwargs["embodiment_tag"] == "LIBERO_PANDA" + assert policy.kwargs["strict"] is False + assert pipeline.weights_sources == () + assert pipeline.load_weights(iter(())) == set() + + +def test_forward_returns_dict_actions_in_output(): + pipeline = _pipeline() + req = OmniDiffusionRequest( + prompt="pick", + request_id="req", + sampling_params=OmniDiffusionSamplingParams( + extra_args={ + "robot_obs": { + "images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)}, + "state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)}, + "prompt": "pick the cube", + "session_id": "session-a", + }, + "reset": True, + } + ), + ) + + output = pipeline.forward(req) + + assert output.error is None + actions = output.output["actions"] + assert set(actions) == {"arm", "gripper"} + assert actions["arm"].dtype == np.float32 + np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32)) + policy = FakeGr00tPolicy.instances[0] + assert "video" in policy.seen_obs + assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]} + assert "images" not in policy.seen_obs + assert "prompt" not in policy.seen_obs + assert "session_id" not in policy.seen_obs + assert policy.reset_calls == 1 + + +def test_dummy_warmup_returns_shape_correct_zero_actions(): + pipeline = _pipeline() + req = OmniDiffusionRequest( + prompt="dummy run", + request_id="dummy_req_id", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + ) + + output = pipeline.forward(req) + + assert output.error is None + actions = output.output["actions"] + assert set(actions) == {"arm", "gripper"} + assert actions["arm"].shape == (1, 2, 2) + assert actions["gripper"].shape == (1, 2, 1) + assert not actions["arm"].any() + assert FakeGr00tPolicy.instances[0].seen_obs is None + + +def test_reset_delegates_to_policy(): + pipeline = _pipeline() + + assert pipeline.reset() == {"reset": True} + assert FakeGr00tPolicy.instances[0].reset_calls == 1 diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py index 626f78eed9c..e8ade2e51e8 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py @@ -50,6 +50,48 @@ def test_noop_when_context_unavailable(self, mocker): mock_get.assert_not_called() +class TestSetForwardContextDPMetadata: + """Test diffusion MoE DP metadata setup for vLLM 0.24 all2all dispatch.""" + + def test_sets_dp_metadata_from_vllm_dp_group(self, mocker): + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe + + mock_ctx = mocker.MagicMock() + mock_ctx.dp_metadata = None + fake_group = mocker.MagicMock() + fake_group.world_size = 2 + fake_group.cpu_group = object() + + def _all_gather_object(output, obj, group): + assert group is fake_group.cpu_group + output[:] = [obj, obj + 1] + + mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True) + mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx) + mocker.patch.object(hunyuan_moe._vllm_ps, "_DP", fake_group, create=True) + mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object", side_effect=_all_gather_object) + + hunyuan_moe._set_forward_context_dp_metadata(1024) + + assert mock_ctx.dp_metadata is not None + assert mock_ctx.dp_metadata.num_tokens_across_dp_cpu.tolist() == [1024, 1025] + + def test_keeps_existing_dp_metadata(self, mocker): + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe + + existing_metadata = object() + mock_ctx = mocker.MagicMock() + mock_ctx.dp_metadata = existing_metadata + mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=True) + mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context", return_value=mock_ctx) + mock_all_gather = mocker.patch.object(hunyuan_moe.torch.distributed, "all_gather_object") + + hunyuan_moe._set_forward_context_dp_metadata(1024) + + assert mock_ctx.dp_metadata is existing_metadata + mock_all_gather.assert_not_called() + + class TestHunyuanFusedMoEPlatformDispatch: """Test platform dispatch via platform qualname hooks.""" diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_rgb.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_rgb.py new file mode 100644 index 00000000000..5dcb682f473 --- /dev/null +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_rgb.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for HunyuanImage3 AR image RGB conversion.""" + +from types import SimpleNamespace + +import pytest +import torch +from PIL import Image + +from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import ( + HunyuanImage3Processor, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FixedResolutionGroup: + def get_target_size(self, width: int, height: int): + return width, height + + def get_base_size_and_ratio_index(self, width: int, height: int): + return 1024, 0 + + +def test_ar_processor_uses_pil_rgb_conversion_for_rgba_images(): + captured: dict[str, object] = {} + + def fake_vit_processor(image: Image.Image): + captured["mode"] = image.mode + captured["pixel"] = image.getpixel((0, 0)) + return { + "pixel_values": torch.zeros((1, 1, 3), dtype=torch.float32), + "pixel_attention_mask": torch.ones((1, 1), dtype=torch.bool), + "spatial_shapes": torch.tensor([[1, 1]], dtype=torch.long), + } + + def fake_vae_processor(_image: Image.Image): + return torch.zeros((1, 3, 1, 1), dtype=torch.float32) + + processor = object.__new__(HunyuanImage3Processor) + processor.hf_config = SimpleNamespace( + vit={"num_channels": 3}, + vae_downsample_factor=(1, 1), + patch_size=1, + ) + processor.reso_group = _FixedResolutionGroup() + processor.vision_encoder_processor = fake_vit_processor + processor.vae_processor = fake_vae_processor + + processor.process_image(Image.new("RGBA", (1, 1), color=(255, 0, 0, 0))) + + assert captured["mode"] == "RGB" + assert captured["pixel"] == (255, 0, 0) diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py index 51f6a85f580..34f86347a8a 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py @@ -20,6 +20,7 @@ RATIO_END = 210 RATIO_OTHER_START = 220 RATIO_OTHER_END = 223 +TIMESTEP = 224 class FakeSamplerModel: @@ -60,6 +61,67 @@ def __init__(self, *, is_comprehension: bool = False): _apply_ratio_restriction = _Real._apply_ratio_restriction +class FakeTokenEmbedding: + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return torch.stack( + [ + input_ids.to(dtype=torch.float32), + input_ids.to(dtype=torch.float32) + 0.5, + ], + dim=-1, + ) + + +class FakeEmbeddingModel: + def __init__(self): + self.model = FakeTokenEmbedding() + self._timestep_token_id = TIMESTEP + + from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import ( + HunyuanImage3ForConditionalGeneration as _Real, + ) + + embed_input_ids = _Real.embed_input_ids + + def _timestep_encode(self, timestep: torch.Tensor) -> torch.Tensor: + return torch.full( + (1, 2), + -7.0, + device=timestep.device, + dtype=timestep.dtype, + ) + + +class TestTimestepEmbedding: + def test_timestep_embedding_avoids_scalar_item_sync(self, monkeypatch): + model = FakeEmbeddingModel() + input_ids = torch.tensor([1, TIMESTEP, 2, TIMESTEP], dtype=torch.long) + + def fail_item(_self): + raise AssertionError("embed_input_ids should not synchronize timestep count") + + monkeypatch.setattr(torch.Tensor, "item", fail_item) + + actual = model.embed_input_ids(input_ids) + expected = model.model.embed_input_ids(input_ids) + expected[input_ids == TIMESTEP] = -7.0 + + torch.testing.assert_close(actual, expected) + + def test_timestep_embedding_preserves_no_timestep_inputs(self, monkeypatch): + model = FakeEmbeddingModel() + input_ids = torch.tensor([1, 2, 3], dtype=torch.long) + + def fail_item(_self): + raise AssertionError("embed_input_ids should not synchronize timestep count") + + monkeypatch.setattr(torch.Tensor, "item", fail_item) + + actual = model.embed_input_ids(input_ids) + + torch.testing.assert_close(actual, model.model.embed_input_ids(input_ids)) + + class TestGetForcedToken: """Tests for the stateless _get_forced_token method.""" diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py index 7b6c94c62f2..b4f1e7493f2 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py @@ -19,6 +19,7 @@ HunyuanImage3Pipeline, ) from vllm_omni.diffusion.worker.input_batch import InputBatch +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.utils import DiffusionRequestState pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -41,7 +42,7 @@ def _state(request_id: str, step_index: int) -> DiffusionRequestState: state = DiffusionRequestState( request_id=request_id, sampling=SimpleNamespace(), - prompts=["prompt"], + prompt="prompt", ) state.step_index = step_index state.timesteps = torch.tensor([1.0, 0.5, 0.25, 0.0]) @@ -135,7 +136,7 @@ def fake_prepare_model_inputs(**kwargs): state = DiffusionRequestState( request_id="req-bot-task", sampling=sampling, - prompts=[prompt_item], + prompt=prompt_item, ) with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"): @@ -160,10 +161,14 @@ def fake_prepare_model_inputs(**kwargs): monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt) pipeline.prepare_model_inputs = fake_prepare_model_inputs - req = SimpleNamespace( - request_id="req-forward-bot-task", - sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"), - prompts=[{"prompt": "prompt", "bot_task": "vanilla"}], + req = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="req-forward-bot-task", + sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"), + prompt={"prompt": "prompt", "bot_task": "vanilla"}, + ) + ] ) with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"): diff --git a/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py b/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py index dc4d49b111b..9609a39baf5 100644 --- a/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py +++ b/tests/diffusion/models/hunyuan_image3/test_multi_resolution.py @@ -524,11 +524,11 @@ def _call_ar2diffusion( self, token_ids: list[int], prompt_kwargs: dict[str, Any] | None = None, - ) -> tuple[Any, Any]: + ) -> dict[str, Any] | None: """Invoke ``ar2diffusion`` with a single fake AR output. - Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can - assert on the resolved height/width. + Returns the diffusion request payload so callers can assert on the + resolved height/width. """ from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import ( ar2diffusion, @@ -601,8 +601,8 @@ def test_ar2diffusion_extracts_height_width(self, ratio_idx): 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 result is not None + diffusion_input = result assert diffusion_input["height"] == expected_h, ( f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}" ) @@ -617,8 +617,8 @@ def test_no_ratio_token_falls_back_to_prompt_size(self): 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 result is not None + diffusion_input = result assert diffusion_input["height"] == 768 assert diffusion_input["width"] == 1024 @@ -633,9 +633,9 @@ def test_out_of_range_ratio_index_falls_back(self): 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 + assert result is not None + assert result["height"] == 512 + assert result["width"] == 512 # -- cache behaviour ----------------------------------------------------- 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 index 524a9b93b33..d4bbb0f8ef0 100644 --- 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 @@ -10,6 +10,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch +import PIL.Image import pytest import torch @@ -19,6 +20,21 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] +def test_i2v_pre_process_uses_single_request_prompt(): + preprocess = i2v_module.get_hunyuan_video_15_i2v_pre_process_func(SimpleNamespace()) + image = PIL.Image.new("RGB", (640, 320)) + request = SimpleNamespace( + prompt={"prompt": "turn this into a video", "multi_modal_data": {"image": image}}, + sampling_params=SimpleNamespace(height=None, width=None), + ) + + result = preprocess(request) + + assert result is request + assert request.sampling_params.height == 448 + assert request.sampling_params.width == 896 + + class TestHunyuanVideoQuantConfigPropagation: """Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines.""" diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py deleted file mode 100644 index bf6a7563623..00000000000 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ /dev/null @@ -1,915 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Unit tests for LTX-2.3 pipeline integration. - -These tests verify: -- Pipeline is properly registered in the diffusion registry -- Post-process function is registered -- Cache-DiT enablers are registered -- Pipeline does NOT inherit from LTX2Pipeline -- Vocoder sample rate detection logic -- Re-export module works correctly -""" - -import json -import os -import tempfile -from types import SimpleNamespace - -import pytest -import torch - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _make_ltx23_pipeline(sequence_parallel_size: int = 1): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - pipeline = object.__new__(LTX23Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.audio_vae_temporal_compression_ratio = 4 - pipeline.audio_vae_mel_compression_ratio = 4 - pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size)) - pipeline.audio_vae = SimpleNamespace( - latents_mean=torch.tensor(0.0), - latents_std=torch.tensor(1.0), - ) - return pipeline - - -class TestPipelineIndependence: - """Verify LTX23Pipeline is fully independent from LTX2Pipeline.""" - - def test_ltx23_pipeline_does_not_inherit_from_ltx2(self): - """LTX23Pipeline must NOT inherit from LTX2Pipeline.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - assert not issubclass(LTX23Pipeline, LTX2Pipeline), ( - "LTX23Pipeline should be fully independent and not inherit from LTX2Pipeline" - ) - - def test_ltx23_pipeline_is_nn_module(self): - """LTX23Pipeline must be an nn.Module.""" - import torch.nn as nn - - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - assert issubclass(LTX23Pipeline, nn.Module) - - def test_ltx23_pipeline_has_progress_bar(self): - """LTX23Pipeline must mix in ProgressBarMixin.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin - - assert issubclass(LTX23Pipeline, ProgressBarMixin) - - def test_ltx23_pipeline_has_cfg_parallel_mixin(self): - """LTX23Pipeline must use the shared CFG parallel implementation.""" - from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - 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.""" - - def test_combine_cfg_noise_matches_x0_space_formula(self): - import torch - - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - pipe = object.__new__(LTX23Pipeline) - video_sample = torch.tensor([[[1.0, -2.0]]]) - audio_sample = torch.tensor([[[0.5, 3.0]]]) - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2]]]) - audio_neg = torch.tensor([[[0.1, 0.4]]]) - video_sigma = torch.tensor(0.25) - audio_sigma = torch.tensor(0.5) - scale = 4.0 - - video_combined, audio_combined = pipe.combine_cfg_noise( - (video_pos, audio_pos), - (video_neg, audio_neg), - scale, - video_latents=video_sample, - audio_latents=audio_sample, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - x0_video_cond = video_sample - video_pos * video_sigma - x0_video_uncond = video_sample - video_neg * video_sigma - x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond) - expected_video = (video_sample - x0_video_guided) / video_sigma - - x0_audio_cond = audio_sample - audio_pos * audio_sigma - x0_audio_uncond = audio_sample - audio_neg * audio_sigma - x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond) - expected_audio = (audio_sample - x0_audio_guided) / audio_sigma - assert torch.allclose(video_combined, expected_video) - assert torch.allclose(audio_combined, expected_audio) - - def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch): - from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 - - pipe = object.__new__(ltx23.LTX23Pipeline) - video_sample = torch.tensor([[[1.0, -2.0]]]) - audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]]) - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]]) - audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]]) - video_sigma = torch.tensor(0.25) - audio_sigma = torch.tensor(0.5) - scale = 4.0 - - class FakeCfgGroup: - def all_gather(self, tensor, separate_tensors=True): - assert separate_tensors - if tensor.shape == video_pos.shape: - return [video_pos, video_neg] - return [audio_pos, audio_neg] - - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) - - expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg( - video_sample, - video_pos, - video_neg, - video_sigma, - scale, - ) - expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg( - audio_sample, - audio_pos, - audio_neg, - audio_sigma, - scale, - ) - - for rank, expected_branch in ((0, "positive"), (1, "negative")): - calls = [] - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank) - - def fake_predict_noise(**kwargs): - calls.append(kwargs["branch"]) - if kwargs["branch"] == "positive": - return video_pos, audio_pos - return video_neg, audio_neg - - object.__setattr__(pipe, "predict_noise", fake_predict_noise) - video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg( - true_cfg_scale=scale, - positive_kwargs={"branch": "positive"}, - negative_kwargs={"branch": "negative"}, - cfg_normalize=False, - video_latents=video_sample, - audio_latents=audio_sample, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - assert calls == [expected_branch] - torch.testing.assert_close(video_combined, expected_video) - torch.testing.assert_close(audio_combined, expected_audio) - - assert "_cfg_video_latents" not in pipe.__dict__ - assert "_cfg_audio_latents" not in pipe.__dict__ - - -class TestCFGParallelForwardPath: - """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" - - @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)]) - def test_forward_cfg_parallel_steps_video_and_audio_scheduler( - self, - monkeypatch, - cfg_rank, - expected_prompt_value, - ): - from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 - from vllm_omni.diffusion.request import OmniDiffusionRequest - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pipe = object.__new__(ltx23.LTX23Pipeline) - torch.nn.Module.__init__(pipe) - pipe.device = torch.device("cpu") - pipe.tokenizer_max_length = 1 - pipe.vae_spatial_compression_ratio = 32 - pipe.vae_temporal_compression_ratio = 1 - pipe.transformer_spatial_patch_size = 1 - pipe.transformer_temporal_patch_size = 1 - pipe.audio_sampling_rate = 1 - pipe.audio_hop_length = 1 - pipe.audio_vae_temporal_compression_ratio = 1 - pipe.audio_vae_mel_compression_ratio = 1 - pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) - pipe.tokenizer = SimpleNamespace(padding_side="left") - pipe.vae = SimpleNamespace( - latents_mean=torch.zeros(2), - latents_std=torch.ones(2), - config=SimpleNamespace(scaling_factor=1.0), - ) - pipe.audio_vae = SimpleNamespace( - latents_mean=torch.zeros(2), - latents_std=torch.ones(2), - config=SimpleNamespace(mel_bins=2, latent_channels=1), - ) - - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2]]]) - audio_neg = torch.tensor([[[0.1, 0.4]]]) - - class FakeCfgGroup: - def all_gather(self, tensor, separate_tensors=True): - assert separate_tensors - if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg): - return [video_pos, video_neg] - if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg): - return [audio_pos, audio_neg] - raise AssertionError(f"Unexpected gathered tensor: {tensor}") - - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank) - monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) - - def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): - scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) - return torch.tensor([1.0, 0.5], device=device), 2 - - monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps) - - class FakeScheduler: - def __init__(self, name="video", calls=None): - self.name = name - self.calls = [] if calls is None else calls - self.config = { - "max_image_seq_len": 4096, - "base_image_seq_len": 1024, - "base_shift": 0.95, - "max_shift": 2.05, - } - self.sigmas = torch.tensor([0.25, 0.25]) - - def __deepcopy__(self, memo): - return FakeScheduler("audio", self.calls) - - def step(self, noise_pred, t, latents, return_dict=False, generator=None): - self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone())) - return (latents - noise_pred,) - - class FakeConnectors: - def to(self, device): - return self - - def __call__(self, prompt_embeds, prompt_attention_mask, padding_side): - assert padding_side == "left" - assert prompt_embeds.shape[0] == 2 - return prompt_embeds, prompt_embeds, prompt_attention_mask - - class FakeRope: - def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps): - return torch.zeros(batch_size, num_frames * height * width, 3, device=device) - - def prepare_audio_coords(self, batch_size, num_frames, device): - return torch.zeros(batch_size, num_frames, 1, device=device) - - class FakeTransformer: - def __init__(self): - self.config = SimpleNamespace(in_channels=2) - self.rope = FakeRope() - self.audio_rope = FakeRope() - self.calls = [] - - def __call__(self, **kwargs): - self.calls.append(kwargs) - expected_prompt = torch.full((1, 1, 1), expected_prompt_value) - torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt) - torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt) - assert kwargs["hidden_states"].shape == (1, 1, 2) - assert kwargs["audio_hidden_states"].shape == (1, 1, 2) - if cfg_rank == 0: - return video_pos, audio_pos - return video_neg, audio_neg - - class DummyProgress: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def update(self): - pass - - pipe.scheduler = FakeScheduler() - pipe.connectors = FakeConnectors() - pipe.transformer = FakeTransformer() - object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress()) - - def fake_encode_prompt(**kwargs): - return ( - torch.ones(1, 1, 1), - torch.ones(1, 1, dtype=torch.bool), - torch.zeros(1, 1, 1), - torch.ones(1, 1, dtype=torch.bool), - ) - - object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) - - video_latents = torch.tensor([[[1.0, -2.0]]]) - audio_latents = torch.tensor([[[0.5, 3.0]]]) - req = OmniDiffusionRequest( - prompts=[{"prompt": "prompt", "negative_prompt": "negative"}], - sampling_params=OmniDiffusionSamplingParams( - height=32, - width=32, - num_frames=1, - frame_rate=1.0, - num_inference_steps=2, - guidance_scale=4.0, - latents=video_latents, - audio_latents=audio_latents, - output_type="latent", - ), - request_id="ltx23-cfg-parallel-forward-test", - ) - - output = pipe.forward(req) - - expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( - video_latents, - video_pos, - video_neg, - pipe.scheduler.sigmas[0], - 4.0, - ) - expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( - audio_latents, - audio_pos, - audio_neg, - pipe.scheduler.sigmas[0], - 4.0, - ) - scheduler_call_names = [call[0] for call in pipe.scheduler.calls] - assert scheduler_call_names == ["video", "audio", "video", "audio"] - assert len(pipe.transformer.calls) == 2 - torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise) - torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise) - torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise) - - video_out, audio_out = output.output - torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) - torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) - - -class TestRegistryIntegration: - """Verify all LTX-2.3 pipeline variants are registered.""" - - def test_pipeline_models_registered(self): - """LTX-2.3 pipeline variants must be in _DIFFUSION_MODELS.""" - from vllm_omni.diffusion.registry import _DIFFUSION_MODELS - - expected = [ - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", - ] - for name in expected: - assert name in _DIFFUSION_MODELS, f"{name} not found in _DIFFUSION_MODELS" - - def test_pipeline_module_paths(self): - """Registry entries must point to the correct modules.""" - from vllm_omni.diffusion.registry import _DIFFUSION_MODELS - - # T2V -> pipeline_ltx2_3 - assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2_3", "LTX23Pipeline") - - # I2V -> pipeline_ltx2_3_image2video - assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == ( - "ltx2", - "pipeline_ltx2_3_image2video", - "LTX23ImageToVideoPipeline", - ) - - def test_post_process_funcs_registered(self): - """Pipeline variants must map to get_ltx2_post_process_func.""" - from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS - - expected = [ - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", - ] - for name in expected: - assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS" - assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func" - - def test_cache_dit_enablers_registered(self): - """Pipeline variants must be registered in CUSTOM_DIT_ENABLERS.""" - from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS - - expected = [ - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", - ] - for name in expected: - assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" - - -class TestVocoderSampleRateDetection: - """Test _detect_vocoder_output_sample_rate logic.""" - - def test_detects_48khz_from_config(self): - """Should detect output_sampling_rate=48000 from vocoder/config.json.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate - - with tempfile.TemporaryDirectory() as tmpdir: - vocoder_dir = os.path.join(tmpdir, "vocoder") - os.makedirs(vocoder_dir) - with open(os.path.join(vocoder_dir, "config.json"), "w") as f: - json.dump({"output_sampling_rate": 48000, "input_sampling_rate": 16000}, f) - - result = _detect_vocoder_output_sample_rate(tmpdir) - assert result == 48000 - - def test_returns_none_for_no_output_sr(self): - """Should return None if vocoder config has no output_sampling_rate.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate - - with tempfile.TemporaryDirectory() as tmpdir: - vocoder_dir = os.path.join(tmpdir, "vocoder") - os.makedirs(vocoder_dir) - with open(os.path.join(vocoder_dir, "config.json"), "w") as f: - json.dump({"sampling_rate": 16000}, f) - - result = _detect_vocoder_output_sample_rate(tmpdir) - assert result is None - - def test_returns_none_for_missing_directory(self): - """Should return None if vocoder directory doesn't exist.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate - - result = _detect_vocoder_output_sample_rate("/nonexistent/path") - assert result is None - - -class TestPostProcessFunction: - """Test the post-process function factory.""" - - def test_post_process_includes_audio_sample_rate(self): - """Post-process func should include audio_sample_rate when detected.""" - import torch - - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func - - with tempfile.TemporaryDirectory() as tmpdir: - vocoder_dir = os.path.join(tmpdir, "vocoder") - os.makedirs(vocoder_dir) - with open(os.path.join(vocoder_dir, "config.json"), "w") as f: - json.dump({"output_sampling_rate": 48000}, f) - - # Create a minimal od_config mock - class MockConfig: - model = tmpdir - - func = get_ltx2_post_process_func(MockConfig()) - - video = torch.zeros(1, 3, 4, 64, 64) - audio = torch.zeros(1, 1, 48000) - result = func((video, audio)) - - assert isinstance(result, dict) - assert "video" in result - assert "audio" in result - assert result["audio_sample_rate"] == 48000 - - def test_post_process_without_vocoder_config(self): - """Post-process func should work without vocoder config (no audio_sample_rate key).""" - import torch - - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func - - class MockConfig: - model = "/nonexistent/path" - - func = get_ltx2_post_process_func(MockConfig()) - - video = torch.zeros(1, 3, 4, 64, 64) - audio = torch.zeros(1, 1, 16000) - result = func((video, audio)) - - assert isinstance(result, dict) - assert "video" in result - assert "audio" in result - assert "audio_sample_rate" not in result - - -class TestReExportModule: - """Test that pipeline_ltx2_3_image2video.py correctly re-exports.""" - - def test_i2v_classes_importable(self): - """I2V classes must be importable from the re-export module.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline - - assert LTX23ImageToVideoPipeline is not None - - def test_post_process_func_importable(self): - """get_ltx2_post_process_func must be importable from re-export module.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func - - assert callable(get_ltx2_post_process_func) - - def test_i2v_classes_are_same_as_direct_import(self): - """Re-exported classes must be the same objects as direct imports.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23ImageToVideoPipeline as Direct - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import ( - LTX23ImageToVideoPipeline as ReExported, - ) - - assert Direct is ReExported - - -class TestInitExports: - """Test that __init__.py exports all LTX-2.3 classes.""" - - def test_all_ltx23_classes_exported(self): - """All LTX-2.3 pipeline classes must be in the ltx2 package __all__.""" - from vllm_omni.diffusion.models import ltx2 - - expected_classes = [ - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", - ] - for name in expected_classes: - assert hasattr(ltx2, name), f"{name} not exported from ltx2 package" - assert name in ltx2.__all__, f"{name} not in ltx2.__all__" - - -class TestAudioLatentSPPadding: - def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self): - pipeline = _make_ltx23_pipeline(sequence_parallel_size=2) - - latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( - batch_size=1, - num_channels_latents=8, - num_mel_bins=64, - audio_latent_length=1, - dtype=torch.float32, - device=torch.device("cpu"), - ) - - assert original_num_frames == 1 - assert padded_num_frames == 2 - assert latents.shape == (1, 2, 128) - - def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self): - pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) - latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4) - - padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( - batch_size=1, - num_channels_latents=2, - num_mel_bins=8, - audio_latent_length=10, - dtype=torch.float32, - device=torch.device("cpu"), - latents=latents, - ) - - assert original_num_frames == 10 - assert padded_num_frames == 12 - assert padded.shape == (1, 12, 4) - torch.testing.assert_close(padded[:, :10], latents) - torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4)) - - def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self): - pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) - latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4) - - audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) - padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( - batch_size=1, - num_channels_latents=2, - num_mel_bins=16, - audio_latent_length=audio_latent_length, - dtype=torch.float32, - device=torch.device("cpu"), - latents=latents, - ) - - assert audio_latent_length == 10 - assert original_num_frames == 10 - assert padded_num_frames == 12 - assert padded.shape == (1, 12, 8) - torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents)) - - def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self): - pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) - latents = torch.zeros(1, 2, 13, 4) - - audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) - - assert audio_latent_length == 13 - - def test_prepare_audio_latents_rejects_incompatible_provided_length(self): - pipeline = _make_ltx23_pipeline(sequence_parallel_size=4) - latents = torch.zeros(1, 11, 4) - - with pytest.raises(ValueError, match="incompatible audio frame count"): - pipeline.prepare_audio_latents( - batch_size=1, - num_channels_latents=2, - num_mel_bins=8, - audio_latent_length=10, - dtype=torch.float32, - device=torch.device("cpu"), - latents=latents, - ) diff --git a/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py b/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py deleted file mode 100644 index bbfe63dfa58..00000000000 --- a/tests/diffusion/models/ltx2/test_ltx2_cfg_parallel_adaptation.py +++ /dev/null @@ -1,58 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _make_pipeline(sequence_parallel_size: int = 1) -> LTX2Pipeline: - pipeline = object.__new__(LTX2Pipeline) - torch.nn.Module.__init__(pipeline) - pipeline.audio_vae_temporal_compression_ratio = 4 - pipeline.audio_vae_mel_compression_ratio = 4 - pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size)) - # Mock audio_vae with identity normalization (mean=0, std=1) so - # _normalize_audio_latents is a no-op and test values are preserved. - pipeline.audio_vae = SimpleNamespace( - latents_mean=torch.tensor(0.0), - latents_std=torch.tensor(1.0), - ) - return pipeline - - -def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents(): - pipeline = _make_pipeline(sequence_parallel_size=4) - latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4) - - padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( - batch_size=1, - num_channels_latents=2, - num_mel_bins=8, - audio_latent_length=10, - dtype=torch.float32, - device=torch.device("cpu"), - latents=latents, - ) - - assert original_num_frames == 10 - assert padded_num_frames == 12 - assert padded.shape == (1, 12, 4) - torch.testing.assert_close(padded[:, :10], latents) - torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4)) - - -def test_unpad_audio_latents_restores_original_frames_before_unpack(): - pipeline = _make_pipeline() - original = torch.arange(40, dtype=torch.float32).view(1, 10, 4) - padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1) - - unpadded = pipeline._unpad_audio_latents(padded, 10) - unpacked = pipeline._unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2) - expected = pipeline._unpack_audio_latents(original, latent_length=10, num_mel_bins=2) - - assert unpacked.shape == (1, 2, 10, 2) - assert not (unpacked == 999.0).any() - torch.testing.assert_close(unpacked, expected) diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py new file mode 100644 index 00000000000..43737f8c102 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -0,0 +1,219 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for LTX image-to-video input and conditioning behavior.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_ltx23_request_pipe(cls): + pipe = object.__new__(cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 99 + return pipe + + +class TestLTXImageToVideoForwardStages: + def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline) + image = torch.zeros(3, 8, 8) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "make the image move", + "negative_prompt": "jitter", + "multi_modal_data": {"image": image}, + }, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + num_inference_steps=2, + ), + request_id="ltx23-i2v-forward-stage-delegation", + ) + ] + ) + seen = {} + + def fake_forward_impl(req_arg, request_inputs, **kwargs): + seen["req"] = req_arg + seen["request_inputs"] = request_inputs + seen["kwargs"] = kwargs + return ["i2v-delegated"] + + object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + + output = pipe.forward(req, noise_scale=0.5) + + assert output == ["i2v-delegated"] + assert seen["req"] is req + assert seen["request_inputs"].prompt == ["make the image move"] + assert seen["request_inputs"].negative_prompt == ["jitter"] + assert seen["kwargs"]["image"] is image + assert seen["kwargs"]["noise_scale"] == 0.5 + + def test_denoise_timestep_kwargs_masks_video_only(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + ts = torch.tensor([2.0, 4.0]) + denoise_ctx = SimpleNamespace( + conditioning_mask=torch.tensor([[1.0, 0.0]]), + conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + ) + + kwargs = pipe._denoise_timestep_kwargs( + ts, + SimpleNamespace(cfg_parallel_ready=False), + denoise_ctx, + ) + + torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]])) + torch.testing.assert_close(kwargs["audio_timestep"], ts) + torch.testing.assert_close(kwargs["sigma"], ts) + + +class TestLTXImageToVideoConditioning: + def test_ltx23_i2v_supports_image_input(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + assert LTX23ImageToVideoPipeline.support_image_input is True + + def test_ltx23_i2v_rejects_multi_image_prompt_list(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + image = object() + + assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image + with pytest.raises(ValueError, match="exactly one image per prompt"): + LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()]) + + def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + image = torch.zeros(1, 3, 4, 4) + additional = { + "preprocessed_image": None, + "pixel_values": image, + "image": torch.ones_like(image), + } + + assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image + + def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning + import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + torch.nn.Module.__init__(pipe) + pipe.vae_spatial_compression_ratio = 1 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + + def fake_randn_tensor(shape, generator=None, device=None, dtype=None): + raise AssertionError("packed I2V latents should not be noised") + + monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor) + monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor) + + latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]]) + + out, conditioning_mask = pipe.prepare_latents( + image=None, + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=3, + noise_scale=1.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) + torch.testing.assert_close(out, latents) + + def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning + import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + torch.nn.Module.__init__(pipe) + pipe.vae_spatial_compression_ratio = 1 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + + def fake_randn_tensor(shape, generator=None, device=None, dtype=None): + return torch.ones(shape, device=device, dtype=dtype) + + monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor) + monkeypatch.setattr(ltx2_latents, "randn_tensor", fake_randn_tensor) + + latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]]) + + out, conditioning_mask = pipe.prepare_latents( + image=None, + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=3, + noise_scale=1.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) + torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]])) + + def test_ltx23_i2v_video_step_preserves_conditioning_frame(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + torch.nn.Module.__init__(pipe) + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + + class FakeScheduler: + def step(self, noise_pred, t, latents, return_dict=False): + return (latents + noise_pred + t,) + + pipe.scheduler = FakeScheduler() + latents = torch.tensor([[[1.0], [2.0], [3.0]]]) + noise_pred = torch.full_like(latents, 10.0) + + out = pipe._step_video_latents_i2v( + noise_pred, + latents, + torch.tensor(0.5), + latent_num_frames=3, + latent_height=1, + latent_width=1, + ) + + torch.testing.assert_close(out[:, :1], latents[:, :1]) + torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5) diff --git a/tests/diffusion/models/ltx2/test_ltx2_hsdp.py b/tests/diffusion/models/ltx2/test_ltx2_hsdp.py deleted file mode 100644 index 4dd07e1bf82..00000000000 --- a/tests/diffusion/models/ltx2/test_ltx2_hsdp.py +++ /dev/null @@ -1,25 +0,0 @@ -import pytest -import torch.nn as nn - -from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def test_ltx2_exposes_hsdp_shard_conditions_for_transformer_blocks(): - model = object.__new__(LTX2VideoTransformer3DModel) - nn.Module.__init__(model) - model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) - model.norm_out = nn.LayerNorm(4) - - conditions = getattr(model, "_hsdp_shard_conditions", None) - - assert conditions is not None - assert len(conditions) == 1 - - matched = [] - for name, module in model.named_modules(): - if any(cond(name, module) for cond in conditions): - matched.append(name) - - assert matched == ["transformer_blocks.0", "transformer_blocks.1"] diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py new file mode 100644 index 00000000000..74cad09d1a9 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1): + pipeline = object.__new__(pipeline_cls) + torch.nn.Module.__init__(pipeline) + pipeline.audio_vae_temporal_compression_ratio = 4 + pipeline.audio_vae_mel_compression_ratio = 4 + pipeline.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=sequence_parallel_size)) + # Mock audio_vae with identity normalization (mean=0, std=1). + pipeline.audio_vae = SimpleNamespace( + latents_mean=torch.tensor(0.0), + latents_std=torch.tensor(1.0), + ) + return pipeline + + +def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(): + pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2) + + latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=8, + num_mel_bins=64, + audio_latent_length=1, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + assert original_num_frames == 1 + assert padded_num_frames == 2 + assert latents.shape == (1, 2, 128) + + +def test_prepare_audio_latents_pads_packed_sequence_dim_for_provided_latents(): + pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4) + latents = torch.arange(40, dtype=torch.float32).view(1, 10, 4) + + padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=10, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + assert original_num_frames == 10 + assert padded_num_frames == 12 + assert padded.shape == (1, 12, 4) + torch.testing.assert_close(padded[:, :10], latents) + torch.testing.assert_close(padded[:, 10:], torch.zeros(1, 2, 4)) + + +def test_unpad_audio_latents_restores_original_frames_before_unpack(): + original = torch.arange(40, dtype=torch.float32).view(1, 10, 4) + padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1) + + unpadded = latent_ops.unpad_audio_latents(padded, 10) + unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2) + expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2) + + assert unpacked.shape == (1, 2, 10, 2) + assert not (unpacked == 999.0).any() + torch.testing.assert_close(unpacked, expected) + + +def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(): + pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4) + latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4) + + audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) + padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=16, + audio_latent_length=audio_latent_length, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + assert audio_latent_length == 10 + assert original_num_frames == 10 + assert padded_num_frames == 12 + assert padded.shape == (1, 12, 8) + torch.testing.assert_close(padded, latent_ops.pack_audio_latents(latents)) + + +def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(): + pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4) + latents = torch.zeros(1, 2, 13, 4) + + audio_latent_length = pipeline._resolve_audio_latent_length(10, latents) + + assert audio_latent_length == 13 + + +def test_prepare_audio_latents_rejects_incompatible_provided_length(): + pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4) + latents = torch.zeros(1, 11, 4) + + with pytest.raises(ValueError, match="incompatible audio frame count"): + pipeline.prepare_audio_latents( + batch_size=1, + num_channels_latents=2, + num_mel_bins=8, + audio_latent_length=10, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py new file mode 100644 index 00000000000..0c6b12fbe48 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -0,0 +1,424 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for LTX guidance and forward-parallel behavior.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestCFGParallelHelpers: + """Test LTX-2.3 CFG helper math without loading model weights.""" + + def test_combine_cfg_noise_matches_x0_space_formula(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + + pipe = object.__new__(LTX23Pipeline) + video_sample = torch.tensor([[[1.0, -2.0]]]) + audio_sample = torch.tensor([[[0.5, 3.0]]]) + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2]]]) + audio_neg = torch.tensor([[[0.1, 0.4]]]) + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + scale = 4.0 + + video_combined, audio_combined = pipe.combine_cfg_noise( + (video_pos, audio_pos), + (video_neg, audio_neg), + scale, + video_latents=video_sample, + audio_latents=audio_sample, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + x0_video_cond = video_sample - video_pos * video_sigma + x0_video_uncond = video_sample - video_neg * video_sigma + x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond) + expected_video = (video_sample - x0_video_guided) / video_sigma + + x0_audio_cond = audio_sample - audio_pos * audio_sigma + x0_audio_uncond = audio_sample - audio_neg * audio_sigma + x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond) + expected_audio = (audio_sample - x0_audio_guided) / audio_sigma + assert torch.allclose(video_combined, expected_video) + assert torch.allclose(audio_combined, expected_audio) + + def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_guidance + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23 + + pipe = object.__new__(ltx23.LTX23Pipeline) + video_sample = torch.tensor([[[1.0, -2.0]]]) + audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]]) + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]]) + audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]]) + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + scale = 4.0 + + class FakeCfgGroup: + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + if tensor.shape == video_pos.shape: + return [video_pos, video_neg] + return [audio_pos, audio_neg] + + monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup()) + + expected_video = ltx2_guidance.combine_velocity_via_x0( + video_sample, + video_pos, + video_neg, + video_sigma, + scale, + ) + expected_audio = ltx2_guidance.combine_velocity_via_x0( + audio_sample, + audio_pos, + audio_neg, + audio_sigma, + scale, + ) + + for rank, expected_branch in ((0, "positive"), (1, "negative")): + calls = [] + monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank) + + def fake_predict_noise(**kwargs): + calls.append(kwargs["branch"]) + if kwargs["branch"] == "positive": + return video_pos, audio_pos + return video_neg, audio_neg + + object.__setattr__(pipe, "predict_noise", fake_predict_noise) + video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg( + true_cfg_scale=scale, + positive_kwargs={"branch": "positive"}, + negative_kwargs={"branch": "negative"}, + cfg_normalize=False, + video_latents=video_sample, + audio_latents=audio_sample, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + assert calls == [expected_branch] + torch.testing.assert_close(video_combined, expected_video) + torch.testing.assert_close(audio_combined, expected_audio) + + assert "_cfg_video_latents" not in pipe.__dict__ + assert "_cfg_audio_latents" not in pipe.__dict__ + + +class TestCFGParallelForwardPath: + """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" + + @pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"]) + def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name): + from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipeline_cls = getattr(ltx23, pipeline_cls_name) + pipe = object.__new__(pipeline_cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 4 + monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1) + + class StopAtEncodePromptError(Exception): + pass + + captured = {} + + def fake_encode_prompt(**kwargs): + captured.update(kwargs) + raise StopAtEncodePromptError + + object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) + + prompt_embeds_a = torch.zeros(2, 3) + prompt_embeds_b = torch.ones(2, 3) + negative_prompt_embeds_a = torch.full((2, 3), 2.0) + negative_prompt_embeds_b = torch.full((2, 3), 3.0) + prompt_attention_mask_a = torch.tensor([True, True]) + prompt_attention_mask_b = torch.tensor([True, False]) + negative_attention_mask_a = torch.tensor([False, True]) + negative_attention_mask_b = torch.tensor([False, False]) + + requests = [ + OmniDiffusionRequest( + prompt={ + "prompt": "prompt-a", + "negative_prompt": "negative-a", + "prompt_embeds": prompt_embeds_a, + "negative_prompt_embeds": negative_prompt_embeds_a, + "prompt_attention_mask": prompt_attention_mask_a, + "negative_prompt_attention_mask": negative_attention_mask_a, + }, + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=1.0, + num_inference_steps=2, + ), + request_id="ltx23-prompt-local-a", + ), + OmniDiffusionRequest( + prompt={ + "prompt": "prompt-b", + "negative_prompt": "negative-b", + "prompt_embeds": prompt_embeds_b, + "negative_prompt_embeds": negative_prompt_embeds_b, + "attention_mask": prompt_attention_mask_b, + "negative_attention_mask": negative_attention_mask_b, + }, + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=1.0, + num_inference_steps=2, + ), + request_id="ltx23-prompt-local-b", + ), + ] + + with pytest.raises(StopAtEncodePromptError): + pipe.forward(DiffusionRequestBatch(requests=requests)) + + assert captured["prompt"] is None + assert captured["negative_prompt"] is None + torch.testing.assert_close( + captured["prompt_embeds"], + torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + captured["negative_prompt_embeds"], + torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + captured["prompt_attention_mask"], + torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0), + ) + torch.testing.assert_close( + captured["negative_prompt_attention_mask"], + torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0), + ) + + @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)]) + @pytest.mark.parametrize( + ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"), + [(1.0, 1, 1.0), (None, 24, 24.0)], + ) + def test_forward_cfg_parallel_steps_video_and_audio_scheduler( + self, + monkeypatch, + cfg_rank, + expected_prompt_value, + frame_rate_input, + audio_sampling_rate, + expected_frame_rate, + ): + from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = object.__new__(ltx23.LTX23Pipeline) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 1 + pipe.vae_spatial_compression_ratio = 32 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.audio_sampling_rate = audio_sampling_rate + pipe.audio_hop_length = 1 + pipe.audio_vae_temporal_compression_ratio = 1 + pipe.audio_vae_mel_compression_ratio = 1 + pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) + pipe.tokenizer = SimpleNamespace(padding_side="left") + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.audio_vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(mel_bins=2, latent_channels=1), + ) + + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2]]]) + audio_neg = torch.tensor([[[0.1, 0.4]]]) + + class FakeCfgGroup: + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg): + return [video_pos, video_neg] + if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg): + return [audio_pos, audio_neg] + raise AssertionError(f"Unexpected gathered tensor: {tensor}") + + monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank) + monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup()) + + def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): + scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) + return torch.tensor([1.0, 0.5], device=device), 2 + + monkeypatch.setattr(ltx2_denoise, "retrieve_timesteps", fake_retrieve_timesteps) + + class FakeScheduler: + def __init__(self, name="video", calls=None): + self.name = name + self.calls = [] if calls is None else calls + self.config = { + "max_image_seq_len": 4096, + "base_image_seq_len": 1024, + "base_shift": 0.95, + "max_shift": 2.05, + } + self.sigmas = torch.tensor([0.25, 0.25]) + + def __deepcopy__(self, memo): + return FakeScheduler("audio", self.calls) + + def step(self, noise_pred, t, latents, return_dict=False, generator=None): + self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone())) + return (latents - noise_pred,) + + class FakeConnectors: + def to(self, device): + return self + + def __call__(self, prompt_embeds, prompt_attention_mask, padding_side): + assert padding_side == "left" + assert prompt_embeds.shape[0] == 2 + return prompt_embeds, prompt_embeds, prompt_attention_mask + + rope_video_fps: list[float] = [] + + class FakeRope: + def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps): + rope_video_fps.append(fps) + return torch.zeros(batch_size, num_frames * height * width, 3, device=device) + + def prepare_audio_coords(self, batch_size, num_frames, device): + return torch.zeros(batch_size, num_frames, 1, device=device) + + class FakeTransformer: + def __init__(self): + self.config = SimpleNamespace(in_channels=2) + self.rope = FakeRope() + self.audio_rope = FakeRope() + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + expected_prompt = torch.full((1, 1, 1), expected_prompt_value) + torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt) + torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt) + assert kwargs["hidden_states"].shape == (1, 1, 2) + assert kwargs["audio_hidden_states"].shape == (1, 1, 2) + if cfg_rank == 0: + return video_pos, audio_pos + return video_neg, audio_neg + + class DummyProgress: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def update(self): + pass + + pipe.scheduler = FakeScheduler() + pipe.connectors = FakeConnectors() + pipe.transformer = FakeTransformer() + object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress()) + + def fake_encode_prompt(**kwargs): + return ( + torch.ones(1, 1, 1), + torch.ones(1, 1, dtype=torch.bool), + torch.zeros(1, 1, 1), + torch.ones(1, 1, dtype=torch.bool), + ) + + object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) + + video_latents = torch.tensor([[[1.0, -2.0]]]) + audio_latents = torch.tensor([[[0.5, 3.0]]]) + req = OmniDiffusionRequest( + prompt={"prompt": "prompt", "negative_prompt": "negative"}, + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=frame_rate_input, + num_inference_steps=2, + guidance_scale=4.0, + latents=video_latents, + audio_latents=audio_latents, + output_type="latent", + ), + request_id="ltx23-cfg-parallel-forward-test", + ) + + output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0] + + expected_video_noise = ltx2_guidance.combine_velocity_via_x0( + video_latents, + video_pos, + video_neg, + pipe.scheduler.sigmas[0], + 4.0, + ) + expected_audio_noise = ltx2_guidance.combine_velocity_via_x0( + audio_latents, + audio_pos, + audio_neg, + pipe.scheduler.sigmas[0], + 4.0, + ) + scheduler_call_names = [call[0] for call in pipe.scheduler.calls] + assert scheduler_call_names == ["video", "audio", "video", "audio"] + assert len(pipe.transformer.calls) == 2 + torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise) + torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise) + torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise) + + video_out, audio_out = output.output + torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) + torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) + + # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve + # to the model's own 24.0 default, not crash on None; a provided rate is passed through. + assert rope_video_fps + assert all(fps == expected_frame_rate for fps in rope_video_fps) diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py new file mode 100644 index 00000000000..2f9155fe0ae --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -0,0 +1,793 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the shared LTX pipeline runtime and public contracts.""" + +import json +import os +import tempfile +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.models.ltx2.ltx2_components import ( + LTX2_COMPONENT_PROFILE, + LTX23_COMPONENT_PROFILE, +) +from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin +from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseExecutor, LTXPhaseResult +from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( + LTX_LEGACY_VELOCITY_GUIDANCE, + LTX_OFFICIAL_X0_GUIDANCE, +) +from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState +from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import LTXPipelineRuntime +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE +from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2ImageToVideoPipeline, + LTX2Pipeline, + LTX23ImageToVideoPipeline, + LTX23Pipeline, + LTXOneStagePipeline, +) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( + LTX2ImageToVideoTwoStagesPipeline, + LTX2TwoStagesPipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_ltx_request_pipe(cls): + pipe = object.__new__(cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 99 + return pipe + + +def _resolve_request_inputs_for_test(pipe, req): + return pipe._resolve_request_inputs( + req, + prompt=None, + negative_prompt=None, + height=None, + width=None, + num_frames=None, + frame_rate=None, + num_inference_steps=None, + timesteps=None, + guidance_scale=4.0, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=None, + ) + + +def test_ltx_versions_share_runtime_without_cross_version_inheritance(): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + + assert issubclass(LTX2Pipeline, LTXPipelineRuntime) + assert issubclass(LTX23Pipeline, LTXPipelineRuntime) + assert issubclass(LTX2TwoStagesPipeline, LTXPipelineRuntime) + assert not issubclass(LTX23Pipeline, LTX2Pipeline) + assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE + assert LTX23Pipeline.component_profile is LTX23_COMPONENT_PROFILE + assert LTX2Pipeline.one_stage_recipe is LTX2_ONE_STAGE_RECIPE + assert LTX23Pipeline.one_stage_recipe is LTX23_ONE_STAGE_RECIPE + assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video + assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video + + +def test_ltx_one_stage_variants_expose_shared_request_batch_and_distributed_decode_contracts(): + one_stage_variants = ( + LTX2Pipeline, + LTX2ImageToVideoPipeline, + LTX23Pipeline, + LTX23ImageToVideoPipeline, + ) + assert all(pipeline_cls.supports_request_batch for pipeline_cls in one_stage_variants) + assert all(pipeline_cls.distributed_video_decode for pipeline_cls in one_stage_variants) + assert not LTX2TwoStagesPipeline.supports_request_batch + assert not LTX2ImageToVideoTwoStagesPipeline.supports_request_batch + + +@pytest.mark.parametrize("pipeline_cls", [LTX2Pipeline, LTX23Pipeline]) +def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cls): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + requests = [ + OmniDiffusionRequest( + prompt=f"prompt-{index}", + sampling_params=OmniDiffusionSamplingParams(), + request_id=f"request-{index}", + ) + for index in range(2) + ] + request_batch = DiffusionRequestBatch(requests) + video = torch.tensor([[1.0], [2.0]]) + audio = torch.tensor([[3.0], [4.0]]) + pipe = object.__new__(pipeline_cls) + torch.nn.Module.__init__(pipe) + object.__setattr__(pipe, "_decode_output", lambda **_kwargs: DiffusionOutput(output=(video, audio))) + forward_ctx = SimpleNamespace( + req=request_batch, + request_inputs=SimpleNamespace( + output_type="np", + generator=None, + decode_timestep=0.0, + decode_noise_scale=None, + ), + prompt_context=SimpleNamespace(connector_prompt_embeds=torch.empty(2, 1)), + device=torch.device("cpu"), + batch_size=2, + num_videos_per_prompt=1, + ) + + outputs = pipe.decode_phase( + LTXPhaseResult( + forward_context=forward_ctx, + video=video, + audio=audio, + ) + ) + + assert len(outputs) == 2 + torch.testing.assert_close(outputs[0].output[0], video[:1]) + torch.testing.assert_close(outputs[0].output[1], audio[:1]) + torch.testing.assert_close(outputs[1].output[0], video[1:]) + torch.testing.assert_close(outputs[1].output[1], audio[1:]) + + +def test_ltx_one_stage_variants_share_forward_template(): + assert LTX2Pipeline._forward_impl is LTXOneStagePipeline._forward_impl + assert LTX2ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl + assert LTX23Pipeline._forward_impl is LTXOneStagePipeline._forward_impl + assert LTX23ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl + assert LTX2ImageToVideoPipeline.forward is LTX23ImageToVideoPipeline.forward + + +def test_ltx_versions_share_request_prompt_and_step_templates(): + shared_methods = ( + "_get_gemma_prompt_embeds", + "encode_prompt", + "check_inputs", + "_resolve_request_inputs", + "_prepare_prompt_context", + "_denoise_step", + ) + for method_name in shared_methods: + base_method = getattr(LTXPipelineRuntime, method_name) + assert getattr(LTX2Pipeline, method_name) is base_method + assert getattr(LTX23Pipeline, method_name) is base_method + + +def test_ltx_versions_select_guidance_without_overriding_control_flow(): + assert LTX2Pipeline.guidance_strategy is LTX_LEGACY_VELOCITY_GUIDANCE + assert LTX23Pipeline.guidance_strategy is LTX_OFFICIAL_X0_GUIDANCE + assert LTX2Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step + assert LTX23Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step + assert LTX2Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise + assert LTX23Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise + + +def test_ltx2_two_stage_variants_share_stage_orchestration(): + assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTX2TwoStagesPipeline) + assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTXI2VConditioningMixin) + assert LTX2ImageToVideoTwoStagesPipeline._run_two_stage is LTX2TwoStagesPipeline._run_two_stage + + +def test_ltx_variants_share_denoise_loop_and_i2v_conditioning(): + assert "_denoise_loop" not in LTX2Pipeline.__dict__ + assert "_denoise_loop" not in LTX23Pipeline.__dict__ + assert issubclass(LTX2ImageToVideoPipeline, LTXI2VConditioningMixin) + assert issubclass(LTX23ImageToVideoPipeline, LTXI2VConditioningMixin) + assert LTX2ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents + assert LTX23ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents + assert LTX2Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents + assert LTX23Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents + assert LTX2Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents + assert LTX23Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents + assert LTX2Pipeline.decode_phase is LTXPipelineRuntime.decode_phase + assert LTX23Pipeline.decode_phase is LTXPipelineRuntime.decode_phase + + +def test_denoise_executor_owns_progress_and_interrupt(): + updates = [] + + class Progress: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, traceback): + return False + + def update(self): + updates.append(True) + + pipeline = SimpleNamespace( + interrupt=False, + progress_bar=lambda total: Progress(), + ) + seen: list[tuple[int, float]] = [] + timesteps = torch.tensor([3.0, 2.0, 1.0]) + initial_state = LTXAVState(video=torch.tensor(0.0), audio=torch.tensor(10.0)) + + def step(index, timestep, state): + seen.append((index, timestep.item())) + pipeline.interrupt = True + return LTXAVState(video=state.video + 1, audio=state.audio + 1) + + state = LTXDenoiseExecutor.run(pipeline, initial_state, timesteps, step) + + assert seen == [(0, 3.0)] + assert updates == [True] + torch.testing.assert_close(state.video, torch.tensor(1.0)) + torch.testing.assert_close(state.audio, torch.tensor(11.0)) + + +def test_ltx2_two_stage_reuses_prompt_context_between_phases(): + request_inputs = LTXRequestInputs( + prompt="prompt", + negative_prompt="negative", + height=32, + width=32, + num_frames=1, + frame_rate=24.0, + num_inference_steps=4, + guidance_scale=4.0, + guidance_rescale=0.0, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.25, + decode_noise_scale=0.5, + output_type="np", + max_sequence_length=16, + ) + prompt_context = object() + phase_calls = [] + + def resolve_request_inputs(req, **kwargs): + return request_inputs + + def run_phase(req, inputs, *, prompt_context=None, **kwargs): + phase_calls.append((inputs, prompt_context)) + if len(phase_calls) == 1: + assert prompt_context is None + assert inputs.decode_timestep == 0.25 + assert inputs.decode_noise_scale == 0.5 + context = prompt_context_sentinel + video = torch.tensor([1.0]) + audio = torch.tensor([2.0]) + else: + assert prompt_context is prompt_context_sentinel + torch.testing.assert_close(inputs.latents, torch.tensor([11.0])) + torch.testing.assert_close(inputs.audio_latents, torch.tensor([2.0])) + assert inputs.guidance_scale == 1.0 + assert inputs.num_inference_steps == 3 + assert inputs.decode_timestep == 0.0 + assert inputs.decode_noise_scale is None + context = prompt_context + video = torch.tensor([3.0]) + audio = torch.tensor([4.0]) + return LTXPhaseResult( + forward_context=SimpleNamespace(prompt_context=context), + video=video, + audio=audio, + ) + + def decode_phase(phase): + return DiffusionOutput(output=(phase.video, phase.audio)) + + class FakeUpsampler(torch.nn.Module): + def forward(self, *, latents, output_type, return_dict): + assert output_type == "latent" + assert not return_dict + return (latents + 10,) + + prompt_context_sentinel = prompt_context + pipeline = object.__new__(LTX2TwoStagesPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.distilled = True + pipeline.upsample_pipe = FakeUpsampler() + object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) + object.__setattr__(pipeline, "run_phase", run_phase) + object.__setattr__(pipeline, "decode_phase", decode_phase) + + output = pipeline.forward(SimpleNamespace()) + + assert len(phase_calls) == 2 + assert phase_calls[1][1] is prompt_context_sentinel + torch.testing.assert_close(output.output[0], torch.tensor([3.0])) + torch.testing.assert_close(output.output[1], torch.tensor([4.0])) + + +class TestLTXRequestParsing: + @pytest.mark.parametrize( + ("t2v_cls", "i2v_cls"), + [ + (LTX2Pipeline, LTX2ImageToVideoPipeline), + (LTX23Pipeline, LTX23ImageToVideoPipeline), + ], + ) + def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + prompt_embeds = torch.tensor([[1.0, 2.0]]) + negative_prompt_embeds = torch.tensor([[3.0, 4.0]]) + prompt_attention_mask = torch.tensor([True, False]) + negative_attention_mask = torch.tensor([False, True]) + video_latents = torch.ones(1, 2, 3) + audio_latents = torch.zeros(1, 2, 3) + generator = torch.Generator().manual_seed(123) + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "shared prompt", + "negative_prompt": "shared negative", + "additional_information": { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "attention_mask": prompt_attention_mask, + "negative_attention_mask": negative_attention_mask, + }, + }, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + frame_rate=12.5, + num_inference_steps=1, + num_outputs_per_prompt=2, + guidance_scale=6.0, + generator=generator, + latents=video_latents, + extra_args={"audio_latents": audio_latents}, + decode_timestep=[0.1], + decode_noise_scale=[0.2], + output_type="latent", + max_sequence_length=17, + ), + request_id="ltx23-shared-request-inputs", + ) + ] + ) + + resolved_t2v = _resolve_request_inputs_for_test( + _make_ltx_request_pipe(t2v_cls), + req, + ) + resolved_i2v = _resolve_request_inputs_for_test( + _make_ltx_request_pipe(i2v_cls), + req, + ) + + assert resolved_i2v.prompt == resolved_t2v.prompt is None + assert resolved_i2v.negative_prompt == resolved_t2v.negative_prompt is None + assert resolved_i2v.height == resolved_t2v.height == 384 + assert resolved_i2v.width == resolved_t2v.width == 512 + assert resolved_i2v.num_frames == resolved_t2v.num_frames == 25 + assert resolved_i2v.frame_rate == resolved_t2v.frame_rate == 12.5 + assert resolved_i2v.num_inference_steps == resolved_t2v.num_inference_steps == 2 + assert resolved_i2v.guidance_scale == resolved_t2v.guidance_scale == 6.0 + assert resolved_i2v.num_videos_per_prompt == resolved_t2v.num_videos_per_prompt == 2 + assert isinstance(resolved_i2v.generator, list) + assert isinstance(resolved_t2v.generator, list) + assert [gen.initial_seed() for gen in resolved_i2v.generator] == [123, 123] + assert [gen.initial_seed() for gen in resolved_t2v.generator] == [123, 123] + assert resolved_i2v.decode_timestep == resolved_t2v.decode_timestep == [0.1] + assert resolved_i2v.decode_noise_scale == resolved_t2v.decode_noise_scale == [0.2] + assert resolved_i2v.output_type == resolved_t2v.output_type == "latent" + assert resolved_i2v.max_sequence_length == resolved_t2v.max_sequence_length == 17 + torch.testing.assert_close(resolved_i2v.latents, video_latents) + torch.testing.assert_close(resolved_t2v.latents, video_latents) + torch.testing.assert_close(resolved_i2v.audio_latents, audio_latents) + torch.testing.assert_close(resolved_t2v.audio_latents, audio_latents) + torch.testing.assert_close(resolved_i2v.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved_t2v.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved_i2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) + torch.testing.assert_close(resolved_t2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) + torch.testing.assert_close(resolved_i2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) + torch.testing.assert_close(resolved_t2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) + torch.testing.assert_close( + resolved_i2v.negative_prompt_attention_mask, + torch.stack([negative_attention_mask]), + ) + torch.testing.assert_close( + resolved_t2v.negative_prompt_attention_mask, + torch.stack([negative_attention_mask]), + ) + + def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx_request_pipe(LTX23Pipeline) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "with embeds", + "additional_information": {"prompt_embeds": torch.tensor([[1.0]])}, + }, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + num_inference_steps=2, + ), + request_id="ltx23-mixed-precomputed-fields-0", + ), + OmniDiffusionRequest( + prompt={"prompt": "without embeds"}, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + num_inference_steps=2, + ), + request_id="ltx23-mixed-precomputed-fields-1", + ), + ] + ) + + with pytest.raises(ValueError, match="mix of provided and missing prompt_embeds"): + _resolve_request_inputs_for_test(pipe, req) + + +class TestLTX23ForwardStages: + def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + + pipe = _make_ltx_request_pipe(LTX23Pipeline) + prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]]) + negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]]) + prompt_attention_mask = torch.tensor([[True, False], [False, True]]) + negative_prompt_attention_mask = torch.tensor([[False, False], [True, True]]) + + ( + repeated_prompt_embeds, + repeated_prompt_attention_mask, + repeated_negative_prompt_embeds, + repeated_negative_prompt_attention_mask, + ) = pipe.encode_prompt( + prompt=None, + do_classifier_free_guidance=True, + num_videos_per_prompt=2, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + device=torch.device("cpu"), + ) + + torch.testing.assert_close( + repeated_prompt_embeds, + torch.tensor([[[1.0], [2.0]], [[1.0], [2.0]], [[3.0], [4.0]], [[3.0], [4.0]]]), + ) + torch.testing.assert_close( + repeated_negative_prompt_embeds, + torch.tensor([[[-1.0], [-2.0]], [[-1.0], [-2.0]], [[-3.0], [-4.0]], [[-3.0], [-4.0]]]), + ) + torch.testing.assert_close( + repeated_prompt_attention_mask, + torch.tensor([[True, False], [True, False], [False, True], [False, True]]), + ) + torch.testing.assert_close( + repeated_negative_prompt_attention_mask, + torch.tensor([[False, False], [False, False], [True, True], [True, True]]), + ) + + def test_t2v_forward_delegates_to_shared_forward_impl(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx_request_pipe(LTX23Pipeline) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="stage refactor prompt", + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + frame_rate=16, + num_inference_steps=3, + guidance_scale=4.5, + output_type="latent", + ), + request_id="ltx23-t2v-forward-stage-delegation", + ) + ] + ) + sigmas = [1.0, 0.5] + timesteps = [1000, 500] + attention_kwargs = {"scale": 1.0} + seen = {} + + def fake_forward_impl(req_arg, request_inputs, **kwargs): + seen["req"] = req_arg + seen["request_inputs"] = request_inputs + seen["kwargs"] = kwargs + return ["delegated"] + + object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + + output = pipe.forward( + req, + sigmas=sigmas, + timesteps=timesteps, + noise_scale=0.25, + attention_kwargs=attention_kwargs, + ) + + assert output == ["delegated"] + assert seen["req"] is req + assert seen["request_inputs"].height == 384 + assert seen["request_inputs"].width == 512 + assert seen["request_inputs"].num_frames == 25 + assert seen["request_inputs"].frame_rate == 16.0 + assert seen["request_inputs"].guidance_scale == 4.5 + assert seen["request_inputs"].output_type == "latent" + assert seen["kwargs"] == { + "noise_scale": 0.25, + "sigmas": sigmas, + "timesteps": timesteps, + "attention_kwargs": attention_kwargs, + } + + +class TestPipelineComponents: + 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 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 + + +class TestLTX23DecodeConditioning: + def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self): + from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value + + assert _expand_per_prompt_decode_value( + [0.1, 0.2], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) == [0.1, 0.1, 0.2, 0.2] + assert _expand_per_prompt_decode_value( + [0.3], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) == [0.3, 0.3, 0.3, 0.3] + assert _expand_per_prompt_decode_value( + [0.1, 0.2, 0.3, 0.4], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) == [0.1, 0.2, 0.3, 0.4] + + def test_decode_conditioning_rejects_ambiguous_lengths(self): + from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value + + with pytest.raises(ValueError, match="decode_timestep"): + _expand_per_prompt_decode_value( + [0.1, 0.2, 0.3], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) + + +class TestRegistryIntegration: + """Verify all LTX-2.3 pipeline variants are registered.""" + + def test_pipeline_module_paths(self): + """Registry entries must point to the correct modules.""" + from vllm_omni.diffusion.registry import _DIFFUSION_MODELS + + assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX23Pipeline") + + assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == ( + "ltx2", + "pipeline_ltx2", + "LTX23ImageToVideoPipeline", + ) + + def test_post_process_funcs_registered(self): + """Pipeline variants must map to get_ltx2_post_process_func.""" + from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS + + expected = [ + "LTX2Pipeline", + "LTX2ImageToVideoPipeline", + "LTX2TwoStagesPipeline", + "LTX2ImageToVideoTwoStagesPipeline", + "LTX2T2VDMD2Pipeline", + "LTX2I2VDMD2Pipeline", + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ] + for name in expected: + assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS" + assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func" + + @pytest.mark.parametrize( + "model_class_name", + [ + "LTX2Pipeline", + "LTX2ImageToVideoPipeline", + "LTX2TwoStagesPipeline", + "LTX2ImageToVideoTwoStagesPipeline", + "LTX2T2VDMD2Pipeline", + "LTX2I2VDMD2Pipeline", + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ], + ) + def test_post_process_func_resolves_from_every_entry_module(self, model_class_name): + from vllm_omni.diffusion.registry import get_diffusion_post_process_func + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as config_file: + json.dump({"output_sampling_rate": 48000}, config_file) + + post_process = get_diffusion_post_process_func( + SimpleNamespace(model_class_name=model_class_name, model=tmpdir) + ) + + result = post_process((torch.zeros(1), torch.zeros(1))) + assert result["audio_sample_rate"] == 48000 + + def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self): + """Pipeline variants are *not* registered in CUSTOM_DIT_ENABLERS.""" + from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS + + # NOTE: We used to have custom enablers for this model, but refactored to handle + # it more generically. Now we only need to ensure it has git cache adapter config. + expected = [ + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ] + for name in expected: + assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" + + +class TestVocoderSampleRateDetection: + """Test _detect_vocoder_output_sample_rate logic.""" + + def test_detects_48khz_from_config(self): + """Should detect output_sampling_rate=48000 from vocoder/config.json.""" + from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as f: + json.dump({"output_sampling_rate": 48000, "input_sampling_rate": 16000}, f) + + result = _detect_vocoder_output_sample_rate(tmpdir) + assert result == 48000 + + def test_returns_none_for_no_output_sr(self): + """Should return None if vocoder config has no output_sampling_rate.""" + from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as f: + json.dump({"sampling_rate": 16000}, f) + + result = _detect_vocoder_output_sample_rate(tmpdir) + assert result is None + + def test_returns_none_for_missing_directory(self): + """Should return None if vocoder directory doesn't exist.""" + from vllm_omni.diffusion.models.ltx2.ltx2_components import _detect_vocoder_output_sample_rate + + result = _detect_vocoder_output_sample_rate("/nonexistent/path") + assert result is None + + +class TestPostProcessFunction: + """Test the post-process function factory.""" + + def test_post_process_includes_audio_sample_rate(self): + """Post-process func should include audio_sample_rate when detected.""" + import torch + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import get_ltx2_post_process_func + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as f: + json.dump({"output_sampling_rate": 48000}, f) + + # Create a minimal od_config mock + class MockConfig: + model = tmpdir + + func = get_ltx2_post_process_func(MockConfig()) + + video = torch.zeros(1, 3, 4, 64, 64) + audio = torch.zeros(1, 1, 48000) + result = func((video, audio)) + + assert isinstance(result, dict) + assert "video" in result + assert "audio" in result + assert result["audio_sample_rate"] == 48000 + + def test_post_process_without_vocoder_config(self): + """Post-process func should work without vocoder config (no audio_sample_rate key).""" + import torch + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import get_ltx2_post_process_func + + class MockConfig: + model = "/nonexistent/path" + + func = get_ltx2_post_process_func(MockConfig()) + + video = torch.zeros(1, 3, 4, 64, 64) + audio = torch.zeros(1, 1, 16000) + result = func((video, audio)) + + assert isinstance(result, dict) + assert "video" in result + assert "audio" in result + assert "audio_sample_rate" not in result diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index e712532976f..01a1fa3cfc3 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -1,9 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest from torch import nn -from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel, _make_rms_norm + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer(): @@ -20,3 +24,29 @@ def test_ltx_rms_norm_affine_weight_remains_parameter(): assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter) assert "weight" not in dict(norm.named_buffers()) assert "weight" in norm.state_dict() + + +def test_ltx_transformer_has_separate_cfg_cache_dit_config(): + adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config") + + assert isinstance(adapter_config, CacheDiTAdapterConfig) + assert adapter_config.has_separate_cfg + + +def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks(): + model = object.__new__(LTX2VideoTransformer3DModel) + nn.Module.__init__(model) + model.transformer_blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(2)]) + model.norm_out = nn.LayerNorm(4) + + conditions = getattr(model, "_hsdp_shard_conditions", None) + + assert conditions is not None + assert len(conditions) == 1 + + matched = [] + for name, module in model.named_modules(): + if any(condition(name, module) for condition in conditions): + matched.append(name) + + assert matched == ["transformer_blocks.0", "transformer_blocks.1"] diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py new file mode 100644 index 00000000000..4ae7a0721d3 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -0,0 +1,314 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestLTXOutputRank: + @pytest.mark.parametrize( + ("distributed_vae_state", "expected_decode_calls"), + [(False, 0), (True, 1), (RuntimeError("unavailable"), 0)], + ) + def test_non_output_rank_only_enters_collective_vae_decode( + self, + monkeypatch, + distributed_vae_state, + expected_decode_calls, + ): + import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx + + class FakeVae: + dtype = torch.float32 + config = SimpleNamespace(timestep_conditioning=False) + + def __init__(self): + self.decode_calls = 0 + + def is_distributed_enabled(self): + if isinstance(distributed_vae_state, Exception): + raise distributed_vae_state + return distributed_vae_state + + def decode(self, *_args, **_kwargs): + self.decode_calls += 1 + return (torch.ones(1, 1),) + + monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1) + pipe = object.__new__(ltx.LTX2Pipeline) + torch.nn.Module.__init__(pipe) + pipe.vae = FakeVae() + + output = pipe._decode_output( + latents=torch.ones(1, 1), + audio_latents=torch.ones(1, 1), + output_type="np", + connector_prompt_embeds=torch.ones(1, 1), + generator=None, + device=torch.device("cpu"), + decode_timestep=0.0, + decode_noise_scale=None, + prompt_batch_size=1, + ) + + assert pipe.vae.decode_calls == expected_decode_calls + assert output.output[0].numel() == 0 + assert output.output[1].numel() == 0 + + +class TestLTX23VaeDistributedDecode: + """Test LTX-2.3 distributed VAE 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_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 TestLTX23VaeTiling: + """Test LTX-2.3 video VAE tile helpers without loading weights.""" + + 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" diff --git a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py index 27d1e757e9d..20024abe799 100644 --- a/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py +++ b/tests/diffusion/models/mistral_encoder/test_mistral_encoder_tp.py @@ -5,6 +5,7 @@ import pytest import torch from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm_omni.diffusion.models.mistral_encoder.mistral_encoder import ( MistralEncoderModel, @@ -491,3 +492,57 @@ def test_logits_shape(self): assert logits.dim() == 3 assert logits.shape[0] == 1 assert logits.shape[1] == 4 + + +class TestQuantConfig: + """Verify quant_config is threaded into weight linears only, with correct prefixes.""" + + def test_quant_config_on_weight_linears(self, mocker): + quant_config = mocker.MagicMock(name="QuantizationConfig") + # Keep construction on the CPU-safe unquantized path while still + # verifying LinearBase consults the supplied quant_config. + quant_config.get_quant_method.return_value = UnquantizedLinearMethod() + + config = _make_config(num_hidden_layers=1) + model = MistralEncoderModel(config, prefix="text_encoder", quant_config=quant_config) + + layer = model.language_model.model.layers[0] + for linear in ( + layer.self_attn.qkv_proj, + layer.self_attn.o_proj, + layer.mlp.gate_up_proj, + layer.mlp.down_proj, + ): + assert linear.quant_config is quant_config + + # get_quant_method should be consulted for the four weight linears. + assert quant_config.get_quant_method.call_count == 4 + seen_prefixes = set() + for call in quant_config.get_quant_method.call_args_list: + if "prefix" in call.kwargs: + seen_prefixes.add(call.kwargs["prefix"]) + elif len(call.args) > 1: + seen_prefixes.add(call.args[1]) + assert any(p.endswith("self_attn.qkv_proj") for p in seen_prefixes) + assert any(p.endswith("self_attn.o_proj") for p in seen_prefixes) + assert any(p.endswith("mlp.gate_up_proj") for p in seen_prefixes) + assert any(p.endswith("mlp.down_proj") for p in seen_prefixes) + + # Embeddings / lm_head stay unquantized (quant_config never passed). + embed = model.language_model.model.embed_tokens + lm_head = model.language_model.lm_head + assert getattr(embed, "quant_config", None) is None + assert getattr(lm_head, "quant_config", None) is None + + def test_layer_prefixes_include_text_encoder(self): + config = _make_config(num_hidden_layers=2) + model = MistralEncoderModel(config, prefix="text_encoder") + + layer0 = model.language_model.model.layers[0] + assert layer0.self_attn.qkv_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + assert layer0.self_attn.o_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + assert layer0.mlp.gate_up_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + assert layer0.mlp.down_proj.prefix.startswith("text_encoder.language_model.model.layers.0") + + layer1 = model.language_model.model.layers[1] + assert "layers.1" in layer1.self_attn.qkv_proj.prefix diff --git a/tests/diffusion/models/ovis_image/test_ovis_image.py b/tests/diffusion/models/ovis_image/test_ovis_image.py index 8e357e6c636..05092ade234 100644 --- a/tests/diffusion/models/ovis_image/test_ovis_image.py +++ b/tests/diffusion/models/ovis_image/test_ovis_image.py @@ -24,6 +24,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.ovis_image.pipeline_ovis_image import OvisImagePipeline from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -166,7 +167,7 @@ def test_basic_generation(ovis_pipeline): """Test the forward pass logic.""" # Setup request req = OmniDiffusionRequest( - prompts=["A photo of a cat"], + prompt="A photo of a cat", request_id="ovis-basic", sampling_params=OmniDiffusionSamplingParams( height=256, @@ -176,7 +177,7 @@ def test_basic_generation(ovis_pipeline): ), ) - output = ovis_pipeline(req) + output = ovis_pipeline(DiffusionRequestBatch(requests=[req])) assert output is not None assert output.output is not None @@ -201,12 +202,10 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch): monkeypatch.setattr(f"{_cfg_parallel}.get_classifier_free_guidance_rank", lambda: 0) req = OmniDiffusionRequest( - prompts=[ - { - "prompt": "A photo of a cat", - "negative_prompt": "bad quality", - } - ], + prompt={ + "prompt": "A photo of a cat", + "negative_prompt": "bad quality", + }, request_id="ovis-guidance", sampling_params=OmniDiffusionSamplingParams( height=256, @@ -216,7 +215,7 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch): ), ) - ovis_pipeline(req) + ovis_pipeline(DiffusionRequestBatch(requests=[req])) assert ovis_pipeline.transformer.call_count >= 2 @@ -226,7 +225,7 @@ def test_resolution_check(ovis_pipeline): """Test resolution divisible validation logic if present.""" # Pass odd resolution req = OmniDiffusionRequest( - prompts=["test"], + prompt="test", request_id="ovis-resolution", sampling_params=OmniDiffusionSamplingParams( height=250, # Not divisible by 16 (8*2) @@ -237,7 +236,7 @@ def test_resolution_check(ovis_pipeline): # Should warn but proceed (as per code I read earlier) or resize? # The code had `logger.warning(...)` - output = ovis_pipeline(req) + output = ovis_pipeline(DiffusionRequestBatch(requests=[req])) assert output is not None diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py b/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py index 873b52bf7a6..6f09817ee7b 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_edit_plus.py @@ -25,12 +25,10 @@ def test_qwen_image_edit_plus_rejects_too_many_input_images(tmp_path: Path): pre_process = get_qwen_image_edit_plus_pre_process_func(SimpleNamespace(model=str(tmp_path))) image = Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)) request = SimpleNamespace( - prompts=[ - { - "prompt": "combine", - "multi_modal_data": {"image": [image, image, image, image, image]}, - } - ], + prompt={ + "prompt": "combine", + "multi_modal_data": {"image": [image, image, image, image, image]}, + }, sampling_params=SimpleNamespace(height=None, width=None), ) diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py index f5676a0056f..7f23b4733ca 100644 --- a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py +++ b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py @@ -1,4 +1,3 @@ -import inspect from types import SimpleNamespace import pytest @@ -17,6 +16,7 @@ from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import ( QwenImageLayeredPipeline, ) +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -160,7 +160,7 @@ def _fake_prepare_generation_context(**kwargs): pipeline._prepare_generation_context = _fake_prepare_generation_context state = SimpleNamespace( - prompts=["prompt"], + prompt="prompt", sampling=SimpleNamespace( height=None, width=None, @@ -179,6 +179,105 @@ def _fake_prepare_generation_context(**kwargs): assert captured["max_sequence_length"] == 1024 +def _make_request_batch_prompt_sampling(**overrides): + values = { + "height": 32, + "width": 32, + "num_inference_steps": 2, + "sigmas": None, + "max_sequence_length": None, + "num_outputs_per_prompt": 0, + "generator": None, + "latents": None, + "true_cfg_scale": None, + "guidance_scale_provided": False, + "guidance_scale": 1.0, + "output_type": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_forward_collates_request_prompt_tensors_for_qwen_image(): + pipeline = object.__new__(QwenImagePipeline) + nn.Module.__init__(pipeline) + pipeline.vae_scale_factor = 8 + pipeline.default_sample_size = 128 + + class StopAfterPrepareContextError(Exception): + pass + + captured = {} + + def _fake_prepare_generation_context(**kwargs): + captured.update(kwargs) + raise StopAfterPrepareContextError + + pipeline._prepare_generation_context = _fake_prepare_generation_context + + prompt_embeds_a = torch.zeros(2, 3) + prompt_embeds_b = torch.ones(2, 3) + prompt_embeds_mask_a = torch.tensor([True, True]) + prompt_embeds_mask_b = torch.tensor([True, False]) + negative_prompt_embeds_a = torch.full((2, 3), 2.0) + negative_prompt_embeds_b = torch.full((2, 3), 3.0) + negative_prompt_embeds_mask_a = torch.tensor([False, True]) + negative_prompt_embeds_mask_b = torch.tensor([False, False]) + + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="qwen-prompt-a", + prompt={ + "prompt": "prompt-a", + "negative_prompt": "negative-a", + "prompt_embeds": prompt_embeds_a, + "prompt_embeds_mask": prompt_embeds_mask_a, + "negative_prompt_embeds": negative_prompt_embeds_a, + "negative_prompt_embeds_mask": negative_prompt_embeds_mask_a, + }, + sampling_params=_make_request_batch_prompt_sampling(), + ), + SimpleNamespace( + request_id="qwen-prompt-b", + prompt={ + "prompt": "prompt-b", + "negative_prompt": "negative-b", + "additional_information": { + "prompt_embeds": [prompt_embeds_b], + "prompt_embeds_mask": [prompt_embeds_mask_b], + "negative_prompt_embeds": [negative_prompt_embeds_b], + "negative_prompt_embeds_mask": [negative_prompt_embeds_mask_b], + }, + }, + sampling_params=_make_request_batch_prompt_sampling(), + ), + ] + ) + + with pytest.raises(StopAfterPrepareContextError): + pipeline.forward(batch) + + assert captured["prompt"] is None + assert captured["negative_prompt"] is None + torch.testing.assert_close( + captured["prompt_embeds"], + torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + captured["prompt_embeds_mask"], + torch.stack([prompt_embeds_mask_a, prompt_embeds_mask_b], dim=0), + ) + torch.testing.assert_close( + captured["negative_prompt_embeds"], + torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + captured["negative_prompt_embeds_mask"], + torch.stack([negative_prompt_embeds_mask_a, negative_prompt_embeds_mask_b], dim=0), + ) + + @pytest.mark.parametrize( ("pipeline_class", "drop_idx"), [ @@ -245,16 +344,3 @@ def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_cl with pytest.raises(AssertionError, match="text encoder should not run"): pipeline.encode_prompt(prompt="short prompt") - - -@pytest.mark.parametrize( - "pipeline_class", - [ - QwenImagePipeline, - QwenImageLayeredPipeline, - QwenImageEditPipeline, - QwenImageEditPlusPipeline, - ], -) -def test_forward_max_sequence_length_default_is_1024(pipeline_class: type): - assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024 diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py new file mode 100644 index 00000000000..62d17f70534 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_postprocess.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as qwen_image + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_qwen_image_postprocess_preserves_output_envelope( + monkeypatch: pytest.MonkeyPatch, + tmp_path, +) -> None: + model_path = tmp_path / "model" + vae_path = model_path / "vae" + vae_path.mkdir(parents=True) + (vae_path / "config.json").write_text(json.dumps({}), encoding="utf-8") + + processed = object() + captured: dict[str, object] = {} + + class FakeImageProcessor: + def __init__(self, *, vae_scale_factor: int) -> None: + captured["vae_scale_factor"] = vae_scale_factor + + def postprocess(self, images: torch.Tensor) -> object: + captured["images"] = images + return processed + + monkeypatch.setattr(qwen_image, "VaeImageProcessor", FakeImageProcessor) + postprocess = qwen_image.get_qwen_image_post_process_func(SimpleNamespace(model=str(model_path))) + + image = torch.zeros(1, 3, 4, 4) + output = postprocess( + { + "payload": {"image": image}, + "metadata": {"prompt_embeddings": {"present": True}}, + } + ) + + assert captured["vae_scale_factor"] == 16 + assert captured["images"] is image + assert output == { + "payload": {"image": processed}, + "metadata": {"prompt_embeddings": {"present": True}}, + } diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py new file mode 100644 index 00000000000..2218db8d65b --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_rope_table_failure_case.py @@ -0,0 +1,44 @@ +"""Regression tests for #4443: padded prompt width vs RoPE table length. + +When ``max(txt_seq_lens) < prompt_embeds.shape[1]``, ``QwenEmbedRope`` builds a +text frequency table that is too short for the padded encoder width. +""" + +import pytest +import torch + +from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenEmbedRope +from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Representative request-batch padding: valid tokens 10/20, shared width 32. +PADDED_WIDTH = 32 +VALID_LENS = [10, 20] + + +def test_pre_pr_txt_seq_lens_is_shorter_than_padded_width(): + mask = torch.zeros(len(VALID_LENS), PADDED_WIDTH, dtype=torch.bool) + mask[0, : VALID_LENS[0]] = True + mask[1, : VALID_LENS[1]] = True + + pre_pr_txt_seq_lens = mask.sum(dim=1).tolist() + assert pre_pr_txt_seq_lens == VALID_LENS + assert max(pre_pr_txt_seq_lens) < PADDED_WIDTH + + +def test_post_pr_txt_seq_lens_matches_padded_width(): + prompt_embeds = torch.zeros(len(VALID_LENS), PADDED_WIDTH, 8) + assert txt_seq_lens_from_embeds(prompt_embeds) == [PADDED_WIDTH, PADDED_WIDTH] + + +def test_rope_table_length_matches_padded_width_after_fix(): + rope = QwenEmbedRope(theta=10000, axes_dim=[16, 56, 56], scale_rope=True) + img_shapes = [[(1, 16, 16)]] + + _, freqs_pre_pr = rope(img_shapes, VALID_LENS, device="cpu") + _, freqs_post_pr = rope(img_shapes, [PADDED_WIDTH, PADDED_WIDTH], device="cpu") + + assert freqs_pre_pr.shape[0] == max(VALID_LENS) + assert freqs_pre_pr.shape[0] < PADDED_WIDTH + assert freqs_post_pr.shape[0] == PADDED_WIDTH diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py new file mode 100644 index 00000000000..336ae2ed43c --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_rope_utils.py @@ -0,0 +1,30 @@ +import pytest +import torch + +from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_txt_seq_lens_from_embeds_uses_padded_width_not_valid_token_count(): + prompt_embeds = torch.zeros(2, 64, 8) + prompt_embeds_mask = torch.zeros(2, 64, dtype=torch.bool) + prompt_embeds_mask[:, :10] = True + + assert prompt_embeds_mask.sum(dim=1).tolist() == [10, 10] + assert txt_seq_lens_from_embeds(prompt_embeds) == [64, 64] + + +def test_txt_seq_lens_from_embeds_supports_2d_embeds(): + prompt_embeds = torch.zeros(48, 16) + + assert txt_seq_lens_from_embeds(prompt_embeds) == [48] + + +def test_txt_seq_lens_from_embeds_returns_none_for_missing_embeds(): + assert txt_seq_lens_from_embeds(None) is None + + +def test_txt_seq_lens_from_embeds_rejects_invalid_rank(): + with pytest.raises(ValueError, match="prompt_embeds must be 2D or 3D"): + txt_seq_lens_from_embeds(torch.zeros(2, 3, 4, 5)) diff --git a/tests/diffusion/models/sd3/test_sd3_pipeline.py b/tests/diffusion/models/sd3/test_sd3_pipeline.py new file mode 100644 index 00000000000..4b0bf843fed --- /dev/null +++ b/tests/diffusion/models/sd3/test_sd3_pipeline.py @@ -0,0 +1,155 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.sd3.pipeline_sd3 import StableDiffusion3Pipeline +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_sd3_sampling(**overrides): + values = { + "height": 32, + "width": 32, + "num_inference_steps": 2, + "sigmas": None, + "max_sequence_length": None, + "num_outputs_per_prompt": 0, + "generator": None, + "latents": None, + "guidance_scale": 4.0, + "output_type": None, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _make_sd3_pipeline(): + pipeline = object.__new__(StableDiffusion3Pipeline) + nn.Module.__init__(pipeline) + pipeline.vae_scale_factor = 8 + pipeline.patch_size = 2 + pipeline.default_sample_size = 128 + pipeline.transformer = SimpleNamespace(in_channels=1) + return pipeline + + +def test_forward_collates_request_prompt_tensors_for_sd3(): + pipeline = _make_sd3_pipeline() + + class StopAfterDiffuseError(Exception): + pass + + encode_calls = [] + diffuse_call = {} + + def _fake_encode_prompt(**kwargs): + encode_calls.append(kwargs) + prompt_embeds = kwargs["prompt_embeds"] + if prompt_embeds is None: + prompt_embeds = torch.empty(2, 2, 3) + return prompt_embeds, kwargs.get("pooled_prompt_embeds") + + def _fake_diffuse(**kwargs): + diffuse_call.update(kwargs) + raise StopAfterDiffuseError + + pipeline.encode_prompt = _fake_encode_prompt + pipeline.prepare_latents = lambda *args, **kwargs: torch.zeros(2, 1, 1, 1) + pipeline.prepare_timesteps = lambda *args, **kwargs: (torch.tensor([1.0]), 1) + pipeline.diffuse = _fake_diffuse + + prompt_embeds_a = torch.zeros(2, 3) + prompt_embeds_b = torch.ones(2, 3) + pooled_prompt_embeds_a = torch.full((4,), 2.0) + pooled_prompt_embeds_b = torch.full((4,), 3.0) + negative_prompt_embeds_a = torch.full((2, 3), 4.0) + negative_prompt_embeds_b = torch.full((2, 3), 5.0) + negative_pooled_prompt_embeds_a = torch.full((4,), 6.0) + negative_pooled_prompt_embeds_b = torch.full((4,), 7.0) + + batch = DiffusionRequestBatch( + requests=[ + SimpleNamespace( + request_id="sd3-prompt-a", + prompt={ + "prompt": "prompt-a", + "negative_prompt": "negative-a", + "prompt_embeds": prompt_embeds_a, + "pooled_prompt_embeds": pooled_prompt_embeds_a, + "negative_prompt_embeds": negative_prompt_embeds_a, + "negative_pooled_prompt_embeds": negative_pooled_prompt_embeds_a, + }, + sampling_params=_make_sd3_sampling(), + ), + SimpleNamespace( + request_id="sd3-prompt-b", + prompt={ + "prompt": "prompt-b", + "negative_prompt": "negative-b", + "additional_information": { + "prompt_embeds": [prompt_embeds_b], + "pooled_prompt_embeds": [pooled_prompt_embeds_b], + "negative_prompt_embeds": [negative_prompt_embeds_b], + "negative_pooled_prompt_embeds": [negative_pooled_prompt_embeds_b], + }, + }, + sampling_params=_make_sd3_sampling(), + ), + ] + ) + + with pytest.raises(StopAfterDiffuseError): + pipeline.forward(batch) + + assert encode_calls[0]["prompt"] is None + assert encode_calls[0]["prompt_2"] is None + assert encode_calls[0]["prompt_3"] is None + torch.testing.assert_close( + encode_calls[0]["prompt_embeds"], + torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + encode_calls[0]["pooled_prompt_embeds"], + torch.stack([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=0), + ) + + assert encode_calls[1]["prompt"] is None + assert encode_calls[1]["prompt_2"] is None + assert encode_calls[1]["prompt_3"] is None + torch.testing.assert_close( + encode_calls[1]["prompt_embeds"], + torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + encode_calls[1]["pooled_prompt_embeds"], + torch.stack([negative_pooled_prompt_embeds_a, negative_pooled_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + diffuse_call["pooled_prompt_embeds"], + torch.stack([pooled_prompt_embeds_a, pooled_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + diffuse_call["negative_pooled_prompt_embeds"], + torch.stack([negative_pooled_prompt_embeds_a, negative_pooled_prompt_embeds_b], dim=0), + ) + + +def test_encode_prompt_preserves_direct_pooled_prompt_embeds(): + pipeline = _make_sd3_pipeline() + prompt_embeds = torch.zeros(1, 2, 3) + pooled_prompt_embeds = torch.ones(1, 4) + + actual_prompt_embeds, actual_pooled_prompt_embeds = pipeline.encode_prompt( + prompt=None, + prompt_2=None, + prompt_3=None, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + ) + + assert actual_prompt_embeds is prompt_embeds + assert actual_pooled_prompt_embeds is pooled_prompt_embeds diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py b/tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py new file mode 100644 index 00000000000..d966e5238c2 --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_triton_kernel.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Correctness tests for Sensenova-U1 Triton kernels. + +Verifies that fused-RMSNorm+3D Rope Triton kernels +produce results numerically equivalent to the PyTorch reference. +All computations follow the kernel's internal float32 promotion rules. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + +SEED = 42 +B = 1 +H_Q = 32 +H_K = 8 +HEAD_DIM = 128 +T_DIM = 64 +HW_DIM = 64 +HW_ROPE_DIM = 32 +QKV_DIM = (H_Q + H_K + H_K) * HEAD_DIM +EPS = 1e-6 +DEVICE = torch.device("cuda:0") + + +@dataclass(frozen=True) +class KernelInput: + seq_len: int + q: torch.Tensor + k: torch.Tensor + q_norm_weight: torch.Tensor + k_norm_weight: torch.Tensor + q_norm_hw_weight: torch.Tensor + k_norm_hw_weight: torch.Tensor + cos_t: torch.Tensor + sin_t: torch.Tensor + cos_h: torch.Tensor + sin_h: torch.Tensor + cos_w: torch.Tensor + sin_w: torch.Tensor + + def args(self) -> tuple[torch.Tensor, ...]: + return ( + self.q, + self.k, + self.q_norm_weight, + self.k_norm_weight, + self.q_norm_hw_weight, + self.k_norm_hw_weight, + self.cos_t, + self.sin_t, + self.cos_h, + self.sin_h, + self.cos_w, + self.sin_w, + ) + + +try: + from vllm_omni.diffusion.models.sensenova_u1.fused_rmsnorm_rope import triton_qk_norm_rope + + _TRITON_AVAILABLE = True +except ImportError: + _TRITON_AVAILABLE = False + +triton_available = pytest.mark.skipif(not _TRITON_AVAILABLE, reason="Triton not available on this platform") + + +# --------------------------------------------------------------------------- +# PyTorch reference implementations (mirror the kernel's internal arithmetic) +# --------------------------------------------------------------------------- + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rope( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) + + +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = EPS) -> torch.Tensor: + input_dtype = x.dtype + xf = x.float() + var = xf.pow(2).mean(dim=-1, keepdim=True) + xf = xf * torch.rsqrt(var + eps) + return xf.to(input_dtype) * weight + + +def assert_close_with_error_stats( + actual: torch.Tensor, + expected: torch.Tensor, + *, + name: str, + atol: float, + rtol: float, +) -> None: + try: + torch.testing.assert_close(actual, expected, atol=atol, rtol=rtol) + except AssertionError as exc: + abs_diff = (actual.float() - expected.float()).abs().flatten() + tolerance = atol + rtol * expected.float().abs().flatten() + mismatch = abs_diff > tolerance + mismatch_count = int(mismatch.sum().item()) + total = abs_diff.numel() + p99 = torch.quantile(abs_diff, 0.99).item() if total else 0.0 + stats = ( + f"{name} error stats: " + f"max={abs_diff.max().item() if total else 0.0:.6g}, " + f"p99={p99:.6g}, " + f"mean={abs_diff.mean().item() if total else 0.0:.6g}, " + f"mismatch={mismatch_count}/{total} ({mismatch_count / total:.4%}), " + f"atol={atol}, rtol={rtol}" + ) + raise AssertionError(f"{stats}\n{exc}") from exc + + +def reference_kernel( + q: torch.Tensor, + k: torch.Tensor, + q_norm_weight: torch.Tensor, + k_norm_weight: torch.Tensor, + q_norm_hw_weight: torch.Tensor, + k_norm_hw_weight: torch.Tensor, + cos_t: torch.Tensor, + sin_t: torch.Tensor, + cos_h: torch.Tensor, + sin_h: torch.Tensor, + cos_w: torch.Tensor, + sin_w: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference matching the model's current math.""" + q_t, q_hw = q.chunk(2, dim=-1) + k_t, k_hw = k.chunk(2, dim=-1) + + q_t = rmsnorm(q_t, q_norm_weight).transpose(1, 2) + k_t = rmsnorm(k_t, k_norm_weight).transpose(1, 2) + q_hw = rmsnorm(q_hw, q_norm_hw_weight).transpose(1, 2) + k_hw = rmsnorm(k_hw, k_norm_hw_weight).transpose(1, 2) + + q_h, q_w = q_hw.chunk(2, dim=-1) + k_h, k_w = k_hw.chunk(2, dim=-1) + + q_t, k_t = apply_rope(q_t, k_t, cos_t, sin_t) + q_h, k_h = apply_rope(q_h, k_h, cos_h, sin_h) + q_w, k_w = apply_rope(q_w, k_w, cos_w, sin_w) + + query = torch.cat([q_t, q_h, q_w], dim=-1) + key = torch.cat([k_t, k_h, k_w], dim=-1) + return query, key + + +def make_input(seq_len: int, dtype: torch.dtype, device: torch.device, seed: int) -> KernelInput: + gen = torch.Generator(device=device) + gen.manual_seed(seed + seq_len) + + # Match real model layout: q/k/v are views split from fused qkv. + qkv = torch.randn(B, seq_len, QKV_DIM, device=device, dtype=dtype, generator=gen) + q, k, _v = qkv.split([H_Q * HEAD_DIM, H_K * HEAD_DIM, H_K * HEAD_DIM], dim=-1) + q = q.view(B, seq_len, H_Q, HEAD_DIM) + k = k.view(B, seq_len, H_K, HEAD_DIM) + + q_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen) + k_norm_weight = torch.randn(T_DIM, device=device, dtype=dtype, generator=gen) + q_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen) + k_norm_hw_weight = torch.randn(HW_DIM, device=device, dtype=dtype, generator=gen) + + # Match Qwen3RotaryEmbedding: emb = torch.cat((freqs, freqs), dim=-1). + def make_rope_pair(dim: int) -> tuple[torch.Tensor, torch.Tensor]: + cos_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen) + sin_half = torch.randn(B, seq_len, dim // 2, device=device, dtype=dtype, generator=gen) + return torch.cat((cos_half, cos_half), dim=-1), torch.cat((sin_half, sin_half), dim=-1) + + cos_t, sin_t = make_rope_pair(T_DIM) + cos_h, sin_h = make_rope_pair(HW_ROPE_DIM) + cos_w, sin_w = make_rope_pair(HW_ROPE_DIM) + + return KernelInput( + seq_len=seq_len, + q=q, + k=k, + q_norm_weight=q_norm_weight, + k_norm_weight=k_norm_weight, + q_norm_hw_weight=q_norm_hw_weight, + k_norm_hw_weight=k_norm_hw_weight, + cos_t=cos_t, + sin_t=sin_t, + cos_h=cos_h, + sin_h=sin_h, + cos_w=cos_w, + sin_w=sin_w, + ) + + +# --------------------------------------------------------------------------- +# Fused RMSNorm + 3D Rope +# --------------------------------------------------------------------------- + + +@triton_available +@pytest.mark.parametrize("seq_len", [1, 2, 4, 8, 9, 16, 32, 64, 128, 256, 260, 271, 512, 1024, 2048, 4096]) +@pytest.mark.parametrize( + ("dtype", "atol", "rtol"), + [ + (torch.float32, 1e-5, 1e-5), + (torch.bfloat16, 3e-2, 3e-2), + ], +) +def test_fused_qk_norm_rope_matches_reference(seq_len: int, dtype: torch.dtype, atol: float, rtol: float): + data = make_input(seq_len, dtype=dtype, device=DEVICE, seed=SEED) + call_args = data.args() + ref_q, ref_k = reference_kernel(*call_args) + out_q, out_k = triton_qk_norm_rope(*call_args, EPS) + + assert_close_with_error_stats(out_q, ref_q, name="query", atol=atol, rtol=rtol) + assert_close_with_error_stats(out_k, ref_k, name="key", atol=atol, rtol=rtol) diff --git a/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py new file mode 100644 index 00000000000..818d3d7c0c1 --- /dev/null +++ b/tests/diffusion/models/sensenova_u1/test_sensenova_u1_cfg.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import SenseNovaU1Pipeline + + +def _pipeline_without_init() -> SenseNovaU1Pipeline: + return object.__new__(SenseNovaU1Pipeline) + + +@pytest.mark.parametrize("kwargs", [None, {"step_i": 0}]) +def test_combine_cfg_noise_requires_is_it2i(kwargs): + pipe = _pipeline_without_init() + out_cond = (torch.ones(1, 2, 3),) + out_uncond = (torch.zeros(1, 2, 3),) + + with pytest.raises(ValueError, match="is_it2i"): + pipe.combine_cfg_noise( + out_cond, + out_uncond, + cfg_scale=4.0, + cfg_norm="cfg_zero_star", + kwargs=kwargs, + ) + + +@pytest.mark.parametrize( + "kwargs", + [{"is_it2i": False}, {"is_it2i": False, "step_i": None}], +) +def test_cfg_zero_star_requires_step_i(kwargs): + pipe = _pipeline_without_init() + out_cond = (torch.ones(1, 2, 3),) + out_uncond = (torch.zeros(1, 2, 3),) + with pytest.raises(ValueError, match="step_i"): + pipe.combine_cfg_noise( + out_cond, + out_uncond, + cfg_scale=4.0, + cfg_norm="cfg_zero_star", + kwargs=kwargs, + ) + + +def test_cfg_zero_star_accepts_step_i(): + pipe = _pipeline_without_init() + out_cond = (torch.ones(1, 2, 3),) + out_uncond = (torch.zeros(1, 2, 3),) + result = pipe.combine_cfg_noise( + out_cond, + out_uncond, + cfg_scale=4.0, + cfg_norm="cfg_zero_star", + kwargs={"is_it2i": False, "step_i": 0}, + ) + + assert result.shape == out_cond[0].shape + assert torch.equal(result, torch.zeros_like(out_cond[0])) + assert torch.isfinite(result).all() diff --git a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py index 576678e2cf0..9d5b11d3ba8 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_i2v_pipeline.py @@ -34,19 +34,19 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline: def test_i2v_preprocess_requires_image_and_resizes_to_480p_aspect() -> None: preprocess = get_wan22_i2v_pre_process_func(SimpleNamespace()) request = SimpleNamespace( - prompts=[{"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "red")}}], + prompt={"prompt": "p", "multi_modal_data": {"image": Image.new("RGB", (320, 160), "red")}}, sampling_params=SimpleNamespace(height=None, width=None), ) result = preprocess(request) - prompt = result.prompts[0] + prompt = result.prompt assert result.sampling_params.height == 432 assert result.sampling_params.width == 880 assert prompt["multi_modal_data"]["image"].size == (880, 432) missing_image = SimpleNamespace( - prompts=[{"prompt": "p", "multi_modal_data": {}}], + prompt={"prompt": "p", "multi_modal_data": {}}, sampling_params=SimpleNamespace(height=None, width=None), ) with pytest.raises(ValueError, match="No image is provided"): @@ -100,7 +100,10 @@ def fake_predict_noise_maybe_with_cfg(**kwargs): timestep_dtype = calls[0]["timestep_values"].dtype torch.testing.assert_close(calls[0]["timestep_values"][0, :4], torch.zeros(4, dtype=timestep_dtype)) torch.testing.assert_close(calls[0]["timestep_values"][0, 4:], torch.full((4,), 900, dtype=timestep_dtype)) - torch.testing.assert_close(calls[0]["hidden_states"][:, :, 0], torch.ones(1, 4, 4, 4)) + torch.testing.assert_close( + calls[0]["hidden_states"][:, :, 0], + torch.ones_like(calls[0]["hidden_states"][:, :, 0]), + ) torch.testing.assert_close(result, torch.full_like(latents, 2.0)) diff --git a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py index 54bb672ef81..198c3190859 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py +++ b/tests/diffusion/models/wan2_2/test_wan22_pipeline_diffuse.py @@ -9,6 +9,7 @@ from torch import nn from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -19,6 +20,12 @@ def dtype(self) -> torch.dtype: return torch.float32 +class _StubTextEncoder(nn.Module): + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + class _StubScheduler: def __init__(self, timesteps: list[int]) -> None: self.timesteps = torch.tensor(timesteps, dtype=torch.int64) @@ -40,12 +47,30 @@ def update(self) -> None: yield _Bar() +def _stub_encode_prompt( + prompt, + negative_prompt=None, + do_classifier_free_guidance=True, + num_videos_per_prompt=1, + max_sequence_length=512, + device=None, + dtype=None, +): + del negative_prompt, do_classifier_free_guidance, device, dtype + batch_size = 1 if isinstance(prompt, str) else len(prompt) + n = batch_size * num_videos_per_prompt + hidden_size = 8 + prompt_embeds = torch.zeros(n, max_sequence_length, hidden_size) + return prompt_embeds, None + + def _make_pipeline() -> Wan22Pipeline: pipeline = object.__new__(Wan22Pipeline) nn.Module.__init__(pipeline) pipeline.device = torch.device("cpu") pipeline.transformer = _StubTransformer() pipeline.transformer_2 = None + pipeline.text_encoder = _StubTextEncoder() pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2), in_channels=4, out_channels=4) pipeline.scheduler = _StubScheduler([9, 5]) pipeline.od_config = SimpleNamespace(flow_shift=5.0) @@ -60,15 +85,14 @@ def _make_pipeline() -> Wan22Pipeline: pipeline._num_timesteps = None pipeline._current_timestep = None pipeline.check_inputs = lambda **kwargs: None + pipeline.encode_prompt = _stub_encode_prompt # type: ignore[method-assign] pipeline.prepare_latents = lambda **kwargs: torch.zeros((1, 4, 1, 8, 8), dtype=torch.float32) pipeline.progress_bar = _noop_progress_bar return pipeline -def test_forward_delegates_denoising_to_diffuse(monkeypatch) -> None: +def test_forward_delegates_denoising_to_diffuse() -> None: pipeline = _make_pipeline() - - prompt_embeds = torch.randn(1, 8) captured: dict[str, object] = {} def _fake_diffuse(**kwargs): @@ -77,15 +101,16 @@ def _fake_diffuse(**kwargs): pipeline.diffuse = _fake_diffuse # type: ignore[method-assign] - req = SimpleNamespace( - prompts=["prompt"], + mock_req = SimpleNamespace( + prompt="prompt", + request_id="test-req", sampling_params=SimpleNamespace( height=None, width=None, num_frames=1, num_inference_steps=2, - guidance_scale_provided=False, - guidance_scale=None, + guidance_scale_provided=True, + guidance_scale=1.0, guidance_scale_2=None, boundary_ratio=None, generator=None, @@ -93,13 +118,16 @@ def _fake_diffuse(**kwargs): num_outputs_per_prompt=1, max_sequence_length=32, latents=None, + output_type="latent", extra_args={}, ), ) + batch = DiffusionRequestBatch(requests=[mock_req]) - output = pipeline.forward(req, prompt_embeds=prompt_embeds, output_type="latent", guidance_scale=1.0) + output = pipeline.forward(batch) assert torch.equal(output.output, torch.ones((1, 4, 1, 8, 8))) + assert torch.equal(captured["prompt_embeds"], torch.zeros(1, 32, 8)) assert torch.equal(captured["timesteps"], pipeline.scheduler.timesteps) assert captured["guidance_low"] == 1.0 assert captured["guidance_high"] == 1.0 diff --git a/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py new file mode 100644 index 00000000000..4440a8f359d --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_s2v_pipeline.py @@ -0,0 +1,335 @@ +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_s2v_exposes_hsdp_shard_conditions_for_transformer_blocks(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.blocks = nn.ModuleList([nn.Linear(4, 4) for _ in range(3)]) + + conditions = getattr(model, "_hsdp_shard_conditions", None) + + assert conditions is not None + assert len(conditions) == 1 + + matched = [] + for name, module in model.named_modules(): + if any(cond(name, module) for cond in conditions): + matched.append(name) + + assert matched == ["blocks.0", "blocks.1", "blocks.2"] + + +def test_s2v_hsdp_shard_condition_does_not_match_non_block_modules(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.blocks = nn.ModuleList([nn.Linear(4, 4)]) + model.head_indicator = nn.Linear(4, 4) + model.casual_audio_encoder = nn.Linear(4, 4) + + conditions = model._hsdp_shard_conditions + non_block_matched = [] + for name, module in model.named_modules(): + if name and "blocks" not in name: + if any(cond(name, module) for cond in conditions): + non_block_matched.append(name) + + assert non_block_matched == [] + + +def test_encode_audio_calls_unshard_reshard_when_fsdp_managed(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.enable_adain = False + model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64)) + + model.unshard = MagicMock() + model.reshard = MagicMock() + + audio_input = torch.randn(1, 1, 64, 5) + motion_frames = [2, 2] + + result = model.encode_audio(audio_input, motion_frames) + + model.unshard.assert_called_once() + model.reshard.assert_called_once() + assert "audio_emb" in result + + +def test_encode_audio_reshard_called_on_exception(): + """Test that reshard() is always called even when encode_audio logic raises.""" + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.enable_adain = False + model.casual_audio_encoder = MagicMock(side_effect=RuntimeError("encoder failed")) + + model.unshard = MagicMock() + model.reshard = MagicMock() + + audio_input = torch.randn(1, 1, 64, 5) + motion_frames = [2, 2] + + with pytest.raises(RuntimeError, match="encoder failed"): + model.encode_audio(audio_input, motion_frames) + + model.unshard.assert_called_once() + model.reshard.assert_called_once() + + +def test_encode_audio_skips_unshard_reshard_when_not_fsdp(): + model = object.__new__(WanS2VTransformer3DModel) + nn.Module.__init__(model) + model.enable_adain = False + model.casual_audio_encoder = MagicMock(return_value=torch.zeros(1, 10, 64)) + + audio_input = torch.randn(1, 1, 64, 5) + motion_frames = [2, 2] + + result = model.encode_audio(audio_input, motion_frames) + + assert not hasattr(model, "unshard") + assert not hasattr(model, "reshard") + assert "audio_emb" in result + + +def test_s2v_pipeline_skips_cpu_offload_when_hsdp_enabled(): + """Test that transformer.to('cpu') is NOT called when HSDP is active.""" + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline + + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + + od_config = MagicMock() + od_config.enable_cpu_offload = True + parallel_config = MagicMock() + parallel_config.use_hsdp = True + od_config.parallel_config = parallel_config + pipeline.od_config = od_config + + mock_transformer = MagicMock() + pipeline.transformer = mock_transformer + + # Simulate the offload decision from the forward loop + if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False): + pipeline.transformer.to("cpu") + + mock_transformer.to.assert_not_called() + + +def test_s2v_pipeline_allows_cpu_offload_when_hsdp_disabled(): + """Test that transformer.to('cpu') IS called when HSDP is not active.""" + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline + + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + + od_config = MagicMock() + od_config.enable_cpu_offload = True + parallel_config = MagicMock() + parallel_config.use_hsdp = False + od_config.parallel_config = parallel_config + pipeline.od_config = od_config + + mock_transformer = MagicMock() + pipeline.transformer = mock_transformer + + # Simulate the offload decision from the forward loop + if pipeline.od_config.enable_cpu_offload and not getattr(pipeline.od_config.parallel_config, "use_hsdp", False): + pipeline.transformer.to("cpu") + + mock_transformer.to.assert_called_once_with("cpu") + + +def test_s2v_pipeline_hsdp_forward_complete_process(): + """Integration-level mock test verifying the complete forward path of the + S2V pipeline under HSDP mode. + + Verifies: + - Text encoding is invoked + - Audio encoding invokes unshard/reshard on the transformer + - Reference image encoding via VAE + - Denoising loop calls the transformer + - VAE decode is invoked + - transformer.to('cpu') is NOT called (HSDP mode) + - Final output is a DiffusionOutput with video + audio + """ + import numpy as np + import PIL.Image + + from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import Wan22S2VPipeline + + pipeline = object.__new__(Wan22S2VPipeline) + nn.Module.__init__(pipeline) + + # -- Config -- + od_config = MagicMock() + od_config.enable_cpu_offload = True + od_config.enable_diffusion_pipeline_profiler = False + parallel_config = MagicMock() + parallel_config.use_hsdp = True + od_config.parallel_config = parallel_config + pipeline.od_config = od_config + + # -- Pipeline attributes -- + pipeline.device = torch.device("cpu") + pipeline._guidance_scale = 4.5 + pipeline._num_timesteps = None + pipeline._current_timestep = None + pipeline.vae_scale_factor_spatial = 8 + pipeline.vae_scale_factor_temporal = 4 + pipeline.resolution_divisor = 16 + pipeline.motion_frames = 73 + pipeline.drop_first_motion = True + pipeline.fps = 16 + pipeline.audio_sample_m = 0 + pipeline._DEFAULT_INFER_FRAMES = 80 + + # -- Mock text encoder -- + pipeline.tokenizer = MagicMock() + pipeline.tokenizer.return_value = MagicMock( + input_ids=torch.zeros(1, 512, dtype=torch.long), + attention_mask=torch.ones(1, 512, dtype=torch.long), + ) + mock_text_encoder = MagicMock() + mock_text_encoder.dtype = torch.bfloat16 + mock_text_encoder.return_value = MagicMock(last_hidden_state=torch.zeros(1, 512, 4096)) + pipeline.text_encoder = mock_text_encoder + + # -- Mock transformer -- + mock_transformer = MagicMock() + mock_transformer.dtype = torch.bfloat16 + mock_transformer.parameters = MagicMock(return_value=iter([torch.zeros(1)])) + + # encode_audio returns dict with audio_emb + mock_transformer.encode_audio = MagicMock(return_value={"audio_emb": torch.zeros(1, 10, 64)}) + # Forward returns noise prediction + mock_transformer.return_value = (torch.zeros(1, 16, 20, 88, 128),) + pipeline.transformer = mock_transformer + + # -- Mock VAE -- + mock_vae = MagicMock() + mock_vae.dtype = torch.bfloat16 + mock_vae.config = MagicMock() + mock_vae.config.scale_factor_temporal = 4 + mock_vae.config.scale_factor_spatial = 8 + mock_vae.config.latents_mean = [0.0] * 16 + mock_vae.config.latents_std = [1.0] * 16 + mock_vae.config.z_dim = 16 + # encode returns mock with latent_dist + mock_encode_result = MagicMock() + mock_encode_result.latent_dist = MagicMock() + mock_encode_result.latent_dist.mode = MagicMock(return_value=torch.zeros(1, 16, 1, 88, 128)) + mock_vae.encode = MagicMock(return_value=mock_encode_result) + # decode returns video + mock_vae.decode = MagicMock(return_value=(torch.zeros(1, 3, 80, 704, 1024),)) + pipeline.vae = mock_vae + + # -- Mock audio model -- + mock_audio_model = MagicMock() + mock_audio_model.device = torch.device("cpu") + mock_audio_param = torch.zeros(1) + mock_audio_model.parameters = MagicMock(return_value=iter([mock_audio_param])) + mock_audio_model.return_value = MagicMock(hidden_states=[torch.zeros(1, 100, 1024)] * 25) + pipeline.audio_model = mock_audio_model + + pipeline.audio_processor = MagicMock() + pipeline.audio_processor.return_value = MagicMock(input_values=torch.zeros(1, 16000)) + + # -- Mock scheduler -- + mock_scheduler = MagicMock() + mock_scheduler.timesteps = torch.linspace(999, 0, 5) + mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),)) + pipeline.scheduler = mock_scheduler + + # -- Bind methods from the real class -- + pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline) + pipeline.encode_ref_image = Wan22S2VPipeline.encode_ref_image.__get__(pipeline) + pipeline.prepare_motion_latents = Wan22S2VPipeline.prepare_motion_latents.__get__(pipeline) + pipeline.prepare_latents = Wan22S2VPipeline.prepare_latents.__get__(pipeline) + pipeline.check_inputs = Wan22S2VPipeline.check_inputs.__get__(pipeline) + pipeline.diffuse = Wan22S2VPipeline.diffuse.__get__(pipeline) + pipeline._normalize_latents = Wan22S2VPipeline._normalize_latents.__get__(pipeline) + pipeline._denormalize_latents = Wan22S2VPipeline._denormalize_latents.__get__(pipeline) + pipeline._prompt_clean = Wan22S2VPipeline._prompt_clean + + # -- Build request -- + from vllm_omni.diffusion.request import OmniDiffusionRequest + + sampling_params = MagicMock() + sampling_params.height = 704 + sampling_params.width = 1024 + sampling_params.num_frames = 80 + sampling_params.num_inference_steps = 5 + sampling_params.guidance_scale = 4.5 + sampling_params.guidance_scale_provided = True + sampling_params.generator = None + sampling_params.seed = 42 + + ref_image = PIL.Image.new("RGB", (1024, 704)) + audio_data = np.zeros(16000, dtype=np.float32) + + req = MagicMock(spec=OmniDiffusionRequest) + req.prompts = [ + { + "prompt": "test prompt", + "negative_prompt": "bad quality", + "multi_modal_data": {"image": ref_image, "audio": audio_data}, + "additional_information": {"audio_path": audio_data, "pose_video": None, "init_first_frame": False}, + } + ] + req.sampling_params = sampling_params + + # -- Mock methods that use complex internal state -- + pipeline.encode_audio = MagicMock(return_value=(torch.zeros(1, 25, 64, 80), 1, 80)) + + # Mock progress bar context + pipeline.progress_bar = MagicMock() + pipeline.progress_bar.return_value.__enter__ = MagicMock(return_value=MagicMock()) + pipeline.progress_bar.return_value.__exit__ = MagicMock(return_value=False) + + # Mock predict_noise_maybe_with_cfg from CFGParallelMixin + pipeline.predict_noise_maybe_with_cfg = MagicMock(return_value=torch.zeros(16, 20, 88, 128)) + pipeline.predict_noise = Wan22S2VPipeline.predict_noise.__get__(pipeline) + + # Mock platform methods + with patch("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.current_omni_platform") as mock_platform: + mock_platform.empty_cache = MagicMock() + mock_platform.is_available = MagicMock(return_value=False) + + with patch( + "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v.load_audio", return_value=(audio_data, 16000) + ): + result = Wan22S2VPipeline.forward(pipeline, req=req) + + # -- Assertions -- + # Text encoder was called + mock_text_encoder.assert_called() + + # Audio encoding was invoked + pipeline.encode_audio.assert_called_once() + + # VAE encode was called (for ref image and motion latents) + mock_vae.encode.assert_called() + + # VAE decode was called + mock_vae.decode.assert_called() + + # transformer.to("cpu") must NOT be called under HSDP + mock_transformer.to.assert_not_called() + + # Output is a DiffusionOutput tuple with (video, audio_waveform, sample_rate) + from vllm_omni.diffusion.data import DiffusionOutput + + assert isinstance(result, DiffusionOutput) + video, audio_waveform, audio_sr = result.output + assert video.shape[0] == 1 # batch + assert video.shape[1] == 3 # channels + assert audio_waveform is not None + assert audio_sr == 16000 diff --git a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py index 9fa9b67c499..714b79d5066 100644 --- a/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py +++ b/tests/diffusion/models/wan2_2/test_wan22_vace_pipeline.py @@ -14,6 +14,7 @@ create_vace_transformer_from_config, get_wan22_vace_pre_process_func, ) +from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -37,21 +38,19 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None: frame = Image.new("RGB", (64, 64), "black") mask = Image.new("L", (64, 64), 255) request = SimpleNamespace( - prompts=[ - { - "prompt": "p", - "multi_modal_data": { - "image": ref, - "video": [frame], - "mask": mask, - }, - } - ], + prompt={ + "prompt": "p", + "multi_modal_data": { + "image": ref, + "video": [frame], + "mask": mask, + }, + }, sampling_params=SimpleNamespace(height=None, width=None), ) result = preprocess(request) - additional_info = result.prompts[0]["additional_information"] + additional_info = result.prompt["additional_information"] assert result.sampling_params.height == 432 assert result.sampling_params.width == 880 @@ -93,6 +92,10 @@ def __init__(self, **kwargs) -> None: } +def test_vace_transformer_layerwise_offloads_conditioning_blocks() -> None: + assert WanVACETransformer3DModel._layerwise_offload_blocks_attrs == ["vace_blocks", "blocks"] + + def test_vace_prepare_masks_encodes_spatial_stride_and_reference_padding() -> None: pipeline = _make_vace_pipeline() mask = torch.ones(1, 3, 5, 16, 16) diff --git a/tests/diffusion/offloader/test_diffusion_cpu_offload.py b/tests/diffusion/offloader/test_diffusion_cpu_offload.py index 92e2abc7e6b..555a62ca0e3 100644 --- a/tests/diffusion/offloader/test_diffusion_cpu_offload.py +++ b/tests/diffusion/offloader/test_diffusion_cpu_offload.py @@ -18,10 +18,19 @@ IMAGE_VIDEO_MODELS = { "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100}, + "Tongyi-MAI/Z-Image-Turbo": {"cuda": 2500, "rocm": 2100}, + "OmniGen2/OmniGen2": {"cuda": 2500, "rocm": 2100}, } MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS} +MODEL_MARKS = { + "riverclouds/qwen_image_random": pytest.mark.core_model, + "Tongyi-MAI/Z-Image-Turbo": pytest.mark.advanced_model, + "stabilityai/stable-audio-open-1.0": pytest.mark.full_model, + "OmniGen2/OmniGen2": pytest.mark.full_model, +} + _GATED_MODELS = {"stabilityai/stable-audio-open-1.0"} # Aliased for backward compatibility (imported by test_diffusion_layerwise_offload.py). @@ -95,11 +104,15 @@ def check_audio_determinism(audio1, audio2, atol=1e-2): return True -@pytest.mark.core_model @pytest.mark.diffusion @hardware_test(res={"cuda": "L4", "rocm": "MI325"}) -@pytest.mark.parametrize("model_name", list(MODELS.keys())) +@pytest.mark.parametrize( + "model_name", + [pytest.param(name, marks=MODEL_MARKS[name]) for name in MODELS], +) def test_cpu_offload_diffusion_model(model_name: str): + if model_name == "OmniGen2/OmniGen2": + pytest.skip("issue #4537") if model_name in _GATED_MODELS: _skip_if_gated_repo_inaccessible(model_name) try: diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py index bc7fd97d1eb..e97c1953ab7 100644 --- a/tests/diffusion/offloader/test_module_collector.py +++ b/tests/diffusion/offloader/test_module_collector.py @@ -149,6 +149,40 @@ def __init__(self): self.audio_vae = nn.Linear(10, 10) +class SiblingDiTsPipeline(nn.Module, SupportsComponentDiscovery): + """Two independent DiTs (e.g. a MoE transformer pair).""" + + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.transformer_2 = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class NestedDiTPipeline(nn.Module, SupportsComponentDiscovery): + """Cosmos3-like: a DiT plus one of its own submodules listed as a DiT. + + The nested entry exists so offloading can treat the inner stack as an + independent ring; the outer DiT's _hsdp_shard_conditions still own all + nested blocks for sharding. Order is nested-first, matching Cosmos3. + """ + + _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Module() + self.transformer.language_model = nn.Linear(10, 10) + self.transformer.gen_layers = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -238,3 +272,35 @@ def test_multiple_vaes(self): assert len(result.vaes) == 2 assert result.vaes[0] is pipeline.vae assert result.vaes[1] is pipeline.audio_vae + + +class TestOutermostDits: + """PipelineModules.outermost_dits() drops DiTs nested inside another DiT. + + HSDP shards each returned DiT via its own _hsdp_shard_conditions, so a DiT + nested inside another discovered DiT must be excluded: it is already covered + by its ancestor's conditions and may not declare conditions of its own. + """ + + def test_keeps_independent_sibling_dits(self): + result = ModuleDiscovery.discover(SiblingDiTsPipeline()) + + names, modules = result.outermost_dits() + + assert names == ["transformer", "transformer_2"] + assert [id(m) for m in modules] == [id(m) for m in result.dits] + + def test_drops_nested_dit_keeping_only_ancestor(self): + pipeline = NestedDiTPipeline() + result = ModuleDiscovery.discover(pipeline) + + # Discovery sees both, nested-first (matches Cosmos3's _dit_modules order). + assert result.dit_names == ["transformer.language_model", "transformer"] + + names, modules = result.outermost_dits() + + # Only the ancestor survives; sharding the nested condition-less stack + # again would double-wrap its blocks and raise in apply_hsdp_to_model. + assert names == ["transformer"] + assert len(modules) == 1 + assert modules[0] is pipeline.transformer diff --git a/tests/diffusion/offloader/test_sequential_backend.py b/tests/diffusion/offloader/test_sequential_backend.py index 2539cc06895..dc74248e54d 100644 --- a/tests/diffusion/offloader/test_sequential_backend.py +++ b/tests/diffusion/offloader/test_sequential_backend.py @@ -7,7 +7,8 @@ import torch from torch import nn -from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook +from vllm_omni.diffusion.offloader.base import OffloadConfig, OffloadStrategy +from vllm_omni.diffusion.offloader.sequential_backend import ModelLevelOffloadBackend, SequentialOffloadHook from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] @@ -41,6 +42,40 @@ def mock(self): return tracker, mock +def test_model_level_backend_delegates_to_custom_pipeline_offload() -> None: + class CustomPipeline(nn.Module): + def __init__(self) -> None: + super().__init__() + self.enable_args = None + self.disable_called = False + + def enable_omni_model_cpu_offload(self, **kwargs) -> None: + self.enable_args = kwargs + + def disable_omni_model_cpu_offload(self) -> None: + self.disable_called = True + + pipeline = CustomPipeline() + backend = ModelLevelOffloadBackend( + OffloadConfig(strategy=OffloadStrategy.MODEL_LEVEL, pin_cpu_memory=False), + torch.device("cpu"), + ) + + backend.enable(pipeline) + + assert backend.enabled is True + assert pipeline.enable_args == { + "device": torch.device("cpu"), + "pin_memory": False, + "use_hsdp": False, + } + + backend.disable() + + assert backend.enabled is False + assert pipeline.disable_called is True + + class TestMoveParamsPinMemory: def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch): """DTensor should skip pin_memory to avoid RuntimeError.""" diff --git a/tests/diffusion/quantization/test_bitsandbytes_config.py b/tests/diffusion/quantization/test_bitsandbytes_config.py new file mode 100644 index 00000000000..8da2cd521ee --- /dev/null +++ b/tests/diffusion/quantization/test_bitsandbytes_config.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for BitsAndBytes quantization config.""" + +import importlib.util +import sys +import types + +import pytest +import torch +from pytest_mock import MockerFixture +from torch.nn import Module, Parameter +from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization import build_quant_config +from vllm_omni.quantization.factory import SUPPORTED_QUANTIZATION_METHODS + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] + +cuda_available = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="GPU platform not available.") + +bitsandbytes_available = pytest.mark.skipif( + importlib.util.find_spec("bitsandbytes") is None, + reason="bitsandbytes package not installed", +) + + +def _ensure_bitsandbytes_importable(monkeypatch: pytest.MonkeyPatch) -> None: + """Inject a stub bitsandbytes package when the optional dep is absent. + + ``mocker.patch("bitsandbytes...")`` imports the target module; without a + stub, core-model CI without bitsandbytes raises ModuleNotFoundError. + """ + if importlib.util.find_spec("bitsandbytes") is not None: + return + bnb = types.ModuleType("bitsandbytes") + bnb_functional = types.ModuleType("bitsandbytes.functional") + bnb.functional = bnb_functional + monkeypatch.setitem(sys.modules, "bitsandbytes", bnb) + monkeypatch.setitem(sys.modules, "bitsandbytes.functional", bnb_functional) + + +def test_bitsandbytes_config_creation(): + """Test that BitsAndBytes config can be created.""" + config = build_quant_config("bitsandbytes") + assert config is not None + assert config.get_name() == "bitsandbytes" + + +def test_bitsandbytes_config_with_custom_params(): + """Test BitsAndBytes config with custom parameters.""" + config = build_quant_config( + "bitsandbytes", + quant_type="fp4", + compress_statistics=False, + ignored_layers=["to_out"], + ) + assert config is not None + assert config.quant_type == "fp4" + assert config.compress_statistics is False + assert "to_out" in config.ignored_layers + + +def test_supported_methods(): + """Test that supported methods list includes bitsandbytes.""" + assert "bitsandbytes" in SUPPORTED_QUANTIZATION_METHODS + + +def test_quantization_integration(): + """Test end-to-end quantization flow through OmniDiffusionConfig.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + config = OmniDiffusionConfig(model="test", quantization_config="bitsandbytes") + assert config.quantization_config is not None + assert config.quantization_config.get_name() == "bitsandbytes" + + config2 = OmniDiffusionConfig( + model="test", + quantization_config={ + "method": "bitsandbytes", + "quant_type": "nf4", + "compress_statistics": True, + }, + ) + assert config2.quantization_config is not None + assert config2.quantization_config.get_name() == "bitsandbytes" + assert config2.quantization_config.quant_type == "nf4" + assert config2.quantization_config.compress_statistics is True + + +def test_quantization_dict_not_mutated(): + """Test that passing a dict to quantization_config doesn't mutate it.""" + from vllm_omni.diffusion.data import OmniDiffusionConfig + + original_dict = {"method": "bitsandbytes", "quant_type": "nf4"} + dict_copy = original_dict.copy() + + OmniDiffusionConfig(model="test", quantization_config=original_dict) + + assert original_dict == dict_copy + + +def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch): + """Test get_quant_method routing for CUDA.""" + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + config = build_quant_config("bitsandbytes") + + def _fake_init(self, quant_config): + pass + + layer = mocker.Mock(spec=LinearBase) + mocker.patch.object(BnBOnlineLinearMethod, "__init__", _fake_init) + + prefix = "test_layer" + + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True) + method = config.get_quant_method(layer, prefix) + assert isinstance(method, BnBOnlineLinearMethod) + + config.ignored_layers = [prefix] + method = config.get_quant_method(layer, prefix) + assert isinstance(method, UnquantizedLinearMethod) + + +class TestBnBOnlineLinearMethod: + @pytest.fixture + def mock_quant_config(self, mocker): + config = mocker.Mock() + config.quant_type = "nf4" + config.compress_statistics = True + return config + + @pytest.fixture + def mock_deps(self, mocker, monkeypatch: pytest.MonkeyPatch): + _ensure_bitsandbytes_importable(monkeypatch) + mock_qweight = torch.ones((64, 32), dtype=torch.uint8) + mock_quant_state = mocker.Mock() + mock_quant = mocker.patch( + "bitsandbytes.functional.quantize_4bit", + return_value=(mock_qweight, mock_quant_state), + create=True, + ) + # x is (2, 16, 32) -> flattened (32, 32); matmul must return (32, 64) + # so apply can reshape back to (2, 16, 64). + mock_matmul = mocker.patch( + "bitsandbytes.matmul_4bit", + return_value=torch.randn(32, 64), + create=True, + ) + return { + "quant": mock_quant, + "matmul": mock_matmul, + "mock_qweight": mock_qweight, + "mock_quant_state": mock_quant_state, + } + + def test_process_weights_after_loading(self, mock_deps, mock_quant_config, mocker): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(mock_quant_config) + layer = Module() + layer.weight = Parameter(torch.randn(64, 32)) + mocker.patch.object(torch.Tensor, "cuda", return_value=layer.weight.data) + + method.process_weights_after_loading(layer) + + mock_deps["quant"].assert_called_once() + assert torch.equal(layer.weight.data, mock_deps["mock_qweight"]) + assert layer.quant_state is mock_deps["mock_quant_state"] + assert layer.bnb_shape == (64, 32) + + def test_apply(self, mock_deps, mock_quant_config): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(mock_quant_config) + layer = Module() + layer.weight = Parameter(torch.ones((64, 32), dtype=torch.uint8), requires_grad=False) + layer.quant_state = mock_deps["mock_quant_state"] + + x = torch.randn(2, 16, 32, dtype=torch.float16) + bias = torch.randn(64) + output = method.apply(layer, x, bias) + + mock_deps["matmul"].assert_called_once() + assert output.shape == (2, 16, 64) + assert output.dtype == torch.float16 + + +@pytest.fixture +def quant_config(): + from vllm_omni.quantization.bitsandbytes_config import DiffusionBitsAndBytesConfig + + return DiffusionBitsAndBytesConfig( + quant_type="nf4", + compress_statistics=True, + ) + + +@cuda_available +@bitsandbytes_available +class TestCudaBnBSmoke: + """Smoke tests using real bitsandbytes CUDA kernels.""" + + @pytest.fixture + def real_layer(self): + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.randn(128, 64, dtype=torch.float16, device="cuda"), + requires_grad=False, + ) + layer.logical_widths = [128] + layer.input_size_per_partition = 64 + layer.output_size_per_partition = 128 + layer.orig_dtype = torch.float16 + return layer + + def test_real_cuda_quantize_4bit_shape_contract(self, quant_config): + import bitsandbytes.functional as bnb_F + + weight = torch.randn(128, 64, dtype=torch.float16, device="cuda") + qweight, quant_state = bnb_F.quantize_4bit( + weight, + quant_type=quant_config.quant_type, + compress_statistics=quant_config.compress_statistics, + ) + + assert qweight.dtype == torch.uint8 + assert quant_state is not None + + def test_real_cuda_online_process_weights_after_loading(self, quant_config, real_layer): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(quant_config) + method.process_weights_after_loading(real_layer) + + assert real_layer.weight.dtype == torch.uint8 + assert hasattr(real_layer, "quant_state") + assert real_layer.bnb_shape == (128, 64) + + def test_real_cuda_bnb_apply_forward(self, quant_config): + from vllm_omni.quantization.bitsandbytes_config import BnBOnlineLinearMethod + + method = BnBOnlineLinearMethod(quant_config) + + layer = torch.nn.Module() + layer.weight = torch.nn.Parameter( + torch.randn(128, 64, dtype=torch.float16, device="cuda"), + requires_grad=False, + ) + method.process_weights_after_loading(layer) + + x = torch.randn(2, 16, 64, dtype=torch.float16, device="cuda") + output = method.apply(layer, x) + + assert output.shape == (2, 16, 128) + assert output.dtype == torch.float16 diff --git a/tests/diffusion/quantization/test_fp8_config.py b/tests/diffusion/quantization/test_fp8_config.py index a1f79030e04..9bad167926d 100644 --- a/tests/diffusion/quantization/test_fp8_config.py +++ b/tests/diffusion/quantization/test_fp8_config.py @@ -154,21 +154,6 @@ def test_component_config_with_default(): assert config.resolve("vae.encoder") is None -def test_gguf_config(): - from vllm_omni.quantization import build_quant_config - from vllm_omni.quantization.gguf_config import DiffusionGGUFConfig - - config = build_quant_config( - { - "method": "gguf", - "gguf_model": "path/to/model.gguf", - } - ) - assert isinstance(config, DiffusionGGUFConfig) - assert config.gguf_model == "path/to/model.gguf" - assert config.get_name() == "gguf" - - def test_integration_string(): from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -236,7 +221,7 @@ def test_transformer_config_auto_detects_modelopt_fp8(): def test_supported_methods_includes_vllm(): from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS - for method in ["fp8", "gguf", "awq", "gptq", "bitsandbytes", "modelopt"]: + for method in ["fp8", "awq", "gptq", "bitsandbytes", "modelopt"]: assert method in SUPPORTED_QUANTIZATION_METHODS, f"{method} missing" diff --git a/tests/diffusion/quantization/test_gguf_config.py b/tests/diffusion/quantization/test_gguf_config.py deleted file mode 100644 index 1dd401d59c5..00000000000 --- a/tests/diffusion/quantization/test_gguf_config.py +++ /dev/null @@ -1,128 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for diffusion GGUF quantization config and linear method.""" - -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch -from vllm.model_executor.layers.linear import LinearBase -from vllm.model_executor.layers.quantization.gguf import UNQUANTIZED_TYPES, UnquantizedLinearMethod - -from vllm_omni.quantization import build_quant_config -from vllm_omni.quantization.gguf_config import ( - DiffusionGGUFConfig, - DiffusionGGUFLinearMethod, - dequant_gemm_gguf, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] - - -def test_gguf_config_creation_and_delegation(): - config = DiffusionGGUFConfig( - gguf_model="weights.gguf", - unquantized_modules=["proj_out"], - ) - - assert config.gguf_model == "weights.gguf" - assert config.unquantized_modules == ["proj_out"] - assert config.get_name() == "gguf" - - -def test_build_quant_config_builds_gguf_config(): - config = build_quant_config( - "gguf", - gguf_model="repo/model-Q4_0.gguf", - unquantized_modules=["norm_out"], - ) - - assert isinstance(config, DiffusionGGUFConfig) - assert config.gguf_model == "repo/model-Q4_0.gguf" - assert config.unquantized_modules == ["norm_out"] - - -def test_gguf_vllm_config_returns_diffusion_linear_method_for_linear_layers(): - linear = object.__new__(LinearBase) - method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(linear, "transformer.img_in") - - assert isinstance(method, DiffusionGGUFLinearMethod) - - -def test_gguf_vllm_config_respects_unquantized_modules(): - linear = object.__new__(LinearBase) - method = DiffusionGGUFConfig(unquantized_modules=["proj_out"]).get_quant_method(linear, "transformer.proj_out") - - assert isinstance(method, UnquantizedLinearMethod) - - -def test_gguf_vllm_config_returns_none_for_non_linear_layers(): - method = DiffusionGGUFConfig(unquantized_modules=[]).get_quant_method(torch.nn.LayerNorm(4), "norm") - - assert method is None - - -def test_dequant_gemm_gguf_uses_plain_matmul_for_unquantized_types(): - qweight_type = next(iter(UNQUANTIZED_TYPES)) - x = torch.tensor([[1.0, 2.0]], dtype=torch.float32) - qweight = torch.tensor([[3.0, 4.0], [5.0, 6.0]], dtype=torch.float32) - - out = dequant_gemm_gguf(x, qweight, qweight_type) - - assert torch.allclose(out, x @ qweight.T) - - -def test_diffusion_gguf_linear_method_applies_bias_on_unquantized_weight(): - qweight_type = next(iter(UNQUANTIZED_TYPES)) - method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[])) - layer = SimpleNamespace( - qweight=torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32), - qweight_type=SimpleNamespace(weight_type=qweight_type), - ) - x = torch.tensor([[[1.0, 0.5]]], dtype=torch.float32) - bias = torch.tensor([0.25, -0.5], dtype=torch.float32) - - out = method.apply(layer, x, bias) - - expected = x @ layer.qweight.T + bias - assert torch.allclose(out, expected) - - -def test_diffusion_gguf_linear_method_concatenates_sharded_outputs(): - qweight_type = next(iter(UNQUANTIZED_TYPES)) - method = DiffusionGGUFLinearMethod(DiffusionGGUFConfig(unquantized_modules=[])) - qweight = torch.nn.Parameter( - torch.tensor( - [ - [1.0, 0.0], - [0.0, 1.0], - [2.0, 0.0], - [0.0, 2.0], - ], - dtype=torch.float32, - ), - requires_grad=False, - ) - qweight.shard_id = ["left", "right"] - qweight.shard_offset_map = { - "left": (0, 2, 2), - "right": (2, 4, 2), - } - layer = SimpleNamespace( - qweight=qweight, - qweight_type=SimpleNamespace( - shard_weight_type={ - "left": qweight_type, - "right": qweight_type, - } - ), - ) - x = torch.tensor([[1.5, 2.0]], dtype=torch.float32) - - out = method.apply(layer, x) - - expected_q = x @ qweight[:2].T - expected_k = x @ qweight[2:4].T - assert torch.allclose(out, torch.cat([expected_q, expected_k], dim=-1)) diff --git a/tests/diffusion/quantization/test_hsdp_fp8.py b/tests/diffusion/quantization/test_hsdp_fp8.py new file mode 100644 index 00000000000..482235ff2d7 --- /dev/null +++ b/tests/diffusion/quantization/test_hsdp_fp8.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for HSDP/FSDP2 compatibility with online FP8 quantization.""" + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.quantization.hsdp_fp8 import ( + _build_transposed_get_layer_params, + prepare_fp8_layers_for_fsdp, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# --- shared test doubles --- + + +class _ToyKernel: + """Minimal kernel that mirrors ``FP8ScaledMMLinearKernel``.""" + + def __init__(self): + self.layer_param_names = ("weight", "weight_scale", "input_scale", "input_scale_ub") + + def _get_layer_params(self, layer): + w, w_s, x_s, x_s_ub = self.layer_param_names + return ( + getattr(layer, w), + getattr(layer, w_s), + getattr(layer, x_s, None), + getattr(layer, x_s_ub, None), + ) + + +# --- helper to build a toy module simulating online-FP8 post-load state --- + + +def _make_fp8_toy_module(out_features: int = 16, in_features: int = 32): + """Return a module whose weight is ``qweight.t()`` (non-contiguous) and whose + ``quant_method`` is an ``Fp8LinearMethod`` with a ``_ToyKernel``. + """ + from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod + + # Create the quant-method instance without calling __init__ (avoids + # heavyweight upstream config dependencies). isinstance() still works. + qm = object.__new__(Fp8LinearMethod) + qm.fp8_linear = _ToyKernel() + + # randn does not support float8 on CPU; zeros does and the test only + # cares about shape/stride, not the actual weight values. + qweight = torch.zeros(out_features, in_features, dtype=torch.float8_e4m3fn) + weight = nn.Parameter(qweight.t()) # (in, out) non-contiguous column-major view + + module = nn.Module() + module.quant_method = qm + module.weight = weight + module.weight_scale = nn.Parameter(torch.ones(1, dtype=torch.float32)) + # Set input_scale / input_scale_ub to verify _get_layer_params + # passes them through unchanged. + module.input_scale = nn.Parameter(torch.ones(1, dtype=torch.float32)) + module.input_scale_ub = nn.Parameter(torch.ones(1, dtype=torch.float32)) + + return module + + +# --- tests --- + + +def test_transposed_get_layer_params(): + import types + + kernel = _ToyKernel() + + # original bound method + patch + re-bind, matching prepare_fp8_layers_for_fsdp + original_bound = kernel._get_layer_params + patched_func = _build_transposed_get_layer_params(original_bound) + kernel._get_layer_params = types.MethodType(patched_func, kernel) + + module = _make_fp8_toy_module(16, 32) + # Reproduce the storage rewrite from prepare_fp8_layers_for_fsdp: + # weight was qweight.t() → .t() recovers qweight (row-major contiguous). + module.weight = nn.Parameter(module.weight.data.t(), requires_grad=False) + assert module.weight.is_contiguous(), "qweight should be row-major contiguous" + + w, w_s, x_s, x_s_ub = kernel._get_layer_params(module) + + # patched method transposes row-major → column-major for Cutlass + assert w.shape == (32, 16) + assert not w.is_contiguous() + # scales are passed through unchanged + assert w_s is module.weight_scale + assert x_s is module.input_scale + assert x_s_ub is module.input_scale_ub + + +def test_rewrites_non_contiguous_weight(): + module = _make_fp8_toy_module(16, 32) + assert not module.weight.is_contiguous() + + n_rewritten_layers = prepare_fp8_layers_for_fsdp(module) + assert n_rewritten_layers == 1 + assert module.weight.is_contiguous() + assert module.weight.shape == (16, 32) + + +def test_patched_get_layer_params_returns_transposed_view(): + module = _make_fp8_toy_module(16, 32) + prepare_fp8_layers_for_fsdp(module) + + kernel = module.quant_method.fp8_linear + w, w_s, x_s, x_s_ub = kernel._get_layer_params(module) + + # Cutlass expects column-major B: (in, out), stride (1, in) + assert w.shape == (32, 16) + assert not w.is_contiguous() + assert w.stride() == (1, 32) + + # scales pass through unchanged + assert w_s is module.weight_scale + assert x_s is module.input_scale + assert x_s_ub is module.input_scale_ub + + +def test_contiguous_weights_skipped(): + module = _make_fp8_toy_module(16, 32) + # make weight contiguous before calling + module.weight = nn.Parameter(module.weight.data.t().contiguous()) + assert module.weight.is_contiguous() + + n_rewritten_layers = prepare_fp8_layers_for_fsdp(module) + assert n_rewritten_layers == 0 + + +def test_kernel_patched_only_once_per_instance(): + module = _make_fp8_toy_module(16, 32) + kernel = module.quant_method.fp8_linear + original = kernel._get_layer_params + + prepare_fp8_layers_for_fsdp(module) + first_patched = kernel._get_layer_params + + # Second call with same kernel id must not stack patches + n_rewritten_layers = prepare_fp8_layers_for_fsdp(module) + assert n_rewritten_layers == 0 + assert kernel._get_layer_params is first_patched + assert kernel._get_layer_params is not original diff --git a/tests/diffusion/quantization/test_mxfp4_config.py b/tests/diffusion/quantization/test_mxfp4_config.py index af3a9c9c435..240ee250941 100644 --- a/tests/diffusion/quantization/test_mxfp4_config.py +++ b/tests/diffusion/quantization/test_mxfp4_config.py @@ -537,3 +537,118 @@ def test_single_scale_row_parallel_tp2_loader_simulation(): for rank in range(_TP2): assert _shard(ckpt_weight_scale, layer.weight_scale, rank, _TP2, "input_dim").shape == layer.weight_scale.shape + + +# --------------------------------------------------------------------------- +# ROCm MXFP4 (gfx950) — get_quant_method dispatch +# +# These run on CPU: the platform check, gcnArchName probe and aiter op +# registration are all mocked (stdlib unittest.mock + monkeypatch, so no +# pytest-mock dependency). They cover only the dispatch + weight-allocation +# logic added by the ROCm PR — the AITER GEMM / quant kernels require real +# gfx950 hardware and are intentionally NOT exercised here. +# --------------------------------------------------------------------------- + + +@pytest.fixture +def _rocm_platform(monkeypatch: pytest.MonkeyPatch): + """Make current_omni_platform report ROCm, and stub the aiter + custom-op registration so ROCmMxfp4*Method can be constructed without aiter.""" + from vllm_omni.platforms import current_omni_platform + from vllm_omni.quantization import mxfp4_config + + monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False) + monkeypatch.setattr(current_omni_platform, "is_rocm", lambda: True) + monkeypatch.setattr(mxfp4_config, "_register_rocm_mxfp4_op", lambda: None) + + +def _patch_gcn_arch(monkeypatch: pytest.MonkeyPatch, arch: str) -> None: + """Patch torch.cuda.get_device_properties(...).gcnArchName to return `arch`.""" + from types import SimpleNamespace + + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda *a, **k: SimpleNamespace(gcnArchName=arch), + ) + + +def _fake_linear_layer(): + """A stand-in that passes isinstance(layer, LinearBase) without a real layer.""" + from unittest.mock import MagicMock + + from vllm.model_executor.layers.linear import LinearBase + + return MagicMock(spec=LinearBase) + + +def test_rocm_online_dispatch_returns_rocm_method(_rocm_platform, monkeypatch): + """ROCm + gfx950 + online checkpoint must return ROCmMxfp4OnlineLinearMethod.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod + + _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False) + + method = cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q") + assert isinstance(method, ROCmMxfp4OnlineLinearMethod) + + +def test_rocm_ignored_layer_returns_unquantized(_rocm_platform, monkeypatch): + """A prefix in ignored_layers must return UnquantizedLinearMethod before the gfx950 probe.""" + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + _patch_gcn_arch(monkeypatch, "gfx950:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False, ignored_layers=["proj_out"]) + + assert isinstance(cfg.get_quant_method(_fake_linear_layer(), "proj_out"), UnquantizedLinearMethod) + + +def test_rocm_non_gfx950_raises(_rocm_platform, monkeypatch): + """MXFP4 on ROCm requires gfx950; any other arch must raise.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config + + _patch_gcn_arch(monkeypatch, "gfx942:sramecc+:xnack-") + cfg = DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=False) + + with pytest.raises(NotImplementedError, match="gfx950"): + cfg.get_quant_method(_fake_linear_layer(), "blocks.0.attn1.to_q") + + +# --------------------------------------------------------------------------- +# ROCm MXFP4 — create_weights (online lazy meta-device placeholder) +# +# The base ROCmMxfp4LinearMethod is abstract (no create_weights); the online +# subclass gets create_weights from _LazyWeightMixin, which registers a BF16 +# weight on the meta device to be materialised at load time. +# --------------------------------------------------------------------------- + + +def test_rocm_create_weights_column_parallel_tp2(_rocm_platform): + """Column-parallel TP=2: meta BF16 weight has the output halved, input full.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod + + method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K, output_partition_sizes=[_TP2_N // _TP2]) + + assert layer.weight.shape == (_TP2_N // _TP2, _TP2_K) + assert layer.weight.dtype == torch.bfloat16 + assert layer.weight.device.type == "meta" + assert layer.weight.input_dim == 1 + assert layer.weight.output_dim == 0 + assert layer.logical_widths == [_TP2_N // _TP2] + assert layer.weight_block_size is None + + +def test_rocm_create_weights_row_parallel_tp2(_rocm_platform): + """Row-parallel TP=2: meta BF16 weight has the input halved, output full.""" + from vllm_omni.quantization.mxfp4_config import DiffusionMXFP4Config, ROCmMxfp4OnlineLinearMethod + + method = ROCmMxfp4OnlineLinearMethod(DiffusionMXFP4Config()) + layer = _create_weights(method, input_size_per_partition=_TP2_K // _TP2, output_partition_sizes=[_TP2_N]) + + assert layer.weight.shape == (_TP2_N, _TP2_K // _TP2) + assert layer.input_size_per_partition == _TP2_K // _TP2 + assert layer.output_size_per_partition == _TP2_N diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index 5f0dc4d620a..9b53d894164 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -82,6 +82,7 @@ class QualityTestConfig: gpu: str = "H100" # minimum GPU requirement negative_prompt: str = "" guidance_scale: float | None = None + enable_cpu_offload: bool = False def baseline_ref(self) -> str: return self.baseline_model or self.model or "" @@ -181,6 +182,18 @@ def validate(self) -> None: max_lpips=0.20, num_inference_steps=10, ), + QualityTestConfig( + id="fp8_flux2_dev_text_encoder", + model="black-forest-labs/FLUX.2-dev", + quantization={"text_encoder": "fp8", "transformer": None, "vae": None}, + task="t2i", + prompt="a cup of coffee on a wooden table, morning light", + max_lpips=0.15, + num_inference_steps=10, + enable_cpu_offload=True, + height=1024, + width=1024, + ), QualityTestConfig( id="fp8_qwen_image", model="Qwen/Qwen-Image", @@ -459,7 +472,10 @@ def test_quantization_quality(config: QualityTestConfig): generate_fn = _generate_video if config.task == "t2v" else _generate_image # --- BF16 baseline --- - omni_bl = Omni(model=config.baseline_ref()) + bl_kwargs: dict = {"model": config.baseline_ref()} + if config.enable_cpu_offload: + bl_kwargs["enable_cpu_offload"] = True + omni_bl = Omni(**bl_kwargs) baseline_out, bl_mem = generate_fn(omni_bl, config) omni_bl.shutdown() del omni_bl @@ -468,10 +484,13 @@ def test_quantization_quality(config: QualityTestConfig): # --- Quantized --- quantization = config.quantization_ref() + qt_kwargs: dict = {"model": config.quantized_ref()} + if config.enable_cpu_offload: + qt_kwargs["enable_cpu_offload"] = True if quantization is None: - omni_qt = Omni(model=config.quantized_ref()) + omni_qt = Omni(**qt_kwargs) else: - omni_qt = Omni(model=config.quantized_ref(), quantization_config=quantization) + omni_qt = Omni(**qt_kwargs, quantization_config=quantization) quant_out, qt_mem = generate_fn(omni_qt, config) omni_qt.shutdown() del omni_qt diff --git a/tests/diffusion/quantization/test_trajectory_similarity_tool.py b/tests/diffusion/quantization/test_trajectory_similarity_tool.py index 45d51d35641..252e6f1c18b 100644 --- a/tests/diffusion/quantization/test_trajectory_similarity_tool.py +++ b/tests/diffusion/quantization/test_trajectory_similarity_tool.py @@ -99,7 +99,6 @@ def test_candidate_model_can_point_to_offline_checkpoint_without_online_quantiza candidate_quantization_config_json=None, candidate_ignored_layers=None, ignored_layers=None, - candidate_gguf_model=None, candidate_load_format="default", ) diff --git a/tests/diffusion/test_attention_compile.py b/tests/diffusion/test_attention_compile.py new file mode 100644 index 00000000000..d397446416c --- /dev/null +++ b/tests/diffusion/test_attention_compile.py @@ -0,0 +1,83 @@ +# 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.diffusion.attention.layer import Attention +from vllm_omni.diffusion.forward_context import set_forward_context + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def test_attention_uses_compile_boundary_for_hsdp(monkeypatch): + attention = object.__new__(Attention) + calls = [] + + def _boundary(query, key, value, attn_metadata=None): + calls.append("boundary") + return query + + def _impl(query, key, value, attn_metadata=None): + calls.append("impl") + return query + + attention._forward_hsdp_compile_boundary = _boundary + attention._forward_impl = _impl + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=True)) + query = torch.empty(1) + with set_forward_context(omni_diffusion_config=config): + assert Attention.forward(attention, query, query, query) is query + + assert calls == ["boundary"] + + +def test_attention_keeps_compiled_impl_without_hsdp(monkeypatch): + attention = object.__new__(Attention) + calls = [] + + def _boundary(query, key, value, attn_metadata=None): + calls.append("boundary") + return query + + def _impl(query, key, value, attn_metadata=None): + calls.append("impl") + return query + + attention._forward_hsdp_compile_boundary = _boundary + attention._forward_impl = _impl + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=False)) + query = torch.empty(1) + with set_forward_context(omni_diffusion_config=config): + assert Attention.forward(attention, query, query, query) is query + + assert calls == ["impl"] + + +def test_attention_keeps_compiled_impl_without_diffusion_config(monkeypatch): + attention = object.__new__(Attention) + calls = [] + + def _boundary(query, key, value, attn_metadata=None): + calls.append("boundary") + return query + + def _impl(query, key, value, attn_metadata=None): + calls.append("impl") + return query + + attention._forward_hsdp_compile_boundary = _boundary + attention._forward_impl = _impl + monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True) + + query = torch.empty(1) + with set_forward_context(vllm_config=None): + assert Attention.forward(attention, query, query, query) is query + + assert calls == ["impl"] diff --git a/tests/diffusion/test_compile.py b/tests/diffusion/test_compile.py new file mode 100644 index 00000000000..ca898da95f6 --- /dev/null +++ b/tests/diffusion/test_compile.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch.nn as nn + +import vllm_omni.diffusion.compile as compile_module +from vllm_omni.diffusion.compile import regionally_compile + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _WrappedBlock(nn.Module): + def __init__(self) -> None: + super().__init__() + self.compile_called = False + self.forward_compiled = False + + def compile(self, *args, **kwargs): + self.compile_called = True + return self + + def forward(self, x): + return x + + +class _ModelWithWrappedRepeatedBlocks(nn.Module): + _repeated_blocks = ["OriginalBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks"] + + def __init__(self) -> None: + super().__init__() + self.transformer_blocks = nn.ModuleList([_WrappedBlock(), _WrappedBlock()]) + self.other_blocks = nn.ModuleList([_WrappedBlock()]) + + +def test_regionally_compile_matches_wrapped_blocks_by_declared_container_attr(monkeypatch): + model = _ModelWithWrappedRepeatedBlocks() + compile_calls = [] + + def _compile(fn, *args, **kwargs): + compile_calls.append((fn, args, kwargs)) + + def _compiled(*fn_args, **fn_kwargs): + return f"compiled:{fn(*fn_args, **fn_kwargs)}" + + return _compiled + + monkeypatch.setattr(compile_module.torch, "compile", _compile) + + regionally_compile(model, dynamic=True) + + assert len(compile_calls) == 2 + assert all(not block.compile_called for block in model.transformer_blocks) + assert not model.other_blocks[0].compile_called + assert model.transformer_blocks[0].forward("ok") == "compiled:ok" diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index eb862a91ae0..5f9a5a5c0dc 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -11,8 +11,20 @@ import pytest import torch - -from vllm_omni.diffusion.diffusion_engine import _move_tensor_tree_to_cpu +from pytest_mock import MockerFixture + +import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _move_tensor_tree_to_cpu +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.sched.interface import ( + CachedRequestData, + NewRequestData, +) +from vllm_omni.diffusion.sched.interface import ( + DiffusionSchedulerOutput as RealDiffusionSchedulerOutput, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams @dataclass @@ -66,6 +78,324 @@ def update_from_output(self, sched_output, runner_output): return [req.request_id for req in sched_output.scheduled_new_reqs] +class _BatchCapablePipeline: + supports_request_batch = True + + +class _SingleRequestPipeline: + pass + + +class _SingleRequestOverridePipeline(_BatchCapablePipeline): + def forward(self, req, prompt_ids=None): + return DiffusionOutput(output=None) + + +def _make_request_mode_sched_output(*request_ids: str) -> RealDiffusionSchedulerOutput: + new_reqs = [ + NewRequestData( + request_id=request_id, + req=OmniDiffusionRequest( + prompt=f"prompt_{request_id}", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id=request_id, + ), + ) + for request_id in request_ids + ] + return RealDiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=new_reqs, + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=len(new_reqs), + num_waiting_reqs=0, + ) + + +class TestRequestBatchCapability: + pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + def test_supports_request_batch_uses_registered_model_class(self, monkeypatch: pytest.MonkeyPatch) -> None: + od_config = SimpleNamespace(model_class_name="BatchPipeline", custom_pipeline_args=None) + + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: _BatchCapablePipeline if model_class_name == "BatchPipeline" else None, + ) + + assert diffusion_engine_module.supports_request_batch(od_config) is True + + def test_supports_request_batch_uses_custom_pipeline_class(self, monkeypatch: pytest.MonkeyPatch) -> None: + od_config = SimpleNamespace( + model_class_name="SinglePipeline", + custom_pipeline_args={"pipeline_class": _BatchCapablePipeline}, + ) + + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: _SingleRequestPipeline, + ) + + assert diffusion_engine_module.supports_request_batch(od_config) is True + + def test_supports_request_batch_uses_custom_pipeline_class_name(self, monkeypatch: pytest.MonkeyPatch) -> None: + od_config = SimpleNamespace( + model_class_name="SinglePipeline", + custom_pipeline_args={"pipeline_class": "test.module.BatchPipeline"}, + ) + + monkeypatch.setattr( + diffusion_engine_module, + "resolve_obj_by_qualname", + lambda qualname: _BatchCapablePipeline if qualname == "test.module.BatchPipeline" else None, + ) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: _SingleRequestPipeline, + ) + + assert diffusion_engine_module.supports_request_batch(od_config) is True + + def test_supports_request_batch_uses_only_explicit_pipeline_attribute_for_custom_override( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + od_config = SimpleNamespace( + model_class_name="BatchPipeline", + custom_pipeline_args={"pipeline_class": _SingleRequestOverridePipeline}, + ) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: None, + ) + + assert diffusion_engine_module.supports_request_batch(od_config) is True + + def test_supports_request_batch_honors_explicit_false_on_custom_override( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + class _ExplicitlyUnsupportedOverride(_BatchCapablePipeline): + supports_request_batch = False + + def forward(self, req, prompt_ids=None): + return DiffusionOutput(output=None) + + od_config = SimpleNamespace( + model_class_name="BatchPipeline", + custom_pipeline_args={"pipeline_class": _ExplicitlyUnsupportedOverride}, + ) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: None, + ) + + assert diffusion_engine_module.supports_request_batch(od_config) is False + + def test_supports_request_batch_rejects_invalid_custom_pipeline_class_name( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace( + model_class_name="BatchPipeline", + custom_pipeline_args={"pipeline_class": "test.module.MissingPipeline"}, + ) + + def fail_resolve(qualname): + raise ImportError(qualname) + + monkeypatch.setattr(diffusion_engine_module, "resolve_obj_by_qualname", fail_resolve) + registry_load = mocker.Mock(return_value=_BatchCapablePipeline) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + registry_load, + ) + + with pytest.raises(ValueError, match="Failed to resolve custom diffusion pipeline class"): + diffusion_engine_module.supports_request_batch(od_config) + registry_load.assert_not_called() + + def test_engine_disables_batch_dispatch_for_single_request_pipeline( + self, + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace( + model_class_name="SinglePipeline", + custom_pipeline_args=None, + streaming_output=False, + ) + fake_executor = SimpleNamespace( + execute_request=mocker.Mock(return_value="per-request"), + execute_batch=mocker.Mock(return_value="batch"), + execute_step=mocker.Mock(return_value="step"), + ) + fake_executor_cls = mocker.Mock(return_value=fake_executor) + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", + lambda *args, **kwargs: fake_executor_cls, + ) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: _SingleRequestPipeline, + ) + monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None) + + engine = DiffusionEngine(od_config) + output = engine.execute_fn(_make_request_mode_sched_output("req-a", "req-b")) + + assert engine.supports_request_batch is False + assert output == "per-request" + fake_executor.execute_request.assert_called_once() + fake_executor.execute_batch.assert_not_called() + + @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")]) + def test_engine_enables_batch_dispatch_for_request_batch_pipeline( + self, + request_ids: tuple[str, ...], + monkeypatch: pytest.MonkeyPatch, + mocker: MockerFixture, + ) -> None: + od_config = SimpleNamespace( + model_class_name="BatchPipeline", + custom_pipeline_args=None, + streaming_output=False, + ) + fake_executor = SimpleNamespace( + execute_request=mocker.Mock(return_value="per-request"), + execute_batch=mocker.Mock(return_value="batch"), + execute_step=mocker.Mock(return_value="step"), + ) + fake_executor_cls = mocker.Mock(return_value=fake_executor) + + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", + lambda *args, **kwargs: fake_executor_cls, + ) + monkeypatch.setattr( + diffusion_engine_module.DiffusionModelRegistry, + "_try_load_model_cls", + lambda model_class_name: _BatchCapablePipeline, + ) + monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None) + + engine = DiffusionEngine(od_config) + output = engine.execute_fn(_make_request_mode_sched_output(*request_ids)) + + assert engine.supports_request_batch is True + assert output == "batch" + fake_executor.execute_batch.assert_called_once() + fake_executor.execute_request.assert_not_called() + + +class TestRequestBatchAdmission: + pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + def test_config_rejects_negative_request_batch_max_wait_ms(self) -> None: + with pytest.raises(ValueError, match="request_batch_max_wait_ms"): + OmniDiffusionConfig(model="test", request_batch_max_wait_ms=-1.0) + + def test_config_normalizes_request_batch_max_wait_ms_to_float(self) -> None: + config = OmniDiffusionConfig(model="test", request_batch_max_wait_ms=5) + + assert config.request_batch_max_wait_ms == 5.0 + assert isinstance(config.request_batch_max_wait_ms, float) + + def test_scheduler_exposes_waiting_and_running_counts(self) -> None: + from vllm_omni.diffusion.sched import RequestScheduler + + od_config = SimpleNamespace(max_num_seqs=4) + scheduler = RequestScheduler() + scheduler.initialize(od_config) + + assert scheduler.num_waiting_requests() == 0 + assert scheduler.num_running_requests() == 0 + + scheduler.add_request( + OmniDiffusionRequest( + prompt="prompt_a", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id="req-a", + ) + ) + scheduler.add_request( + OmniDiffusionRequest( + prompt="prompt_b", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id="req-b", + ) + ) + assert scheduler.num_waiting_requests() == 2 + assert scheduler.num_running_requests() == 0 + + scheduler.schedule() + assert scheduler.num_waiting_requests() == 0 + assert scheduler.num_running_requests() == 2 + + def test_request_batch_admission_exits_early_when_waiting_queue_stable(self) -> None: + from vllm_omni.diffusion.sched import RequestScheduler + + od_config = SimpleNamespace( + max_num_seqs=32, + request_batch_max_wait_ms=1000.0, + step_execution=False, + ) + scheduler = RequestScheduler() + scheduler.initialize(od_config) + for idx in range(2): + scheduler.add_request( + OmniDiffusionRequest( + prompt=f"prompt_{idx}", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id=f"req-{idx}", + ) + ) + + engine = object.__new__(DiffusionEngine) + engine.od_config = od_config + engine.scheduler = scheduler + engine.step_execution = False + engine.supports_request_batch = True + engine.stop_event = threading.Event() + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + + start = time.monotonic() + with engine._cv: + engine._wait_for_request_batch_admission_locked() + waited_s = time.monotonic() - start + + # Stable-window exit (~50ms), not the full 1000ms deadline. + assert waited_s < 0.5 + assert waited_s >= 0.04 + assert scheduler.num_waiting_requests() == 2 + assert scheduler.num_running_requests() == 0 + + @pytest.mark.core_model @pytest.mark.diffusion @pytest.mark.cpu @@ -142,8 +472,6 @@ def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None: @pytest.mark.asyncio async def test_async_add_req_and_wait_for_response(): - from vllm_omni.diffusion.diffusion_engine import DiffusionEngine - engine = object.__new__(DiffusionEngine) engine.scheduler = MockScheduler() engine._out_queue = {} @@ -153,8 +481,10 @@ async def test_async_add_req_and_wait_for_response(): engine._cv = threading.Condition(engine._rpc_lock) engine._init_lock = asyncio.Lock() engine._closed = False + engine.od_config = SimpleNamespace(streaming_output=False) engine._loop_started = False engine.main_loop = None + engine.supports_request_batch = False engine._finalize_finished_request = lambda rid, out, err: out.result diff --git a/tests/diffusion/test_diffusion_engine_cleanup.py b/tests/diffusion/test_diffusion_engine_cleanup.py index a8985f3723e..3cc3d48f7a5 100644 --- a/tests/diffusion/test_diffusion_engine_cleanup.py +++ b/tests/diffusion/test_diffusion_engine_cleanup.py @@ -9,6 +9,7 @@ import pytest +from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler @@ -19,7 +20,7 @@ def _make_request(request_id: str) -> OmniDiffusionRequest: return OmniDiffusionRequest( - prompts=[f"prompt_{request_id}"], + prompt=f"prompt_{request_id}", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_id=request_id, ) @@ -33,6 +34,7 @@ def _make_engine() -> DiffusionEngine: engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) engine._out_queue = {} + engine._out_queue_streaming = {} engine._closed = False engine._shutdown_complete = False engine.abort_queue = queue.Queue() @@ -57,6 +59,23 @@ def test_close_completes_pending_async_waiters() -> None: event_loop.close() +def test_close_completes_pending_streaming_waiters() -> None: + engine = _make_engine() + event_loop = asyncio.new_event_loop() + try: + engine.main_loop = event_loop + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_queue_streaming["pending-stream"] = queue + + engine.close() + + output = queue.get_nowait() + assert output.error == "DiffusionEngine is closed." + assert output.finished is True + finally: + event_loop.close() + + def test_handle_finished_requests_ignores_already_drained_waiter() -> None: class RacingOutQueue(dict): def __contains__(self, key): diff --git a/tests/diffusion/test_diffusion_engine_dummy_run.py b/tests/diffusion/test_diffusion_engine_dummy_run.py index 71cb475125c..50dde30839d 100644 --- a/tests/diffusion/test_diffusion_engine_dummy_run.py +++ b/tests/diffusion/test_diffusion_engine_dummy_run.py @@ -3,7 +3,7 @@ import pytest -from vllm_omni.diffusion import diffusion_engine +from vllm_omni.diffusion import io_support pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @@ -13,12 +13,12 @@ class JointAudioVideoModel: dummy_run_num_frames = 2 monkeypatch.setattr( - diffusion_engine.DiffusionModelRegistry, + io_support.DiffusionModelRegistry, "_try_load_model_cls", lambda model_class_name: JointAudioVideoModel, ) - assert diffusion_engine.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2 + assert io_support.get_dummy_run_num_frames("joint_audio_video", supports_audio_input=False) == 2 def test_dummy_run_num_frames_keeps_audio_output_default(monkeypatch: pytest.MonkeyPatch) -> None: @@ -26,12 +26,12 @@ class AudioOutputModel: support_audio_output = True monkeypatch.setattr( - diffusion_engine.DiffusionModelRegistry, + io_support.DiffusionModelRegistry, "_try_load_model_cls", lambda model_class_name: AudioOutputModel, ) - assert diffusion_engine.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2 + assert io_support.get_dummy_run_num_frames("audio_output", supports_audio_input=False) == 2 def test_dummy_run_num_frames_defaults_to_single_frame(monkeypatch: pytest.MonkeyPatch) -> None: @@ -39,19 +39,19 @@ class VideoOnlyModel: pass monkeypatch.setattr( - diffusion_engine.DiffusionModelRegistry, + io_support.DiffusionModelRegistry, "_try_load_model_cls", lambda model_class_name: VideoOnlyModel, ) - assert diffusion_engine.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1 + assert io_support.get_dummy_run_num_frames("video_only", supports_audio_input=False) == 1 def test_dummy_run_num_frames_uses_audio_input_fallback(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( - diffusion_engine.DiffusionModelRegistry, + io_support.DiffusionModelRegistry, "_try_load_model_cls", lambda model_class_name: None, ) - assert diffusion_engine.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2 + assert io_support.get_dummy_run_num_frames("unknown", supports_audio_input=True) == 2 diff --git a/tests/diffusion/test_diffusion_engine_metrics.py b/tests/diffusion/test_diffusion_engine_metrics.py index dd8db9fdf53..2e8eae24286 100644 --- a/tests/diffusion/test_diffusion_engine_metrics.py +++ b/tests/diffusion/test_diffusion_engine_metrics.py @@ -1,11 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Source-level regression tests for DiffusionEngine. +"""Source-level regression tests for diffusion output/engine helpers. These tests verify naming conventions and patterns by inspecting source code at the function level using AST. They are intentionally coupled to the source -layout and should be updated whenever the metrics construction code is -refactored. +layout and should be updated whenever the inspected helper code is refactored. """ from __future__ import annotations @@ -23,10 +22,20 @@ "diffusion_engine.py", ) ) +_FORMATTER_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "vllm_omni", + "diffusion", + "output_formatter.py", + ) +) -def _read_engine_source() -> str: - with open(_ENGINE_PATH) as f: +def _read_source(path: str) -> str: + with open(path) as f: return f.read() @@ -60,37 +69,38 @@ def _get_function_source(source: str, class_name: str | None, func_name: str) -> class TestMetricKeys: - """Verify metric naming conventions in DiffusionEngine.step().""" + """Verify metric naming conventions in diffusion output formatting.""" def test_no_duplicate_preprocess_key(self) -> None: - """step() should not contain 'preprocessing_time_ms' (duplicate of 'preprocess_time_ms').""" - source = _read_engine_source() - step_source = _get_function_source(source, "DiffusionEngine", "step") - assert "preprocessing_time_ms" not in step_source, ( - "Found duplicate key 'preprocessing_time_ms' in step() — should only use 'preprocess_time_ms'" + """format_diffusion_outputs() should not duplicate 'preprocess_time_ms'.""" + source = _read_source(_FORMATTER_PATH) + formatter_source = _get_function_source(source, None, "format_diffusion_outputs") + assert "preprocessing_time_ms" not in formatter_source, ( + "Found duplicate key 'preprocessing_time_ms' in " + "format_diffusion_outputs() — should only use 'preprocess_time_ms'" ) def test_metric_key_naming_consistency(self) -> None: - """In step(): exec_time should use exec_total_time, total_time should use step_total_ms.""" - source = _read_engine_source() - step_source = _get_function_source(source, "DiffusionEngine", "step") - lines = step_source.split("\n") + """Metric keys should map to the explicit timing fields.""" + source = _read_source(_FORMATTER_PATH) + formatter_source = _get_function_source(source, None, "format_diffusion_outputs") + lines = formatter_source.split("\n") found_exec = False found_total = False for line in lines: if '"diffusion_engine_exec_time_ms"' in line: found_exec = True - assert "exec_total_time" in line, ( - "diffusion_engine_exec_time_ms should measure executor time only (exec_total_time)" + assert "timings.exec_time_s" in line, ( + "diffusion_engine_exec_time_ms should measure executor time only (timings.exec_time_s)" ) if '"diffusion_engine_total_time_ms"' in line: found_total = True - assert "step_total_ms" in line, ( - "diffusion_engine_total_time_ms should measure full step time (step_total_ms)" + assert "timings.total_time_ms" in line, ( + "diffusion_engine_total_time_ms should measure full step time (timings.total_time_ms)" ) - assert found_exec, "diffusion_engine_exec_time_ms key not found in step()" - assert found_total, "diffusion_engine_total_time_ms key not found in step()" + assert found_exec, "diffusion_engine_exec_time_ms key not found in format_diffusion_outputs()" + assert found_total, "diffusion_engine_total_time_ms key not found in format_diffusion_outputs()" class TestDummyRunAllocation: @@ -98,7 +108,7 @@ class TestDummyRunAllocation: def test_no_oversized_allocation(self) -> None: """_dummy_run should not allocate more audio than needed.""" - source = _read_engine_source() + source = _read_source(_ENGINE_PATH) dummy_source = _get_function_source(source, "DiffusionEngine", "_dummy_run") assert "audio_sr * audio_duration_sec" not in dummy_source, ( "_dummy_run should generate exact-sized audio, not allocate and slice" diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index 852f6dc2ee2..4b7d0efbb1b 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -32,12 +32,12 @@ from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _RpcTask from vllm_omni.diffusion.sched import RequestScheduler -from vllm_omni.diffusion.sched.interface import SamplingParamsKey +from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey from vllm_omni.diffusion.worker.utils import RunnerOutput # Default values for every batch-key field, so SimpleNamespace-based # sampling_params satisfy ``get_sampling_params_key``'s attribute lookups. -_SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(SamplingParamsKey)} +_SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(RequestBatchSamplingParamsKey)} pytestmark = [pytest.mark.diffusion, pytest.mark.cpu] @@ -59,6 +59,7 @@ def __init__(self, rpc_delay: float = 0.0): self.rpc_delay = rpc_delay self.is_failed = False self._closed = False + self.od_config = SimpleNamespace() def collective_rpc( self, @@ -101,14 +102,15 @@ def execute_request(self, scheduler_output) -> RunnerOutput: # Mimic the real MultiprocDiffusionExecutor.execute_request: it # forwards a single request through collective_rpc. new_req = scheduler_output.scheduled_new_reqs[0] + req = new_req.req result = self.collective_rpc( "execute_model", - args=(new_req.req,), + args=(req, self.od_config), unique_reply_rank=0, exec_all_ranks=True, ) return RunnerOutput( - request_id=new_req.request_id, + request_id=req.request_id, step_index=None, finished=True, result=result, @@ -119,10 +121,12 @@ def shutdown(self) -> None: def _make_request(tag: str): + sampling_params = dict(_SAMPLING_KEY_DEFAULTS) + sampling_params["num_inference_steps"] = 1 return SimpleNamespace( request_id=tag, - prompts=[f"prompt_{tag}"], - sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS), + prompt=f"prompt_{tag}", + sampling_params=SimpleNamespace(**sampling_params), ) @@ -137,17 +141,20 @@ def _make_engine_with_loop( """ engine = DiffusionEngine.__new__(DiffusionEngine) engine._closed = False + engine.od_config = SimpleNamespace(streaming_output=False) engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay) sched = RequestScheduler() sched.initialize(SimpleNamespace(max_num_seqs=1)) engine.scheduler = sched engine.step_execution = False + engine.supports_request_batch = False engine.execute_fn = engine.executor.execute_request engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) engine._out_queue = {} + engine._out_queue_streaming = {} engine._closed = False engine.abort_queue = queue.Queue() engine._rpc_queue = queue.Queue() diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py index 57c653ef104..f7cc39d9430 100644 --- a/tests/diffusion/test_diffusion_ipc.py +++ b/tests/diffusion/test_diffusion_ipc.py @@ -15,6 +15,7 @@ pack_diffusion_output_shm, unpack_diffusion_output_shm, ) +from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] @@ -107,6 +108,30 @@ def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None: assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}] +def test_batch_runner_output_round_trips_nested_results_through_shm() -> None: + first = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + second = torch.arange(_large_numel(torch.float32), dtype=torch.float32) + 1 + output = BatchRunnerOutput.from_list( + [ + RunnerOutput(request_id="req-0", finished=True, result=DiffusionOutput(output=first)), + RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output={"image": second})), + RunnerOutput(request_id="req-error", finished=True, result=DiffusionOutput(error="boom")), + ] + ) + + pack_diffusion_output_shm(output) + + assert output.runner_outputs[0].result.output["__tensor_shm__"] is True + assert output.runner_outputs[1].result.output["image"]["__tensor_shm__"] is True + assert output.runner_outputs[2].result.error == "boom" + + unpack_diffusion_output_shm(output) + + torch.testing.assert_close(output["req-0"].result.output, first) + torch.testing.assert_close(output["req-1"].result.output["image"], second) + assert output["req-error"].result.error == "boom" + + def test_pack_value_keeps_tensor_at_threshold_inline() -> None: tensor = torch.arange( _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(), diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index 960dd072b18..5d4ef5d72a8 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -9,7 +9,9 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module from tests.helpers.mark import hardware_test +from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request pytestmark = [pytest.mark.diffusion] @@ -21,14 +23,79 @@ def _noop_forward_context(*args, **kwargs): class _DummyPipeline: + supports_request_batch = True + def __init__(self, output): self._output = output self.forward_calls = 0 + self.last_req = None def forward(self, req): - del req + self.last_req = req self.forward_calls += 1 - return self._output + return [self._output] + + +class _SingleRequestBatchPipeline: + supports_request_batch = False + + def __init__(self): + self.last_req = None + + def forward(self, req): + self.last_req = req + return [DiffusionOutput(output=req.prompts[0])] + + +class _SingleRequestDiffusionOutputPipeline: + supports_request_batch = False + + def __init__(self): + self.last_req = None + + def forward(self, req): + self.last_req = req + return DiffusionOutput(output=req.prompts[0]) + + +class _ChunkStepPipeline: + device = torch.device("cpu") + supports_step_execution = True + + def __init__(self, outputs): + self._outputs = outputs + self.prepare_calls = 0 + self.decode_calls = 0 + + def prepare_encode(self, state): + self.prepare_calls += 1 + state.prompt_embeds = torch.zeros(1, 1, 1) + state.latents = torch.zeros(1, 1) + state.timesteps = torch.tensor([1.0, 0.0]) + state.step_index = 0 + state.step_in_chunk = 0 + state.chunk_num_steps = 2 + state.total_chunks = len(self._outputs) + return state + + def denoise_step(self, input_batch, states): + del states + return torch.ones_like(input_batch.latents) + + def step_scheduler(self, state, noise_pred): + state.latents = noise_pred + state.step_index += 1 + state.step_in_chunk += 1 + + def post_decode(self, state): + output = self._outputs[self.decode_calls] + self.decode_calls += 1 + state.chunk_index += 1 + state.step_index = 0 + state.step_in_chunk = 0 + if not state.request_denoise_completed: + state.latents = torch.zeros(1, 1) + return output def _make_request(skip_cache_refresh: bool = True): @@ -39,9 +106,29 @@ def _make_request(skip_cache_refresh: bool = True): num_inference_steps=4, ) return SimpleNamespace( - prompts=["a prompt"], + request_id="req-test", + prompt="a prompt", sampling_params=sampling_params, skip_cache_refresh=skip_cache_refresh, + kv_sender_info=None, + ) + + +def _make_request_with_params(req_id: str, sampling_params): + return SimpleNamespace( + request_id=req_id, + prompt=f"prompt-{req_id}", + prompts=[f"prompt-{req_id}"], + sampling_params=sampling_params, + skip_cache_refresh=True, + ) + + +def _fake_platform_for_peak_memory(): + return SimpleNamespace( + reset_peak_memory_stats=lambda: None, + max_memory_reserved=lambda: 0, + max_memory_allocated=lambda: 0, ) @@ -49,23 +136,100 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar runner = object.__new__(DiffusionModelRunner) runner.vllm_config = object() runner.device = torch.device("cpu") - runner.pipeline = _DummyPipeline(output=SimpleNamespace(output="ok")) + runner.pipeline = _DummyPipeline(output=DiffusionOutput(output="ok")) runner.cache_backend = cache_backend runner.offload_backend = None + runner.state_cache = {} runner.prompt_embed_cache = None runner.od_config = SimpleNamespace( cache_backend=cache_backend_name, enable_cache_dit_summary=enable_cache_dit_summary, parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.kv_transfer_manager = SimpleNamespace( receive_kv_cache=lambda req, target_device=None: None, receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None, - receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None, + receive_multi_kv_cache_distributed=lambda *a, **k: None, ) + runner._kv_prefetch_enabled = False return runner +def _make_compile_runner(*, use_hsdp: bool): + runner = object.__new__(DiffusionModelRunner) + runner.pipeline = SimpleNamespace(transformer=SimpleNamespace()) + runner.od_config = SimpleNamespace(parallel_config=SimpleNamespace(use_hsdp=use_hsdp)) + return runner + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("use_hsdp", [False, True]) +def test_compile_transformer_regionally_compiles_blocks(monkeypatch, use_hsdp): + runner = _make_compile_runner(use_hsdp=use_hsdp) + compile_calls = [] + + def _regionally_compile(model, *args, **kwargs): + compile_calls.append((model, args, kwargs)) + return model + + monkeypatch.setattr(model_runner_module, "regionally_compile", _regionally_compile) + + DiffusionModelRunner._compile_transformer(runner, "transformer") + + assert compile_calls == [ + ( + runner.pipeline.transformer, + (), + {"dynamic": True}, + ) + ] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch): + """Step streaming returns empty step results until a chunk decode boundary.""" + chunks = [ + DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), + DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), + ] + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.pipeline = _ChunkStepPipeline(chunks) + runner.od_config.streaming_output = True + runner.od_config.step_execution = True + req = _make_request(skip_cache_refresh=True) + req.request_id = "req" + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req)], + scheduled_cached_reqs=SimpleNamespace(request_ids=[]), + ) + + first = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert first.get_request_output("req").result is None + + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[], + scheduled_cached_reqs=SimpleNamespace(request_ids=["req"]), + ) + second = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert second.get_request_output("req").result == chunks[0] + assert second.get_request_output("req").finished is False + + DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + fourth = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) + assert fourth.get_request_output("req").result == chunks[1] + assert fourth.get_request_output("req").finished is True + + @pytest.mark.core_model @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch): @@ -76,6 +240,9 @@ def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypa cache_summary_calls = [] monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) monkeypatch.setattr( model_runner_module, "cache_summary", @@ -101,6 +268,9 @@ def is_enabled(self): cache_summary_calls = [] monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) monkeypatch.setattr( model_runner_module, "cache_summary", @@ -113,6 +283,251 @@ def is_enabled(self): assert cache_summary_calls == [(runner.pipeline, True)] +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_passes_single_request_batch_to_non_admission_pipeline(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name="none") + runner.pipeline = _SingleRequestBatchPipeline() + req = _make_request(skip_cache_refresh=True) + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + output = DiffusionModelRunner.execute_model(runner, req) + + assert output.output == "a prompt" + assert isinstance(runner.pipeline.last_req, DiffusionRequestBatch) + assert runner.pipeline.last_req.num_reqs == 1 + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_accepts_bare_diffusion_output_from_single_request_pipeline(monkeypatch): + runner = _make_runner(cache_backend=None, cache_backend_name="none") + runner.pipeline = _SingleRequestDiffusionOutputPipeline() + req = _make_request(skip_cache_refresh=True) + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + + output = DiffusionModelRunner.execute_model(runner, req) + + assert output.output == "a prompt" + assert isinstance(runner.pipeline.last_req, DiffusionRequestBatch) + assert runner.pipeline.last_req.num_reqs == 1 + + +class _BatchPipeline: + """Pipeline returning a configurable list of outputs from forward().""" + + supports_request_batch = True + + def __init__(self, outputs): + self._outputs = outputs + self.last_batch = None + + def forward(self, batch): + self.last_batch = batch + return list(self._outputs) + + +class _SingleRequestPipeline: + def forward(self, batch): + del batch + return [DiffusionOutput(output="single")] + + +class _BatchSingleOutputPipeline: + supports_request_batch = True + + def forward(self, batch): + return DiffusionOutput(output=batch.prompts[0]) + + +def _make_batch_runner(pipeline): + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = object() + runner.device = torch.device("cpu") + runner.pipeline = pipeline + runner.cache_backend = None + runner.offload_backend = None + runner.od_config = SimpleNamespace( + cache_backend="none", + enable_cache_dit_summary=False, + parallel_config=SimpleNamespace(use_hsdp=False), + ) + runner.kv_transfer_manager = SimpleNamespace( + receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None, + ) + return runner + + +def _make_scheduler_output(num_reqs: int): + reqs = [_make_request() for _ in range(num_reqs)] + for i, req in enumerate(reqs): + req.request_id = f"req-{i}" + return SimpleNamespace(scheduled_new_reqs=[SimpleNamespace(req=req) for req in reqs]) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_rejects_output_count_mismatch(monkeypatch): + """A pipeline returning the wrong number of outputs must fail loudly, + not silently drop requests or IndexError on the per-request mapping.""" + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + # forward returns 1 output for 2 scheduled requests + runner = _make_batch_runner(_BatchPipeline(outputs=[DiffusionOutput(output="only-one")])) + sched = _make_scheduler_output(num_reqs=2) + + with pytest.raises(RuntimeError, match="returned 1 outputs for 2 requests"): + DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_routes_one_output_per_request(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + outs = [DiffusionOutput(output="a"), DiffusionOutput(output="b")] + runner = _make_batch_runner(_BatchPipeline(outputs=outs)) + sched = _make_scheduler_output(num_reqs=2) + + result = DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + assert len(result.runner_outputs) == 2 + assert result.runner_outputs[0].request_id == "req-0" + assert result.runner_outputs[0].result.output == "a" + assert result.runner_outputs[1].request_id == "req-1" + assert result.runner_outputs[1].result.output == "b" + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_preserves_per_request_sampling_and_seeds_generators(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + outputs = [DiffusionOutput(output="a"), DiffusionOutput(output="b")] + pipeline = _BatchPipeline(outputs=outputs) + runner = _make_batch_runner(pipeline) + sched = _make_scheduler_output(num_reqs=2) + sched.scheduled_new_reqs[0].req.sampling_params.seed = 111 + sched.scheduled_new_reqs[1].req.sampling_params.seed = 222 + + DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + assert pipeline.last_batch is not None + assert [sp.seed for sp in pipeline.last_batch.sampling_params_list] == [111, 222] + assert [sp.generator.initial_seed() for sp in pipeline.last_batch.sampling_params_list] == [111, 222] + with pytest.raises(AssertionError, match="multiple requests"): + _ = pipeline.last_batch.sampling_params + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_rejects_pipeline_without_request_batch_support(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + runner = _make_batch_runner(_SingleRequestPipeline()) + sched = _make_scheduler_output(num_reqs=2) + + with pytest.raises(RuntimeError, match="does not support request-batch forward"): + DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_rejects_single_diffusion_output(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + runner = _make_batch_runner(_BatchSingleOutputPipeline()) + sched = _make_scheduler_output(num_reqs=1) + + with pytest.raises(RuntimeError, match="request-batch forward must return list\\[DiffusionOutput\\]"): + DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_batch_uses_runner_output_helper(monkeypatch): + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", _fake_platform_for_peak_memory()) + outs = [DiffusionOutput(output="a"), DiffusionOutput(output="b")] + runner = _make_batch_runner(_BatchPipeline(outputs=outs)) + sched = _make_scheduler_output(num_reqs=2) + helper_calls = [] + original_from_outputs = DiffusionModelRunner._runner_output_from_outputs + + def _recording_from_outputs(self, reqs, outputs): + helper_calls.append(([req.request_id for req in reqs], [output.output for output in outputs])) + return original_from_outputs(self, reqs, outputs) + + monkeypatch.setattr(DiffusionModelRunner, "_runner_output_from_outputs", _recording_from_outputs) + + result = DiffusionModelRunner.execute_model_batch(runner, sched, runner.od_config) + + assert [runner_output.result.output for runner_output in result.runner_outputs] == ["a", "b"] + assert helper_calls == [(["req-0", "req-1"], ["a", "b"])] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_diffusion_output_by_request_slices_single_and_multi_request_outputs(): + reqs = [_make_request(), _make_request()] + reqs[0].request_id = "req-0" + reqs[1].request_id = "req-1" + batch = DiffusionRequestBatch(requests=reqs) + result = DiffusionOutput(output=["img-0a", "img-0b", "img-1a", "img-1b"], stage_durations={"decode": 1.0}) + + outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=2) + + assert [output.output for output in outputs] == [["img-0a", "img-0b"], ["img-1a", "img-1b"]] + assert [output.stage_durations for output in outputs] == [{"decode": 1.0}, {"decode": 1.0}] + + single = split_diffusion_output_by_request( + result, DiffusionRequestBatch(requests=reqs[:1]), num_outputs_per_prompt=2 + ) + + assert single[0].output == ["img-0a", "img-0b"] + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_split_diffusion_output_by_request_slices_tuple_outputs(): + reqs = [_make_request(), _make_request()] + batch = DiffusionRequestBatch(requests=reqs) + result = DiffusionOutput( + output=( + ["video-0", "video-1"], + torch.tensor([10, 20]), + ) + ) + + outputs = split_diffusion_output_by_request(result, batch, num_outputs_per_prompt=1) + + assert outputs[0].output[0] == ["video-0"] + assert torch.equal(outputs[0].output[1], torch.tensor([10])) + assert outputs[1].output[0] == ["video-1"] + assert torch.equal(outputs[1].output[1], torch.tensor([20])) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_execute_model_runs_forward_after_kv_receive(monkeypatch): + """execute_model runs the pipeline forward after the KV receive step.""" + runner = _make_runner(cache_backend=None, cache_backend_name=None) + runner.kv_transfer_manager.receive_multi_kv_cache_distributed = lambda *a, **k: True + req = _make_request(skip_cache_refresh=True) + + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0) + monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0) + + output = DiffusionModelRunner.execute_model(runner, req) + + assert output.output == "ok" + assert runner.pipeline.forward_calls == 1 + + @pytest.mark.core_model @pytest.mark.cpu def test_load_model_clears_cache_backend_for_unsupported_pipeline(monkeypatch): @@ -157,6 +572,7 @@ def enable(self, pipeline): cache_config={}, model_class_name="NextStep11Pipeline", enforce_eager=True, + streaming_output=False, ) monkeypatch.setattr(model_runner_module, "LoadConfig", lambda: object()) diff --git a/tests/diffusion/test_diffusion_output_formatter.py b/tests/diffusion/test_diffusion_output_formatter.py new file mode 100644 index 00000000000..22b305eabe1 --- /dev/null +++ b/tests/diffusion/test_diffusion_output_formatter.py @@ -0,0 +1,412 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion import output_formatter +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.output_formatter import ( + DiffusionStepTimings, + format_diffusion_outputs, + format_empty_diffusion_outputs, + normalize_diffusion_postprocess_output, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] + + +def _request( + prompt: str | dict | None = None, + *, + num_outputs_per_prompt: int = 1, +) -> OmniDiffusionRequest: + return OmniDiffusionRequest( + prompt=prompt or "prompt", + request_id="req-1", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=1, + num_outputs_per_prompt=num_outputs_per_prompt, + resolution=512, + ), + ) + + +def _config(model_class_name: str = "mock_model") -> SimpleNamespace: + return SimpleNamespace(model_class_name=model_class_name) + + +def _timings() -> DiffusionStepTimings: + return DiffusionStepTimings( + preprocess_time_s=0.01, + exec_time_s=0.02, + postprocess_time_s=0.03, + total_time_ms=60.0, + ) + + +def test_formatter_preserves_single_video_audio_actions_and_metadata( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + postprocess_output = normalize_diffusion_postprocess_output( + { + "payload": { + "video": ["frame-0"], + "audio": "audio-0", + "actions": "action-0", + }, + "metadata": { + "audio": {"sample_rate": 48000}, + "video": {"fps": 24.0}, + "actions": {"action_mode": "policy"}, + }, + } + ) + + results = format_diffusion_outputs( + request=_request("prompt-0"), + od_config=_config(), + diffusion_output=DiffusionOutput( + output=None, + stage_durations={"execute": 1.25}, + peak_memory_mb=321.0, + ), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert len(results) == 1 + result = results[0] + assert result.images == ["frame-0"] + assert result.prompt == "prompt-0" + assert result.final_output_type == "image" + assert result.multimodal_output == { + "metadata": { + "audio": {"sample_rate": 48000}, + "video": {"fps": 24.0}, + "actions": {"action_mode": "policy"}, + }, + "audio": "audio-0", + "audio_sample_rate": 48000, + "fps": 24.0, + "actions": "action-0", + } + assert result.stage_durations == {"execute": 1.25} + assert result.peak_memory_mb == 321.0 + assert result.metrics == { + "preprocess_time_ms": 10.0, + "diffusion_engine_exec_time_ms": 20.0, + "diffusion_engine_total_time_ms": 60.0, + "image_num": 1, + "resolution": 512, + "postprocess_time_ms": 30.0, + } + + +def test_formatter_normalizes_payload_metadata_envelope( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + postprocess_output = normalize_diffusion_postprocess_output( + { + "payload": { + "video": ["frame-0"], + "actions": "action-0", + }, + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + "internal": { + "robolab_action_postprocess": object(), + }, + }, + } + ) + + assert postprocess_output.outputs == { + "video": ["frame-0"], + "actions": "action-0", + } + assert postprocess_output.primary_key == "video" + assert postprocess_output.metadata == { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + "common": { + "action_only_output": True, + }, + } + assert "internal" not in postprocess_output.metadata + + [result] = format_diffusion_outputs( + request=_request("prompt-0"), + od_config=_config(), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.multimodal_output == { + "metadata": postprocess_output.metadata, + "actions": "action-0", + } + + +def test_formatter_maps_trajectory_payload_to_request_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + latents = torch.zeros(2, 1, 4) + timesteps = torch.tensor([1.0, 0.5]) + log_probs = torch.tensor([-0.1, -0.2]) + postprocess_output = normalize_diffusion_postprocess_output( + { + "payload": { + "image": "image-0", + "trajectory": { + "latents": latents, + "timesteps": timesteps, + "log_probs": log_probs, + }, + }, + "metadata": {"trajectory": {"type": "denoising"}}, + } + ) + + [result] = format_diffusion_outputs( + request=_request("prompt-0"), + od_config=_config(), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == ["image-0"] + assert result.latents is latents + assert result.trajectory_latents is latents + assert result.trajectory_timesteps is timesteps + assert result.trajectory_log_probs is log_probs + assert result.multimodal_output == { + "metadata": {"trajectory": {"type": "denoising"}}, + "trajectory": { + "latents": latents, + "timesteps": timesteps, + "log_probs": log_probs, + }, + } + + +def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + postprocess_output = normalize_diffusion_postprocess_output( + { + "payload": {"text": "caption"}, + "metadata": {"text": {"text_output": "caption"}}, + } + ) + + [result] = format_diffusion_outputs( + request=_request("describe this"), + od_config=_config(), + diffusion_output=DiffusionOutput(output=None), + output_data={"payload": {"text": "caption"}}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == [] + assert result.prompt == "describe this" + assert result.final_output_type == "text" + assert result.multimodal_output == { + "metadata": {"text": {"text_output": "caption"}}, + "text": "caption", + } + + +def test_formatter_preserves_audio_output_with_model_sample_rate_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class AudioModel: + audio_sample_rate = 44100 + + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True) + monkeypatch.setattr( + output_formatter.DiffusionModelRegistry, + "_try_load_model_cls", + lambda _: AudioModel, + ) + postprocess_output = normalize_diffusion_postprocess_output(["waveform"]) + + [result] = format_diffusion_outputs( + request=_request("speak"), + od_config=_config("audio_model"), + diffusion_output=DiffusionOutput(output=["waveform"]), + output_data=["waveform"], + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == [] + assert result.prompt == "speak" + assert result.final_output_type == "audio" + assert result.multimodal_output == { + "audio": "waveform", + "audio_sample_rate": 44100, + } + + +def test_formatter_preserves_audio_model_video_audio_and_actions( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True) + postprocess_output = normalize_diffusion_postprocess_output( + { + "video": ["frame-0"], + "audio": "audio-0", + "actions": "action-0", + "audio_sample_rate": 16000, + "fps": 30.0, + } + ) + + [result] = format_diffusion_outputs( + request=_request("watch and listen"), + od_config=_config("audio_video_model"), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == ["frame-0"] + assert result.prompt == "watch and listen" + assert result.final_output_type == "image" + assert result.multimodal_output == { + "metadata": { + "audio": {"sample_rate": 16000}, + "video": {"fps": 30.0}, + }, + "audio": "audio-0", + "audio_sample_rate": 16000, + "fps": 30.0, + "actions": "action-0", + } + + +def test_formatter_preserves_audio_only_postprocess_dict( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True) + postprocess_output = normalize_diffusion_postprocess_output( + { + "audio": "waveform", + "audio_sample_rate": 24000, + } + ) + + [result] = format_diffusion_outputs( + request=_request("speak"), + od_config=_config("audio_model"), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == [] + assert result.prompt == "speak" + assert result.final_output_type == "audio" + assert result.multimodal_output == { + "metadata": { + "audio": {"sample_rate": 24000}, + }, + "audio": "waveform", + "audio_sample_rate": 24000, + } + + +def test_formatter_preserves_single_prompt_multiple_audio_outputs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: True) + monkeypatch.setattr( + output_formatter.DiffusionModelRegistry, + "_try_load_model_cls", + lambda _: None, + ) + postprocess_output = normalize_diffusion_postprocess_output(["waveform-0", "waveform-1"]) + + [result] = format_diffusion_outputs( + request=_request("speak", num_outputs_per_prompt=2), + od_config=_config("audio_model"), + diffusion_output=DiffusionOutput(output=["waveform-0", "waveform-1"]), + output_data=["waveform-0", "waveform-1"], + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert result.images == [] + assert result.prompt == "speak" + assert result.final_output_type == "audio" + assert result.multimodal_output == {"audio": ["waveform-0", "waveform-1"]} + + +def test_formatter_preserves_single_prompt_audio_and_action_payloads( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False) + postprocess_output = normalize_diffusion_postprocess_output( + { + "video": ["frame-0", "frame-1"], + "audio": ["audio-0", "audio-1"], + "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "fps": 12.5, + } + ) + + results = format_diffusion_outputs( + request=_request("prompt-0"), + od_config=_config(), + diffusion_output=DiffusionOutput(output=None), + output_data={"raw": "output"}, + postprocess_output=postprocess_output, + timings=_timings(), + ) + + assert len(results) == 1 + assert results[0].images == ["frame-0", "frame-1"] + assert results[0].prompt == "prompt-0" + assert results[0].multimodal_output["metadata"] == {"video": {"fps": 12.5}} + assert results[0].multimodal_output["audio"] == ["audio-0", "audio-1"] + assert results[0].multimodal_output["fps"] == 12.5 + torch.testing.assert_close( + results[0].multimodal_output["actions"], + torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + ) + + +def test_format_empty_diffusion_outputs_preserves_empty_response_shape() -> None: + results = format_empty_diffusion_outputs(_request("prompt-0")) + + assert len(results) == 1 + assert [result.prompt for result in results] == ["prompt-0"] + assert [result.images for result in results] == [[]] + assert [result.metrics for result in results] == [{}] diff --git a/tests/diffusion/test_diffusion_output_metadata.py b/tests/diffusion/test_diffusion_output_metadata.py new file mode 100644 index 00000000000..ee3559f4755 --- /dev/null +++ b/tests/diffusion/test_diffusion_output_metadata.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest +import torch + +from vllm_omni.outputs.output_metadata import ( + strip_internal_metadata, + validate_diffusion_metadata, + validate_public_diffusion_metadata, +) + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] + + +def test_validate_accepts_canonical_metadata_fields() -> None: + validate_diffusion_metadata( + { + "video": { + "fps": 24.0, + "video_fps_multiplier": 2, + "shape": [1, 3, 16, 64, 64], + }, + "audio": { + "sample_rate": 48000, + "audiox_task": "tts", + }, + "actions": { + "raw_action_dim": 7, + "action_mode": "policy", + "domain_id": 3, + }, + "transfer": { + "controls": {"edge": torch.zeros(1, 3, 4, 8, 8)}, + "hints": ["edge"], + }, + "text": { + "text_output": "caption", + "think_text": "reasoning", + }, + "common": { + "action_only_output": True, + }, + } + ) + + +@pytest.mark.parametrize( + ("metadata", "error_type"), + [ + ({"video": {"fps": 0}}, ValueError), + ({"video": {"fps": "24"}}, TypeError), + ({"video": {"video_fps_multiplier": 0}}, ValueError), + ({"video": {"shape": [1, 0, 64]}}, ValueError), + ({"audio": {"sample_rate": 0}}, ValueError), + ({"audio": {"audiox_task": ""}}, ValueError), + ({"actions": {"raw_action_dim": 0}}, ValueError), + ({"actions": {"action_mode": ""}}, ValueError), + ({"actions": {"domain_id": "7"}}, TypeError), + ({"transfer": {"controls": {"edge": "not-a-tensor"}}}, TypeError), + ({"transfer": {"hints": ["edge", 1]}}, TypeError), + ({"text": {"text_output": 1}}, TypeError), + ({"text": {"think_text": 1}}, TypeError), + ({"common": {"action_only_output": "true"}}, TypeError), + ], +) +def test_validate_rejects_invalid_canonical_metadata(metadata: dict[str, object], error_type: type[Exception]) -> None: + with pytest.raises(error_type): + validate_diffusion_metadata(metadata) + + +def test_validate_allows_unknown_model_metadata() -> None: + validate_diffusion_metadata( + { + "model_specific": {"opaque": object()}, + "video": {"model_key": object()}, + } + ) + + +def test_public_metadata_rejects_internal_group() -> None: + with pytest.raises(ValueError, match="internal"): + validate_public_diffusion_metadata({"internal": {"secret": object()}}) + + +def test_strip_internal_metadata_removes_private_group() -> None: + metadata = { + "actions": {"action_mode": "policy"}, + "internal": {"robolab_action_postprocess": object()}, + } + + public_metadata = strip_internal_metadata(metadata) + + assert public_metadata == {"actions": {"action_mode": "policy"}} + assert "internal" not in public_metadata diff --git a/tests/diffusion/test_diffusion_plugin_hooks.py b/tests/diffusion/test_diffusion_plugin_hooks.py index 6804d550d53..a437a572a49 100644 --- a/tests/diffusion/test_diffusion_plugin_hooks.py +++ b/tests/diffusion/test_diffusion_plugin_hooks.py @@ -16,7 +16,6 @@ import pytest from vllm_omni.diffusion.registry import ( - _DIFFUSION_ACTION_POST_PROCESS_FUNCS, _DIFFUSION_IR_OP_PRIORITY_FUNCS, _DIFFUSION_MODELS, _DIFFUSION_POST_PROCESS_FUNCS, @@ -62,7 +61,6 @@ def cleanup_registry(self): original_models = _DIFFUSION_MODELS.copy() original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy() original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy() - original_action_post = _DIFFUSION_ACTION_POST_PROCESS_FUNCS.copy() original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy() yield _DIFFUSION_MODELS.clear() @@ -71,8 +69,6 @@ def cleanup_registry(self): _DIFFUSION_PRE_PROCESS_FUNCS.update(original_pre) _DIFFUSION_POST_PROCESS_FUNCS.clear() _DIFFUSION_POST_PROCESS_FUNCS.update(original_post) - _DIFFUSION_ACTION_POST_PROCESS_FUNCS.clear() - _DIFFUSION_ACTION_POST_PROCESS_FUNCS.update(original_action_post) _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear() _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority) @@ -84,7 +80,6 @@ def test_register_new_model(self): class_name="TestPipeline", pre_process_func_name="test_pre_process", post_process_func_name="test_post_process", - action_post_process_func_name="test_action_post_process", ir_op_priority_func_name="test_ir_op_priority", ) assert "TestPipeline" in _DIFFUSION_MODELS @@ -95,9 +90,21 @@ def test_register_new_model(self): ) assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process" assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process" - assert _DIFFUSION_ACTION_POST_PROCESS_FUNCS["TestPipeline"] == "test_action_post_process" assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority" + def test_register_model_accepts_deprecated_action_postprocess_keyword(self): + """Deprecated action hook keyword is accepted but not registered.""" + register_diffusion_model( + model_arch="LegacyActionPipeline", + module_name="test_plugin.diffusion.pipeline", + class_name="LegacyActionPipeline", + post_process_func_name="test_post_process", + action_post_process_func_name="test_action_post_process", + ) + + assert "LegacyActionPipeline" in _DIFFUSION_MODELS + assert _DIFFUSION_POST_PROCESS_FUNCS["LegacyActionPipeline"] == "test_post_process" + class TestWorkerUsesHook: """Test that DiffusionWorker resolves model runner via platform hook.""" diff --git a/tests/diffusion/test_diffusion_request.py b/tests/diffusion/test_diffusion_request.py index 9087a6e2f1c..ffaab3c530f 100644 --- a/tests/diffusion/test_diffusion_request.py +++ b/tests/diffusion/test_diffusion_request.py @@ -13,7 +13,7 @@ def _make_request() -> OmniDiffusionRequest: return OmniDiffusionRequest( - prompts=[{"prompt": "a cup of coffee on a table"}], + prompt={"prompt": "a cup of coffee on a table"}, sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_id="request-test", ) @@ -22,7 +22,7 @@ def _make_request() -> OmniDiffusionRequest: def test_request_id_is_required(): with pytest.raises(TypeError, match="request_id"): OmniDiffusionRequest( - prompts=[{"prompt": "a cup of coffee on a table"}], + prompt={"prompt": "a cup of coffee on a table"}, sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), ) @@ -37,7 +37,7 @@ def test_request_ids_identity_list_is_removed(): def test_request_id_must_be_non_empty(): with pytest.raises(ValueError, match="request_id must be a non-empty string"): OmniDiffusionRequest( - prompts=[{"prompt": "a cup of coffee on a table"}], + prompt={"prompt": "a cup of coffee on a table"}, sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_id="", ) @@ -45,7 +45,7 @@ def test_request_id_must_be_non_empty(): def test_dummy_run_request_is_identified_by_reserved_request_id(): req = OmniDiffusionRequest( - prompts=[{"prompt": "dummy run"}], + prompt={"prompt": "dummy run"}, sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_id=DUMMY_DIFFUSION_REQUEST_ID, ) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index f5a4f7f7717..7d6d2d11bff 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio import queue import threading from types import SimpleNamespace @@ -28,7 +29,7 @@ def _make_request(req_id: str) -> OmniDiffusionRequest: return OmniDiffusionRequest( - prompts=[f"prompt_{req_id}"], + prompt=f"prompt_{req_id}", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_id=req_id, ) @@ -66,7 +67,7 @@ def _make_step_request( sampling_params: OmniDiffusionSamplingParams | None = None, ) -> OmniDiffusionRequest: return OmniDiffusionRequest( - prompts=[f"prompt_{req_id}"], + prompt=f"prompt_{req_id}", sampling_params=sampling_params or OmniDiffusionSamplingParams( num_inference_steps=num_inference_steps, @@ -127,6 +128,12 @@ def update_from_output(self, sched_output, output) -> set[str]: def has_requests(self) -> bool: return not self._scheduled + def num_waiting_requests(self) -> int: + return 0 if self._scheduled else 1 + + def num_running_requests(self) -> int: + return 1 if self._scheduled else 0 + def get_request_state(self, request_id: str): del request_id return self._state @@ -163,7 +170,7 @@ def _make(lora_int_id: int | None = None, lora_scale: float = 1.0) -> OmniDiffus ) sp.lora_scale = lora_scale return OmniDiffusionRequest( - prompts=["prompt"], + prompt="prompt", sampling_params=sp, request_id=f"req-{lora_int_id}-{lora_scale}", ) @@ -195,6 +202,41 @@ def test_equal_for_same_lora_identity(self) -> None: assert a == b +class TestGetRequestBatchSamplingParamsKey: + """Pure-function tests for the request-batch compatibility key builder.""" + + @staticmethod + def _make( + *, + num_inference_steps: int = 2, + seed: int | None = 123, + generator: torch.Generator | None = None, + ) -> OmniDiffusionRequest: + sp = OmniDiffusionSamplingParams( + num_inference_steps=num_inference_steps, + seed=seed, + generator=generator, + ) + return OmniDiffusionRequest(prompt="prompt", sampling_params=sp, request_id=f"req-{num_inference_steps}") + + def test_distinguishes_num_inference_steps(self) -> None: + from vllm_omni.diffusion.sched.base_scheduler import get_request_batch_sampling_params_key + + assert get_request_batch_sampling_params_key( + self._make(num_inference_steps=2) + ) != get_request_batch_sampling_params_key(self._make(num_inference_steps=4)) + + def test_ignores_seed_and_generator(self) -> None: + from vllm_omni.diffusion.sched.base_scheduler import get_request_batch_sampling_params_key + + gen_a = torch.Generator(device="cpu").manual_seed(1) + gen_b = torch.Generator(device="cpu").manual_seed(2) + + assert get_request_batch_sampling_params_key( + self._make(seed=1, generator=gen_a) + ) == get_request_batch_sampling_params_key(self._make(seed=2, generator=gen_b)) + + class TestRequestScheduler: def setup_method(self) -> None: self.scheduler: RequestScheduler = RequestScheduler() @@ -238,6 +280,35 @@ def test_empty_output_without_error_marks_completed(self) -> None: assert finished == {req_id} assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED + def test_streaming_output_keeps_request_running_until_final_chunk(self) -> None: + scheduler = StepScheduler() + scheduler.initialize(SimpleNamespace()) + req_id = scheduler.add_request(_make_request("stream")) + + sched_output = scheduler.schedule() + chunk = RunnerOutput( + request_id=req_id, + step_index=1, + finished=False, + result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), + ) + finished = scheduler.update_from_output(sched_output, chunk) + + assert finished == set() + assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING + assert scheduler.has_requests() is True + + final_chunk = RunnerOutput( + request_id=req_id, + step_index=2, + finished=True, + result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), + ) + finished = scheduler.update_from_output(sched_output, final_chunk) + + assert finished == {req_id} + assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED + def test_fifo_single_request_scheduling(self) -> None: req_id_a = self.scheduler.add_request(_make_request("a")) req_id_b = self.scheduler.add_request(_make_request("b")) @@ -267,8 +338,18 @@ def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=2)) - req_id_a = scheduler.add_request(_make_request("a")) - req_id_b = scheduler.add_request(_make_request("b")) + req_id_a = scheduler.add_request( + _make_step_request( + "a", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1, seed=123), + ) + ) + req_id_b = scheduler.add_request( + _make_step_request( + "b", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1, seed=123), + ) + ) sched_output = scheduler.schedule() @@ -276,6 +357,50 @@ def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None: assert sched_output.num_running_reqs == 2 assert sched_output.num_waiting_reqs == 0 + def test_batches_incompatible_request_sampling_params_separately(self) -> None: + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + + req_id_a = scheduler.add_request( + _make_step_request( + "a", num_inference_steps=2, sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123) + ) + ) + scheduler.add_request( + _make_step_request( + "b", num_inference_steps=4, sampling_params=OmniDiffusionSamplingParams(num_inference_steps=4, seed=123) + ) + ) + + first = scheduler.schedule() + + assert _new_ids(first) == [req_id_a] + assert first.num_running_reqs == 1 + assert first.num_waiting_reqs == 1 + + def test_batches_different_request_local_seed_together(self) -> None: + scheduler = RequestScheduler() + scheduler.initialize(SimpleNamespace(max_num_seqs=2)) + + req_id_a = scheduler.add_request( + _make_step_request( + "a", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=123), + ) + ) + req_id_b = scheduler.add_request( + _make_step_request( + "b", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, seed=456), + ) + ) + + first = scheduler.schedule() + + assert _new_ids(first) == [req_id_a, req_id_b] + assert first.num_running_reqs == 2 + assert first.num_waiting_reqs == 0 + def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None: scheduler = RequestScheduler() scheduler.initialize(SimpleNamespace(max_num_seqs=3)) @@ -283,7 +408,7 @@ def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None req_id_a = scheduler.add_request(_make_request("a")) req_id_b = scheduler.add_request( OmniDiffusionRequest( - prompts=["prompt_b"], + prompt="prompt_b", sampling_params=OmniDiffusionSamplingParams(width=768), request_id="b", ) @@ -344,7 +469,7 @@ def test_has_requests_state_transition(self) -> None: def test_request_id_is_scheduler_key(self) -> None: request = OmniDiffusionRequest( - prompts=["prompt_map_a", "prompt_map_b"], + prompt="prompt_map_a", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), request_id="map-parent", ) @@ -369,6 +494,7 @@ def test_duplicate_request_id_is_rejected(self) -> None: class TestDiffusionEngine: def test_add_req_and_wait_for_response_single_path(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(streaming_output=False) engine.scheduler = RequestScheduler() engine.scheduler.initialize(SimpleNamespace()) engine._rpc_lock = threading.RLock() @@ -391,6 +517,7 @@ def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> scheduler = _StubScheduler(request, runner_output) engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(streaming_output=False) engine.scheduler = scheduler engine._rpc_lock = threading.RLock() engine._cv = threading.Condition(engine._rpc_lock) @@ -410,7 +537,7 @@ def test_initializes_injected_scheduler( ) -> None: request = _make_request("init") scheduler = _StubScheduler(request, DiffusionOutput(output=None)) - od_config = SimpleNamespace(model_class_name="mock_model") + od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False) fake_executor_cls = mocker.Mock(return_value=mocker.Mock()) monkeypatch.setattr( @@ -472,7 +599,7 @@ def test_abort_queue_marks_request_finished_aborted(self) -> None: def test_finalize_finished_request_returns_aborted_output(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) - engine.scheduler = RequestScheduler() + engine.scheduler = StepScheduler() engine.scheduler.initialize(SimpleNamespace()) req_id = engine.scheduler.add_request(_make_request("req-finalize")) @@ -483,12 +610,77 @@ def test_finalize_finished_request_returns_aborted_output(self) -> None: assert output.aborted is True assert output.abort_message == "Request req-finalize aborted." + @pytest.mark.asyncio + async def test_streaming_runner_output_notifies_each_chunk(self) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = StepScheduler() + engine.scheduler.initialize(SimpleNamespace()) + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._out_queue_streaming = {} + engine.main_loop = asyncio.get_running_loop() + + req_id = engine.scheduler.add_request(_make_request("stream-engine")) + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_queue_streaming[req_id] = queue + sched_output = engine.scheduler.schedule() + + chunk = RunnerOutput( + request_id=req_id, + step_index=1, + finished=False, + result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2), + ) + finished_req_ids = engine.scheduler.update_from_output(sched_output, chunk) + engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, chunk) + + notified_chunk = await asyncio.wait_for(queue.get(), timeout=1) + assert notified_chunk.output == "chunk-0" + assert notified_chunk.finished is False + assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING + + final_chunk = RunnerOutput( + request_id=req_id, + step_index=2, + finished=True, + result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2), + ) + finished_req_ids = engine.scheduler.update_from_output(sched_output, final_chunk) + engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, final_chunk) + + notified_final = await asyncio.wait_for(queue.get(), timeout=1) + assert notified_final.output == "chunk-1" + assert notified_final.finished is True + assert engine.scheduler.get_request_state(req_id) is None + + @pytest.mark.asyncio + async def test_finished_streaming_request_without_runner_output_notifies_waiter(self) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.scheduler = RequestScheduler() + engine.scheduler.initialize(SimpleNamespace()) + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._out_queue_streaming = {} + engine.main_loop = asyncio.get_running_loop() + + req_id = engine.scheduler.add_request(_make_request("stream-abort")) + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + engine._out_queue_streaming[req_id] = queue + engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED) + + engine._handle_empty_streaming_requests({req_id}) + + output = await asyncio.wait_for(queue.get(), timeout=1) + assert output.aborted is True + assert output.finished is True + assert engine.scheduler.get_request_state(req_id) is None + def test_initializes_step_scheduler_when_step_execution_enabled( self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, ) -> None: - od_config = SimpleNamespace(model_class_name="mock_model") + od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False) od_config.step_execution = True fake_executor = mocker.Mock() fake_executor_cls = mocker.Mock(return_value=fake_executor) @@ -521,124 +713,6 @@ def test_dummy_run_raises_on_output_error(self, mocker: MockerFixture) -> None: with pytest.raises(RuntimeError, match="Dummy run failed: boom"): engine._dummy_run() - @pytest.mark.asyncio - async def test_step_multi_request_reuses_multimodal_slice_logic(self, mocker: MockerFixture) -> None: - engine = DiffusionEngine.__new__(DiffusionEngine) - engine.od_config = SimpleNamespace( - model_class_name="mock_model", - enable_cpu_offload=False, - ) - engine.pre_process_func = None - engine.post_process_func = None - engine._check_and_start_background_loop = mocker.AsyncMock() - engine.async_add_req_and_wait_for_response = mocker.AsyncMock( - return_value=DiffusionOutput( - output={ - "video": ["frame-0", "frame-1"], - "audio": ["audio-0", "audio-1"], - "actions": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), - } - ) - ) - - request = OmniDiffusionRequest( - prompts=["prompt-0", "prompt-1"], - sampling_params=OmniDiffusionSamplingParams( - num_inference_steps=1, - num_outputs_per_prompt=1, - ), - request_id="req-batch", - ) - - mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False) - outputs = await engine.step(request) - - assert len(outputs) == 2 - assert outputs[0].images == ["frame-0"] - assert outputs[1].images == ["frame-1"] - assert outputs[0].multimodal_output["audio"] == "audio-0" - assert outputs[1].multimodal_output["audio"] == "audio-1" - torch.testing.assert_close( - outputs[0].multimodal_output["actions"], - torch.tensor([1.0, 2.0]), - ) - torch.testing.assert_close( - outputs[1].multimodal_output["actions"], - torch.tensor([3.0, 4.0]), - ) - - @pytest.mark.asyncio - async def test_step_empty_dict_output_still_runs_postprocess(self, mocker: MockerFixture) -> None: - engine = DiffusionEngine.__new__(DiffusionEngine) - engine.od_config = SimpleNamespace( - model_class_name="mock_model", - enable_cpu_offload=False, - ) - engine.pre_process_func = None - engine.post_process_func = mocker.Mock(return_value={"video": ["processed"]}) - engine._post_process_accepts_sampling_params = False - engine._check_and_start_background_loop = mocker.AsyncMock() - engine.async_add_req_and_wait_for_response = mocker.AsyncMock( - return_value=DiffusionOutput( - output={}, - custom_output={"actions": torch.tensor([[1.0, 2.0]])}, - ) - ) - - request = OmniDiffusionRequest( - prompts=["prompt"], - sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), - request_id="req-action", - ) - - mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False) - outputs = await engine.step(request) - - engine.post_process_func.assert_called_once_with({}) - assert outputs[0].images == ["processed"] - torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[1.0, 2.0]])) - - @pytest.mark.asyncio - async def test_step_action_only_flag_skips_postprocess(self, mocker: MockerFixture) -> None: - engine = DiffusionEngine.__new__(DiffusionEngine) - engine.od_config = SimpleNamespace( - model_class_name="mock_model", - enable_cpu_offload=False, - ) - engine.pre_process_func = None - engine.post_process_func = mocker.Mock(side_effect=AssertionError("postprocess should be skipped")) - engine.action_post_process_func = mocker.Mock(return_value=torch.tensor([[3.0, 4.0]])) - engine._post_process_accepts_sampling_params = False - engine._action_post_process_accepts_custom_output = True - engine._action_post_process_accepts_sampling_params = False - engine._check_and_start_background_loop = mocker.AsyncMock() - raw_action = torch.tensor([[1.0, 2.0]]) - engine.async_add_req_and_wait_for_response = mocker.AsyncMock( - return_value=DiffusionOutput( - output={}, - custom_output={ - "action": raw_action, - "action_only_output": True, - }, - ) - ) - - request = OmniDiffusionRequest( - prompts=["prompt"], - sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), - request_id="req-action", - ) - - mocker.patch("vllm_omni.diffusion.diffusion_engine.supports_audio_output", return_value=False) - outputs = await engine.step(request) - - engine.post_process_func.assert_not_called() - engine.action_post_process_func.assert_called_once() - assert engine.action_post_process_func.call_args.args[0] is raw_action - assert "custom_output" in engine.action_post_process_func.call_args.kwargs - assert outputs[0].images == [] - torch.testing.assert_close(outputs[0].multimodal_output["actions"], torch.tensor([[3.0, 4.0]])) - class TestStepScheduler: def setup_method(self) -> None: diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index a21813d28c6..7e18996ede1 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -30,6 +30,9 @@ pack_diffusion_output_shm, unpack_diffusion_output_shm, ) +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, +) from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import StepScheduler from vllm_omni.diffusion.sched.interface import ( @@ -97,6 +100,51 @@ def post_decode(self, state, **kwargs): return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32)) +class _ProfilingStepPipeline(_StepPipeline): + enable_diffusion_pipeline_profiler = True + + def __init__(self): + super().__init__() + self._stage_durations: dict[str, float] = {} + + @property + def stage_durations(self) -> dict[str, float]: + return dict(self._stage_durations) + + def clear_profiler_records(self) -> None: + self._stage_durations.clear() + + def prepare_encode(self, state, **kwargs): + result = super().prepare_encode(state, **kwargs) + self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0 + return result + + def denoise_step(self, input_batch, **kwargs): + result = super().denoise_step(input_batch, **kwargs) + self._stage_durations["QwenImagePipeline.diffuse"] = 2.0 + return result + + def post_decode(self, state, **kwargs): + result = super().post_decode(state, **kwargs) + self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0 + return result + + +class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin): + _PROFILER_TARGETS: list[str] = [] + + def __init__(self): + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=True, + ) + + def forward(self): + return None + + def denoise_step(self): + return "ok" + + class _InterruptingStepPipeline(_StepPipeline): interrupt = True @@ -114,6 +162,26 @@ def post_decode(self, state, **kwargs): raise AssertionError("post_decode should not run after interrupt") +class _FakePeakMemoryPlatform: + def __init__(self, reserved_mb: list[float]): + self._reserved_mb = reserved_mb + self.reset_calls = 0 + + def reset_peak_memory_stats(self): + self.reset_calls += 1 + + def max_memory_reserved(self): + index = min(self.reset_calls - 1, len(self._reserved_mb) - 1) + return int(self._reserved_mb[index] * 1024**2) + + def max_memory_allocated(self): + index = min(self.reset_calls - 1, len(self._reserved_mb) - 1) + return int((self._reserved_mb[index] - 100) * 1024**2) + + def is_available(self) -> bool: + return True + + class _IdentityNoiseTransformer(torch.nn.Module): def forward(self, x: torch.Tensor, **kwargs): del kwargs @@ -205,15 +273,10 @@ def post_decode(self, state, **kwargs): def _make_step_request(num_inference_steps: int = 2): - return SimpleNamespace( - prompts=["a prompt"], + return OmniDiffusionRequest( + prompt="a prompt", request_id="req-1", - sampling_params=SimpleNamespace( - generator=None, - seed=None, - generator_device=None, - num_inference_steps=num_inference_steps, - ), + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps), ) @@ -226,7 +289,7 @@ def _assert_aborted_output(output: DiffusionOutput, request_id: str) -> None: def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) -> OmniDiffusionRequest: return OmniDiffusionRequest( - prompts=[f"prompt-{req_id}"], + prompt=f"prompt-{req_id}", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps), request_id=req_id, ) @@ -249,6 +312,7 @@ def _make_runner(): runner.od_config = SimpleNamespace( cache_backend=None, parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.device = torch.device("cpu") runner.pipeline = _StepPipeline() @@ -267,6 +331,7 @@ def _make_distributed_runner(mode: str, device: torch.device): runner.od_config = SimpleNamespace( cache_backend=None, parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.device = device runner.pipeline = _DistributedStepPipeline(mode=mode, device=device) @@ -280,6 +345,7 @@ def _make_distributed_runner(mode: str, device: torch.device): def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None): + req.request_id = request_id return DiffusionSchedulerOutput( step_id=step_id, scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)], @@ -307,7 +373,7 @@ def _make_input_batch_state(request_id: str, latent_value: float) -> DiffusionRe state = DiffusionRequestState( request_id=request_id, sampling=SimpleNamespace(), - prompts=None, + prompt=None, ) state.latents = torch.tensor([[latent_value]]) state.timesteps = torch.tensor([1.0]) @@ -327,7 +393,7 @@ def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_id def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine: engine = object.__new__(DiffusionEngine) - engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline") + engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline", streaming_output=False) engine.pre_process_func = None engine.post_process_func = None engine.scheduler = scheduler @@ -422,6 +488,16 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3)) +@pytest.mark.cpu +def test_step_profiler_reports_denoise_step_as_diffuse(): + pipeline = _AutoDenoiseProfilerPipeline() + + assert pipeline.denoise_step() == "ok" + + assert any(key.endswith(".diffuse") for key in pipeline.stage_durations) + assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations) + + @pytest.mark.cpu class TestRunner: """DiffusionModelRunner.execute_stepwise""" @@ -454,6 +530,71 @@ def test_completes_request_and_clears_state(self, monkeypatch): assert runner.pipeline.scheduler_calls == 2 assert runner.pipeline.decode_calls == 1 + def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch): + runner = _make_runner() + runner.pipeline = _ProfilingStepPipeline() + req = _make_step_request() + reset_calls = [] + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "is_available", + lambda: True, + ) + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "reset_peak_memory_stats", + lambda: reset_calls.append(True), + ) + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "max_memory_reserved", + lambda: 2 * 1024**2, + ) + monkeypatch.setattr( + model_runner_module.current_omni_platform, + "max_memory_allocated", + lambda: 1024**2, + ) + + DiffusionModelRunner.execute_stepwise( + runner, + _make_scheduler_output(req, step_id=0), + ) + result = DiffusionModelRunner.execute_stepwise( + runner, + _make_cached_scheduler_output(step_id=1), + ) + + output = result.get_request_output("req-1") + assert output.finished is True + assert output.result is not None + assert output.result.peak_memory_mb == 2 + assert output.result.stage_durations == { + "QwenImagePipeline.text_encoder.forward": 1.0, + "QwenImagePipeline.diffuse": 4.0, + "QwenImagePipeline.vae.decode": 3.0, + } + assert reset_calls == [True] + + def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch): + runner = _make_runner() + req = _make_step_request() + fake_platform = _FakePeakMemoryPlatform([1500.0, 1200.0]) + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + monkeypatch.setattr(model_runner_module, "current_omni_platform", fake_platform) + + first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) + first_output = first.get_request_output("req-1") + assert first_output.finished is False + assert first_output.result is None + + second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1)) + second_output = second.get_request_output("req-1") + assert second_output.finished is True + assert second_output.result is not None + assert second_output.result.peak_memory_mb == pytest.approx(1500.0) + def test_rejects_multi_request_step_batch(self): runner = _make_runner() req_1 = _make_step_request() @@ -563,6 +704,7 @@ def __exit__(self, exc_type, exc, tb): step_execution=True, model_class_name="RequestOnlyPipeline", parallel_config=SimpleNamespace(use_hsdp=False), + streaming_output=False, ) runner.device = torch.device("cpu") runner.pipeline = None @@ -710,6 +852,7 @@ class TestExecutor: def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture): executor = object.__new__(MultiprocDiffusionExecutor) + executor.od_config = SimpleNamespace(streaming_output=False) executor._ensure_open = lambda: None expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None) executor.collective_rpc = mocker.Mock(return_value=expected) diff --git a/tests/diffusion/test_diffusion_streaming_output.py b/tests/diffusion/test_diffusion_streaming_output.py new file mode 100644 index 00000000000..ae5f9eee8bc --- /dev/null +++ b/tests/diffusion/test_diffusion_streaming_output.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Integration tests for diffusion pipeline streaming output mode. + +Covers: +- Mock pipeline -> StageDiffusionClient (ZMQ subprocess path) +- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni +- Mock pipeline -> InlineStageDiffusionClient -> Orchestrator -> AsyncOmni -> `/v1/realtime/video` +""" + +import asyncio +import contextlib +import queue +import time +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +import zmq +from fastapi import FastAPI, WebSocket +from pytest_mock import MockerFixture +from starlette.testclient import TestClient +from vllm.utils.network_utils import get_open_zmq_ipc_path + +import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module +from tests.engine.test_orchestrator import OrchestratorFixture, _build_harness, _wait_for +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient +from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient +from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.distributed.omni_connectors.utils.serialization import ( + OmniMsgpackDecoder, + OmniMsgpackEncoder, +) +from vllm_omni.engine.async_omni_engine import StageRuntimeInfo +from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni +from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +class _StepStreamingPipeline: + supports_step_execution = True + + def __init__(self, outputs: list[DiffusionOutput]) -> None: + self.outputs = outputs + self.requests: list[Any] = [] + + def step_outputs(self, request): + self.requests.append(request) + return list(self.outputs) + + +class _FailingStreamingPipeline: + supports_step_execution = True + + def __init__(self) -> None: + self.requests: list[Any] = [] + + def step_outputs(self, request): + self.requests.append(request) + raise RuntimeError("stream failed midway") + + +class _PipelineBackedEngine: + """DiffusionEngine stand-in that exposes step-streaming chunks from a mock pipeline.""" + + def __init__(self, pipeline: _StepStreamingPipeline | _FailingStreamingPipeline) -> None: + self.pipeline = pipeline + self.executor = SimpleNamespace( + register_failure_callback=MagicMock(), + check_health=MagicMock(), + ) + + async def step_streaming(self, request): + try: + outputs = self.pipeline.step_outputs(request) + except Exception as exc: + outputs = [ + _streaming_diffusion_output(chunk=0, finished=False), + DiffusionOutput(error=str(exc), finished=True), + ] + for output in outputs: + if output.error is not None: + yield [OmniRequestOutput.from_error(request_id=request.request_id, error_message=output.error)] + continue + payload, metadata = _streaming_payload_and_metadata(output) + stream_metadata = metadata.get("stream", {}) + yield [ + OmniRequestOutput.from_diffusion( + request_id=request.request_id, + images=list(payload.get("image") or []), + multimodal_output={"metadata": {"stream": stream_metadata}}, + finished=output.finished, + ) + ] + + def abort(self, request_id: str) -> None: + del request_id + + +def _streaming_diffusion_output( + *, + chunk: int, + finished: bool, + images: list | None = None, +) -> DiffusionOutput: + payload = {"image": images or []} + metadata = {"stream": {"chunk": chunk}} + return DiffusionOutput( + output={ + "payload": payload, + "metadata": metadata, + }, + finished=finished, + ) + + +def _streaming_payload_and_metadata(output: DiffusionOutput) -> tuple[dict, dict]: + envelope = output.output if isinstance(output.output, dict) else {} + payload = envelope.get("payload") if isinstance(envelope.get("payload"), dict) else {} + metadata = envelope.get("metadata") if isinstance(envelope.get("metadata"), dict) else {} + return payload, metadata + + +class TestPipelineStreamingOutputToStageDiffusionClient: + """Streaming pipeline output over ZMQ into ``StageDiffusionClient``.""" + + @pytest.mark.asyncio + async def test_streaming_output_reaches_stage_client_from_mock_pipeline(self) -> None: + """Mock pipeline chunks over ZMQ reach StageDiffusionClient with correct finished flags.""" + pipeline = _StepStreamingPipeline( + [ + _streaming_diffusion_output(chunk=0, finished=False), + _streaming_diffusion_output(chunk=1, finished=True), + ] + ) + + outputs = await self._run_streaming_client_proc(pipeline, request_id="req-stream") + + assert [output.request_id for output in outputs] == ["req-stream", "req-stream"] + assert [output.finished for output in outputs] == [False, True] + assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1] + assert pipeline.requests[0].request_id == "req-stream" + + @pytest.mark.asyncio + async def test_streaming_midway_error_reaches_stage_client_after_prior_chunk(self) -> None: + """A pipeline error after the first chunk still delivers that chunk then a terminal error output.""" + pipeline = _FailingStreamingPipeline() + + outputs = await self._run_streaming_client_proc(pipeline, request_id="req-error") + + assert len(outputs) == 2 + assert outputs[0].request_id == "req-error" + assert outputs[0].finished is False + assert outputs[0].multimodal_output["metadata"]["stream"]["chunk"] == 0 + assert outputs[1].request_id == "req-error" + assert outputs[1].finished is True + assert outputs[1].error is not None + assert "stream failed midway" in outputs[1].error + + @classmethod + async def _run_streaming_client_proc( + cls, + pipeline: _StepStreamingPipeline | _FailingStreamingPipeline, + *, + request_id: str, + ) -> list[OmniRequestOutput]: + request_address = get_open_zmq_ipc_path() + response_address = get_open_zmq_ipc_path() + proc = object.__new__(StageDiffusionProc) + proc._od_config = SimpleNamespace(streaming_output=True) + proc._engine = _PipelineBackedEngine(pipeline) + proc._closed = False + client = cls._make_client(request_address, response_address) + proc_task = asyncio.create_task(proc.run_loop(request_address, response_address)) + try: + await asyncio.sleep(0.05) + await client.add_request_async( + request_id, + "prompt", + OmniDiffusionSamplingParams(), + ) + deadline = asyncio.get_running_loop().time() + 2.0 + outputs: list[OmniRequestOutput] = [] + while (not outputs or not outputs[-1].finished) and asyncio.get_running_loop().time() < deadline: + output = client.get_diffusion_output_nowait() + if output is not None: + outputs.append(output) + else: + await asyncio.sleep(0.01) + return outputs + finally: + try: + client._request_socket.send(client._encoder.encode({"type": "shutdown"})) + except Exception: + pass + try: + await asyncio.wait_for(proc_task, timeout=2) + finally: + cls._close_client(client) + + @staticmethod + def _make_client(request_address: str, response_address: str) -> StageDiffusionClient: + client = object.__new__(StageDiffusionClient) + client.stage_id = 0 + client.replica_id = 0 + client.final_output = True + client.final_output_type = "image" + client.default_sampling_params = {} + client.requires_multimodal_data = False + client.custom_process_input_func = None + client.engine_input_source = [] + client._proc = None + client._proc_manager = None + client._owns_process = False + client._zmq_ctx = zmq.Context() + client._request_socket = client._zmq_ctx.socket(zmq.PUSH) + client._request_socket.bind(request_address) + client._response_socket = client._zmq_ctx.socket(zmq.PULL) + client._response_socket.bind(response_address) + client._encoder = OmniMsgpackEncoder() + client._decoder = OmniMsgpackDecoder() + client._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() + client._rpc_results = {} + client._pending_rpcs = set() + client._tasks = {} + client._shutting_down = False + client._engine_dead = False + return client + + @staticmethod + def _close_client(client: StageDiffusionClient) -> None: + client._request_socket.close(linger=0) + client._response_socket.close(linger=0) + client._zmq_ctx.term() + + +class TestPipelineStreamingOutputToEntrypoint: + """Streaming pipeline output through inline stage, orchestrator, and entrypoints.""" + + @pytest.mark.asyncio + async def test_streaming_output_reaches_async_omni_from_mock_pipeline(self) -> None: + """Mock pipeline streaming chunks reach AsyncOmni.generate() via inline stage and orchestrator.""" + pipeline = _StepStreamingPipeline( + [ + _streaming_diffusion_output(chunk=0, finished=False), + _streaming_diffusion_output(chunk=1, finished=True), + ] + ) + inline_client = self._make_inline_pipeline_client(pipeline) + fixture = _build_harness([inline_client]) + omni = self._make_async_omni(self._OrchestratorBridgeEngine(fixture)) + + try: + outputs: list[OmniRequestOutput] = [] + async for output in omni.generate( + prompt={"prompt": "a cat"}, + request_id="req-omni", + sampling_params_list=[OmniDiffusionSamplingParams()], + output_modalities=["image"], + ): + outputs.append(output) + + await _wait_for(lambda: len(pipeline.requests) == 1) + assert pipeline.requests[0].request_id.startswith("req-omni-") + assert [output.multimodal_output["metadata"]["stream"]["chunk"] for output in outputs] == [0, 1] + assert [getattr(output.request_output, "finished", output.finished) for output in outputs] == [False, True] + finally: + await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client) + + @pytest.mark.asyncio + async def test_streaming_output_reaches_api_websocket_from_mock_pipeline(self, mocker: MockerFixture) -> None: + """Mock pipeline streaming chunks reach the video WebSocket handler through AsyncOmni.generate().""" + frames = [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(2)] + pipeline = _StepStreamingPipeline( + [ + _streaming_diffusion_output(chunk=0, images=[frames], finished=False), + _streaming_diffusion_output(chunk=1, images=[frames], finished=True), + ] + ) + inline_client = self._make_inline_pipeline_client(pipeline) + fixture = _build_harness([inline_client]) + bridge_engine = self._OrchestratorBridgeEngine(fixture) + omni = self._make_async_omni(bridge_engine) + + encoded = [b"pipeline-fmp4-0", b"pipeline-fmp4-1"] + + class FakeStreamingVideoEncoder: + encode_calls: list[int] = [] + + def encode(self, video): + del video + idx = len(self.encode_calls) + self.encode_calls.append(idx) + return encoded[idx] + + def close(self): + return b"" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", + lambda *, output_format, fps, video_codec_options=None: FakeStreamingVideoEncoder(), + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", + return_value="diffusion", + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + + handler = OmniStreamingVideoOutputHandler( + engine_client=omni, + model_name="test-model", + stage_configs=bridge_engine.stage_configs, + stall_timeout=5.0, + start_timeout=5.0, + ) + app = FastAPI() + + @app.websocket("/v1/realtime/video") + async def ws_endpoint(websocket: WebSocket): + await handler.handle_session(websocket) + + try: + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "integration test"}) + assert ws.receive_json()["type"] == "video.start" + assert ws.receive_bytes() == b"pipeline-fmp4-0" + assert ws.receive_bytes() == b"pipeline-fmp4-1" + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["chunks"] == 2 + + await _wait_for(lambda: len(pipeline.requests) == 1) + assert pipeline.requests[0].request_id + finally: + await self._shutdown_pipeline_omni_harness(omni, fixture, inline_client) + + class _OrchestratorBridgeEngine: + """Minimal AsyncOmni engine facade backed by a live Orchestrator harness.""" + + def __init__(self, fixture: OrchestratorFixture) -> None: + self._fixture = fixture + self.stage_metadata = [ + StageRuntimeInfo( + stage_type="diffusion", + final_output=True, + final_output_type="image", + ) + ] + self.stage_configs: list[Any] = [SimpleNamespace(stage_type="diffusion")] + self.default_sampling_params_list = [OmniDiffusionSamplingParams()] + self.num_stages = 1 + self.supported_tasks = ("generate",) + self._alive = True + + async def add_request_async( + self, + *, + request_id: str, + prompt: Any, + sampling_params_list: list[Any], + final_stage_id: int, + **kwargs: Any, + ) -> None: + self._fixture.request_sync_q.put_nowait( + StageSubmissionMessage( + type="add_request", + request_id=request_id, + prompt=prompt, + original_prompt=prompt, + output_prompt_text=None, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=kwargs.get("final_output_stage_ids"), + preprocess_ms=0.0, + request_timestamp=kwargs.get("arrival_time", time.time()), + enqueue_ts=time.perf_counter(), + ) + ) + + async def try_get_output_async(self) -> Any | None: + try: + return self._fixture.output_sync_q.get_nowait() + except queue.Empty: + return None + + def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: + return self.stage_metadata[stage_id] + + def is_alive(self) -> bool: + return self._fixture.thread.is_alive() + + async def abort_async(self, request_ids: list[str]) -> None: + del request_ids + + @classmethod + def _make_inline_pipeline_client(cls, pipeline: _StepStreamingPipeline) -> InlineStageDiffusionClient: + """Build an inline diffusion stage client whose engine runs the given mock pipeline.""" + metadata = StageMetadata( + stage_id=0, + stage_type="diffusion", + engine_output_type="image", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source=[], + final_output=True, + final_output_type="image", + default_sampling_params=OmniDiffusionSamplingParams(), + custom_process_input_func=None, + model_stage=None, + runtime_cfg=None, + ) + pipeline_engine = _PipelineBackedEngine(pipeline) + with patch.object(InlineStageDiffusionClient, "_enrich_config"): + with patch( + "vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine.make_engine", + return_value=pipeline_engine, + ): + od_config = MagicMock(spec=OmniDiffusionConfig) + od_config.streaming_output = True + return InlineStageDiffusionClient( + model="test_model", + od_config=od_config, + metadata=metadata, + batch_size=1, + ) + + @staticmethod + def _make_async_omni(engine: _OrchestratorBridgeEngine) -> AsyncOmni: + omni = object.__new__(AsyncOmni) + omni.engine = engine + omni.log_stats = False + omni._pause_cond = asyncio.Condition() + omni._paused = False + omni.request_states = {} + omni.final_output_task = None + omni.event_resolver = AsyncEventResolver() + omni._enable_ar_profiler = False + omni._is_sleeping = False + omni.prom_metrics = MagicMock() + omni.mod_metrics = MagicMock() + omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params + omni._compute_final_stage_id = lambda output_modalities: 0 + omni._compute_final_output_stage_ids = lambda output_modalities: [0] + omni.default_sampling_params_list = engine.default_sampling_params_list + omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None) + return omni + + @staticmethod + async def _shutdown_pipeline_omni_harness( + omni: AsyncOmni, + fixture: OrchestratorFixture, + inline_client: InlineStageDiffusionClient, + ) -> None: + if omni.final_output_task is not None: + omni.final_output_task.cancel() + await asyncio.gather(omni.final_output_task, return_exceptions=True) + inline_client.shutdown() + fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + await asyncio.to_thread(fixture.thread.join, 5) + + +def _make_vllm_config(): + @contextlib.contextmanager + def set_priority(*args, **kwargs): + del args, kwargs + yield + + return SimpleNamespace( + kernel_config=SimpleNamespace(ir_op_priority=SimpleNamespace(set_priority=set_priority)), + compilation_config=SimpleNamespace(ir_enable_torch_wrap=True), + ) + + +class TestSupportedPipelines: + """Streaming-output protocol checks for supported pipelines.""" + + def test_helios_supports_step_execution_for_streaming_output(self) -> None: + from vllm_omni.diffusion.models.helios.pipeline_helios import HeliosPipeline + from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution + + # Avoid loading model weights; protocol membership depends on the class contract. + pipeline = object.__new__(HeliosPipeline) + + assert pipeline.supports_step_execution is True + assert supports_step_execution(pipeline) is True + assert isinstance(pipeline, SupportsStepExecution) is True + + def test_load_model_rejects_streaming_output_without_step_execution(self, monkeypatch) -> None: + class _NoStepPipeline: + def forward(self): ... + + class _FakeDiffusersPipelineLoader: + def __init__(self, *args, **kwargs): ... + + def load_model(self, **kwargs): + return _NoStepPipeline() + + class _FakeDeviceMemoryProfiler: + consumed_memory = 0 + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + runner = object.__new__(DiffusionModelRunner) + runner.vllm_config = _make_vllm_config() + runner.od_config = SimpleNamespace( + enable_cpu_offload=False, + enable_layerwise_offload=False, + enforce_eager=True, + cache_backend=None, + cache_config=None, + streaming_output=True, + step_execution=False, + model_class_name="NoStepPipeline", + parallel_config=SimpleNamespace(use_hsdp=False), + ) + runner.device = torch.device("cpu") + runner.pipeline = None + runner.cache_backend = None + runner.offload_backend = None + runner.kv_transfer_manager = SimpleNamespace() + + monkeypatch.setattr(model_runner_module, "DiffusersPipelineLoader", _FakeDiffusersPipelineLoader) + monkeypatch.setattr(model_runner_module, "DeviceMemoryProfiler", _FakeDeviceMemoryProfiler) + monkeypatch.setattr(model_runner_module, "get_offload_backend", lambda *args, **kwargs: None) + monkeypatch.setattr(model_runner_module, "get_cache_backend", lambda *args, **kwargs: None) + + with pytest.raises(ValueError, match="NoStepPipeline"): + DiffusionModelRunner.load_model(runner) diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 8ad6c273416..fb45cd41d0d 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -23,6 +23,15 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] +def patch_cumem_allocator(mocker: MockerFixture): + mock_allocator_class = mocker.Mock() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class", + return_value=mock_allocator_class, + ) + return mock_allocator_class + + @pytest.fixture def mock_od_config(mocker: MockerFixture): """Create a mock OmniDiffusionConfig.""" @@ -118,7 +127,7 @@ def setup_allocator(self, mocker: MockerFixture): """ Unified interception of Allocators, and provision of default security values. """ - self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + self.mock_allocator_class = patch_cumem_allocator(mocker) self.mock_allocator = mocker.Mock() self.mock_allocator_class.get_instance.return_value = self.mock_allocator self.mock_allocator.get_current_usage.return_value = 4 * 1024**3 @@ -126,7 +135,7 @@ def setup_allocator(self, mocker: MockerFixture): def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3] mock_platform.get_device_total_memory.return_value = 80 * 1024**3 @@ -158,7 +167,7 @@ def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3] mock_platform.get_device_total_memory.return_value = 80 * 1024**3 @@ -201,7 +210,7 @@ def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker): def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker): """Test that sleep validates memory was actually freed.""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.return_value = 10 * 1024**3 mock_platform.get_device_total_memory.return_value = 80 * 1024**3 @@ -226,7 +235,7 @@ def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_wor def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep uses device-scoped fallback when NVML is unavailable.""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") mock_get_process_memory.side_effect = [None, None] @@ -252,7 +261,7 @@ class TestDiffusionWorkerWakeUp: def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up without saved buffers (level 1 sleep).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) # Setup allocator mock mock_allocator = mocker.Mock() @@ -272,7 +281,7 @@ def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker): def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up with saved buffers (level 2 sleep).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) # Setup allocator mock mock_allocator = mocker.Mock() @@ -317,7 +326,7 @@ def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker): def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up only restores buffers that were saved.""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) # Setup allocator mock mock_allocator = mocker.Mock() diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py index 9c10d46b41f..b94f86db0e6 100644 --- a/tests/diffusion/test_inline_stage_diffusion_client.py +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -46,6 +46,7 @@ def client(mock_engine): ) with patch.object(InlineStageDiffusionClient, "_enrich_config"): od_config = MagicMock(spec=OmniDiffusionConfig) + od_config.streaming_output = False c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) yield c c.shutdown() @@ -72,6 +73,35 @@ async def test_inline_dispatch_request_success(client, mock_engine): mock_engine.step.assert_called_once() +@pytest.mark.asyncio +async def test_inline_dispatch_request_streaming_success(client, mock_engine): + """Inline StageDiffusionClient correctly receives streaming chunk output from Diffusion Engine""" + chunks = [ + OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=False), + OmniRequestOutput.from_diffusion(request_id="req-stream", images=[], finished=True), + ] + + async def _step_streaming(_request): + for chunk in chunks: + yield [chunk] + + client.od_config.streaming_output = True + mock_engine.step_streaming = _step_streaming + + await client.add_request_async("req-stream", "A test prompt", OmniDiffusionSamplingParams()) + + outputs = [] + for _ in range(20): + output = client.get_diffusion_output_nowait() + if output is not None: + outputs.append(output) + if output.finished: + break + await asyncio.sleep(0.01) + + assert outputs == chunks + + @pytest.mark.asyncio async def test_inline_dispatch_request_error(client, mock_engine): # Setup mock engine step to raise an exception diff --git a/tests/diffusion/test_input_batch_txt_seq_lens.py b/tests/diffusion/test_input_batch_txt_seq_lens.py new file mode 100644 index 00000000000..315f0f703f6 --- /dev/null +++ b/tests/diffusion/test_input_batch_txt_seq_lens.py @@ -0,0 +1,31 @@ +import pytest +import torch + +from vllm_omni.diffusion.worker.input_batch import _prepare_request_prompt_field +from vllm_omni.diffusion.worker.utils import DiffusionRequestState +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.cpu] + + +def test_prepare_request_prompt_field_refreshes_txt_seq_lens_after_padding(): + state = DiffusionRequestState( + request_id="req-0", + sampling=OmniDiffusionSamplingParams(), + prompt_embeds=torch.randn(1, 10, 4), + prompt_embeds_mask=torch.ones(1, 10, dtype=torch.bool), + txt_seq_lens=[10], + ) + + embeds, mask = _prepare_request_prompt_field( + state, + embeds_attr="prompt_embeds", + mask_attr="prompt_embeds_mask", + seq_lens_attr="txt_seq_lens", + target_seq_len=32, + ) + + assert embeds.shape == (1, 32, 4) + assert mask is not None and mask.shape == (1, 32) + assert state.txt_seq_lens == [32] + assert mask.sum(dim=1).tolist() == [10] diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index e5add76213c..070d899e6ee 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -18,9 +18,16 @@ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE +from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import RequestScheduler +from vllm_omni.diffusion.sched.interface import ( + CachedRequestData, + DiffusionSchedulerOutput, + NewRequestData, +) from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc +from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -39,7 +46,7 @@ def _mock_request(tag: str): """Return a lightweight request object identifiable by *tag*.""" return SimpleNamespace( request_id=tag, - prompts=[f"prompt_{tag}"], + prompt=f"prompt_{tag}", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), ) @@ -49,7 +56,7 @@ def _make_executor(num_gpus: int = 1): Returns ``(executor, request_queue, result_queue)``. """ - od_cfg = SimpleNamespace(num_gpus=num_gpus) + od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None) executor = MultiprocDiffusionExecutor(od_cfg) @@ -75,6 +82,7 @@ def _make_engine(num_gpus: int = 1): """Create a lightweight ``DiffusionEngine`` wired to mocked executor.""" executor, req_q, res_q = _make_executor(num_gpus) engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(streaming_output=False) sched = RequestScheduler() sched.initialize(SimpleNamespace()) engine.scheduler = sched @@ -99,13 +107,24 @@ def _run(): req = req_q.get(timeout=10) method = req.get("method", "") args = req.get("args", ()) - if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"): + if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput): + sched_output = args[0] + runner_outputs = [] + for nr in sched_output.scheduled_new_reqs: + tag = f"result_for_{nr.request_id}" + runner_outputs.append( + RunnerOutput(request_id=nr.request_id, finished=True, result=_tagged_output(tag)) + ) + res_q.put(BatchRunnerOutput.from_list(runner_outputs)) + elif method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"): tag = f"result_for_{args[0].request_id}" + res_q.put(_tagged_output(tag)) elif args: tag = f"result_for_{args[0]}" + res_q.put(_tagged_output(tag)) else: tag = f"result_for_{method}" - res_q.put(_tagged_output(tag)) + res_q.put(_tagged_output(tag)) t = threading.Thread(target=_run, daemon=True) t.start() @@ -170,6 +189,64 @@ def _b(): assert results["B"].error == "result_for_B" +# ───────────────── request-mode dispatch (per-request vs batch) ───────────── + + +def _make_sched_output(*request_ids: str) -> DiffusionSchedulerOutput: + """Build a request-mode scheduler output with the given new requests.""" + new_reqs = [ + NewRequestData( + request_id=rid, + req=OmniDiffusionRequest( + prompt=f"prompt_{rid}", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id=rid, + ), + ) + for rid in request_ids + ] + return DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=new_reqs, + scheduled_cached_reqs=CachedRequestData.make_empty(), + finished_req_ids=set(), + num_running_reqs=len(new_reqs), + num_waiting_reqs=0, + ) + + +class TestRequestModeDispatch: + """Request-batch-capable dispatch uses ``execute_batch`` for request-mode cycles.""" + + @pytest.mark.parametrize("request_ids", [("solo",), ("A", "B", "C")]) + def test_request_batch_capable_pipeline_uses_execute_batch(self, request_ids): + engine, executor, _, _ = _make_engine() + executor.execute_request = Mock(return_value="per-request") + executor.execute_batch = Mock(return_value="batch") + engine.execute_fn = executor.execute_batch + + out = engine.execute_fn(_make_sched_output(*request_ids)) + + executor.execute_batch.assert_called_once() + assert out == "batch" + executor.execute_request.assert_not_called() + + @pytest.mark.parametrize("request_ids", [("solo",), ("A", "B")]) + def test_batch_path_routes_results_through_worker(self, request_ids): + """End-to-end: a request-batch cycle goes out as one ``execute_model_batch`` + RPC and comes back as a per-request-routed ``BatchRunnerOutput``.""" + engine, executor, req_q, res_q = _make_engine() + engine.execute_fn = executor.execute_batch + wt = _start_worker(req_q, res_q, count=1) + + out = engine.execute_fn(_make_sched_output(*request_ids)) + wt.join(5) + + assert isinstance(out, BatchRunnerOutput) + results = {ro.request_id: ro.result.error for ro in out.runner_outputs} + assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids} + + # ───────────────── concurrent collective RPC ───────────────── @@ -581,6 +658,58 @@ def _dying_dequeue(timeout=None): ) +class TestMultiprocExecutorStepStreamingOutput: + """Streaming output uses step execution and one worker reply per step.""" + + def test_execute_step_allows_streaming_output_mode(self): + executor, req_q, res_q = _make_executor() + executor.od_config = SimpleNamespace(streaming_output=True) # pyright: ignore[reportAttributeAccessIssue] + runner_outputs = [ + RunnerOutput( + request_id="sched-stream", + step_index=1, + finished=False, + result=DiffusionOutput(output={"chunk": 0}, finished=False, chunk_index=0, total_chunks=2), + ), + RunnerOutput( + request_id="sched-stream", + step_index=2, + finished=True, + result=DiffusionOutput(output={"chunk": 1}, finished=True, chunk_index=1, total_chunks=2), + ), + ] + scheduler_output = SimpleNamespace( + scheduled_request_ids=["sched-stream"], + ) + + def _worker(): + for runner_output in runner_outputs: + req_q.get(timeout=10) + res_q.put(runner_output) + + thread = threading.Thread(target=_worker, daemon=True) + thread.start() + + first: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType] + executor, + scheduler_output, # pyright: ignore[reportArgumentType] + ) + second: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType] + executor, + scheduler_output, # pyright: ignore[reportArgumentType] + ) + + assert first is runner_outputs[0] + assert first.result is not None + assert first.result.output == {"chunk": 0} + assert first.finished is False + assert second is runner_outputs[1] + assert second.result is not None + assert second.result.output == {"chunk": 1} + assert second.finished is True + thread.join(timeout=2) + + class TestDiffusionEngineDeadErrorPassthrough: """``DiffusionEngine.add_req_and_wait_for_response`` re-raises ``EngineDeadError`` from executor and wraps other errors.""" @@ -625,11 +754,11 @@ def _make_client(self, *, engine_dead=False, proc_alive=True): client._tasks = {} client._shutting_down = False client._engine_dead = engine_dead - client._owns_process = True - client._proc = MagicMock( + proc = MagicMock( is_alive=MagicMock(return_value=proc_alive), exitcode=1, ) + client._proc_manager = SimpleNamespace(proc=proc) client._request_socket = MagicMock() client._response_socket = MagicMock() client._encoder = MagicMock() @@ -672,7 +801,7 @@ def test_get_output_returns_none_when_alive_and_empty(self): assert client.get_diffusion_output_nowait() is None def test_check_health_raises_when_proc_dead(self): - """``check_health`` detects a dead subprocess via ``_proc.is_alive()`` + """``check_health`` detects a dead subprocess via the manager's proc and raises ``EngineDeadError``, setting ``_engine_dead`` as a side effect.""" client = self._make_client(proc_alive=False) @@ -699,7 +828,7 @@ def test_get_output_returns_none_on_signal_death(self): ``get_diffusion_output_nowait`` returns ``None`` and sets ``_shutting_down`` instead of raising.""" client = self._make_client(proc_alive=False) - client._proc.exitcode = 137 # SIGKILL (128 + 9) + client._proc_manager.proc.exitcode = 137 # SIGKILL (128 + 9) client._response_socket.recv.side_effect = zmq.Again result = client.get_diffusion_output_nowait() @@ -727,10 +856,42 @@ def test_initialize_client_requires_replica_id(self): metadata, "tcp://req", "tcp://resp", - proc=None, batch_size=1, ) + @pytest.mark.asyncio + async def test_collective_rpc_async_returns_none_result(self, monkeypatch): + client = self._make_client() + client._owns_process = False + client._proc = None + client._encoder.encode.return_value = b"encoded-rpc" + + async def _unexpected_poll(*_, **__): + raise AssertionError("collective_rpc_async should not keep polling after a None rpc_result arrives") + + client._response_poller = SimpleNamespace(poll=_unexpected_poll) + + rpc_id = "rpc-none" + monkeypatch.setattr( + "vllm_omni.diffusion.stage_diffusion_client.uuid.uuid4", + lambda: SimpleNamespace(hex=rpc_id), + ) + + def _drain() -> None: + client._rpc_results[rpc_id] = None + + client._drain_responses = _drain + + result = await client.collective_rpc_async( + method="profile", + timeout=0.01, + args=(False, None), + ) + + assert result is None + client._request_socket.send.assert_called_once_with(b"encoded-rpc") + assert rpc_id not in client._pending_rpcs + # ───────── monitor thread & death sentinel integration tests ───────── @@ -836,7 +997,7 @@ def test_proc_monitor_sets_engine_dead_on_process_death(self): client._engine_dead = False proc = _make_short_lived_process() - client._proc = proc + client._proc_manager = SimpleNamespace(proc=proc) client._start_proc_monitor() proc.join(5) diff --git a/tests/diffusion/test_stage_diffusion_proc.py b/tests/diffusion/test_stage_diffusion_proc.py index dcefafb8d40..c96aa8ba118 100644 --- a/tests/diffusion/test_stage_diffusion_proc.py +++ b/tests/diffusion/test_stage_diffusion_proc.py @@ -3,79 +3,17 @@ import asyncio import time -from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass -from types import SimpleNamespace import pytest from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] -def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info(): - async def run_test(): - captured = {} - - async def step(request): - captured["request"] = request - return [ - SimpleNamespace( - images=["img-1"], - _multimodal_output={}, - _custom_output={}, - metrics={}, - stage_durations={}, - peak_memory_mb=0.0, - latents=None, - trajectory_latents=None, - trajectory_timesteps=None, - trajectory_log_probs=None, - trajectory_decoded=None, - final_output_type="image", - ), - SimpleNamespace( - images=["img-2"], - _multimodal_output={}, - _custom_output={}, - metrics={}, - stage_durations={}, - peak_memory_mb=0.0, - latents=None, - trajectory_latents=None, - trajectory_timesteps=None, - trajectory_log_probs=None, - trajectory_decoded=None, - final_output_type="image", - ), - ] - - proc = object.__new__(StageDiffusionProc) - proc._engine = SimpleNamespace(step=step) - proc._executor = ThreadPoolExecutor(max_workers=1) - - try: - result = await proc._process_batch_request( - request_id="req-parent", - prompts=["hello", "world"], - sampling_params_dict=asdict(OmniDiffusionSamplingParams()), - kv_sender_info={0: {"host": "10.0.0.2", "zmq_port": 50151}}, - ) - finally: - proc._executor.shutdown(wait=True) - - request = captured["request"] - assert request.request_id == "req-parent" - assert not hasattr(request, "request_ids") - assert request.kv_sender_info == {0: {"host": "10.0.0.2", "zmq_port": 50151}} - assert result.request_id == "req-parent" - assert result.images == ["img-1", "img-2"] - - asyncio.run(run_test()) - - @dataclass class MockOmniRequestOutput: request_id: str = "" @@ -102,6 +40,40 @@ def simulate_step_delay(height, width, num_inference_steps) -> float: return [MockOmniRequestOutput(request_id=request.request_id)] +@pytest.mark.asyncio +async def test_proc_streaming_request_yields_each_engine_chunk(): + """Ensure that the streaming output chunks from DiffusionEngine reaches StageDiffusionProc""" + captured = {} + chunks = [ + OmniRequestOutput.from_diffusion(request_id="", images=[], finished=False), + OmniRequestOutput.from_diffusion(request_id="", images=[], finished=True), + ] + + class _StreamingEngine: + async def step_streaming(self, request): + captured["request"] = request + for chunk in chunks: + yield [chunk] + + stage_proc = object.__new__(StageDiffusionProc) + stage_proc._engine = _StreamingEngine() + + outputs = [ + output + async for output in stage_proc._process_streaming_request( + request_id="req-stream", + prompt="prompt", + sampling_params_dict=asdict(OmniDiffusionSamplingParams()), + kv_sender_info={0: {"host": "127.0.0.1"}}, + ) + ] + + assert outputs == chunks + assert [output.request_id for output in outputs] == ["req-stream", "req-stream"] + assert [output.finished for output in outputs] == [False, True] + assert captured["request"].kv_sender_info == {0: {"host": "127.0.0.1"}} + + @pytest.mark.asyncio async def test_proc_process_request_with_batching_async_output(): stage_proc = object.__new__(StageDiffusionProc) diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index fa26ae2763a..ba22743fc70 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -184,7 +184,7 @@ def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config): mock_reqs = [mocker.Mock()] result = wrapper.execute_model(mock_reqs, mock_od_config) - wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config) + wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config, kv_prefetch_jobs=None) assert result == mock_output def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config): diff --git a/tests/diffusion/utils/__init__.py b/tests/diffusion/utils/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/diffusion/utils/test_kv_utils.py b/tests/diffusion/utils/test_kv_utils.py new file mode 100644 index 00000000000..6b523b92033 --- /dev/null +++ b/tests/diffusion/utils/test_kv_utils.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for kv utils.""" + +import pytest +import torch + +from vllm_omni.diffusion.utils.kv_utils import left_pad_stack + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_uniform_lengths(): + """Ensure uniform 3D tensors correctly stack to 4D and have no mask.""" + tensors = [torch.randn(10, 4, 128) for _ in range(3)] + stacked, mask = left_pad_stack(tensors) + assert stacked.shape == (3, 10, 4, 128) + assert mask is None + for i in range(3): + assert torch.equal(stacked[i], tensors[i]) + + +def test_variable_lengths(): + """Ensure variable 3D tensors correctly stack to 4D with a mask.""" + t1 = torch.ones(5, 2, 4) + t2 = torch.ones(8, 2, 4) + t3 = torch.ones(3, 2, 4) + stacked, mask = left_pad_stack([t1, t2, t3]) + + assert stacked.shape == (3, 8, 2, 4) + assert mask is not None + assert mask.shape == (3, 8) + # Ensure summing over dim 1 gives our seq lens back + assert mask.sum(dim=1).tolist() == [5, 8, 3] + + +def test_single_tensor_is_4d(): + """Ensure a single 3D tensor is expanded to 4D.""" + t = torch.randn(7, 4, 128) + stacked, mask = left_pad_stack([t]) + assert stacked.shape == (1, 7, 4, 128) + assert mask is None + assert torch.equal(stacked[0], t) + + +def test_preserves_device_and_dtype(): + """Ensure device/dtype is preserved.""" + t1 = torch.randn(3, 2, dtype=torch.bfloat16) + t2 = torch.randn(5, 2, dtype=torch.bfloat16) + stacked, mask = left_pad_stack([t1, t2]) + assert stacked.dtype == torch.bfloat16 + assert stacked.device == t1.device + assert mask.device == t1.device + + +def test_mismatched_trailing_shapes_raises(): + """Ensure that mismatched dims outside of 0 explodes.""" + t1 = torch.randn(5, 4, 128) + t2 = torch.randn(5, 8, 128) + with pytest.raises(ValueError): + left_pad_stack([t1, t2]) + + +def test_mismatched_ndim_raises(): + """Ensure that mismatched ndims explodes.""" + t1 = torch.randn(5, 4) + t2 = torch.randn(5, 4, 128) + with pytest.raises(ValueError): + left_pad_stack([t1, t2]) + + +def test_empty_list_raises(): + """Ensure that tensors must be nonempty.""" + with pytest.raises(ValueError): + left_pad_stack([]) diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index 106507114b1..fef63f4f8e5 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py @@ -32,7 +32,7 @@ # Reference pixel data extracted from the known-good output image # Each entry contains (x, y) position and expected (R, G, B) values -# "Generated with seed=52, num_inference_steps=14, +# "Generated with seed=52, num_inference_steps=15, # prompt='A cute cat'" REFERENCE_PIXELS = [ {"position": (100, 100), "rgb": (115, 113, 94)}, @@ -63,7 +63,7 @@ def _find_free_port() -> int: return port -def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list: +def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list: """Configure sampling parameters for Bagel text2img generation. Args: @@ -266,7 +266,6 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str: return temp_file.name -@pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index eb5a79d6aff..c21af394a06 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py @@ -7,7 +7,7 @@ - img2img: validates output vs reference pixels within a ±10 tolerance. - text2img: validates output vs reference pixels within a ±5 tolerance (equivalent to `examples/offline_inference/bagel/end2end.py` with - text2img modality and 15 steps). + text2img modality and 14 steps). """ import os @@ -30,7 +30,7 @@ BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml") # Reference pixel data extracted from the known-good output image -# Generated with seed=52, num_inference_steps=14, +# Generated with seed=52, num_inference_steps=15, # prompt='Change the grass color to red', # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg REFERENCE_PIXELS = [ @@ -46,9 +46,8 @@ {"position": (256, 256), "rgb": (181, 201, 221)}, ] - # text2img reference pixels (aligned with offline `bagel/end2end.py` text2img, 15 steps) -# "Generated with seed=52, num_inference_steps=14, +# "Generated with seed=52, num_inference_steps=15, # prompt='A cute cat'" TEXT2IMG_REFERENCE_PIXELS = [ {"position": (100, 100), "rgb": (115, 113, 94)}, @@ -88,7 +87,7 @@ def _find_free_port() -> int: return port -def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list: +def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list: """Configure sampling parameters for Bagel img2img generation. Args: @@ -228,7 +227,6 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str: return config_path -@pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @@ -245,7 +243,6 @@ def test_bagel_img2img_shared_memory_connector(run_level): _validate_pixels(generated_image) -@pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) diff --git a/tests/distributed/omni_connectors/test_kv_async_prefetch.py b/tests/distributed/omni_connectors/test_kv_async_prefetch.py new file mode 100644 index 00000000000..6ef5a3f8cc3 --- /dev/null +++ b/tests/distributed/omni_connectors/test_kv_async_prefetch.py @@ -0,0 +1,356 @@ +"""Background KV prefetch — CPU-testable mechanism. + +Covers opt-in gating, the start_prefetch / consume_prefetched_kv round trip over a +mock connector, dedup, per-call sender_info isolation, role classification, and +miss/abort paths. GPU H2D and D2D paths need CUDA (integration tests). +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + OmniKVCacheConfig, + OmniKVTransferManager, + ReceiveRole, +) +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# start_prefetch() skips stubs without a sender endpoint, so submissions carry it. +_SENDER_INFO = {"host": "127.0.0.1", "zmq_port": 50051} + + +class MockConnector: + def __init__(self): + self.store = {} + self.get_calls = [] + + def put(self, from_stage, to_stage, put_key, data): + self.store[f"{from_stage}->{to_stage}:{put_key}"] = data + return True, len(str(data)), None + + def get(self, from_stage, to_stage, get_key, metadata=None): + self.get_calls.append({"get_key": get_key, "metadata": metadata}) + key = f"{from_stage}->{to_stage}:{get_key}" + if key in self.store: + return self.store[key], len(str(self.store[key])) + return None + + +def _make_sender_receiver(*, async_prefetch=True): + """Build sender+receiver managers sharing one mock connector, seed one KV.""" + connector = MockConnector() + + sender = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={"type": "mock"}, from_stage="sender", to_stage="receiver", need_send_cache=True + ) + ) + sender._connector = connector + + receiver = OmniKVTransferManager( + OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="sender", + stage_id="receiver", + need_recv_cache=True, + recv_timeout=1.0, + ), + async_prefetch=async_prefetch, + ) + receiver._connector = connector + return sender, receiver, connector + + +def _seed_payload(sender, req_id, *, num_layers=2, block_size=4, num_heads=4, head_dim=8): + kv_caches = [torch.randn(2, 5, block_size, num_heads, head_dim) for _ in range(num_layers)] + sender.handle_finished_requests_kv_transfer( + {req_id: {"block_ids": [0, 1], "seq_len": 10}}, kv_caches, block_size, "float32" + ) + + +def _req(rid): + """A minimal request carrying just ``request_id`` (consume_prefetched_kv only needs that).""" + return SimpleNamespace(request_id=rid) + + +# --------------------------------------------------------------------------- # +# Opt-in gating +# --------------------------------------------------------------------------- # + + +def test_async_prefetch_defaults_false(): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + assert mgr._async_prefetch is False + + +def test_from_od_config_enables_prefetch_flag(): + od = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": True}) + assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True + od2 = SimpleNamespace(omni_kv_config=None) + assert OmniKVTransferManager.from_od_config(od2)._async_prefetch is False + od3 = SimpleNamespace(omni_kv_config={"enable_kv_async_prefetch": False}) + assert OmniKVTransferManager.from_od_config(od3)._async_prefetch is False + + +def test_prefetch_preserved_with_connector_config(): + # The flag must survive merging with a populated connector config. + od = SimpleNamespace( + omni_kv_config={ + "connector_config": {"type": "mock"}, + "need_recv_cache": True, + "enable_kv_async_prefetch": True, + }, + ) + assert OmniKVTransferManager.from_od_config(od)._async_prefetch is True + + +def test_prefetch_disabled_with_cfg_companion_collector(): + # A CFG companion collector forces the synchronous path (see _resolve_async_prefetch). + od = SimpleNamespace( + omni_kv_config={"enable_kv_async_prefetch": True}, + cfg_kv_collect_func=lambda *a, **k: {}, + ) + assert OmniKVTransferManager.from_od_config(od)._async_prefetch is False + + +def test_from_vllm_config_ar_path_is_always_sync(): + # The AR entry point never enables prefetch. + model_config = SimpleNamespace(omni_kv_config=None) + vllm_config = SimpleNamespace(kv_transfer_config=None) + assert OmniKVTransferManager.from_vllm_config(vllm_config, model_config)._async_prefetch is False + + +def test_start_prefetch_noop_when_disabled(): + _, receiver, _ = _make_sender_receiver(async_prefetch=False) + receiver.start_prefetch({"request_id": "r1", "kv_sender_info": None}) + assert receiver._prefetch_futures == {} + + +# --------------------------------------------------------------------------- # +# Round trip +# --------------------------------------------------------------------------- # + + +def test_prefetch_round_trip_returns_data(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-1") + + receiver.start_prefetch({"request_id": "rid-1", "kv_sender_info": _SENDER_INFO}) + assert "rid-1" in receiver._prefetch_futures + + data, size = receiver.consume_prefetched_kv(_req("rid-1")) + assert data is not None + assert "layer_blocks" in data + assert data["metadata"]["seq_len"] == 10 + assert size > 0 + # Future is consumed (popped) on retrieval. + assert "rid-1" not in receiver._prefetch_futures + + +def test_consume_prefetched_kv_miss_without_prefetch_returns_none(): + _, receiver, _ = _make_sender_receiver(async_prefetch=False) + assert receiver.consume_prefetched_kv(_req("never-prefetched")) == (None, 0) + + +def test_start_prefetch_dedups_same_request(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-dup") + receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO}) + fut1 = receiver._prefetch_futures["rid-dup"] + receiver.start_prefetch({"request_id": "rid-dup", "kv_sender_info": _SENDER_INFO}) + assert receiver._prefetch_futures["rid-dup"] is fut1 + + +def test_start_prefetch_skips_without_sender_info(): + # No endpoint => no submission; the sync path handles it. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-nosi") + receiver.start_prefetch({"request_id": "rid-nosi", "kv_sender_info": None}) + assert receiver._prefetch_futures == {} + + +def test_consume_prefetched_kv_sweeps_stale_entries(): + # Consuming a slot drops every other tracked prefetch. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-stale") + receiver.start_prefetch({"request_id": "rid-stale", "kv_sender_info": _SENDER_INFO}) + assert "rid-stale" in receiver._prefetch_futures + # Miss on the current request still sweeps stale entries. + receiver._async_prefetch = False + assert receiver.consume_prefetched_kv(_req("rid-current")) == (None, 0) + assert receiver._prefetch_futures == {} + + +# --------------------------------------------------------------------------- # +# consume_prefetched_kv: miss returns None, stale sweep +# --------------------------------------------------------------------------- # + + +def test_start_prefetch_sweeps_orphan(): + # An orphan prefetch (e.g. aborted) is dropped when the next prefetch starts. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-orphan") + _seed_payload(sender, "rid-next") + receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO}) + receiver.consume_prefetched_kv(_req("rid-orphan")) # complete the future + receiver.start_prefetch({"request_id": "rid-orphan", "kv_sender_info": _SENDER_INFO}) # re-add as leftover + receiver.start_prefetch({"request_id": "rid-next", "kv_sender_info": _SENDER_INFO}) + assert "rid-orphan" not in receiver._prefetch_futures + assert "rid-next" in receiver._prefetch_futures + + +def test_shutdown_prefetch_clears_state(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-sd") + receiver.start_prefetch({"request_id": "rid-sd", "kv_sender_info": _SENDER_INFO}) + receiver.shutdown_prefetch() + assert receiver._prefetch_futures == {} + assert receiver._prefetch_executor is None + + +# --------------------------------------------------------------------------- # +# Correctness: per-call sender_info must not touch instance state +# --------------------------------------------------------------------------- # + + +def test_receive_with_sender_info_does_not_mutate_instance_state(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-si") + assert receiver._sender_base_host is None + assert receiver._sender_base_zmq_port is None + + receiver.receive_kv_cache_for_request( + "rid-si", + target_device=None, + sender_info={"host": "1.2.3.4", "zmq_port": 50051}, + ) + # The per-call endpoint must not leak into instance fields. + assert receiver._sender_base_host is None + assert receiver._sender_base_zmq_port is None + + +# --------------------------------------------------------------------------- # +# Receive-role classification (recv_role) and the follower prefetch guard +# --------------------------------------------------------------------------- # + +import vllm_omni.diffusion.distributed.parallel_state as ps # noqa: E402 + + +def _patch_topo(monkeypatch, *, world_size, world_rank=0, cfg_size=1, cfg_rank=0, sp_size=1, sp_rank=0): + monkeypatch.setattr(ps, "get_world_group", lambda: SimpleNamespace(world_size=world_size, rank_in_group=world_rank)) + monkeypatch.setattr(ps, "get_classifier_free_guidance_world_size", lambda: cfg_size) + monkeypatch.setattr(ps, "get_classifier_free_guidance_rank", lambda: cfg_rank) + monkeypatch.setattr(ps, "get_cfg_group", lambda: None) + monkeypatch.setattr(ps, "get_sequence_parallel_world_size", lambda: sp_size) + monkeypatch.setattr(ps, "get_sequence_parallel_rank", lambda: sp_rank) + monkeypatch.setattr(ps, "get_sp_group", lambda: None) + + +def _set_tp(mgr, src, tgt): + mgr._tp_topo = SimpleNamespace(source_tp_size=src, target_tp_size=tgt, local_rank=0) + + +def _role(mgr): + """Recompute the receive role from the (patched) parallel state, bypassing the cache.""" + mgr._topo_config = None + return mgr.topo_config.role + + +def test_recv_role_world1_is_local(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 1, 1) + _patch_topo(monkeypatch, world_size=1) + assert _role(mgr) is ReceiveRole.LOCAL + + +def test_recv_role_pure_tp_is_local(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=2) + assert _role(mgr) is ReceiveRole.LOCAL + + +def test_recv_role_cfg_leader_and_follower(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_sp_leader_and_follower(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_tp_and_sp_leader_per_tp_rank(monkeypatch): + # TP=2 + SP=2: sp_rank 0 is leader (pulls), sp_rank 1 is follower (gets broadcast). + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 2, 2) + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=4, sp_size=2, sp_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_hsdp_world_broadcast(monkeypatch): + # TP inactive, world > 1 (HSDP): rank-0 leader, others follower. + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + _set_tp(mgr, 1, 1) + _patch_topo(monkeypatch, world_size=2, world_rank=0) + assert _role(mgr) is ReceiveRole.LEADER + _patch_topo(monkeypatch, world_size=2, world_rank=1) + assert _role(mgr) is ReceiveRole.FOLLOWER + + +def test_recv_role_uninitialized_defaults_local(monkeypatch): + mgr = OmniKVTransferManager(OmniKVCacheConfig()) + + def _boom(): + raise AssertionError("world group is not initialized") + + monkeypatch.setattr(ps, "get_world_group", _boom) + assert _role(mgr) is ReceiveRole.LOCAL + + +def test_start_prefetch_skips_follower(monkeypatch): + # A follower never pulls, so start_prefetch must not submit. + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-fol") + _set_tp(receiver, 2, 2) + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=1) + receiver.start_prefetch({"request_id": "rid-fol", "kv_sender_info": _SENDER_INFO}) + assert receiver._prefetch_futures == {} + + +def test_start_prefetch_submits_for_owner(monkeypatch): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-own") + _set_tp(receiver, 2, 2) + _patch_topo(monkeypatch, world_size=4, cfg_size=2, cfg_rank=0) + receiver.start_prefetch({"request_id": "rid-own", "kv_sender_info": _SENDER_INFO}) + assert "rid-own" in receiver._prefetch_futures + + +def test_consume_then_apply_attaches_payload(): + sender, receiver, _ = _make_sender_receiver() + _seed_payload(sender, "rid-apply") + receiver.start_prefetch({"request_id": "rid-apply", "kv_sender_info": _SENDER_INFO}) + data, _ = receiver.consume_prefetched_kv(_req("rid-apply")) + assert data is not None + + req = OmniDiffusionRequest(prompt="p", sampling_params=OmniDiffusionSamplingParams(), request_id="rid-apply") + # Mirror consume_and_distribute_kv_cache's LOCAL apply path (CPU: record_stream no-op). + receiver._record_stream_for_prefetched(data) + receiver.apply_kv_cache_to_request(req, data) + assert req.past_key_values is not None diff --git a/tests/distributed/omni_connectors/test_kv_flow.py b/tests/distributed/omni_connectors/test_kv_flow.py index 2cc19f1b366..83206fefbda 100644 --- a/tests/distributed/omni_connectors/test_kv_flow.py +++ b/tests/distributed/omni_connectors/test_kv_flow.py @@ -1,7 +1,6 @@ import json import struct -import numpy as np import pytest import torch @@ -206,17 +205,6 @@ def test_update_sender_info_uses_configured_source_stage(): assert manager.config.connector_config["sender_zmq_port"] == 50152 -def test_clone_received_payload_tensors_breaks_buffer_alias(): - payload, key_tensor = _make_serialized_payload() - raw = np.frombuffer(bytearray(payload), dtype=np.uint8) - data = KVCacheTransferData.from_bytes(memoryview(raw)) - - OmniKVTransferManager._clone_received_payload_tensors(data) - raw[:] = 0 - - assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor) - - def test_receive_kv_cache_uses_exponential_backoff(monkeypatch): config = OmniKVCacheConfig( connector_config={"type": "mock"}, @@ -382,7 +370,7 @@ def test_manager_reception(kv_config, mock_connector, common_constants): mock_connector.store[store_key] = data_to_receive req = OmniDiffusionRequest( - prompts=["test_recv"], + prompt="test_recv", sampling_params=OmniDiffusionSamplingParams(), request_id=req_id, ) @@ -426,7 +414,7 @@ def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_conf mock_connector.store[store_key] = data_to_receive req = OmniDiffusionRequest( - prompts=["prompt-a", "prompt-b"], + prompt="prompt-a", sampling_params=OmniDiffusionSamplingParams(), request_id=parent_req_id, ) @@ -451,7 +439,7 @@ def collect_cfg(request_id, cfg_request_ids, kv_transfer_manager, target_device) return {"cfg_text_kv_metadata": {"ok": True}} req = OmniDiffusionRequest( - prompts=["prompt-a", "prompt-b"], + prompt="prompt-a", sampling_params=OmniDiffusionSamplingParams(), request_id="req-parent", ) @@ -512,7 +500,7 @@ def test_integration_flow(common_constants): receiver_manager._connector = connector req = OmniDiffusionRequest( - prompts=["test_integ"], + prompt="test_integ", sampling_params=OmniDiffusionSamplingParams(), request_id=req_id, ) @@ -526,6 +514,57 @@ def test_integration_flow(common_constants): assert req.kv_metadata["seq_len"] == 10 +def test_deferred_release_isolates_data_from_pool_buffer(): + payload, key_tensor = _make_serialized_payload() + + # ManagedBuffer-like object whose release() zeros storage (simulates reuse). + raw_array = bytearray(payload) + buf_tensor = torch.frombuffer(raw_array, dtype=torch.uint8) + + class _CorruptingManagedBuffer: + def __init__(self, storage, tensor): + self._storage = storage + self._tensor = tensor + + @property + def tensor(self): + return self._tensor + + def release(self): + for i in range(len(self._storage)): + self._storage[i] = 0 + + managed_buf = _CorruptingManagedBuffer(raw_array, buf_tensor) + + class _ManagedBufConnector: + def __init__(self, buf, from_stage, to_stage, req_id): + self._buf = buf + self._key = f"{from_stage}->{to_stage}:omni_{from_stage}_to_{to_stage}_kv_cache_{req_id}" + + def get(self, from_stage, to_stage, get_key, metadata=None): + if f"{from_stage}->{to_stage}:{get_key}" == self._key: + return self._buf, len(self._buf._storage) + return None + + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="sender", + stage_id="receiver", + need_recv_cache=True, + recv_timeout=1.0, + ) + manager = OmniKVTransferManager(config) + req_id = "req-payload" + manager._connector = _ManagedBufConnector(managed_buf, "sender", "receiver", req_id) + + data, size = manager.receive_kv_cache_for_request(req_id, target_device=torch.device("cpu")) + + assert data is not None + assert size > 0 + # Returned data must survive the pool buffer being zeroed by release(). + assert torch.equal(data["layer_blocks"]["key_cache"][0], key_tensor) + + def test_manager_extraction_no_connector(kv_config, common_constants): """Test extraction when connector is unavailable (should still return IDs).""" block_size = common_constants["block_size"] diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index 73928aaf6f4..572376aef46 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -41,6 +41,19 @@ def test_get_nonexistent_key_returns_none(self, connector): result = connector.get("s0", "s1", "no_such_key_xyz", metadata=None) assert result is None + def test_get_empty_shm_race_returns_none(self, connector, monkeypatch): + def raise_empty_file(*args, **kwargs): + raise ValueError("cannot mmap an empty file") + + monkeypatch.setattr( + "vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_pkg.SharedMemory", + raise_empty_file, + ) + + result = connector.get("s0", "s1", "not_ready_yet", metadata=None) + + assert result is None + def test_rank_aware_keys_independent(self, connector): """Each TP rank writes/reads its own key — simulates homogeneous TP.""" payloads = {} diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py index 495c10a657e..261de0b32d9 100644 --- a/tests/distributed/omni_connectors/test_tp_rank_aware.py +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -586,8 +586,17 @@ def _receive(req_obj, cfg_func, target_device): return_value=0, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), ): - assert mgr.receive_multi_kv_cache_distributed(req) is True + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True mgr.receive_multi_kv_cache.assert_called_once() assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") @@ -642,6 +651,15 @@ def test_tp_cfg_follower_receives_local_payload_without_receiving(self): return_value=1, ), patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), ): assert mgr.receive_multi_kv_cache_distributed(req) is True @@ -666,7 +684,27 @@ def test_tp_without_cfg_keeps_independent_receive_path(self): world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1) mgr.receive_multi_kv_cache = MagicMock(return_value=True) - with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group): + with ( + patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=None), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_world_size", + return_value=1, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_sequence_parallel_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=None), + ): assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu")) @@ -722,7 +760,7 @@ def _receive(req_obj, cfg_func, target_device): ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): - assert mgr.receive_multi_kv_cache_distributed(req) is True + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True # Verify behavior based on role if is_owner: @@ -817,7 +855,7 @@ def _receive(req_obj, cfg_func, target_device): ), patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group", return_value=sp_group), ): - assert mgr.receive_multi_kv_cache_distributed(req) is True + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True # Verify behavior based on role if role == "owner": @@ -899,6 +937,76 @@ def test_owner_receive_failure_scenarios(self, has_cfg, cfg_size): # ── TP auto-detect ─────────────────────────────────────────────────── +class TestBuildCfgRankLocalPayloads: + """Guard _build_cfg_rank_local_payloads padding logic for the + cfg_size==2 / no branch_roles scenario (hunyuan-image shared-prefix reuse).""" + + def test_cfg_size2_no_branch_roles_copies_main_payload_to_follower(self): + """When AR doesn't split KV by branch, cfg follower (rank 1) should + receive a copy of the positive KV, not None.""" + req = SimpleNamespace( + request_id="req-1", + past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]), + kv_metadata={"seq_len": 10}, + sampling_params=SimpleNamespace(), + ) + payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2) + + assert len(payloads) == 2 + # Rank 0: main payload (always present) + assert payloads[0] is not None + assert "past_key_values" in payloads[0] + # Rank 1: must also be non-None — same positive KV for shared-prefix reuse + assert payloads[1] is not None + assert payloads[1]["past_key_values"] is req.past_key_values + assert payloads[1]["kv_metadata"] == {"seq_len": 10} + + def test_cfg_size3_no_branch_roles_keeps_none_padding(self): + """When cfg_size > 2 and no branch_roles, padding stays None (safety guard).""" + req = SimpleNamespace( + request_id="req-1", + past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]), + kv_metadata={"seq_len": 10}, + sampling_params=SimpleNamespace(), + ) + payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=3) + + assert len(payloads) == 3 + assert payloads[0] is not None + assert payloads[1] is None + assert payloads[2] is None + + def test_empty_main_payload_keeps_none_padding(self): + """When main_payload is empty, padding stays None even with cfg_size==2.""" + req = SimpleNamespace( + request_id="req-1", + sampling_params=SimpleNamespace(), + ) + payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2) + + assert len(payloads) == 2 + # No past_key_values → main_payload is empty dict → falsy → padding stays None + assert payloads[1] is None + + def test_with_branch_roles_uses_original_behavior(self): + """When branch_roles exist, follower payload comes from branch KV, not padding.""" + req = SimpleNamespace( + request_id="req-1", + past_key_values=SimpleNamespace(key_cache=[torch.tensor([1.0])]), + kv_metadata={"seq_len": 10}, + sampling_params=SimpleNamespace( + cfg_text_past_key_values=SimpleNamespace(key_cache=[torch.tensor([2.0])]), + cfg_text_kv_metadata={"source": "cfg_text"}, + ), + ) + payloads = OmniKVTransferManager._build_cfg_rank_local_payloads(req, cfg_size=2) + + assert len(payloads) == 2 + assert payloads[0] is not None + assert payloads[1] is not None + assert payloads[1]["sp.cfg_active_branch"] == "cfg_text" + + class TestAutoDetectTP: def test_auto_detect_when_config_defaults(self): """When config from_tp/to_tp == 1 (default), manager should auto-detect.""" diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py index c7b1b597e70..dfb291ccf24 100644 --- a/tests/distributed/omni_coordinator/test_omni_coordinator.py +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import json +import threading import time import pytest @@ -48,6 +49,33 @@ def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None: _recv_replica_list(sub, timeout_ms=50) +def test_omni_coordinator_wait_for_shutdown_unblocks_on_close(): + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=1000.0, + ) + + waiter = threading.Thread(target=coordinator.wait_for_shutdown) + waiter.start() + time.sleep(0.05) + assert waiter.is_alive() + + coordinator.close() + waiter.join(timeout=1.0) + assert not waiter.is_alive() + + def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): """Rapid updates should be coalesced into fewer PUB messages.""" router_addr = get_engine_client_zmq_addr( diff --git a/tests/docs/test_generate_examples.py b/tests/docs/test_generate_examples.py new file mode 100644 index 00000000000..0b29837a998 --- /dev/null +++ b/tests/docs/test_generate_examples.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import importlib.util +from pathlib import Path + +import pytest + +ROOT_DIR = Path(__file__).parents[2] +HOOK_PATH = ROOT_DIR / "docs/mkdocs/hooks/generate_examples.py" +SPEC = importlib.util.spec_from_file_location("generate_examples", HOOK_PATH) +assert SPEC and SPEC.loader +generate_examples = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(generate_examples) + + +def test_model_example_titles_match_shared_display_names(): + for slug, display_name in generate_examples.MODEL_DISPLAY_NAMES.items(): + for category, mode_title in generate_examples.SERVING_MODE_TITLES.items(): + example_path = ROOT_DIR / "examples" / category / slug + example = generate_examples.Example(example_path, category) + + assert example.title == f"{display_name}: {mode_title}" + assert example.nav_title == display_name + + +def test_model_example_title_mismatch_is_rejected(tmp_path): + example_path = tmp_path / "mimo_audio" + example_path.mkdir() + (example_path / "README.md").write_text("# A different title\n", encoding="utf-8") + + with pytest.raises(ValueError, match="Model example title mismatch"): + generate_examples.Example(example_path, "offline_inference") + + +def test_model_display_names_require_string_mapping(tmp_path, monkeypatch): + display_names_path = tmp_path / "model_display_names.yml" + display_names_path.write_text("- not a mapping\n", encoding="utf-8") + monkeypatch.setattr(generate_examples, "MODEL_DISPLAY_NAMES_FILE", display_names_path) + + with pytest.raises(ValueError, match="MODEL_DISPLAY_NAMES_FILE"): + generate_examples.load_model_display_names() diff --git a/tests/dreamzero/test_crossattn_cache.py b/tests/dreamzero/test_crossattn_cache.py new file mode 100644 index 00000000000..94361c090a0 --- /dev/null +++ b/tests/dreamzero/test_crossattn_cache.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests that DreamZero's cross-attention KV cache is actually init'd and reused. + +DreamZero's I2V cross-attention caches the text k/v and image k_img/v_img in a +per-layer ``crossattn_cache`` dict: they are computed on the first forward of a +session (``is_init`` False -> True) and reused on every later step, because the +text/image context is session-invariant. Only the query (which depends on the +per-step hidden state) is recomputed. + +These tests lock in that ``WanI2VCrossAttention.forward`` initialises the cache +on the first call and reuses it (not recompute) on later calls, and recomputes +every call when no cache is provided. + +CPU-only: ``WanI2VCrossAttention.__init__`` needs a tensor-parallel group (it +builds ColumnParallel/RowParallel linears), so we bypass ``__init__`` and stub +the projections/attention to count calls and exercise only the caching logic. +""" + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm_omni.diffusion.models.dreamzero.causal_wan_model import WanI2VCrossAttention + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_CLIP_LEN = 257 # image tokens the I2V cross-attn splits off the front of context + + +def _make_cross_attn(n_heads: int = 2, head_dim: int = 4) -> WanI2VCrossAttention: + """A WanI2VCrossAttention whose submodules are call-counting stubs. + + ``__init__`` is bypassed (it needs a TP group); we only wire what + ``forward`` touches. Each projection returns a [B, S, n*head_dim] tensor so + the subsequent ``unflatten(2, (n, head_dim))`` works, and is a MagicMock so + we can assert how many times it was invoked. + """ + inner = n_heads * head_dim + module = WanI2VCrossAttention.__new__(WanI2VCrossAttention) + module.tp_num_heads = n_heads + module.head_dim = head_dim + + def _proj() -> MagicMock: + return MagicMock(side_effect=lambda t: torch.zeros(t.shape[0], t.shape[1], inner)) + + module.q = _proj() + module.k = _proj() + module.v = _proj() + module.k_img = _proj() + module.v_img = _proj() + module.norm_q = lambda t: t + module.norm_k = lambda t: t + module.norm_k_img = lambda t: t + # attn returns the query unchanged ([B, S, n, d]); flatten(2) downstream is happy. + module.attn = lambda q, k, v: q + module.o = lambda t: t + return module + + +def _fresh_cache() -> dict: + return {"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} + + +def _inputs(inner: int = 8): + x = torch.randn(1, 3, inner) # B=1, query_len=3 + context = torch.randn(1, _CLIP_LEN + 5, inner) # 257 image tokens + 5 text tokens + return x, context + + +def test_cache_inits_on_first_call_and_reuses_on_second() -> None: + module = _make_cross_attn() + cache = _fresh_cache() + x, context = _inputs() + + # First call: cache is initialised and every k/v projection runs once. + module.forward(x, context, crossattn_cache=cache) + assert cache["is_init"] is True + for key in ("k", "v", "k_img", "v_img"): + assert cache[key] is not None, f"{key} should be cached after first call" + assert module.k.call_count == 1 + assert module.v.call_count == 1 + assert module.k_img.call_count == 1 + assert module.v_img.call_count == 1 + + cached = {key: cache[key] for key in ("k", "v", "k_img", "v_img")} + + # Second call: k/v/k_img/v_img are read from cache, NOT recomputed. + module.forward(x, context, crossattn_cache=cache) + assert module.k.call_count == 1 + assert module.v.call_count == 1 + assert module.k_img.call_count == 1 + assert module.v_img.call_count == 1 + for key, tensor in cached.items(): + assert cache[key] is tensor, f"{key} should be the same cached object" + + # The query always depends on the per-step hidden state, so it is recomputed. + assert module.q.call_count == 2 + + +def test_no_cache_recomputes_every_call() -> None: + module = _make_cross_attn() + x, context = _inputs() + + module.forward(x, context, crossattn_cache=None) + module.forward(x, context, crossattn_cache=None) + + # Without a cache, the image/text k/v are recomputed on every call. + assert module.k_img.call_count == 2 + assert module.v_img.call_count == 2 + assert module.k.call_count == 2 + assert module.v.call_count == 2 diff --git a/tests/dreamzero/test_fused_qk_rms_norm.py b/tests/dreamzero/test_fused_qk_rms_norm.py new file mode 100644 index 00000000000..dbbb439400b --- /dev/null +++ b/tests/dreamzero/test_fused_qk_rms_norm.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Equivalence tests for DreamZero's fused q/k qk-norm TP all-reduce. + +`fused_qk_rms_norm` packs the q and k per-token sum-of-squares into one tensor +and issues a single tensor-parallel all-reduce instead of two. These tests lock +in that it is numerically identical to applying `norm_q`/`norm_k` separately, +and that it really issues exactly one (fused) all-reduce. + +CPU-only: the TP world size and the collective are mocked, so a single process +simulates one rank (all-reduce is identity). This exercises the cat/split and +count-scaling logic that a real multi-rank run depends on. +""" + +from unittest.mock import patch + +import pytest +import torch + +from vllm_omni.diffusion.models.dreamzero.causal_wan_model import ( + DistributedRMSNorm, + fused_qk_rms_norm, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_MODULE = "vllm_omni.diffusion.models.dreamzero.causal_wan_model" + + +def _make_norms(dim: int) -> tuple[DistributedRMSNorm, DistributedRMSNorm]: + """Two norms with distinct random weights so a swapped q/k slice is caught.""" + torch.manual_seed(0) + norm_q = DistributedRMSNorm(dim) + norm_k = DistributedRMSNorm(dim) + norm_q.weight.data.normal_() + norm_k.weight.data.normal_() + return norm_q, norm_k + + +def _identity_all_reduce(tensor: torch.Tensor) -> torch.Tensor: + """Single-rank all-reduce: the reduction over one rank is the identity.""" + return tensor + + +def test_fused_matches_separate_tp1() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + with patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=1): + ref_q, ref_k = norm_q(q), norm_k(k) + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0) + + +def test_fused_matches_separate_tp_gt1() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + # tp_size=2 exercises the reduce branch and the `global_count *= tp_size` + # path in both the per-norm forward and the fused helper. + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2), + patch(f"{_MODULE}.tensor_model_parallel_all_reduce", side_effect=_identity_all_reduce), + ): + ref_q, ref_k = norm_q(q), norm_k(k) + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, ref_q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, ref_k, rtol=0, atol=0) + + +def test_fused_issues_single_all_reduce_over_packed_pair() -> None: + dim = 64 + norm_q, norm_k = _make_norms(dim) + q = torch.randn(2, 5, dim) + k = torch.randn(2, 5, dim) + + with ( + patch(f"{_MODULE}.get_tensor_model_parallel_world_size", return_value=2), + patch( + f"{_MODULE}.tensor_model_parallel_all_reduce", + side_effect=_identity_all_reduce, + ) as mock_ar, + ): + fused_qk_rms_norm(norm_q, norm_k, q, k) + + # Exactly one collective, over a tensor whose last dim packs the q+k pair. + assert mock_ar.call_count == 1 + packed = mock_ar.call_args[0][0] + assert packed.shape[-1] == 2 + + +def test_fused_identity_fallback_when_qk_norm_disabled() -> None: + # qk_norm=False makes the norms nn.Identity; the helper must fall back to + # plain passthrough without touching the collective. + norm_q, norm_k = torch.nn.Identity(), torch.nn.Identity() + q = torch.randn(2, 5, 64) + k = torch.randn(2, 5, 64) + + with patch(f"{_MODULE}.tensor_model_parallel_all_reduce") as mock_ar: + fused_q, fused_k = fused_qk_rms_norm(norm_q, norm_k, q, k) + + torch.testing.assert_close(fused_q, q, rtol=0, atol=0) + torch.testing.assert_close(fused_k, k, rtol=0, atol=0) + mock_ar.assert_not_called() diff --git a/tests/dreamzero/test_pipeline_state.py b/tests/dreamzero/test_pipeline_state.py index a3a5aaec03c..f00caf53ca7 100644 --- a/tests/dreamzero/test_pipeline_state.py +++ b/tests/dreamzero/test_pipeline_state.py @@ -4,7 +4,6 @@ from collections import OrderedDict import pytest -import torch from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState @@ -46,14 +45,11 @@ def test_dreamzero_pipeline_state_lru_caps_retained_sessions() -> None: assert "session-b" not in pipeline._states -def test_dreamzero_state_cache_access_requires_initialization() -> None: +def test_dreamzero_state_owns_no_kv_caches() -> None: + """KV (self- and cross-attention) is engine-owned since the AR-Diffusion + paged backend: the model-local cache accessors must be gone so nothing can + silently bypass the engine pool.""" state = DreamZeroState() - with pytest.raises(RuntimeError, match="KV caches not initialized"): - state.get_kv_caches() - - with pytest.raises(RuntimeError, match="Cross-attn caches not initialized"): - state.get_crossattn_caches() - - with pytest.raises(RuntimeError, match="create_kv_caches first"): - state.update_kv_cache(0, torch.empty(0)) + for removed in ("get_kv_caches", "create_kv_caches", "update_kv_cache", "get_crossattn_caches"): + assert not hasattr(state, removed) diff --git a/tests/dreamzero/test_qkv_fusion.py b/tests/dreamzero/test_qkv_fusion.py new file mode 100644 index 00000000000..16b78693f69 --- /dev/null +++ b/tests/dreamzero/test_qkv_fusion.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Equivalence test for DreamZero self-attn QKV fusion. + +CausalWanSelfAttention now projects q/k/v with a single QKVParallelLinear and +splits the result, instead of three separate ColumnParallelLinear. The DreamZero +checkpoint still ships separate q/k/v weights, so load_weights routes each into +the packed `qkv` param with its shard id. This test proves the chain +"load separate q/k/v via shard id -> single GEMM -> split" reproduces the three +independent projections, given identical weights. + +CPU-only: a single-rank (world_size=1) tensor-parallel group is initialized and +the GEMM dispatcher is forced to the default (CPU-compatible) path, following +tests/diffusion/layers/test_adalayernorm.py. +""" + +import os + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture(autouse=True) +def _init_distributed(): + """Minimal world_size=1 TP group required to build QKVParallelLinear.""" + from vllm.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + init_distributed_environment, + initialize_model_parallel, + ) + + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29507") + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method="env://", + ) + initialize_model_parallel() + yield + cleanup_dist_env_and_memory() + + +@pytest.fixture(autouse=True) +def _force_default_gemm(monkeypatch): + """Force CPU-compatible GEMM dispatch for CPU test tensors.""" + from vllm.model_executor.layers.utils import default_unquantized_gemm + + monkeypatch.setattr( + "vllm.model_executor.layers.linear.dispatch_unquantized_gemm", + lambda: default_unquantized_gemm, + ) + + +def test_fused_qkv_load_and_split_matches_separate_projections(): + from vllm.model_executor.layers.linear import QKVParallelLinear + + torch.manual_seed(0) + dim, head_dim, num_heads = 16, 4, 4 # inner_dim == dim, no GQA + + qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, + ) + + # Distinct random weight + bias per projection, loaded the exact way + # DreamZero.load_weights does: param.weight_loader(param, tensor, shard_id). + refs: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for shard_id in ("q", "k", "v"): + w = torch.randn(dim, dim) + b = torch.randn(dim) + refs[shard_id] = (w, b) + qkv.weight.weight_loader(qkv.weight, w, shard_id) + qkv.bias.weight_loader(qkv.bias, b, shard_id) + + x = torch.randn(2, 3, dim) + fused, _ = qkv(x) + q, k, v = fused.split([dim, dim, dim], dim=-1) + + for got, shard_id in zip((q, k, v), ("q", "k", "v"), strict=True): + w, b = refs[shard_id] + ref = torch.nn.functional.linear(x, w, b) + torch.testing.assert_close(got, ref) + + +def test_fused_qkv_shard_ids_are_not_transposed(): + # Guard against q/k/v being routed to the wrong slice: give each projection a + # constant, distinct weight so a swapped shard id would surface immediately. + from vllm.model_executor.layers.linear import QKVParallelLinear + + dim, head_dim, num_heads = 16, 4, 4 + qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + bias=True, + ) + + scale = {"q": 1.0, "k": 2.0, "v": 3.0} + for shard_id, s in scale.items(): + qkv.weight.weight_loader(qkv.weight, torch.full((dim, dim), s), shard_id) + qkv.bias.weight_loader(qkv.bias, torch.zeros(dim), shard_id) + + x = torch.ones(1, dim) + fused, _ = qkv(x) + q, k, v = fused.split([dim, dim, dim], dim=-1) + + # F.linear(ones, full(s)) == s * dim per element. + torch.testing.assert_close(q, torch.full((1, dim), scale["q"] * dim)) + torch.testing.assert_close(k, torch.full((1, dim), scale["k"] * dim)) + torch.testing.assert_close(v, torch.full((1, dim), scale["v"] * dim)) diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index 593a9960a38..ebe635ea811 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/tests/e2e/accuracy/helpers.py @@ -305,7 +305,7 @@ def assert_video_similarity_metrics( offline_path: Path, ssim_threshold: float, psnr_threshold: float, -) -> None: +) -> tuple[float, float]: ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path) psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path) ssim_score = parse_ssim_score(ssim_output) @@ -339,6 +339,7 @@ def assert_video_similarity_metrics( f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. " f"online={online_path} offline={offline_path}" ) + return ssim_score, psnr_score def assert_similarity( 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 8ce6f19888b..76fa67fdb82 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 @@ -33,7 +33,7 @@ * **Seed-TTS** optional extras for WER/SIM/UTMOS:: - pip install 'vllm-omni[seed-tts-eval]' + pip install 'vllm-omni[dev]' Examples -------- diff --git a/tests/e2e/accuracy/test_cosmos3_similarity.py b/tests/e2e/accuracy/test_cosmos3_similarity.py index ff2350096e0..efa095fb1cb 100644 --- a/tests/e2e/accuracy/test_cosmos3_similarity.py +++ b/tests/e2e/accuracy/test_cosmos3_similarity.py @@ -42,8 +42,6 @@ def _server_args() -> list[str]: return [ "--num-gpus", "1", - "--model-class-name", - "Cosmos3OmniDiffusersPipeline", "--stage-init-timeout", "900", "--init-timeout", diff --git a/tests/e2e/accuracy/test_diffusers_backend_similarity.py b/tests/e2e/accuracy/test_diffusers_backend_similarity.py index ad4dd57260c..52915229088 100644 --- a/tests/e2e/accuracy/test_diffusers_backend_similarity.py +++ b/tests/e2e/accuracy/test_diffusers_backend_similarity.py @@ -338,7 +338,7 @@ def test_diffusers_backend_i2v_matches_diffusers( def _diffusers_dummy_run(pipe: DiffusionPipeline) -> None: - from vllm_omni.diffusion.diffusion_engine import supports_multimodal_input + from vllm_omni.diffusion.io_support import supports_multimodal_input supports_image_input, supports_audio_input = supports_multimodal_input( SimpleNamespace( diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 620254cf9ad..2214b734ddc 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/tests/e2e/accuracy/test_hunyuan_image3.py @@ -71,6 +71,7 @@ def _assert_cot_structural_markers(cot_text: str, label: str) -> None: os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA" +os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG" pytestmark = [pytest.mark.local_model, pytest.mark.diffusion] @@ -297,7 +298,7 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag system_prompt_type = result.system_prompt_type ar_stop_token_ids = resolve_stop_token_ids(task="it2i", bot_task="think_recaption", tokenizer=tokenizer) - with OmniRunner(MODEL_PATH, deploy_config=deploy_config_path) as runner: + with OmniRunner(MODEL_PATH, deploy_config=deploy_config_path, trust_remote_code=True) as runner: params_list = list(runner.omni.default_sampling_params_list) for sp in params_list: if isinstance(sp, OmniDiffusionSamplingParams): @@ -330,7 +331,10 @@ def _run_offline(deploy_config_path: str, output_path: Path) -> tuple[Image.Imag if ro and getattr(ro, "outputs", None): cot_text = "".join(getattr(o, "text", "") or "" for o in ro.outputs) if not cot_text: - ar_text = getattr(out, "custom_output", {}).get("ar_generated_text") + multimodal_output = getattr(out, "multimodal_output", {}) or {} + metadata = multimodal_output.get("metadata", {}) if isinstance(multimodal_output, dict) else {} + text_metadata = metadata.get("text", {}) if isinstance(metadata, dict) else {} + ar_text = text_metadata.get("ar_generated_text") if isinstance(text_metadata, dict) else None if isinstance(ar_text, list): cot_text = "\n".join(text for text in ar_text if text) else: @@ -364,6 +368,7 @@ def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image "300", "--init-timeout", "900", + "--trust-remote-code", ] try: with OmniServer(MODEL_PATH, server_args, use_omni=True) as omni_server: @@ -487,7 +492,12 @@ def _run_dit_model( os.environ["VLLM_NVFP4_GEMM_BACKEND"] = nvfp4_backend try: - with OmniRunner(model, deploy_config=deploy_config_path, mode="text-to-image") as runner: + with OmniRunner( + model, + deploy_config=deploy_config_path, + mode="text-to-image", + trust_remote_code=True, + ) as runner: generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED) params = OmniDiffusionSamplingParams( height=QUANT_HEIGHT, diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index 6ac930f4611..3f799c038e0 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py @@ -48,6 +48,7 @@ "trust_remote_code": True, "devices": "0,1,2,3", # set dynamically by _write_deploy_config "vae_use_slicing": False, + "moe_backend": "flashinfer_cutlass", "vae_use_tiling": False, "parallel_config": { "pipeline_parallel_size": 1, diff --git a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py index 7fb2297a2c1..a169817a20f 100644 --- a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py +++ b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py @@ -40,7 +40,7 @@ _DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit()) _DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38) -MODEL_ID = "dg845/LTX-2.3-Diffusers" +MODEL_ID = "diffusers/LTX-2.3-Diffusers" MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL" PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting" NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark" @@ -216,7 +216,10 @@ def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Ima from vllm.config import VllmConfig, set_current_vllm_config from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import create_transformer_from_config, load_transformer_config + from vllm_omni.diffusion.models.ltx2 import ( + create_transformer_from_config, + load_transformer_config, + ) vllm_config = VllmConfig() ctx = set_current_vllm_config(vllm_config) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index 2b513ecda9d..7495bf08f1f 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -29,7 +29,7 @@ NUM_INFERENCE_STEPS = 20 TRUE_CFG_SCALE = 4.0 SEED = 42 -SSIM_THRESHOLD = 0.97 +SSIM_THRESHOLD = 0.94 PSNR_THRESHOLD = 30.0 MODEL_2512_ID = "Qwen/Qwen-Image-2512" diff --git a/tests/e2e/accuracy/test_video_streaming_output_similarity.py b/tests/e2e/accuracy/test_video_streaming_output_similarity.py new file mode 100644 index 00000000000..799e6e1258e --- /dev/null +++ b/tests/e2e/accuracy/test_video_streaming_output_similarity.py @@ -0,0 +1,143 @@ +""" +Online serving smoke test for streaming video output similarity. +""" + +import json +from pathlib import Path + +import pytest + +from tests.e2e.accuracy.helpers import assert_video_similarity_metrics, probe_video +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +MODEL = "BestWishYsh/Helios-Distilled" +PROMPT = "A serene lakeside sunrise with mist over the water." +WIDTH = 640 +HEIGHT = 384 +FPS = 16 +NUM_FRAMES = 65 +NUM_INFERENCE_STEPS = 5 +SEED = 42 +MIN_PSNR_DB = 28.0 +MIN_SSIM = 0.95 + +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + +HELIOS_DISTILLED_EXTRA_PARAMS = { + "is_enable_stage2": True, + "pyramid_num_stages": 3, + "pyramid_num_inference_steps_list": [1, 1, 1], + "is_amplify_first_chunk": True, +} + + +def _server_args(*, streaming_output: bool) -> list[str]: + args = [ + "--stage-init-timeout", + "600", + "--init-timeout", + "900", + "--log-stats", + ] + if streaming_output: + args.append("--diffusion-streaming-output") + return args + + +def _request_config(model: str, *, serialize_extra_params: bool) -> dict: + extra_params = ( + json.dumps(HELIOS_DISTILLED_EXTRA_PARAMS) if serialize_extra_params else HELIOS_DISTILLED_EXTRA_PARAMS + ) + return { + "model": model, + "form_data": { + "prompt": PROMPT, + "width": WIDTH, + "height": HEIGHT, + "fps": FPS, + "num_frames": NUM_FRAMES, + "num_inference_steps": NUM_INFERENCE_STEPS, + "guidance_scale": 1.0, + "seed": SEED, + "extra_params": extra_params, + }, + } + + +def _client_for(server: OmniServer, *, run_level: str) -> OpenAIClientHandler: + return OpenAIClientHandler( + host=server.host, + port=server.port, + api_key="EMPTY", + run_level=run_level, + log_stats=server.log_stats, + ) + + +def _write_video_artifact(response: DiffusionResponse, output_path: Path) -> Path: + assert response.videos, "No video artifact returned" + output_path.write_bytes(response.videos[0]) + return output_path + + +@pytest.mark.parametrize( + "model", + [ + pytest.param( + MODEL, + id="helios_distilled", + marks=SINGLE_CARD_MARKS, + ) + ], +) +def test_helios_streaming_output_matches_non_streaming( + model: str, + model_prefix: str, + run_level: str, + tmp_path: Path, + record_property, +): + model = model_prefix + model + + with OmniServer(model, _server_args(streaming_output=True), use_omni=True) as server: + client = _client_for(server, run_level=run_level) + streaming_response = client.send_streaming_video_diffusion_request( + _request_config(server.model, serialize_extra_params=False) + )[0] + + streaming_path = _write_video_artifact(streaming_response, tmp_path / "helios_streaming.mp4") + record_property("helios_streaming_video", str(streaming_path)) + record_property("helios_streaming_latency_s", streaming_response.e2e_latency) + + with OmniServer(model, _server_args(streaming_output=False), use_omni=True) as server: + client = _client_for(server, run_level=run_level) + non_streaming_response = client.send_video_diffusion_request( + _request_config(server.model, serialize_extra_params=True) + )[0] + + non_streaming_path = _write_video_artifact(non_streaming_response, tmp_path / "helios_non_streaming.mp4") + record_property("helios_non_streaming_video", str(non_streaming_path)) + record_property("helios_non_streaming_latency_s", non_streaming_response.e2e_latency) + + streaming_metadata = probe_video(streaming_path) + non_streaming_metadata = probe_video(non_streaming_path) + assert streaming_metadata == non_streaming_metadata, ( + f"Video metadata mismatch for helios_streaming_output:\n" + f"streaming={streaming_metadata}\n" + f"non_streaming={non_streaming_metadata}\n" + f"streaming_path={streaming_path}\n" + f"non_streaming_path={non_streaming_path}" + ) + + ssim, psnr = assert_video_similarity_metrics( + label="helios_streaming_output", + online_path=streaming_path, + offline_path=non_streaming_path, + ssim_threshold=MIN_SSIM, + psnr_threshold=MIN_PSNR_DB, + ) + record_property("helios_streaming_vs_non_streaming_psnr_db", psnr) + record_property("helios_streaming_vs_non_streaming_ssim_percent", ssim * 100.0) diff --git a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py index 709c6655565..193d1b1e2bd 100644 --- a/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py +++ b/tests/e2e/offline_inference/custom_pipeline/qwen_image_pipeline_with_logprob.py @@ -21,7 +21,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.qwen_image import QwenImagePipeline -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch def _maybe_to_cpu(v): @@ -33,6 +33,8 @@ def _maybe_to_cpu(v): # Custom pipeline class for QwenImage that returns log probabilities during the diffusion process. # This is for test class QwenImagePipelineWithLogProbForTest(QwenImagePipeline): + supports_request_batch = False + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) self.device = get_local_device() @@ -211,7 +213,7 @@ def diffuse( def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt_ids: torch.Tensor | list[int] | None = None, prompt_mask: torch.Tensor | None = None, negative_prompt_ids: torch.Tensor | list[int] | None = None, @@ -239,8 +241,12 @@ def forward( sde_type: Literal["sde", "cps"] = "sde", logprobs: bool = True, ) -> DiffusionOutput: - # Extract prompt data from OmniCustomPrompt in req.prompts[0] - custom_prompt = req.prompts[0] if req.prompts else {} + if req.num_reqs != 1: + raise ValueError("QwenImagePipelineWithLogProbForTest supports only single-request forward.") + request = req.requests[0] + + # Extract prompt data from OmniCustomPrompt in req.prompt + custom_prompt = request.prompt if request.prompt is not None else {} if isinstance(custom_prompt, dict): prompt_ids = custom_prompt.get("prompt_ids", prompt_ids) prompt_mask = custom_prompt.get("prompt_mask", prompt_mask) @@ -248,7 +254,7 @@ def forward( negative_prompt_mask = custom_prompt.get("negative_prompt_mask", negative_prompt_mask) # Read sampling params from req.sampling_params - sp = req.sampling_params + sp = request.sampling_params height = sp.height or self.default_sample_size * self.vae_scale_factor width = sp.width or self.default_sample_size * self.vae_scale_factor num_inference_steps = sp.num_inference_steps or num_inference_steps @@ -282,7 +288,7 @@ def forward( else: # Both prompt_ids and prompt_embeds are None (e.g. during warmup/dummy run). # Return a minimal dummy output to avoid crashing. - return DiffusionOutput(output=None, custom_output={}) + return DiffusionOutput(output=None) if isinstance(negative_prompt_ids, list): negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device) @@ -393,14 +399,20 @@ def forward( image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] return DiffusionOutput( - output=_maybe_to_cpu(image), + output={ + "payload": { + "image": _maybe_to_cpu(image), + }, + "metadata": { + "prompt_embeddings": { + "prompt_embeds": _maybe_to_cpu(prompt_embeds), + "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask), + "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds), + "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask), + }, + }, + }, trajectory_latents=_maybe_to_cpu(all_latents), trajectory_log_probs=_maybe_to_cpu(all_log_probs), trajectory_timesteps=_maybe_to_cpu(all_timesteps), - custom_output={ - "prompt_embeds": _maybe_to_cpu(prompt_embeds), - "prompt_embeds_mask": _maybe_to_cpu(prompt_embeds_mask), - "negative_prompt_embeds": _maybe_to_cpu(negative_prompt_embeds), - "negative_prompt_embeds_mask": _maybe_to_cpu(negative_prompt_embeds_mask), - }, ) diff --git a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py index 9bcaf0e4079..69952435e9c 100644 --- a/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py +++ b/tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py @@ -390,20 +390,23 @@ async def generate( diffusion_output = self._to_tensor(final_res.images[0]).float() / 255.0 - mm_output = final_res.custom_output or {} + metadata = ( + final_res.multimodal_output.get("metadata", {}) if isinstance(final_res.multimodal_output, dict) else {} + ) + prompt_embeddings = metadata.get("prompt_embeddings", {}) if isinstance(metadata, dict) else {} if sampling_params.get("logprobs", False): - all_log_probs = mm_output.get("all_log_probs") + all_log_probs = final_res.trajectory_log_probs log_probs = all_log_probs[0] if all_log_probs is not None else None else: log_probs = None - all_latents = mm_output.get("all_latents") - all_timesteps = mm_output.get("all_timesteps") - prompt_embeds = mm_output.get("prompt_embeds") - prompt_embeds_mask = mm_output.get("prompt_embeds_mask") - negative_prompt_embeds = mm_output.get("negative_prompt_embeds") - negative_prompt_embeds_mask = mm_output.get("negative_prompt_embeds_mask") + all_latents = final_res.trajectory_latents + all_timesteps = final_res.trajectory_timesteps + prompt_embeds = prompt_embeddings.get("prompt_embeds") + prompt_embeds_mask = prompt_embeddings.get("prompt_embeds_mask") + negative_prompt_embeds = prompt_embeddings.get("negative_prompt_embeds") + negative_prompt_embeds_mask = prompt_embeddings.get("negative_prompt_embeds_mask") extra_fields = { "all_latents": all_latents[0] if all_latents is not None else None, diff --git a/tests/e2e/offline_inference/test_audiox_model_expansion.py b/tests/e2e/offline_inference/test_audiox_model_expansion.py index b78dc1a5eef..d4b2f5223c8 100644 --- a/tests/e2e/offline_inference/test_audiox_model_expansion.py +++ b/tests/e2e/offline_inference/test_audiox_model_expansion.py @@ -21,7 +21,7 @@ _OMNI_RUNNER_PARAMS = [(m, None, {"model_class_name": "AudioXPipeline"}) for m in models] pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True, ids=models), ] diff --git a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py index ea6ac8dd3ce..49079266d6c 100644 --- a/tests/e2e/offline_inference/test_cosyvoice3_expansion.py +++ b/tests/e2e/offline_inference/test_cosyvoice3_expansion.py @@ -141,7 +141,7 @@ def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[s ), ] pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.tts, pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True), ] diff --git a/tests/e2e/offline_inference/test_covo_audio_expansion.py b/tests/e2e/offline_inference/test_covo_audio_expansion.py index 3c6e5efb752..11672686e40 100644 --- a/tests/e2e/offline_inference/test_covo_audio_expansion.py +++ b/tests/e2e/offline_inference/test_covo_audio_expansion.py @@ -14,14 +14,14 @@ models = ["tencent/Covo-Audio-Chat"] stage_config = get_deploy_config_path("covo_audio.yaml") -test_params = [(model, stage_config) for model in models] +test_params = [(model, stage_config, {"trust_remote_code": True}) for model in models] def get_question(): return "请回答这段音频里的问题。" -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards={"cuda": 1}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) diff --git a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py index a3f57064f35..e22aa117470 100644 --- a/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py +++ b/tests/e2e/offline_inference/test_dreamid_omni_cache_dit.py @@ -30,7 +30,7 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow] +pytestmark = [pytest.mark.diffusion, pytest.mark.cache, pytest.mark.slow] REPO_ROOT = Path(__file__).resolve().parents[3] DEFAULT_MODEL_PATH = REPO_ROOT / "dreamid_omni" diff --git a/tests/e2e/offline_inference/test_dynin_omni_expansion.py b/tests/e2e/offline_inference/test_dynin_omni_expansion.py index 7aae111f058..19a84733276 100644 --- a/tests/e2e/offline_inference/test_dynin_omni_expansion.py +++ b/tests/e2e/offline_inference/test_dynin_omni_expansion.py @@ -29,12 +29,12 @@ models = ["snu-aidas/Dynin-Omni"] stage_configs = [str(_DEFAULT_STAGE_CONFIG_PATH)] -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +test_params = [(model, stage_config, {"trust_remote_code": True}) for model in models for stage_config in stage_configs] DYNIN_CONFIG_PATH = str(_DEFAULT_DYNIN_CONFIG_PATH) if _DEFAULT_DYNIN_CONFIG_PATH is not None else None pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.omni, pytest.mark.parametrize("omni_runner", test_params, indirect=True), ] @@ -290,7 +290,7 @@ def _numel(value: Any) -> int: return 0 -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_dynin_t2i_decode_to_image(omni_runner) -> None: _configure_dynin_config_env() prompt = _build_t2i_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH) @@ -304,7 +304,7 @@ def test_dynin_t2i_decode_to_image(omni_runner) -> None: assert _numel(image_value) > 0 -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_dynin_mmu_to_text(omni_runner) -> None: _configure_dynin_config_env() tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) @@ -322,7 +322,7 @@ def test_dynin_mmu_to_text(omni_runner) -> None: assert text_content -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_dynin_image_to_text(omni_runner) -> None: _configure_dynin_config_env() tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) @@ -341,7 +341,7 @@ def test_dynin_image_to_text(omni_runner) -> None: assert text_content -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_dynin_speech_to_text(omni_runner) -> None: _configure_dynin_config_env() tokenizer = AutoTokenizer.from_pretrained(omni_runner.model_name, trust_remote_code=True) @@ -360,7 +360,7 @@ def test_dynin_speech_to_text(omni_runner) -> None: assert text_content -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_dynin_t2s_decode_to_audio(omni_runner) -> None: _configure_dynin_config_env() prompt = _build_t2s_decode_prompt(dynin_config_path=DYNIN_CONFIG_PATH) diff --git a/tests/e2e/offline_inference/test_ernie_image_expansion.py b/tests/e2e/offline_inference/test_ernie_image_expansion.py index 2375ab38bd9..19793ff0007 100644 --- a/tests/e2e/offline_inference/test_ernie_image_expansion.py +++ b/tests/e2e/offline_inference/test_ernie_image_expansion.py @@ -30,7 +30,7 @@ ) pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] diff --git a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py index a9b24ef9d4a..4e9e3071289 100644 --- a/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_dev_autoround_w4a16_expansion.py @@ -105,7 +105,7 @@ def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]: @pytest.mark.skip(reason="BF16 baseline model cannot be loaded correctly yet.") -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}) @pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) @@ -124,7 +124,7 @@ def test_flux_autoround_w4a16_generates_image(omni_runner: OmniRunner): @pytest.mark.skip(reason="BF16 baseline model cannot be loaded correctly yet.") -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}) def test_flux_autoround_w4a16_memory_savings(): diff --git a/tests/e2e/offline_inference/test_flux1_kontext_expansion.py b/tests/e2e/offline_inference/test_flux1_kontext_expansion.py index 3ab1aa40136..abf7bbc2c50 100644 --- a/tests/e2e/offline_inference/test_flux1_kontext_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_kontext_expansion.py @@ -14,6 +14,7 @@ import pytest from vllm.assets.image import ImageAsset +from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunnerHandler from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -30,7 +31,7 @@ ) pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] @@ -45,6 +46,7 @@ def _sampling_512() -> OmniDiffusionSamplingParams: ) +@hardware_test(res={"cuda": "H100"}, num_cards=2) def test_flux_kontext_text_to_image(omni_runner_handler: OmniRunnerHandler): """Test FluxKontext text-to-image generation with real model.""" request_config = { @@ -55,6 +57,7 @@ def test_flux_kontext_text_to_image(omni_runner_handler: OmniRunnerHandler): omni_runner_handler.send_diffusion_request(request_config) +@hardware_test(res={"cuda": "H100"}, num_cards=2) def test_flux_kontext_image_edit(omni_runner_handler: OmniRunnerHandler): """Test FluxKontext image-to-image editing with real model.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") diff --git a/tests/e2e/offline_inference/test_flux1_schnell_expansion.py b/tests/e2e/offline_inference/test_flux1_schnell_expansion.py index 583ceaf5979..de6a758f690 100644 --- a/tests/e2e/offline_inference/test_flux1_schnell_expansion.py +++ b/tests/e2e/offline_inference/test_flux1_schnell_expansion.py @@ -5,6 +5,7 @@ import pytest +from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunnerHandler from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -13,12 +14,13 @@ _OMNI_RUNNER_PARAM = (MODEL, None) pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] +@hardware_test(res={"cuda": "H100"}, num_cards=1) def test_flux_schnell_text_to_image(omni_runner_handler: OmniRunnerHandler): """Test FLUX.1-schnell text-to-image generation.""" request_config = { diff --git a/tests/e2e/offline_inference/test_flux2_klein.py b/tests/e2e/offline_inference/test_flux2_klein.py index 3bd8152291c..9656772c9a3 100644 --- a/tests/e2e/offline_inference/test_flux2_klein.py +++ b/tests/e2e/offline_inference/test_flux2_klein.py @@ -18,6 +18,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform +# NOTE: Most of the flux2klein tests have migrated to the common diffusion tests. +# This model will also use the tiny model builder when run at the core model level. MODEL = "black-forest-labs/FLUX.2-klein-4B" _OMNI_RUNNER_PARAM = (MODEL, None) @@ -70,41 +72,6 @@ def _send_inpaint_with_generator( ) -# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097 -@pytest.mark.advanced_model -@pytest.mark.diffusion -def test_flux2_klein_can_accept_text_inputs(omni_runner_handler: OmniRunnerHandler): - omni_runner_handler.send_diffusion_request( - { - "model": MODEL, - "prompt": "a cup of coffee on the table", - "sampling_params": OmniDiffusionSamplingParams(num_inference_steps=2, seed=42), - } - ) - - -@pytest.mark.advanced_model -@pytest.mark.diffusion -@hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_flux2_klein_inpaint_basic(omni_runner_handler: OmniRunnerHandler): - input_image, mask_image = _create_test_inputs() - omni_runner_handler.send_diffusion_request( - { - "model": MODEL, - "prompt": "Fill in the masked area with a beautiful garden", - "multi_modal_data": {"image": input_image, "mask_image": mask_image}, - "sampling_params": OmniDiffusionSamplingParams( - height=_HEIGHT, - width=_WIDTH, - num_inference_steps=_NUM_INFERENCE_STEPS, - guidance_scale=0.0, - generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), - num_outputs_per_prompt=1, - ), - } - ) - - @pytest.mark.advanced_model @pytest.mark.diffusion @hardware_test(res={"cuda": "L4"}, num_cards=1) diff --git a/tests/e2e/offline_inference/test_glm_tts.py b/tests/e2e/offline_inference/test_glm_tts_expansion.py similarity index 90% rename from tests/e2e/offline_inference/test_glm_tts.py rename to tests/e2e/offline_inference/test_glm_tts_expansion.py index 0b902c6bc2f..e54dd5480f1 100644 --- a/tests/e2e/offline_inference/test_glm_tts.py +++ b/tests/e2e/offline_inference/test_glm_tts_expansion.py @@ -1,10 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E Offline inference tests for GLM-TTS voice cloning. - -GLM-TTS is a zero-shot voice cloning model that always requires ref_audio. -There is no text-only / non-clone inference path in the official implementation. -""" +"""E2E expansion tests for GLM-TTS voice cloning offline inference (nightly CI).""" import os import tempfile @@ -23,11 +19,12 @@ from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.model_executor.models.glm_tts.glm_tts import build_glm_tts_prefill_metadata +pytestmark = [pytest.mark.slow, pytest.mark.tts] + MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS") REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。" DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml") - REFERENCE_PROMPT_WAV_PATH = get_asset_path("glm_tts/jiayan_zh.wav") ASYNC_CHUNK_MODES = [ @@ -84,9 +81,7 @@ def _get_deploy_config(*, async_chunk: bool) -> str: ) -@pytest.mark.advanced_model -@pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("async_chunk", ASYNC_CHUNK_MODES) def test_offline_voice_clone_zh(async_chunk: bool) -> None: """ @@ -94,9 +89,6 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None: Deploy Setting: glm_tts.yaml with explicit sync/async mode, enforce_eager Input Modal: text + ref_audio + ref_text Output Modal: audio - - Uses the official jiayan_zh.wav reference audio from the upstream - GLM-TTS repository (real human speech matching REF_TEXT). """ synth_text = "今天天气真不错，适合出去散散步。" prompt_audio = _load_ref_audio() @@ -106,6 +98,7 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None: MODEL, stage_configs_path=_get_deploy_config(async_chunk=async_chunk), stage_init_timeout=600, + trust_remote_code=True, ) as omni_runner: outputs = omni_runner.omni.generate( [ @@ -144,5 +137,5 @@ def test_offline_voice_clone_zh(async_chunk: bool) -> None: audio_format="audio/wav", ), {"input": synth_text, "response_format": "wav"}, - run_level="advanced_model", + run_level="full_model", ) diff --git a/tests/e2e/offline_inference/test_higgs_audio_v2.py b/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py similarity index 96% rename from tests/e2e/offline_inference/test_higgs_audio_v2.py rename to tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py index fdeccc9243d..00482a132a4 100644 --- a/tests/e2e/offline_inference/test_higgs_audio_v2.py +++ b/tests/e2e/offline_inference/test_higgs_audio_v2_expansion.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -End-to-end offline inference tests for higgs-audio v2. +E2E Offline expansion tests for higgs-audio v2. Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py: build Stage-0 prompt_token_ids with the model's own processor + plain-text @@ -34,6 +34,8 @@ from tests.helpers.runtime import OmniRunner from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.slow, pytest.mark.tts] + MODEL = "bosonai/higgs-audio-v2-generation-3B-base" STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml") SAMPLE_RATE = 24_000 @@ -89,9 +91,7 @@ def _extract_pcm(outputs) -> torch.Tensor: raise AssertionError(f"no audio payload in any stage output (got {len(outputs)} stages)") -@pytest.mark.core_model -@pytest.mark.tts -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> None: """ @@ -131,9 +131,7 @@ def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> assert np.isfinite(arr).all(), "audio contains non-finite samples (nan/inf)" -@pytest.mark.core_model -@pytest.mark.tts -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) def test_higgs_audio_v2_offline_batch_two_prompts(omni_runner_function: OmniRunner) -> None: """ @@ -199,9 +197,7 @@ def _compose_voice_clone_request(model_name: str, text: str, ref_text: str, ref_ } -@pytest.mark.advanced_model -@pytest.mark.tts -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) def test_higgs_audio_v2_offline_voice_clone(omni_runner_function: OmniRunner) -> None: """ diff --git a/tests/e2e/offline_inference/test_hunyuanimage3.py b/tests/e2e/offline_inference/test_hunyuanimage3.py index ac1cb13cba7..ebd338fe8fd 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3.py @@ -275,6 +275,7 @@ def omni() -> Generator[Omni, None, None]: MODEL_NAME, deploy_config=str(DEPLOY_CONFIG_PATH), mode="text-to-image", + trust_remote_code=True, ) as runner: yield runner.omni diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py index 5178a40a6ce..008f0d2a9be 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_i2t_expansion.py @@ -53,6 +53,7 @@ def omni() -> Generator[Omni, None, None]: with OmniRunner( MODEL_NAME, stage_configs_path=str(STAGE_CONFIG_PATH), + trust_remote_code=True, ) as runner: yield runner.omni diff --git a/tests/e2e/offline_inference/test_indextts2_expansion.py b/tests/e2e/offline_inference/test_indextts2_expansion.py new file mode 100644 index 00000000000..ba5b984f926 --- /dev/null +++ b/tests/e2e/offline_inference/test_indextts2_expansion.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for IndexTTS2 two-stage pipeline. + +Stage 0 (GPT AR) → Stage 1 (S2Mel + BigVGAN). Output is 22050 Hz mono WAV. +Covers: basic TTS. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch + +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 Omni +from vllm_omni.model_executor.models.indextts2.prompt_utils import ( + build_indextts2_prefill_prompt_ids, +) + +MODEL_NAME = "IndexTeam/IndexTTS-2" +STAGE_CONFIG = get_deploy_config_path("indextts2.yaml") + +_OMNI_RUNNER_PARAM = ( + MODEL_NAME, + STAGE_CONFIG, +) + +ASSETS_DIR = Path(__file__).resolve().parents[2] / "assets" / "indextts2" +REF_AUDIO = str(ASSETS_DIR / "ref_audio.wav") + +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), +] + +SAMPLE_RATE = 22050 + + +def _audio_from_mm(mm: dict) -> torch.Tensor | None: + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if isinstance(audio, list): + chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0] + return torch.cat(chunks, dim=0) if chunks else None + return audio if isinstance(audio, torch.Tensor) else None + + +def _sample_rate_from_mm(mm: dict) -> int: + sr = mm.get("sr") + if isinstance(sr, list): + sr = sr[-1] if sr else None + if hasattr(sr, "item"): + return int(sr.item()) + return int(sr) if sr is not None else SAMPLE_RATE + + +def _build_request(text: str, ref_audio: str, **extra) -> dict: + info = { + "text": [text], + "voice": [ref_audio], + } + info.update(extra) + return { + "prompt_token_ids": build_indextts2_prefill_prompt_ids(MODEL_NAME, text), + "additional_information": info, + } + + +def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: + for omni_out in omni.generate(request): + mm = omni_out.multimodal_output + assert mm is not None, "Expected multimodal_output" + audio = _audio_from_mm(mm) + assert audio is not None, "Expected audio output" + return audio.cpu(), _sample_rate_from_mm(mm) + raise AssertionError("No outputs received") + + +def _assert_valid_audio(audio: torch.Tensor, sr: int) -> None: + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert torch.any(audio != 0).item() + + +@hardware_test(res={"cuda": "L4"}) +def test_basic_english(omni_runner: OmniRunner) -> None: + """English TTS: text + ref_audio, no emotion conditioning.""" + req = _build_request("Hello, this is a voice synthesis test.", REF_AUDIO) + audio, sr = _collect_audio(omni_runner.omni, req) + _assert_valid_audio(audio, sr) diff --git a/tests/e2e/offline_inference/test_krea2_expansion.py b/tests/e2e/offline_inference/test_krea2_expansion.py new file mode 100644 index 00000000000..f00b4bb0be3 --- /dev/null +++ b/tests/e2e/offline_inference/test_krea2_expansion.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""L4 functionality tests for the Krea 2 text-to-image diffusion pipeline (``Krea2Pipeline``). + +Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default (override with the +``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the Raw checkpoint or a local diffusers +directory). They cover a basic functional smoke plus the layerwise-CPU-offload path that the pipeline declares +via ``SupportsComponentDiscovery`` / ``_layerwise_offload_blocks_attrs``. +""" + +import os + +import numpy as np +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunnerHandler +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 = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo") +# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32). +LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA") +PROMPT = "a fox in the snow, photorealistic" +# darkbrush is a trigger-word LoRA: the trained monochrome ink-wash style only manifests when +# its trigger phrase is present, so the LoRA test prompt appends it for a strong, stable signal. +LORA_TRIGGER = "monochrome ink wash style" +LORA_PROMPT = f"a fox in the snow, {LORA_TRIGGER}" + +pytestmark = [ + pytest.mark.diffusion, + pytest.mark.full_model, +] + + +def _sampling() -> OmniDiffusionSamplingParams: + # Small resolution + few steps to keep the L4 case light. guidance_scale resolution is checkpoint-aware inside + # the pipeline (distilled -> no-CFG, Raw -> CFG), so this stays agnostic to which checkpoint KREA2_MODEL points at. + return OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=8, + guidance_scale=0.0, + seed=42, + ) + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) +def test_krea2_text_to_image_001(omni_runner_handler: OmniRunnerHandler) -> None: + """Basic functional smoke: a single prompt produces a decoded image.""" + omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize( + "omni_runner", + [(MODEL, None, {"enable_layerwise_offload": True})], + indirect=True, +) +def test_krea2_layerwise_offload(omni_runner_handler: OmniRunnerHandler) -> None: + """Exercise layerwise CPU offload on the DiT (SupportsComponentDiscovery + _layerwise_offload_blocks_attrs).""" + omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.cache +@pytest.mark.parametrize("omni_runner", [(MODEL, None, {"cache_backend": "cache_dit"})], indirect=True) +def test_krea2_cache_dit(omni_runner_handler: OmniRunnerHandler) -> None: + """Exercise Cache-DiT on Krea 2 via the custom Krea2Pipeline enabler (``enable_cache_for_krea2``). + + Validates the docs' Cache-DiT support claim for Krea 2. ``has_separate_cfg`` is checkpoint-aware + (False for the distilled Turbo no-CFG path, True for the Raw CFG path); the default Turbo checkpoint + exercises the no-CFG branch. + """ + omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize( + "omni_runner", + # Standalone HSDP: use_hsdp=True with an explicit hsdp_shard_size shards the DiT weights + # across 2 ranks (replicate_size defaults to 1 -> world_size = 1 * 2 = 2). HSDP must keep + # tensor_parallel_size/data_parallel_size at 1 (see vllm_omni/diffusion/data.py __post_init__). + # Also validated locally at 8 GPUs (hsdp_shard_size=8); CI shard provisions 2. + [(MODEL, None, {"use_hsdp": True, "hsdp_shard_size": 2})], + indirect=True, +) +def test_krea2_hsdp(omni_runner_handler: OmniRunnerHandler) -> None: + """Exercise Hybrid Sharded Data Parallel (HSDP) weight sharding on the DiT across 2 GPUs. + + Validates the docs' HSDP-support claim for Krea 2. Scheduled on the multi-GPU (distributed_cuda) + shard by ``num_cards=2``; skipped automatically on boxes with fewer than 2 CUDA devices. + """ + omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize( + "omni_runner", + # VAE patch-parallel decode: partition the VAE decode across the TP world. Mirrors the online + # "tp_vae_patch" case (tensor_parallel_size=2 + vae_patch_parallel_size=2) from the wan2.2/z-image + # expansion suites, so vae_patch_parallel_size must equal the world size established by TP. + # Also validated locally at 8 GPUs (tp=8, vae_patch=8); CI shard provisions 2. + [(MODEL, None, {"tensor_parallel_size": 2, "vae_patch_parallel_size": 2})], + indirect=True, +) +def test_krea2_vae_patch_parallel(omni_runner_handler: OmniRunnerHandler) -> None: + """Exercise VAE patch-parallel (decode) across 2 GPUs. + + Validates the docs' VAE-patch-parallel (decode) support claim for Krea 2. Scheduled on the + multi-GPU (distributed_cuda) shard by ``num_cards=2``; skipped automatically on boxes with + fewer than 2 CUDA devices. + """ + omni_runner_handler.send_diffusion_request({"model": MODEL, "prompt": PROMPT, "sampling_params": _sampling()}) + + +def _generate(handler: OmniRunnerHandler, lora_request: LoRARequest | None, lora_scale: float = 1.0) -> np.ndarray: + sp = _sampling() + sp.lora_request = lora_request + sp.lora_scale = lora_scale + resp = handler.send_diffusion_request({"model": MODEL, "prompt": LORA_PROMPT, "sampling_params": sp}) + return np.asarray(resp.images[0].convert("RGB"), dtype=np.int16) + + +@hardware_test(res={"cuda": "H100"}) +@pytest.mark.parametrize("omni_runner", [(MODEL, None)], indirect=True) +def test_krea2_lora(omni_runner_handler: OmniRunnerHandler) -> None: + """Validate diffusion LoRA on Krea 2: visible effect, scale sensitivity, clean deactivation. + + Uses ``LORA`` (a vLLM-Omni-compatible PEFT repackaging of ``krea/Krea-2-LoRA-darkbrush``, + 264 modules matched via the ReplicatedLinear projections) passed per-request through + ``OmniDiffusionSamplingParams.lora_request`` / ``lora_scale`` — the same fields the + diffusion LoRA manager reads to activate/deactivate the adapter. + """ + lora_request = LoRARequest(lora_name="darkbrush", lora_int_id=stable_lora_int_id(LORA), lora_path=LORA) + + baseline = _generate(omni_runner_handler, None) + img_1x = _generate(omni_runner_handler, lora_request, lora_scale=1.0) + img_2x = _generate(omni_runner_handler, lora_request, lora_scale=2.0) + restored = _generate(omni_runner_handler, None) + + diff_1x = np.abs(baseline - img_1x).mean() + diff_2x = np.abs(baseline - img_2x).mean() + diff_restored = np.abs(baseline - restored).mean() + + # (a) Adapter has a 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) Scale changes the output (stronger adapter -> larger perturbation). + assert not np.isclose(diff_1x, diff_2x, atol=1.0), f"LoRA scale had no effect: 1x={diff_1x:.2f}, 2x={diff_2x:.2f}" + # (c) Passing lora_request=None cleanly deactivates: same seed -> back to baseline (modulo fp drift). + assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}" diff --git a/tests/e2e/offline_inference/test_magi_human_expansion.py b/tests/e2e/offline_inference/test_magi_human_expansion.py index 9dd4a1d19cd..e6b3c6a429a 100644 --- a/tests/e2e/offline_inference/test_magi_human_expansion.py +++ b/tests/e2e/offline_inference/test_magi_human_expansion.py @@ -22,19 +22,31 @@ { "init_timeout": 1200, "tensor_parallel_size": 2, + "vae_use_tiling": True, + }, +) + +_OMNI_RUNNER_PARAM_OFFLOAD = ( + MODEL, + None, + { + "init_timeout": 1200, + "tensor_parallel_size": 2, + "enable_cpu_offload": True, + "vae_use_tiling": True, }, ) pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, - pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), + pytest.mark.parametrize( + "omni_runner", + [_OMNI_RUNNER_PARAM, _OMNI_RUNNER_PARAM_OFFLOAD], + indirect=True, + ), ] -_SKIP_ISSUE_3236 = pytest.mark.skip( - reason="https://github.com/vllm-project/vllm-omni/issues/3236", -) - def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None: """Validate that the MP4 contains meaningful video and audio tracks.""" @@ -60,7 +72,6 @@ def _validate_mp4(video_bytes: bytes, min_frames: int = 10) -> None: container.close() -@_SKIP_ISSUE_3236 @hardware_test(res={"cuda": "H100"}, num_cards=2) def test_magi_human_e2e(omni_runner: OmniRunner): """End-to-end test for MagiHuman generating video and audio.""" diff --git a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py index 931ce5931e9..04308c4a796 100644 --- a/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py +++ b/tests/e2e/offline_inference/test_mammoth_moda2_expansion.py @@ -5,7 +5,7 @@ match a golden reference. Model Hub repo id: ``bytedance-research/MammothModa2-Preview``. -Stage config: ``get_deploy_config_path("mammoth_moda2.yaml")`` → ``vllm_omni/deploy/mammoth_moda2.yaml`` +Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml`` Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json`` Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py`` @@ -36,9 +36,9 @@ _AR_PATCH_SIZE = 16 MODEL_PATH = "bytedance-research/MammothModa2-Preview" -T2I_STAGE_CONFIG = get_deploy_config_path("mammoth_moda2.yaml") +T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml") -_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_STAGE_CONFIG) +_OMNI_RUNNER_PARAM = (MODEL_PATH, None, {"deploy_config": T2I_DEPLOY_CONFIG}) # Golden pixel reference file. Set UPDATE_GOLDEN=1 to regenerate. _GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json" @@ -99,7 +99,7 @@ def _sample_pixels(img_tensor: torch.Tensor) -> list[float]: # End-to-end test # --------------------------------------------------------------------------- pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.diffusion, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] 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 19cf6fa4ce6..6d4bd197a25 100644 --- a/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py +++ b/tests/e2e/offline_inference/test_ming_flash_omni_expansion.py @@ -49,7 +49,7 @@ def get_eager_config(): test_params = [(m, c) for m in models for c in [get_eager_config()]] -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) @@ -65,7 +65,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) @@ -82,7 +82,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) @@ -101,7 +101,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) @@ -118,7 +118,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True) @@ -138,7 +138,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) @@ -154,7 +154,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) diff --git a/tests/e2e/offline_inference/test_ming_tts.py b/tests/e2e/offline_inference/test_ming_tts_expansion.py similarity index 97% rename from tests/e2e/offline_inference/test_ming_tts.py rename to tests/e2e/offline_inference/test_ming_tts_expansion.py index 775942274ac..82da742273f 100644 --- a/tests/e2e/offline_inference/test_ming_tts.py +++ b/tests/e2e/offline_inference/test_ming_tts_expansion.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""End-to-end offline inference tests for Ming-omni-tts.""" +"""E2E expansion tests for Ming-omni-tts offline inference (nightly CI).""" import asyncio import uuid @@ -32,6 +32,12 @@ TEST_INSTRUCTION = "轻柔的ASMR耳语，慢速，贴近麦克风" MIN_AUDIO_SAMPLES = 1000 +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4704"), +] + @pytest.fixture(scope="module") def ming_tokenizer(): @@ -163,8 +169,6 @@ def _extract_multimodal_output(output) -> Mapping[str, Any]: 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.""" @@ -185,8 +189,6 @@ def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None: 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.""" @@ -220,8 +222,6 @@ def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> N ) -@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.""" @@ -251,8 +251,6 @@ def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None: 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.""" diff --git a/tests/e2e/offline_inference/test_moss_tts.py b/tests/e2e/offline_inference/test_moss_tts.py deleted file mode 100644 index 505db1af718..00000000000 --- a/tests/e2e/offline_inference/test_moss_tts.py +++ /dev/null @@ -1,225 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E offline inference tests for MOSS-TTS two-stage pipeline. - -Tests the MossTTSDelayModel path using MOSS-VoiceGenerator (1.7B) — the -smallest MossTTSDelayModel variant — so the test runs on a single A10G or L4 -without requiring an 80 GB GPU. - -A second test class covers MOSS-TTS-Realtime (1.7B, MossTTSRealtime) to -exercise the local-transformer generation path. - -Both classes use a module-scoped engine so the model is loaded once per test -module execution. Do NOT add a second engine fixture in this file — that -triggers mid-module teardown races (see skill invariant I4). -""" - -from __future__ import annotations - -import os -import urllib.request - -import pytest -import torch -from vllm import SamplingParams - -from tests.helpers.mark import hardware_test -from vllm_omni import Omni - -# --------------------------------------------------------------------------- -# Shared constants -# --------------------------------------------------------------------------- - -# Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts"). -# Without the tts + run-level marks these tests are collected but never run. -# TODO: Fix this test -# pytestmark = [pytest.mark.full_model, pytest.mark.tts] - -SAMPLE_RATE = 24_000 # All MOSS-TTS full variants output 24 kHz - -REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" - -_DEFAULT_SAMPLING = SamplingParams( - temperature=1.7, - top_p=0.8, - top_k=25, - max_tokens=512, # short for tests (~20 s at 24 kHz, ~12.5 tokens/s) - seed=42, - detokenize=False, -) - -# --------------------------------------------------------------------------- -# Session fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture(scope="session") -def ref_audio_path(tmp_path_factory) -> str: - """Download the upstream reference clip once per session. - - Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments. - """ - cache_dir = tmp_path_factory.mktemp("moss_tts_ref") - target = cache_dir / "zh_1.wav" - try: - with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: - target.write_bytes(resp.read()) - except Exception as exc: - msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" - if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): - pytest.skip(msg) - pytest.fail(msg) - if not target.exists() or target.stat().st_size == 0: - pytest.fail(f"Reference clip empty after download: {target}") - return str(target) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict: - return { - "prompt": "<|im_start|>assistant\n", - "additional_information": { - "task_type": ["voice_clone"], - "text": [text], - "mode": ["voice_clone"], - "prompt_audio_path": [ref_audio_path], - "seed": [seed], - }, - } - - -def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: - """Run one request; return (waveform_cpu, sample_rate).""" - for stage_outputs in omni.generate(request, _DEFAULT_SAMPLING): - for req_output in stage_outputs.request_output: - mm = req_output.outputs[0].multimodal_output - assert mm is not None, "Expected non-None multimodal_output" - # Never use ``a or b`` on tensor values — a multi-element tensor - # raises on truthiness. Check ``is None`` explicitly and handle the - # pre-consolidation list form. - audio = mm.get("audio") - if audio is None: - audio = mm.get("model_outputs") - sr = mm.get("sr") - assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output" - if isinstance(audio, list): - audio = torch.cat( - [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], - dim=0, - ) - assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}" - return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE - raise AssertionError("No stage outputs received") - - -# --------------------------------------------------------------------------- -# MossTTSDelayModel — MOSS-VoiceGenerator 1.7B -# (tests the delay-pattern generation path on a small GPU) -# --------------------------------------------------------------------------- - -_DELAY_MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator" - - -@pytest.fixture(scope="module") -def delay_engine(): - return Omni(model=_DELAY_MODEL, stage_init_timeout=300) - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_english(delay_engine, ref_audio_path): - """MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio.""" - req = _build_request("Hello, this is a MOSS-TTS voice cloning test.", ref_audio_path) - audio, sr = _collect_audio(delay_engine, req) - - assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" - assert audio.numel() > 0, "Audio tensor is empty" - assert not torch.all(audio == 0), "Audio is silence" - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_chinese(delay_engine, ref_audio_path): - """MossTTSDelayModel: Chinese input produces non-empty audio.""" - req = _build_request("你好，这是语音合成测试。", ref_audio_path) - audio, sr = _collect_audio(delay_engine, req) - - assert sr == SAMPLE_RATE - assert audio.numel() > 0 - assert not torch.all(audio == 0) - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path): - """MossTTSDelayModel: same seed yields identical waveforms.""" - req = _build_request("Reproducibility check.", ref_audio_path, seed=99) - audio1, _ = _collect_audio(delay_engine, req) - audio2, _ = _collect_audio(delay_engine, req) - - assert audio1.shape == audio2.shape, "Shapes differ across identical seeds" - assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds" - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_delay_batch(delay_engine, ref_audio_path): - """MossTTSDelayModel: batch of two requests each returns non-empty audio.""" - requests = [ - _build_request("First sentence.", ref_audio_path), - _build_request("Second sentence.", ref_audio_path), - ] - results: list[torch.Tensor] = [] - for stage_outputs in delay_engine.generate(requests, [_DEFAULT_SAMPLING] * 2): - for req_output in stage_outputs.request_output: - mm = req_output.outputs[0].multimodal_output - assert mm is not None - audio = mm.get("audio") - if audio is None: - audio = mm.get("model_outputs") - assert audio is not None - if isinstance(audio, list): - audio = torch.cat( - [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], - dim=0, - ) - results.append(audio.reshape(-1).cpu()) - - assert len(results) == 2, f"Expected 2 outputs, got {len(results)}" - for i, audio in enumerate(results): - assert audio.numel() > 0, f"Audio[{i}] is empty" - - -# --------------------------------------------------------------------------- -# MossTTSRealtime — MOSS-TTS-Realtime 1.7B -# (tests the local-transformer generation path) -# --------------------------------------------------------------------------- - -_REALTIME_MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" - - -@pytest.fixture(scope="module") -def realtime_engine(): - return Omni(model=_REALTIME_MODEL, stage_init_timeout=300) - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_realtime_english(realtime_engine, ref_audio_path): - """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" - req = _build_request("This is a real-time TTS streaming test.", ref_audio_path) - audio, sr = _collect_audio(realtime_engine, req) - - assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" - assert audio.numel() > 0, "Audio tensor is empty" - assert not torch.all(audio == 0), "Audio is silence" - - -@hardware_test(res={"cuda": "L4"}) -def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path): - """MossTTSRealtime: same seed yields identical waveforms.""" - req = _build_request("Determinism check for realtime model.", ref_audio_path, seed=7) - audio1, _ = _collect_audio(realtime_engine, req) - audio2, _ = _collect_audio(realtime_engine, req) - - assert audio1.shape == audio2.shape - assert torch.allclose(audio1, audio2, atol=1e-4) diff --git a/tests/e2e/offline_inference/test_moss_tts_expansion.py b/tests/e2e/offline_inference/test_moss_tts_expansion.py new file mode 100644 index 00000000000..81a64fcd9d8 --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_expansion.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-VoiceGenerator (MossTTSDelayModel, 1.7B). + +Weekly expansion coverage (``pytest.mark.slow``) for the delay-model path. +Upstream keeps the same tests in ``test_moss_tts.py``; this branch uses the +``*_expansion.py`` naming convention and deletes the base file on merge. + +MOSS-VoiceGenerator synthesizes speech from a text *instruction* that describes +the desired voice (e.g. "a warm female voice with an American accent"). It does +NOT accept a reference audio clip. The request format requires running +AutoProcessor.from_pretrained once per call to encode the (text, instruction) +pair into the (prompt_token_ids, codes.ref) grid that the MossTTSDelayModel +talker expects — same as examples/offline_inference/text_to_speech/moss_tts/ +end2end.py:_build_unified_codes. + +Realtime coverage: ``test_moss_tts_realtime_expansion.py``. + +No determinism test: VoiceGenerator produces variable-length output even with +a fixed seed; waveform length reproducibility is not guaranteed. +""" + +from __future__ import annotations + +import gc +import os + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +MODEL = "OpenMOSS-Team/MOSS-VoiceGenerator" +SAMPLE_RATE = 24_000 + +# Stage 0 = AR talker; max_tokens=512 keeps tests fast (~20 s of audio). +# Stage 1 = codec decoder; greedy, large context to hold all codec tokens. +_STAGE0_PARAMS = SamplingParams( + temperature=1.7, + top_p=0.8, + top_k=25, + max_tokens=512, + seed=42, + detokenize=False, +) +_STAGE1_PARAMS = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + seed=42, + detokenize=False, +) +_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS] + + +def _get_test_config() -> str: + """Derive a CI-friendly config from moss_voice_generator.yaml.""" + return modify_stage_config( + get_deploy_config_path("moss_voice_generator.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + }, + }, + }, + ) + + +pytestmark = [ + pytest.mark.slow, + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), + pytest.mark.tts, + pytest.mark.parametrize( + "omni_runner", + [(MODEL, _get_test_config(), {"trust_remote_code": True})], + indirect=True, + ), +] + + +def _build_request(text: str, instruction: str) -> dict: + """Build a VoiceGenerator request using AutoProcessor.""" + from transformers import AutoProcessor + + try: + proc = AutoProcessor.from_pretrained(MODEL, trust_remote_code=True) + except Exception as exc: + msg = f"Cannot load AutoProcessor for {MODEL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + user_msg = proc.build_user_message(text=text, instruction=instruction) + batch = proc(conversations=[[user_msg]], mode="generation") + unified = batch["input_ids"][0] + text_ids = unified[:, 0].tolist() + audio_codes = unified[:, 1:].contiguous().to(torch.int64) + del proc + gc.collect() + + return { + "prompt_token_ids": text_ids, + "additional_information": {"codes": {"ref": audio_codes}}, + } + + +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + chunks: list[torch.Tensor] = [] + sr_final = SAMPLE_RATE + for out in omni_runner.omni.generate(request, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + sr = mm.get("sr") + if sr is not None: + sr_final = int(sr.item() if hasattr(sr, "item") else sr) + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + chunks.append(audio.reshape(-1).cpu()) + if not chunks: + raise AssertionError("No audio output received from generate()") + return torch.cat(chunks, dim=0), sr_final + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_delay_english(omni_runner: OmniRunner) -> None: + """VoiceGenerator: English instruction produces non-empty 24 kHz audio.""" + req = _build_request( + "Hello, this is a MOSS voice design test.", + "a warm female voice with an American accent", + ) + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" + assert audio.numel() > 0, "Audio tensor is empty" + assert not torch.all(audio == 0), "Audio is silence" + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_delay_chinese(omni_runner: OmniRunner) -> None: + """VoiceGenerator: Chinese input produces non-empty audio.""" + req = _build_request( + "你好，这是语音合成测试。", + "一个清晰温暖的女声", + ) + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert not torch.all(audio == 0) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_delay_batch(omni_runner: OmniRunner) -> None: + """VoiceGenerator: batch of two requests each returns non-empty audio.""" + requests = [ + _build_request("First sentence.", "a warm female voice"), + _build_request("Second sentence.", "a young male voice"), + ] + results: list[torch.Tensor] = [] + for out in omni_runner.omni.generate(requests, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + results.append(audio.reshape(-1).cpu()) + + assert len(results) >= 2, f"Expected at least 2 audio outputs, got {len(results)}" + for i, audio in enumerate(results): + assert audio.numel() > 0, f"Audio chunk[{i}] is empty" diff --git a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py index 0cc8deea2d0..d3e7fcb0f34 100644 --- a/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py +++ b/tests/e2e/offline_inference/test_moss_tts_nano_expansion.py @@ -29,11 +29,11 @@ _OMNI_RUNNER_PARAM = ( MODEL_NAME, STAGE_CONFIG, + {"trust_remote_code": True}, ) pytestmark = [ - pytest.mark.skip(reason="issue#4361"), - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.tts, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime.py b/tests/e2e/offline_inference/test_moss_tts_realtime.py new file mode 100644 index 00000000000..892a395ba9f --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_realtime.py @@ -0,0 +1,295 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). + +Uses the standard omni_runner + pytestmark pattern (one module-scoped engine +per file, skill invariant I4). The delay-model variant lives in the sibling +file test_moss_tts.py. + +The request format mirrors end2end.py:_build_realtime_prompt exactly: + 1. snapshot_download to locate the MossTTSRealtimeProcessor module. + 2. MOSS-Audio-Tokenizer (separate HF repo) to encode the reference clip. + 3. Build the (L, 17) int grid; split into prompt_token_ids + codes.ref + + ids.all (remaining text tokens fed one-per-step during decode). + +Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments. + +No determinism test: MossTTSRealtime uses a local depth transformer that is +sensitive to async scheduling timing; waveform reproducibility is not +guaranteed across back-to-back calls. +""" + +from __future__ import annotations + +import gc +import importlib.util +import os +import sys +import urllib.request +from pathlib import Path + +import numpy as np +import pytest +import soundfile as sf +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" +SAMPLE_RATE = 24_000 + +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" + +_STAGE0_PARAMS = SamplingParams( + temperature=1.0, + top_p=0.95, + top_k=50, + max_tokens=512, + seed=42, + detokenize=False, +) +_STAGE1_PARAMS = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=65536, + seed=42, + detokenize=False, +) +_SAMPLING = [_STAGE0_PARAMS, _STAGE1_PARAMS] + +# --------------------------------------------------------------------------- +# Deploy config +# --------------------------------------------------------------------------- + + +def _get_test_config() -> str: + """Derive a CI-friendly config from moss_tts_realtime.yaml. + + Reduces Stage 0 gpu_memory_utilization from 0.60 → 0.45 to leave headroom + for Stage 1 (0.12) on a shared L4/A10G. max_num_seqs=1 keeps per-test + peak memory predictable. + """ + return modify_stage_config( + get_deploy_config_path("moss_tts_realtime.yaml"), + updates={ + "stages": { + 0: { + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + }, + }, + }, + ) + + +# --------------------------------------------------------------------------- +# pytestmark — one engine for the whole module +# --------------------------------------------------------------------------- + +pytestmark = [ + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), + pytest.mark.full_model, + pytest.mark.tts, + pytest.mark.parametrize( + "omni_runner", + [(MODEL, _get_test_config(), {"trust_remote_code": True})], + indirect=True, + ), +] + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def ref_audio_path(tmp_path_factory) -> str: + """Download the upstream reference clip once per test module. + + Set MOSS_TTS_SKIP_ON_NET_FAIL=1 to skip in air-gapped environments. + """ + target = tmp_path_factory.mktemp("moss_tts_realtime_ref") / "zh_1.wav" + try: + with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: + target.write_bytes(resp.read()) + except Exception as exc: + msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + if not target.exists() or target.stat().st_size == 0: + pytest.fail(f"Reference clip empty after download: {target}") + return str(target) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor: + """Load and resample reference audio to (1, T) float32 at target_sr.""" + data, sr = sf.read(path, always_2d=True) # (T, C) + wav = torch.from_numpy(data.T.astype("float32")) # (C, T) + if wav.shape[0] > 1: + wav = wav.mean(dim=0, keepdim=True) + if sr != target_sr: + import torchaudio + + wav = torchaudio.functional.resample(wav, sr, target_sr) + return wav # (1, T) + + +def _build_request(ref_audio_path: str, text: str) -> dict: + """Build a MOSS-TTS-Realtime request. + + Mirrors end2end.py:_build_realtime_prompt exactly: + 1. Locate MossTTSRealtimeProcessor in the HF snapshot. + 2. Encode ref audio via MOSS-Audio-Tokenizer (separate repo). + 3. Build the (L, 17) int grid and return prompt_token_ids + + additional_information: {codes.ref, ids.all}. + + Frees the codec model before returning to avoid competing with the + running vllm engine for GPU/CPU memory. + """ + from huggingface_hub import snapshot_download + from transformers import AutoModel, AutoTokenizer + + # Step 1: locate the realtime processor module in the snapshot. + try: + snap_dir = Path(snapshot_download(repo_id=MODEL)) + except Exception as exc: + msg = f"Cannot locate snapshot for {MODEL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + proc_module_path = snap_dir / "processing_mossttsrealtime.py" + if not proc_module_path.exists(): + pytest.fail(f"Realtime processor module missing at {proc_module_path}") + + mod_name = "_moss_tts_realtime_proc" + if mod_name not in sys.modules: + spec = importlib.util.spec_from_file_location(mod_name, proc_module_path) + mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type] + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) # type: ignore[union-attr] + else: + mod = sys.modules[mod_name] + + tokenizer = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True) + realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer) + + # Step 2: encode the reference audio clip via MOSS-Audio-Tokenizer. + try: + codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).to("cpu").eval() + except Exception as exc: + msg = f"Cannot load MOSS-Audio-Tokenizer: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + + ref_wav = _load_ref_audio(ref_audio_path) + wav_2d = ref_wav if ref_wav.dim() == 2 else ref_wav.unsqueeze(0) + with torch.no_grad(): + enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16) + codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())] + audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16) + del codec + gc.collect() + + # Step 3: build the (L, 17) prefill grid. + system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) + + assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False) + assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64) + assistant_grid[:, 0] = assistant_header_ids + + text_ids_only = tokenizer.encode(text, add_special_tokens=False) + if not text_ids_only: + pytest.fail(f"Empty text after tokenization: {text!r}") + + PREFILL_MAX_TEXT = 12 + cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT) + text_grid = np.full((cur_len, 17), 1024, dtype=np.int64) + text_grid[:, 0] = text_ids_only[:cur_len] + text_grid[-1, 1] = 1025 # audio_bos at the last prefilled text token + + grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0) + prompt_token_ids = grid[:, 0].tolist() + audio_codes_tensor = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16) + remaining_text_ids = list(text_ids_only[cur_len:]) + + additional_info: dict = {"codes": {"ref": audio_codes_tensor}} + if remaining_text_ids: + additional_info["ids"] = {"all": remaining_text_ids} + + return { + "prompt_token_ids": prompt_token_ids, + "additional_information": additional_info, + } + + +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + """Run one request and return (waveform_cpu, sample_rate).""" + chunks: list[torch.Tensor] = [] + sr_final = SAMPLE_RATE + for out in omni_runner.omni.generate(request, _SAMPLING): + mm = out.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + sr = mm.get("sr") + if sr is not None: + sr_final = int(sr.item() if hasattr(sr, "item") else sr) + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + chunks.append(audio.reshape(-1).cpu()) + if not chunks: + raise AssertionError("No audio output received from generate()") + return torch.cat(chunks, dim=0), sr_final + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path: str) -> None: + """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" + req = _build_request(ref_audio_path, "This is a real-time TTS streaming test.") + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" + assert audio.numel() > 0, "Audio tensor is empty" + assert not torch.all(audio == 0), "Audio is silence" + + +@pytest.mark.advanced_model +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_moss_tts_realtime_chinese(omni_runner: OmniRunner, ref_audio_path: str) -> None: + """MossTTSRealtime: Chinese input produces non-empty audio.""" + req = _build_request(ref_audio_path, "你好，这是语音合成测试。") + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE + assert audio.numel() > 0 + assert not torch.all(audio == 0) diff --git a/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py new file mode 100644 index 00000000000..b84e5a0ff1e --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_realtime_expansion.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference tests for MOSS-TTS-Realtime (MossTTSRealtime, 1.7B). + +Exercises the local-transformer generation path on L4. Delay-model coverage is in +``test_moss_tts_expansion.py``. + +No determinism test: MossTTSRealtime is sensitive to async scheduling timing; +waveform reproducibility is not guaranteed across sequential ``generate`` calls. +""" + +from __future__ import annotations + +import os +import urllib.request + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" +DEPLOY_CONFIG = get_deploy_config_path("moss_tts_realtime.yaml") +_OMNI_RUNNER_PARAM = ( + MODEL, + DEPLOY_CONFIG, + {"stage_init_timeout": 300, "trust_remote_code": True}, +) + +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), + pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), +] + +SAMPLE_RATE = 24_000 +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" + +_DEFAULT_SAMPLING = SamplingParams( + temperature=1.7, + top_p=0.8, + top_k=25, + max_tokens=512, + seed=42, + detokenize=False, +) + + +@pytest.fixture(scope="session") +def ref_audio_path(tmp_path_factory) -> str: + cache_dir = tmp_path_factory.mktemp("moss_tts_realtime_ref") + target = cache_dir / "zh_1.wav" + try: + with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: + target.write_bytes(resp.read()) + except Exception as exc: + msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + if not target.exists() or target.stat().st_size == 0: + pytest.fail(f"Reference clip empty after download: {target}") + return str(target) + + +def _build_request(text: str, ref_audio_path: str, seed: int = 42) -> dict: + return { + "prompt": "<|im_start|>assistant\n", + "additional_information": { + "task_type": ["voice_clone"], + "text": [text], + "mode": ["voice_clone"], + "prompt_audio_path": [ref_audio_path], + "seed": [seed], + }, + } + + +def _sampling_for(omni_runner: OmniRunner) -> SamplingParams | list[SamplingParams]: + omni = omni_runner.omni + if omni.num_stages == 1: + return _DEFAULT_SAMPLING + params = omni_runner.get_default_sampling_params_list() + params[0] = _DEFAULT_SAMPLING + return params + + +def _audio_from_stage(stage_outputs) -> tuple[torch.Tensor, int] | None: + mm = stage_outputs.multimodal_output + if not mm: + return None + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + return None + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if not isinstance(audio, torch.Tensor) or audio.numel() == 0: + return None + sr = mm.get("sr") + sample_rate = int(sr.item()) if sr is not None else SAMPLE_RATE + return audio.reshape(-1).cpu(), sample_rate + + +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + for stage_outputs in omni_runner.omni.generate(request, _sampling_for(omni_runner)): + parsed = _audio_from_stage(stage_outputs) + if parsed is not None: + return parsed + raise AssertionError("No stage outputs received") + + +@hardware_test(res={"cuda": "L4"}) +def test_moss_tts_realtime_english(omni_runner: OmniRunner, ref_audio_path) -> None: + """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio.""" + req = _build_request("This is a real-time TTS streaming test.", ref_audio_path) + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" + assert audio.numel() > 0, "Audio tensor is empty" + assert not torch.all(audio == 0), "Audio is silence" diff --git a/tests/e2e/offline_inference/test_moss_tts_v1_5.py b/tests/e2e/offline_inference/test_moss_tts_v1_5.py deleted file mode 100644 index 99c3b77604b..00000000000 --- a/tests/e2e/offline_inference/test_moss_tts_v1_5.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""E2E offline inference test for MOSS-TTS-v1.5 (MossTTSDelay-8B). - -MOSS-TTS-v1.5 is a continued-training upgrade of MOSS-TTS 1.0 with the SAME -``MossTTSDelay`` architecture and API, so it reuses the existing talker/codec -path unchanged — this test just pins that the 8B v1.5 checkpoint loads and -produces audio. - -The model is 8B, so this lives in its own file (one module-scoped engine per -file, per skill invariant I4) and is gated behind an 80 GB GPU; the small-GPU -MossTTSDelay coverage stays in ``test_moss_tts.py`` (1.7B VoiceGenerator). -""" - -from __future__ import annotations - -import os -import urllib.request - -import pytest -import torch -from vllm import SamplingParams - -from tests.helpers.mark import hardware_test -from vllm_omni import Omni - -# H100-gated (8B). Picked up by the nightly "TTS · Function Test with H100" -# job (-m "full_model and H100 and tts"); the L4 MOSS-TTS tests land in the -# sibling L4 tts lane instead. -# TODO: Fix this test -# pytestmark = [pytest.mark.full_model, pytest.mark.tts] - -SAMPLE_RATE = 24_000 -REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" -_MODEL = "OpenMOSS-Team/MOSS-TTS-v1.5" - -_DEFAULT_SAMPLING = SamplingParams( - temperature=1.7, - top_p=0.8, - top_k=25, - max_tokens=512, # short for tests (~20 s at 24 kHz, ~12.5 tokens/s) - seed=42, - detokenize=False, -) - - -@pytest.fixture(scope="session") -def ref_audio_path(tmp_path_factory) -> str: - """Download the upstream reference clip once per session. - - Set ``MOSS_TTS_SKIP_ON_NET_FAIL=1`` to skip in air-gapped environments. - """ - target = tmp_path_factory.mktemp("moss_tts_v15_ref") / "zh_1.wav" - try: - with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: - target.write_bytes(resp.read()) - except Exception as exc: # noqa: BLE001 - msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" - if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): - pytest.skip(msg) - pytest.fail(msg) - if not target.exists() or target.stat().st_size == 0: - pytest.fail(f"Reference clip empty after download: {target}") - return str(target) - - -@pytest.fixture(scope="module") -def v15_engine(): - return Omni(model=_MODEL, stage_init_timeout=600) - - -def _build_request(text: str, ref_audio_path: str) -> dict: - return { - "prompt": "<|im_start|>assistant\n", - "additional_information": { - "task_type": ["voice_clone"], - "text": [text], - "mode": ["voice_clone"], - "prompt_audio_path": [ref_audio_path], - "seed": [42], - }, - } - - -def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]: - for stage_outputs in omni.generate(request, _DEFAULT_SAMPLING): - for req_output in stage_outputs.request_output: - mm = req_output.outputs[0].multimodal_output - assert mm is not None, "Expected non-None multimodal_output" - audio = mm.get("audio") - if audio is None: - audio = mm.get("model_outputs") - sr = mm.get("sr") - assert audio is not None, "Expected 'audio' or 'model_outputs' in multimodal_output" - if isinstance(audio, list): - audio = torch.cat( - [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], - dim=0, - ) - assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}" - return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE - raise AssertionError("No stage outputs received") - - -@hardware_test(res={"cuda": "H100"}) -def test_moss_tts_v15_voice_clone(v15_engine, ref_audio_path): - """MOSS-TTS-v1.5 (8B): voice_clone produces non-empty 24 kHz audio.""" - req = _build_request("Hello, this is a MOSS-TTS v1.5 voice cloning test.", ref_audio_path) - audio, sr = _collect_audio(v15_engine, req) - - assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" - assert audio.numel() > 0, "Audio tensor is empty" - assert not torch.all(audio == 0), "Audio is silence" diff --git a/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py new file mode 100644 index 00000000000..25c671e0540 --- /dev/null +++ b/tests/e2e/offline_inference/test_moss_tts_v1_5_expansion.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E offline inference test for MOSS-TTS-v1.5 (MossTTSDelay-8B). + +MOSS-TTS-v1.5 is a continued-training upgrade of MOSS-TTS 1.0 with the SAME +``MossTTSDelay`` architecture and API. The 8B checkpoint is H100-gated; small-GPU +MossTTSDelay coverage stays in ``test_moss_tts_expansion.py``. +""" + +from __future__ import annotations + +import os +import urllib.request + +import pytest +import torch +from vllm import SamplingParams + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = "OpenMOSS-Team/MOSS-TTS-v1.5" +DEPLOY_CONFIG = get_deploy_config_path("moss_tts.yaml") +_OMNI_RUNNER_PARAM = ( + MODEL, + DEPLOY_CONFIG, + {"stage_init_timeout": 600, "trust_remote_code": True}, +) + +pytestmark = [ + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4643"), + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), +] + +SAMPLE_RATE = 24_000 +REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav" + +_DEFAULT_SAMPLING = SamplingParams( + temperature=1.7, + top_p=0.8, + top_k=25, + max_tokens=512, + seed=42, + detokenize=False, +) + + +@pytest.fixture(scope="session") +def ref_audio_path(tmp_path_factory) -> str: + target = tmp_path_factory.mktemp("moss_tts_v15_ref") / "zh_1.wav" + try: + with urllib.request.urlopen(REF_AUDIO_URL, timeout=30) as resp: + target.write_bytes(resp.read()) + except Exception as exc: # noqa: BLE001 + msg = f"Cannot fetch reference clip {REF_AUDIO_URL}: {exc}" + if os.environ.get("MOSS_TTS_SKIP_ON_NET_FAIL"): + pytest.skip(msg) + pytest.fail(msg) + if not target.exists() or target.stat().st_size == 0: + pytest.fail(f"Reference clip empty after download: {target}") + return str(target) + + +def _build_request(text: str, ref_audio_path: str) -> dict: + return { + "prompt": "<|im_start|>assistant\n", + "additional_information": { + "task_type": ["voice_clone"], + "text": [text], + "mode": ["voice_clone"], + "prompt_audio_path": [ref_audio_path], + "seed": [42], + }, + } + + +def _sampling_for(omni_runner: OmniRunner) -> SamplingParams | list[SamplingParams]: + omni = omni_runner.omni + if omni.num_stages == 1: + return _DEFAULT_SAMPLING + params = omni_runner.get_default_sampling_params_list() + params[0] = _DEFAULT_SAMPLING + return params + + +def _collect_audio(omni_runner: OmniRunner, request: dict) -> tuple[torch.Tensor, int]: + for stage_outputs in omni_runner.omni.generate(request, _sampling_for(omni_runner)): + mm = stage_outputs.multimodal_output + if not mm: + continue + audio = mm.get("audio") + if audio is None: + audio = mm.get("model_outputs") + if audio is None: + continue + if isinstance(audio, list): + audio = torch.cat( + [t.reshape(-1) for t in audio if isinstance(t, torch.Tensor) and t.numel() > 0], + dim=0, + ) + if not isinstance(audio, torch.Tensor) or audio.numel() == 0: + continue + sr = mm.get("sr") + return audio.reshape(-1).cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE + raise AssertionError("No stage outputs received") + + +@hardware_test(res={"cuda": "H100"}) +def test_moss_tts_v15_voice_clone(omni_runner: OmniRunner, ref_audio_path) -> None: + """MOSS-TTS-v1.5 (8B): voice_clone produces non-empty 24 kHz audio.""" + req = _build_request("Hello, this is a MOSS-TTS v1.5 voice cloning test.", ref_audio_path) + audio, sr = _collect_audio(omni_runner, req) + + assert sr == SAMPLE_RATE, f"Expected {SAMPLE_RATE} Hz, got {sr}" + assert audio.numel() > 0, "Audio tensor is empty" + assert not torch.all(audio == 0), "Audio is silence" diff --git a/tests/e2e/offline_inference/test_omnivoice_expansion.py b/tests/e2e/offline_inference/test_omnivoice_expansion.py index 7ea646a8555..e3a13811293 100644 --- a/tests/e2e/offline_inference/test_omnivoice_expansion.py +++ b/tests/e2e/offline_inference/test_omnivoice_expansion.py @@ -18,20 +18,22 @@ from tests.helpers.stage_config import get_deploy_config_path MODEL = "k2-fsa/OmniVoice" -STAGE_CONFIG = get_deploy_config_path("omnivoice.yaml") +DEPLOY_CONFIG = get_deploy_config_path("omnivoice.yaml") -# (model, stage_config_path, extra_omni_kwargs) — see ``omni_runner`` in tests.helpers.fixtures.runtime +# OmniRunner tuple: model, legacy stage config path, extra Omni kwargs. +# The migrated test passes deploy_config through extra Omni kwargs. _OMNI_RUNNER_PARAM = ( MODEL, - STAGE_CONFIG, + None, { + "deploy_config": DEPLOY_CONFIG, "trust_remote_code": True, "log_stats": True, }, ) pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.tts, pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True), ] diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py index 190857e3b13..6be54969ca5 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16_expansion.py @@ -23,7 +23,7 @@ from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config pytestmark = [ - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.omni, ] @@ -36,6 +36,10 @@ _CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml") +_SKIP_ISSUE_4663 = pytest.mark.skip( + reason="https://github.com/vllm-project/vllm-omni/issues/4663", +) + def _get_stage_config(): """Build a CI-friendly stage config with eager mode.""" @@ -61,6 +65,7 @@ def _get_stage_config(): # ------------------------------------------------------------------ +@_SKIP_ISSUE_4663 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) def test_text_to_text(omni_runner, omni_runner_handler): @@ -79,6 +84,7 @@ def test_text_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ +@_SKIP_ISSUE_4663 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) def test_audio_to_text(omni_runner, omni_runner_handler): @@ -102,6 +108,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ +@_SKIP_ISSUE_4663 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) def test_image_to_text(omni_runner, omni_runner_handler): @@ -123,6 +130,7 @@ def test_image_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ +@_SKIP_ISSUE_4663 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) def test_video_to_text(omni_runner, omni_runner_handler): @@ -144,6 +152,7 @@ def test_video_to_text(omni_runner, omni_runner_handler): # ------------------------------------------------------------------ +@_SKIP_ISSUE_4663 @hardware_test(res={"cuda": "L4"}, num_cards=4) @pytest.mark.parametrize("omni_runner", quant_params, indirect=True) def test_mix_to_audio(omni_runner, omni_runner_handler): diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py b/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py index 34a87845632..9d368a6882d 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py @@ -51,7 +51,7 @@ def get_question(prompt_type="mix"): return prompts.get(prompt_type, prompts["mix"]) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) @@ -82,7 +82,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None: omni_runner_handler.send_omni_request(request_config) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3}) @pytest.mark.parametrize("omni_runner", test_params, indirect=True) diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py new file mode 100644 index 00000000000..5d1b3dda7c7 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Qwen3-Omni ModelOpt NVFP4 W4A4 — Marlin fallback path. + +Loads the published full-thinker NVFP4 W4A4 checkpoint on H100 (sm_90), +where ``init_nvfp4_linear_kernel()`` selects the ``MarlinNvFp4LinearKernel`` +weight-only fallback. The Marlin kernel's PWAL casts ``weight_scale`` +FP8 -> bf16 and permutes; the patched +``ModelOptNvFp4LinearMethod.process_weights_after_loading`` in +``vllm_omni.patch`` therefore clamps NaN bytes BEFORE delegating to the +original PWAL. This test asserts the served response is coherent, i.e. +no `!!!!` token cascade caused by NaN bytes in per-block ``weight_scale``. + +Hardware coverage caveat: this DOES NOT exercise the native FlashInfer +FP4 GEMM (Cutlass / TRT-LLM / cuDNN) that is selected on Blackwell +sm_100+ — those kernel PWALs swizzle/shuffle FP8 bytes without dtype +casting, a different code path from Marlin. TODO(B200): add a companion +Blackwell smoke once a Blackwell runner is wired into ``tests.helpers.mark``. +""" + +import os +import re + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.omni, +] + +QUANTIZED_MODEL = os.environ.get( + "QWEN3_OMNI_W4A4_MODEL", + "YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip", +) + +_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +# `!!!!` (4+ exclamation marks) is the signature degenerate output the +# NaN-clamp patch defends against. Any test that loads a W4A4 checkpoint +# must reject this. +_NAN_COLLAPSE_RE = re.compile(r"!{4,}") + + +@pytest.fixture(scope="module", autouse=True) +def _qwen3_omni_env(): + with pytest.MonkeyPatch.context() as mp: + mp.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + yield + + +def _stage_config(): + # enforce_eager=True here to keep CI memory + warm-up time bounded for the + # smoke test, matching the existing Qwen3-Omni e2e pattern (see + # test_qwen3_omni_autoround_w4a16_expansion.py). The throughput claims in + # the PR body are from explicit no-enforce-eager runs, not from this test. + # + # TODO: the production code path uses FULL_AND_PIECEWISE CUDA graphs. + # The NaN clamp itself is graph-mode-independent so this enforce-eager + # smoke is still a valid regression guard against `!!!!`, but a future + # CUDA-graph-specific regression would not be caught here. If/when CI + # gains enough VRAM headroom for a no-eager run, add a companion test. + return modify_stage_config( + _CI_DEPLOY, + updates={"stages": {0: {"enforce_eager": True}, 1: {"enforce_eager": True}}}, + ) + + +quant_params = [(QUANTIZED_MODEL, _stage_config())] + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_marlin_fallback_no_nan_collapse(omni_runner, omni_runner_handler): + """H100 Marlin FP4 fallback: text in -> coherent text out. + + On H100 the upstream NVFP4 PWAL routes through + ``MarlinNvFp4LinearKernel.process_weights_after_loading`` -> + ``prepare_fp4_layer_for_marlin``, which casts FP8 weight_scale to + bf16 and permutes. The patched PWAL's NaN clamp runs BEFORE this + transform, so this test is the regression guard for the Marlin + code path specifically (Blackwell sm_100+ takes the FlashInfer + Cutlass/TRT-LLM/cuDNN path and is NOT covered by this test). + + Asserts the response is non-empty and does NOT match the `!!!!` + degenerate-output signature that indicates NaN propagation through + the FP4 GEMM. + """ + response = omni_runner_handler.send_omni_request({"prompts": "Why is the sky blue?", "modalities": ["text"]}) + assert response.success, "request failed" + text = (response.text_content or "").strip() + assert text, "empty response — likely server returned only EOS" + assert not _NAN_COLLAPSE_RE.search(text), ( + f"degenerate `!!!!` output detected — NaN propagated through FP4 GEMM. " + f"This means vllm_omni.patch's NVFP4 weight_scale NaN clamp did not install or " + f"did not run. Full response: {text!r}" + ) diff --git a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py index 416a5009cc6..4fa2d4b6fe0 100644 --- a/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py +++ b/tests/e2e/offline_inference/test_qwen3_tts_customvoice.py @@ -34,12 +34,10 @@ def get_cuda_graph_config(): "stages": { 0: { "max_num_seqs": 1, - "gpu_memory_utilization": 0.2, "enforce_eager": True, "async_scheduling": False, }, 1: { - "gpu_memory_utilization": 0.2, "enforce_eager": True, "async_scheduling": False, }, diff --git a/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py new file mode 100644 index 00000000000..8d7d3188281 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen_image_autoround_w4a16_expansion.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E test for Qwen-Image AutoRound W4A16 quantized inference. + +Verifies that the W4A16 quantized Qwen-Image checkpoint loads end-to-end +(quant_config propagated correctly) and produces a valid image. +""" + +import os +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + +QUANTIZED_MODEL = os.environ.get( + "QWEN_IMAGE_AUTOROUND_MODEL", + "INC4AI/Qwen-Image-AutoRound-W4A16", +) + +HEIGHT = 512 +WIDTH = 512 +NUM_STEPS = 2 + + +def _sampling_params() -> OmniDiffusionSamplingParams: + return OmniDiffusionSamplingParams( + height=HEIGHT, + width=WIDTH, + num_inference_steps=NUM_STEPS, + true_cfg_scale=1.0, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + + +def _first_request_images(outputs) -> list: + first_output = outputs[0] + assert first_output.final_output_type == "image" + req_out = first_output.request_output + assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") + return req_out.images + + +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) +def test_qwen_image_autoround_w4a16_load(): + """Load the W4A16 quantized Qwen-Image model and run a minimal generation. + + Verifies that quant_config is propagated to all transformer blocks + and produces a valid, non-blank image. + """ + with OmniRunner(QUANTIZED_MODEL, enforce_eager=True) as runner: + outputs = runner.omni.generate( + "a cup of coffee on a table", + _sampling_params(), + ) + images = _first_request_images(outputs) + assert len(images) >= 1, "Expected at least one generated image" + img = images[0] + assert img.width == WIDTH and img.height == HEIGHT + arr = np.array(img) + assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)" + + +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}) +def test_qwen_image_autoround_w4a16_generate(): + """Full generation: 512×512, 20 steps, CFG=5.0. + + Validates that realistic quantized inference (multi-step with CFG + negative prompts and longer denoising) produces a valid non-blank image. + """ + params = OmniDiffusionSamplingParams( + height=HEIGHT, + width=WIDTH, + num_inference_steps=20, + true_cfg_scale=5.0, + generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42), + ) + with OmniRunner(QUANTIZED_MODEL, enforce_eager=True) as runner: + outputs = runner.omni.generate( + "a cup of coffee on a wooden table, morning sunlight, photorealistic", + params, + ) + images = _first_request_images(outputs) + assert len(images) >= 1, "Expected at least one generated image" + img = images[0] + + # Save for manual inspection + output_path = Path(__file__).resolve().parents[4] / "ar_coffee_full.png" + img.save(str(output_path)) + + assert img.width == WIDTH and img.height == HEIGHT + arr = np.array(img) + assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)" diff --git a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py index 1f9769f8199..4f5c97d1148 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_expansion.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_expansion.py @@ -124,7 +124,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest: # --------------------------------------------------------------------------- -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}) def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None: diff --git a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py similarity index 96% rename from tests/e2e/offline_inference/test_sensenova_u1_img2img.py rename to tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py index fedd8e90e51..242d6a12eb2 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_img2img.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_img2img_expansion.py @@ -153,9 +153,9 @@ def _generate_sensenova_u1_img2img( return generated_image -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion +@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4636") @hardware_test(res={"cuda": "H100"}) def test_sensenova_u1_img2img(run_level): """Test SenseNova-U1 img2img (single-stage diffusion, no deploy YAML).""" @@ -165,11 +165,11 @@ def test_sensenova_u1_img2img(run_level): stage_configs_path=None, ) as runner: generated_image = _generate_sensenova_u1_img2img(runner.omni, input_image) - if run_level == "advanced_model": + if run_level == "full_model" or run_level == "advanced_model": _validate_pixels(generated_image) -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.cache @hardware_test(res={"cuda": "H100"}) diff --git a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py similarity index 96% rename from tests/e2e/offline_inference/test_sensenova_u1_text2img.py rename to tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py index fc9b5411e22..65b11d3ee5c 100644 --- a/tests/e2e/offline_inference/test_sensenova_u1_text2img.py +++ b/tests/e2e/offline_inference/test_sensenova_u1_text2img_expansion.py @@ -143,9 +143,9 @@ def _generate_sensenova_u1_image( return generated_image -@pytest.mark.core_model -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion +@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4636") @hardware_test(res={"cuda": "H100"}) def test_sensenova_u1_text2img(run_level): """Test SenseNova-U1 text2img (single-stage diffusion, no deploy YAML).""" @@ -154,11 +154,11 @@ def test_sensenova_u1_text2img(run_level): stage_configs_path=None, ) as runner: generated_image = _generate_sensenova_u1_image(runner.omni) - if run_level == "advanced_model": + if run_level == "full_model" or run_level == "advanced_model": _validate_pixels(generated_image) -@pytest.mark.core_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.cache @hardware_test(res={"cuda": "H100"}) diff --git a/tests/e2e/offline_inference/test_soulxsinger.py b/tests/e2e/offline_inference/test_soulxsinger.py new file mode 100644 index 00000000000..465f5547b62 --- /dev/null +++ b/tests/e2e/offline_inference/test_soulxsinger.py @@ -0,0 +1,241 @@ +"""E2E offline inference tests for SoulX-Singer (single-stage, preprocess inline).""" + +import functools +import importlib +import json +import os +import shutil +from pathlib import Path + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import numpy as np +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import get_asset_path +from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +PROMPT_AUDIO = get_asset_path("soulxsinger/zh_prompt.mp3") +TARGET_AUDIO = get_asset_path("soulxsinger/music.mp3") +PROMPT_METADATA = get_asset_path("soulxsinger/zh_prompt.json") +TARGET_METADATA = get_asset_path("soulxsinger/music.json") +SAMPLE_RATE = 24_000 + +# phone_set.json (phoneme vocab) is not shipped on HuggingFace; SVS loads it +# from the model dir. Stage it from this pinned upstream commit so the SVS DiT +# can run in CI; the SVS tests skip gracefully if it cannot be fetched. +_PHONE_SET_URL = ( + "https://raw.githubusercontent.com/Soul-AILab/SoulX-Singer/" + "81aeb3ae772c70093c3de74dc23c92d983801ae4/soulxsinger/utils/phoneme/phone_set.json" +) + +if not PROMPT_AUDIO.is_file() or not TARGET_AUDIO.is_file(): + pytest.skip( + f"Missing SoulX-Singer audio assets: {PROMPT_AUDIO.name}, {TARGET_AUDIO.name}", + allow_module_level=True, + ) + +pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.tts] + +_CASES = ( + pytest.param( + "SoulXSingerPipeline", + "soulxsinger_svs.yaml", + { + "language": "Mandarin", + "vocal_sep": False, + "control": "score", + "auto_shift": False, + "pitch_shift": 0, + }, + ("g2pM", "g2p_en"), + id="svs", + ), + pytest.param( + "SoulXSingerSVCPipeline", + "soulxsinger_svc.yaml", + {"vocal_sep": False, "auto_shift": False, "pitch_shift": 0}, + (), + id="svc", + ), +) + + +@functools.lru_cache(maxsize=1) +def _resolve_weights() -> tuple[Path, Path, Path]: + from huggingface_hub import snapshot_download + + base = Path(snapshot_download("Soul-AILab/SoulX-Singer", allow_patterns=["*"])) + + # phone_set.json is not on HF; best-effort stage it from pinned upstream so + # SVS can load. SVS tests skip (not fail) if both this and a manual copy are absent. + phone_set = base / "phoneme" / "phone_set.json" + if not phone_set.is_file() and not (base / "phone_set.json").is_file(): + try: + from urllib.request import urlretrieve + + phone_set.parent.mkdir(exist_ok=True) + urlretrieve(_PHONE_SET_URL, phone_set) + except Exception: + pass + + # make a temporary svc directory + svc_dir = base.parent / "SoulX-Singer-SVC" + svc_dir.mkdir(parents=True, exist_ok=True) + shutil.copyfile(base / "config.yaml", svc_dir / "config.yaml") + shutil.move(base / "model-svc.pt", svc_dir / "model-svc.pt") + + if raw := os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR"): + pre = Path(raw).expanduser().resolve() + if (pre / "rmvpe" / "rmvpe.pt").is_file(): + return base, svc_dir, pre + + from huggingface_hub import snapshot_download + + pre = Path(snapshot_download("Soul-AILab/SoulX-Singer-Preprocess", allow_patterns=["*"])) + return base, svc_dir, pre + + +@pytest.fixture(scope="session") +def soulx_weights() -> tuple[Path, Path, Path]: + try: + return _resolve_weights() + except Exception as exc: + pytest.skip(f"Set SOULXSINGER_MODEL_DIR / SOULX_PREPROCESS_WEIGHTS_DIR. ({exc})") + + +def _flatten_audio(audio_val) -> np.ndarray: + import torch + + if isinstance(audio_val, list): + chunks = [c.detach().cpu().float().numpy().reshape(-1) for c in audio_val if c is not None] + return np.concatenate(chunks) if chunks else np.array([], dtype=np.float32) + if isinstance(audio_val, torch.Tensor): + return audio_val.detach().cpu().float().numpy().reshape(-1) + return np.asarray(audio_val, dtype=np.float32).reshape(-1) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("architecture,deploy_yaml,extra_args,py_deps", _CASES) +def test_soulxsinger_multistage_from_audio( + soulx_weights: tuple[Path, Path, Path], + architecture: str, + deploy_yaml: str, + extra_args: dict, + py_deps: tuple[str, ...], +) -> None: + for mod in py_deps: + try: + importlib.import_module(mod) + except ImportError as exc: + pytest.fail(f"SoulX SVS requires {mod}: {exc}") + + base_dir, svc_dir, preprocess_dir = soulx_weights + + # SVS mode requires phone_set.json in the model directory + if architecture == "SoulXSingerPipeline": + json.dump( + {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1}, + open(base_dir / "config.json", "w"), + indent=4, + ) + model = str(base_dir) + if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file(): + pytest.skip( + "SoulX-Singer SVS test requires phoneme/phone_set.json. " + "Copy it from github.com/Soul-AILab/SoulX-Singer into the model dir. " + "See `examples/offline_inference/text_to_speech/README.md` for details." + ) + if architecture == "SoulXSingerSVCPipeline": + json.dump( + {"model_type": "soulxsinger", "architectures": ["SoulXSingerSVCPipeline"], "max_num_seqs": 1}, + open(svc_dir / "config.json", "w"), + indent=4, + ) + model = str(svc_dir) + + with OmniRunner( + model, + stage_configs_path=get_deploy_config_path(deploy_yaml), + async_chunk=False, + ) as runner: + sampling = OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=3.0, + seed=42, + extra_args={ + "prompt_audio": str(PROMPT_AUDIO), + "target_audio": str(TARGET_AUDIO), + "preprocess_weights_dir": str(preprocess_dir), + **extra_args, + }, + ) + prompt = {"prompt_token_ids": [0]} + outputs = runner.generate([prompt], sampling) + + assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output" + mm = outputs[0].multimodal_output + assert isinstance(mm, dict) and "audio" in mm + audio = _flatten_audio(mm["audio"]) + assert 12_000 <= audio.size + assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4 + duration_s = audio.size / SAMPLE_RATE + assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s" + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_soulxsinger_svs_precomputed(soulx_weights: tuple[Path, Path, Path]) -> None: + """SVS from precomputed note/lyric metadata. + + Exercises the SVS DiT + vocoder serving path while skipping the heavy + upstream preprocess (funasr ASR + ROSVOT note transcription), so it runs in + the merge gate without those non-pip dependencies. + """ + base_dir, _svc_dir, preprocess_dir = soulx_weights + + if not PROMPT_METADATA.is_file() or not TARGET_METADATA.is_file(): + pytest.skip(f"Missing SoulX-Singer metadata fixtures: {PROMPT_METADATA.name}, {TARGET_METADATA.name}") + if not (base_dir / "phoneme" / "phone_set.json").is_file() and not (base_dir / "phone_set.json").is_file(): + pytest.skip( + "SoulX-Singer SVS requires phoneme/phone_set.json (upstream fetch unavailable). " + "See `examples/offline_inference/text_to_speech/README.md`." + ) + + json.dump( + {"model_type": "soulxsinger", "architectures": ["SoulXSingerPipeline"], "max_num_seqs": 1}, + open(base_dir / "config.json", "w"), + indent=4, + ) + + with OmniRunner( + str(base_dir), + stage_configs_path=get_deploy_config_path("soulxsinger_svs.yaml"), + async_chunk=False, + ) as runner: + sampling = OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=3.0, + seed=42, + extra_args={ + "language": "Mandarin", + "control": "score", + "prompt_metadata_path": str(PROMPT_METADATA), + "target_metadata_path": str(TARGET_METADATA), + "audio_path": str(PROMPT_AUDIO), + "preprocess_weights_dir": str(preprocess_dir), + }, + ) + prompt = {"prompt_token_ids": [0]} + outputs = runner.generate([prompt], sampling) + + assert outputs and outputs[0].error is None, outputs[0].error if outputs else "no output" + mm = outputs[0].multimodal_output + assert isinstance(mm, dict) and "audio" in mm + audio = _flatten_audio(mm["audio"]) + assert 12_000 <= audio.size + assert np.isfinite(audio).all() and float(np.max(np.abs(audio))) > 1e-4 + duration_s = audio.size / SAMPLE_RATE + assert 50.0 <= duration_s <= 52.0, f"duration={duration_s:.1f}s" diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 96365c7d381..fc3b5596d47 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -23,7 +23,7 @@ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] _MODEL_REPO = "stabilityai/stable-audio-open-1.0" @@ -103,3 +103,23 @@ def test_stable_audio_quantization_and_teacache() -> None: ) finally: m.close() + + +@hardware_test(res={"cuda": "L4", "xpu": "B60"}) +def test_stable_audio_cpu_offload() -> None: + """Stable Audio Open with FP8 + CPU offload.""" + m = Omni( + model="stabilityai/stable-audio-open-1.0", + quantization="fp8", + enable_cpu_offload=True, + ) + try: + audio = generate_stable_audio_short_clip(m) + assert_audio_valid( + audio, + sample_rate=_SAMPLE_RATE, + channels=2, + duration_s=_CLIP_DURATION_S, + ) + finally: + m.close() diff --git a/tests/e2e/offline_inference/test_step_audio2.py b/tests/e2e/offline_inference/test_step_audio2.py deleted file mode 100644 index cd779a6def9..00000000000 --- a/tests/e2e/offline_inference/test_step_audio2.py +++ /dev/null @@ -1,303 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E tests for Step-Audio2 model with audio input and audio output. - -These tests verify: -1. Processor loads correctly and handles audio input -2. Field configuration for vLLM batching (flat_from_sizes) -3. Tensor shapes are correct (3D not 4D for profiling) -4. Basic E2E inference works with audio input -""" - -from pathlib import Path - -import numpy as np -import pytest -import torch - -from tests.helpers.runtime import OmniRunner - -models = ["stepfun-ai/Step-Audio-2-mini"] - -stage_configs = [str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")] - -test_params = [(model, stage_config) for model in models for stage_config in stage_configs] - - -def create_dummy_audio(sample_rate: int = 16000, duration_sec: float = 1.0) -> tuple[np.ndarray, int]: - """Create a dummy audio signal (sine wave) for testing.""" - t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), dtype=np.float32) - audio = 0.5 * np.sin(2 * np.pi * 440 * t) - return audio, sample_rate - - -class StepAudio2OmniRunner(OmniRunner): - """Test runner specialized for Step-Audio2 model.""" - - def get_omni_inputs( - self, - prompts: list[str] | str, - system_prompt: str | None = None, - audios=None, - **kwargs, - ) -> list[dict]: - """Construct Omni input format for Step-Audio2.""" - if system_prompt is None: - system_prompt = "你是一个语音对话助手，能够理解音频输入并生成语音回复。" - - audio_placeholder = "<audio_patch>" - - if isinstance(prompts, str): - prompts = [prompts] - - num_prompts = len(prompts) - - if audios is None: - audios_list = [None] * num_prompts - elif isinstance(audios, list): - audios_list = audios if len(audios) == num_prompts else [audios[0]] * num_prompts - else: - audios_list = [audios] * num_prompts - - omni_inputs = [] - for i, prompt_text in enumerate(prompts): - user_content = "" - multi_modal_data = {} - - audio = audios_list[i] - if audio is not None: - user_content += audio_placeholder - multi_modal_data["audio"] = audio - - user_content += prompt_text - - full_prompt = ( - f"<|im_start|>system\n{system_prompt}<|im_end|>\n" - f"<|im_start|>user\n{user_content}<|im_end|>\n" - f"<|im_start|>assistant\n" - ) - - input_dict = {"prompt": full_prompt} - if multi_modal_data: - input_dict["multi_modal_data"] = multi_modal_data - - omni_inputs.append(input_dict) - - return omni_inputs - - -@pytest.fixture(scope="session") -def step_audio2_runner(): - return StepAudio2OmniRunner - - -@pytest.mark.core_model -def test_processor_loads_correctly(): - """Test that Step-Audio2 processor can be imported and initialized.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( - StepAudio2MultiModalProcessor, - StepAudio2ProcessingInfo, - StepAudio2Processor, - ) - - assert StepAudio2MultiModalProcessor is not None - assert StepAudio2ProcessingInfo is not None - assert StepAudio2Processor is not None - - -@pytest.mark.core_model -def test_audio_preprocessing(): - """Test that audio preprocessing produces correct tensor shapes.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( - log_mel_spectrogram, - padding_mels, - ) - - audio, sample_rate = create_dummy_audio(sample_rate=16000, duration_sec=1.0) - - mel = log_mel_spectrogram(audio) - - assert mel.ndim == 2, f"Expected 2D tensor, got {mel.ndim}D" - assert mel.shape[0] == 128, f"Expected 128 mel bins, got {mel.shape[0]}" - - mels = [mel, mel] - padded_mels, mel_lens = padding_mels(mels) - - assert padded_mels.ndim == 3, f"Expected 3D tensor, got {padded_mels.ndim}D" - assert padded_mels.shape[0] == 2, f"Expected batch size 2, got {padded_mels.shape[0]}" - assert padded_mels.shape[1] == 128, f"Expected 128 mel bins, got {padded_mels.shape[1]}" - - -@pytest.mark.core_model -def test_feature_length_calculation(): - """Test that audio feature length calculation is correct.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( - calculate_audio_feature_length, - ) - - assert calculate_audio_feature_length(1000) == 125 - - assert calculate_audio_feature_length(100) > 0 - - assert calculate_audio_feature_length(1) >= 1 - - -@pytest.mark.core_model -def test_token_config_constants(): - """Test that token configuration constants are set correctly.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( - DEFAULT_TOKEN_CONFIG, - STEP_AUDIO2_AUDIO_END, - STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID, - STEP_AUDIO2_AUDIO_START, - STEP_AUDIO2_AUDIO_VOCAB_SIZE, - STEP_AUDIO2_TEXT_MAX, - ) - - assert STEP_AUDIO2_TEXT_MAX < STEP_AUDIO2_AUDIO_START, "Text tokens should come before audio tokens" - assert STEP_AUDIO2_AUDIO_END == STEP_AUDIO2_AUDIO_START + STEP_AUDIO2_AUDIO_VOCAB_SIZE - 1 - assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID > STEP_AUDIO2_TEXT_MAX - assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID < STEP_AUDIO2_AUDIO_START - - assert DEFAULT_TOKEN_CONFIG.text_max == STEP_AUDIO2_TEXT_MAX - assert DEFAULT_TOKEN_CONFIG.audio_start == STEP_AUDIO2_AUDIO_START - - -@pytest.mark.core_model -def test_mm_field_config_structure(): - """Test that multimodal field config is properly structured for vLLM batching.""" - from vllm.multimodal.inputs import MultiModalFieldConfig - - audio_lens = torch.tensor([100, 200, 150], dtype=torch.int32) - - field_config = MultiModalFieldConfig.flat_from_sizes("audio", audio_lens) - - assert field_config is not None - - -@pytest.mark.core_model -def test_audio_encoder_output_shape(): - """Test that audio encoder produces correct output shapes.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( - Adaptor, - AudioEncoder, - ) - - encoder = AudioEncoder(n_mels=128, n_ctx=1500, n_state=512, n_head=8, n_layer=6) - - batch_size = 2 - n_mels = 128 - time_steps = 400 - x = torch.randn(batch_size, n_mels, time_steps) - x_len = torch.tensor([time_steps, time_steps // 2], dtype=torch.int32) - - encoded, encoded_lens = encoder(x, x_len) - - assert encoded.ndim == 3, f"Expected 3D tensor, got {encoded.ndim}D" - assert encoded.shape[0] == batch_size - assert encoded.shape[2] == 512 - - adapter = Adaptor(n_state=512, n_hidden=4096, kernel_size=3, stride=2) - adapted = adapter(encoded) - - assert adapted.ndim == 3 - assert adapted.shape[0] == batch_size - assert adapted.shape[2] == 4096 - - -@pytest.mark.core_model -def test_token_separation(): - """Test that token separation works correctly.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( - StepAudio2ThinkerForConditionalGeneration, - ) - - token_ids = [100, 200, 151700, 300, 151800, 400] - - text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(token_ids) - - assert text_tokens == [100, 200, 300, 400] - assert audio_tokens == [4, 104] - - -@pytest.mark.core_model -def test_has_audio_output(): - """Test detection of audio tokens in output.""" - from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( - StepAudio2ThinkerForConditionalGeneration, - ) - - text_only = [100, 200, 300] - assert not StepAudio2ThinkerForConditionalGeneration.has_audio_output(text_only) - - with_audio = [100, 200, 151700, 300] - assert StepAudio2ThinkerForConditionalGeneration.has_audio_output(with_audio) - - -@pytest.mark.advanced_model -@pytest.mark.parametrize("test_config", test_params) -def test_audio_to_text_and_audio(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None: - """Test processing audio input and generating text + audio output.""" - model, stage_config_path = test_config - - with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner: - audio = create_dummy_audio(sample_rate=16000, duration_sec=2.0) - - outputs = runner.generate_multimodal( - prompts="请复述这段音频的内容。", - audios=audio, - ) - - assert len(outputs) > 0 - - text_output = None - for stage_output in outputs: - if stage_output.final_output_type == "text": - text_output = stage_output - break - - assert text_output is not None - assert text_output.request_output is not None - text_content = text_output.request_output.outputs[0].text - assert text_content is not None - assert len(text_content.strip()) > 0 - - audio_output = None - for stage_output in outputs: - if stage_output.final_output_type == "audio": - audio_output = stage_output - break - - if audio_output is not None: - assert audio_output.request_output is not None - audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio") - if audio_tensor is not None: - assert audio_tensor.numel() > 0 - - -@pytest.mark.advanced_model -@pytest.mark.parametrize("test_config", test_params) -def test_text_only_input(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None: - """Test processing text-only input (no audio).""" - model, stage_config_path = test_config - - with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner: - outputs = runner.generate_multimodal( - prompts="你好，请用中文回答。", - audios=None, - ) - - assert len(outputs) > 0 - - text_output = None - for stage_output in outputs: - if stage_output.final_output_type == "text": - text_output = stage_output - break - - assert text_output is not None - assert text_output.request_output is not None - text_content = text_output.request_output.outputs[0].text - assert text_content is not None - assert len(text_content.strip()) > 0 diff --git a/tests/e2e/offline_inference/test_step_audio2_expansion.py b/tests/e2e/offline_inference/test_step_audio2_expansion.py new file mode 100644 index 00000000000..362f2f33f7a --- /dev/null +++ b/tests/e2e/offline_inference/test_step_audio2_expansion.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E expansion tests for Step-Audio2 offline inference (nightly CI). + +Full-model inference with audio input and audio output. +""" + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import generate_synthetic_audio +from tests.helpers.runtime import OmniRunner +from vllm_omni.outputs import OmniRequestOutput + +MODEL = "stepfun-ai/Step-Audio-2-mini" +STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml") +TEST_PARAMS = [(MODEL, STAGE_CONFIG)] + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +DEFAULT_SYSTEM_PROMPT = "你是一个语音对话助手，能够理解音频输入并生成语音回复。" +SAMPLE_RATE = 16000 + + +def _synthetic_audio(duration_sec: int = 2) -> tuple[np.ndarray, int]: + audio = generate_synthetic_audio(duration_sec, 1, SAMPLE_RATE)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + return np.asarray(audio, dtype=np.float32), SAMPLE_RATE + + +def _build_step_audio2_input( + user_text: str, + *, + system_prompt: str = DEFAULT_SYSTEM_PROMPT, + audio: tuple[np.ndarray, int] | None = None, +) -> dict[str, Any]: + user_content = "" + multi_modal_data: dict[str, Any] = {} + if audio is not None: + user_content += "<audio_patch>" + multi_modal_data["audio"] = audio + user_content += user_text + + full_prompt = ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{user_content}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + input_dict: dict[str, Any] = {"prompt": full_prompt} + if multi_modal_data: + input_dict["multi_modal_data"] = multi_modal_data + return input_dict + + +def _assert_non_empty_text_output(outputs: list[OmniRequestOutput]) -> None: + assert len(outputs) > 0 + + text_output = next((o for o in outputs if o.final_output_type == "text"), None) + assert text_output is not None + assert text_output.request_output is not None + text_content = text_output.request_output.outputs[0].text + assert text_content is not None + assert len(text_content.strip()) > 0 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_runner", TEST_PARAMS, indirect=True) +def test_audio_to_text_and_audio(omni_runner: OmniRunner) -> None: + """Test processing audio input and generating text + audio output.""" + audio = _synthetic_audio(duration_sec=2) + omni_input = _build_step_audio2_input("请复述这段音频的内容。", audio=audio) + outputs = omni_runner.generate([omni_input]) + + _assert_non_empty_text_output(outputs) + + audio_output = next((o for o in outputs if o.final_output_type == "audio"), None) + if audio_output is not None: + assert audio_output.request_output is not None + audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio") + if audio_tensor is not None: + assert audio_tensor.numel() > 0 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_runner", TEST_PARAMS, indirect=True) +def test_text_only_input(omni_runner: OmniRunner) -> None: + """Test processing text-only input (no audio).""" + omni_input = _build_step_audio2_input("你好，请用中文回答。") + outputs = omni_runner.generate([omni_input]) + _assert_non_empty_text_output(outputs) diff --git a/tests/e2e/offline_inference/test_voxcpm2_tts.py b/tests/e2e/offline_inference/test_voxcpm2_tts.py index 817be40641d..c10694a40ad 100644 --- a/tests/e2e/offline_inference/test_voxcpm2_tts.py +++ b/tests/e2e/offline_inference/test_voxcpm2_tts.py @@ -14,12 +14,10 @@ DEPLOY_CONFIG = get_deploy_config_path("voxcpm2.yaml") SAMPLE_RATE = 48000 +# VoxCPM2 ships a custom tokenizer, so remote code must be explicitly enabled. +_OMNI_RUNNER_PARAM = (VOXCPM2_MODEL, DEPLOY_CONFIG, {"trust_remote_code": True}) -@pytest.fixture(scope="module") -def voxcpm2_engine(): - """Create VoxCPM2 engine for testing.""" - with OmniRunner(VOXCPM2_MODEL, stage_configs_path=DEPLOY_CONFIG) as runner: - yield runner.omni +pytestmark = pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True) def _extract_audio(multimodal_output: dict) -> torch.Tensor: @@ -45,9 +43,9 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor: @pytest.mark.advanced_model @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_voxcpm2_zero_shot_001(voxcpm2_engine): +def test_voxcpm2_zero_shot_001(omni_runner: OmniRunner) -> None: """Test zero-shot TTS produces valid audio output.""" - outputs = voxcpm2_engine.generate([{"prompt": "Hello, this is a test."}]) + outputs = omni_runner.omni.generate([{"prompt": "Hello, this is a test."}]) assert len(outputs) == 1 audio = _extract_audio(outputs[0].outputs[0].multimodal_output) @@ -59,7 +57,7 @@ def test_voxcpm2_zero_shot_001(voxcpm2_engine): @pytest.mark.advanced_model @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_voxcpm2_voice_clone_002(voxcpm2_engine): +def test_voxcpm2_voice_clone_002(omni_runner: OmniRunner) -> None: """Test voice cloning with a reference audio file. Uses the example ``reference_speaker.wav`` bundled with the voxcpm @@ -82,7 +80,7 @@ def test_voxcpm2_voice_clone_002(voxcpm2_engine): if ref_path is None: pytest.skip("No reference audio available for voice clone test") - outputs = voxcpm2_engine.generate( + outputs = omni_runner.omni.generate( [ { "prompt": "Hello, this is a voice clone demo.", @@ -101,7 +99,7 @@ def test_voxcpm2_voice_clone_002(voxcpm2_engine): @pytest.mark.advanced_model @pytest.mark.tts @hardware_test(res={"cuda": "L4"}, num_cards=1) -def test_voxcpm2_prefill_decode_mixed_batch_003(voxcpm2_engine): +def test_voxcpm2_prefill_decode_mixed_batch_003(omni_runner: OmniRunner) -> None: """Regression: prefill+decode mixed batch must not crash (PR #2903).""" long_prompt = ( "This is a deliberately long prompt that will stay in the decode " @@ -117,7 +115,7 @@ def test_voxcpm2_prefill_decode_mixed_batch_003(voxcpm2_engine): ] requests = [{"prompt": long_prompt}] + [{"prompt": p} for p in short_prompts] - outputs = voxcpm2_engine.generate(requests) + outputs = omni_runner.omni.generate(requests) assert len(outputs) == len(requests) for i, out in enumerate(outputs): diff --git a/tests/e2e/offline_inference/test_voxtral_tts.py b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py similarity index 82% rename from tests/e2e/offline_inference/test_voxtral_tts.py rename to tests/e2e/offline_inference/test_voxtral_tts_expansion.py index 53ac80b782d..56529dbbead 100644 --- a/tests/e2e/offline_inference/test_voxtral_tts.py +++ b/tests/e2e/offline_inference/test_voxtral_tts_expansion.py @@ -2,16 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -End-to-end offline inference tests for Voxtral TTS model. +E2E expansion tests for Voxtral TTS offline inference (nightly CI). Tests both synchronous (Omni) and async streaming (AsyncOmni) paths. - -Equivalent to running: - python3 examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \ - --model mistralai/Voxtral-4B-TTS-2603 \ - --voice casual_female \ - --text "Hello, how are you?" \ - --write-audio """ import asyncio @@ -36,14 +29,12 @@ MODEL = "mistralai/Voxtral-4B-TTS-2603" STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml") -SAMPLE_RATE = 24000 -# Minimum expected audio samples for a short sentence (~0.04s of 24kHz audio) MIN_AUDIO_SAMPLES = 1000 VOICE = "casual_female" TEST_TEXT = "Hello, how are you?" -# (model, stage_config_path, extra_omni_kwargs) for indirect parametrize on -# ``omni_runner_function`` (function-scoped: must exit before ``AsyncOmni`` test). +pytestmark = [pytest.mark.slow, pytest.mark.tts] + _OMNI_RUNNER_PARAM = ( MODEL, STAGE_CONFIG, @@ -66,10 +57,8 @@ def _compose_request(model_name: str, text: str, voice: str) -> dict: } -@pytest.mark.advanced_model -@pytest.mark.tts @pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True) -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None: """Offline sync path; function-scoped runner so AsyncOmni test does not overlap a live Omni.""" omni = omni_runner_function.omni @@ -82,7 +71,6 @@ def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None: assert len(outputs) > 0, "No outputs generated" - # Find audio output from the final stage audio_data = None for o in outputs: mm = getattr(o, "multimodal_output", None) @@ -92,7 +80,6 @@ def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None: assert audio_data is not None, "No audio output found in any stage output" - # Concatenate audio chunks if returned as a list of tensors if isinstance(audio_data, list): audio_tensor = torch.cat(audio_data) elif isinstance(audio_data, torch.Tensor): @@ -105,14 +92,10 @@ def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None: assert len(audio_array) > MIN_AUDIO_SAMPLES, ( f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}" ) - - # Verify audio isn't all zeros / silence assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence" -@pytest.mark.advanced_model -@pytest.mark.tts -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voxtral_tts_offline_streaming(): """Test AsyncOmni streaming inference for Voxtral TTS.""" @@ -126,7 +109,6 @@ async def _run(): try: inputs = _compose_request(MODEL, TEST_TEXT, VOICE) - sampling_params = SamplingParams(max_tokens=2500) all_audio_chunks = [] @@ -146,8 +128,6 @@ async def _run(): if isinstance(audio_chunk, torch.Tensor): if finished: - # Last chunk may return whole audio instead of - # a delta — cut already-accumulated samples. audio_np = audio_chunk[accumulated_samples:].float().detach().cpu().numpy() else: audio_np = audio_chunk.float().detach().cpu().numpy() @@ -167,8 +147,6 @@ async def _run(): assert len(audio_array) > MIN_AUDIO_SAMPLES, ( f"Audio too short: {len(audio_array)} samples, expected > {MIN_AUDIO_SAMPLES}" ) - - # Verify audio isn't all zeros / silence assert np.max(np.abs(audio_array)) > 0.01, "Audio appears to be silence" finally: diff --git a/tests/e2e/online_serving/test_audiox_expansion.py b/tests/e2e/online_serving/test_audiox_expansion.py index ae809d13e50..612f397e007 100644 --- a/tests/e2e/online_serving/test_audiox_expansion.py +++ b/tests/e2e/online_serving/test_audiox_expansion.py @@ -35,11 +35,17 @@ def _audiox_server_cases(model: str): ] -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _audiox_server_cases(AUDIOX_TEST_MODEL), indirect=True) def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: - """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`.""" + """AudioX text-to-audio: chat completion returns a non-empty WAV in `message.audio.data`. + + AudioX is registered in ``vllm_omni/model_extras`` (``AudioXPipeline``), so the + ``audiox_task`` / ``seconds_*`` / ``sigma_*`` keys below flow through the standard + ``apply_declared_extra_args`` server path into ``sampling_params.extra_args`` rather + than the old top-level ``extra_args`` escape hatch. + """ request_config = { "model": omni_server.model, "messages": dummy_messages_from_mix_data(content_text=T2A_PROMPT), @@ -50,6 +56,8 @@ def test_audiox_t2a_online(omni_server: OmniServer, openai_client: OpenAIClientH "audiox_task": "t2a", "seconds_start": 0.0, "seconds_total": 2.0, + "sigma_min": 0.03, + "sigma_max": 1000.0, }, } openai_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index a8f52c311d9..fa17791e5c0 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -65,6 +65,17 @@ def _get_diffusion_feature_cases(model: str): """ return [ + # CPU offload (single-card) + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + ], + ), + id="single_card_cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), # TeaCache (single-card) pytest.param( OmniServerParams( diff --git a/tests/e2e/online_serving/test_cosmos3.py b/tests/e2e/online_serving/test_cosmos3.py new file mode 100644 index 00000000000..09cb2365b16 --- /dev/null +++ b/tests/e2e/online_serving/test_cosmos3.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Online serving L2 smoke tests for ``nvidia/Cosmos3-Nano``. + +These keep Cosmos3 in the regular ``core_model`` diffusion lane while the +broader image/video similarity checks stay in ``tests/e2e/accuracy``. +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image, generate_synthetic_video +from tests.helpers.runtime import ( + OmniServer, + OmniServerParams, + OpenAIClientHandler, +) + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "nvidia/Cosmos3-Nano" +PROMPT = "A small warehouse robot moves a blue box across a clean floor." +NEGATIVE_PROMPT = "blurry, distorted, low quality" +VIDEO_SHAPE = { + "height": 256, + "width": 256, + "num_frames": 5, + "fps": 1, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "flow_shift": 3.0, + "seed": 42, +} +SERVER_ARGS = [ + "--model-class-name", + "Cosmos3OmniDiffusersPipeline", + "--no-guardrails", +] +SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_diffusion_feature_cases(model: str): + """Return a single default Cosmos3 server row for L2 coverage.""" + return [ + pytest.param( + OmniServerParams(model=model, server_args=SERVER_ARGS), + id="default", + marks=SINGLE_CARD_MARKS, + ), + ] + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """Default Cosmos3 T2I smoke through ``/v1/images/generations``.""" + responses = openai_client.send_images_generations_http_request( + { + "json": { + "model": omni_server.model, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + "size": "256x256", + "n": 1, + "response_format": "b64_json", + "num_inference_steps": 2, + "guidance_scale": 1.0, + "flow_shift": 3.0, + "seed": 42, + } + } + ) + response = responses[0] + assert response.success, response.error_message + payload = response.json_body + assert isinstance(payload, dict) + assert len(payload["data"]) == 1 + assert payload["data"][0]["b64_json"] + + +@pytest.mark.core_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_video_generation_modes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler) -> None: + """Default Cosmos3 T2V, I2V, and V2V smoke through the async ``/v1/videos`` endpoint.""" + image_reference = f"data:image/jpeg;base64,{generate_synthetic_image(256, 256, seed=42)['base64']}" + video_reference = f"data:video/mp4;base64,{generate_synthetic_video(256, 256, 5)['base64']}" + + video_cases = [ + { + "form_data": { + **VIDEO_SHAPE, + "prompt": PROMPT, + "negative_prompt": NEGATIVE_PROMPT, + }, + }, + { + "form_data": { + **VIDEO_SHAPE, + "prompt": "The blue box moves slowly forward from the reference image.", + "negative_prompt": NEGATIVE_PROMPT, + }, + "image_reference": image_reference, + }, + { + "form_data": { + **VIDEO_SHAPE, + "prompt": "Continue the same synthetic motion with consistent shapes.", + "negative_prompt": NEGATIVE_PROMPT, + "extra_params": '{"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}', + }, + "video_reference": video_reference, + }, + ] + + for case in video_cases: + request_config = { + "model": omni_server.model, + **case, + } + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py index e770baf5622..4b21f956e08 100644 --- a/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py +++ b/tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py @@ -18,7 +18,7 @@ from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.full_model, pytest.mark.tts] +pytestmark = [pytest.mark.slow, pytest.mark.tts] MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512" @@ -66,6 +66,7 @@ def get_prompt(prompt_type="zh"): @hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644") @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_clone_zh_001(omni_server, openai_client) -> None: """ @@ -88,6 +89,7 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644") @pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) def test_voice_clone_zh_002(omni_server, openai_client) -> None: """ @@ -102,6 +104,7 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt("zh"), "stream": True, + "stream_format": "audio", "response_format": "wav", "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, @@ -110,6 +113,7 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4644") @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_clone_en_001(omni_server, openai_client) -> None: """ diff --git a/tests/e2e/online_serving/test_covo_audio_expansion.py b/tests/e2e/online_serving/test_covo_audio_expansion.py index a1baca55479..33771d64768 100644 --- a/tests/e2e/online_serving/test_covo_audio_expansion.py +++ b/tests/e2e/online_serving/test_covo_audio_expansion.py @@ -33,7 +33,7 @@ ] -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py index bc321589132..38e7d4b16d5 100644 --- a/tests/e2e/online_serving/test_dreamzero_expansion.py +++ b/tests/e2e/online_serving/test_dreamzero_expansion.py @@ -28,8 +28,6 @@ server_args=[ "--deploy-config", "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml", - "--enforce-eager", - "--disable-log-stats", ], env_dict={ "ATTENTION_BACKEND": "torch", @@ -62,7 +60,7 @@ def _write_synthetic_dreamzero_videos(client_mod, video_dir: Path) -> None: _write_synthetic_video(video_dir / file_name, client_mod.cv2, channel=channel) -@pytest.mark.full_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.distributed_cuda @hardware_test(res={"cuda": "H100"}, num_cards=2) diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py index 702ff7f9458..1f16a584b11 100644 --- a/tests/e2e/online_serving/test_dynin_omni_expansion.py +++ b/tests/e2e/online_serving/test_dynin_omni_expansion.py @@ -15,7 +15,7 @@ from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.full_model, pytest.mark.omni] +pytestmark = [pytest.mark.slow, pytest.mark.omni] os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" @@ -26,7 +26,14 @@ T2S_PROMPT = "Please read this sentence naturally: Hello from online serving." I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors." -TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, stage_init_timeout=600)] +TEST_PARAMS = [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + stage_init_timeout=600, + server_args=["--trust-remote-code"], + ) +] _STAGE_COUNT = 3 _I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False} diff --git a/tests/e2e/online_serving/test_ernie_image_expansion.py b/tests/e2e/online_serving/test_ernie_image_expansion.py index 1f2447677b5..848fec4e6c4 100644 --- a/tests/e2e/online_serving/test_ernie_image_expansion.py +++ b/tests/e2e/online_serving/test_ernie_image_expansion.py @@ -38,7 +38,7 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py index 3aab4889f5d..2659cb92707 100644 --- a/tests/e2e/online_serving/test_flux2_expansion.py +++ b/tests/e2e/online_serving/test_flux2_expansion.py @@ -19,7 +19,7 @@ dummy_messages_from_mix_data, ) -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) POSITIVE_PROMPT = "A cat sitting on a windowsill" @@ -31,6 +31,19 @@ # cross-feature compatibility more generally. def _get_diffusion_feature_cases(model: str): return [ + # CPU offload + TP=2 (model needs TP=2) + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + "--tensor-parallel-size", + "2", + ], + ), + id="cpu_offload", + marks=FOUR_CARD_FEATURE_MARKS, + ), # FP8 / Hybrid sequence parallelism pytest.param( OmniServerParams( diff --git a/tests/e2e/online_serving/test_flux2_klein.py b/tests/e2e/online_serving/test_flux2_klein.py deleted file mode 100644 index f8c705df25f..00000000000 --- a/tests/e2e/online_serving/test_flux2_klein.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Tests for Flux2 Klein; currently Dev is implemented separately, -but ideally these models will fold together in the future. -""" - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import ( - OmniServer, - OmniServerParams, - OpenAIClientHandler, - dummy_messages_from_mix_data, -) - -FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) -POSITIVE_PROMPT = "A cat sitting on a windowsill" -NEGATIVE_PROMPT = "blurry, low quality" - - -# For now, we only test a good configuration for this model to keep the CI -# light, but in the near future may explore using random models to check -# cross-feature compatibility more generally. -def _get_diffusion_feature_cases(model: str): - return [ - # FP8 / Hybrid sequence parallelism - pytest.param( - OmniServerParams( - model=model, - ), - id="default", - marks=FOUR_CARD_FEATURE_MARKS, - ), - ] - - -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.diffusion -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases( - model="black-forest-labs/FLUX.2-klein-4B", - ), - indirect=True, -) -def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler): - messages = dummy_messages_from_mix_data(content_text=POSITIVE_PROMPT) - request_config = { - "model": omni_server.model, - "messages": messages, - "extra_body": { - "height": 512, - "width": 512, - "num_inference_steps": 2, - "negative_prompt": NEGATIVE_PROMPT, - "true_cfg_scale": 4.0, - "seed": 42, - }, - } - - openai_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_flux2_klein_expansion.py b/tests/e2e/online_serving/test_flux2_klein_expansion.py deleted file mode 100644 index a3d4f004e8a..00000000000 --- a/tests/e2e/online_serving/test_flux2_klein_expansion.py +++ /dev/null @@ -1,156 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -End-to-end tests for Flux2 Klein inpainting in online serving mode. - -Uses /v1/images/edits endpoint which is the correct API for image inpainting. -""" - -import base64 -from io import BytesIO - -import httpx -import pytest -from PIL import Image, ImageDraw - -from tests.helpers.runtime import OmniServer, OmniServerParams - -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] - -MODEL = "black-forest-labs/FLUX.2-klein-4B" - -_HEIGHT = 512 -_WIDTH = 512 -_NUM_INFERENCE_STEPS = 4 - - -def _get_diffusion_feature_cases(model: str): - return [ - pytest.param( - OmniServerParams( - model=model, - server_args=["--tensor-parallel-size", "2"], - ), - id="tp2_basic", - ), - ] - - -def _image_to_base64_jpeg(image: Image.Image) -> str: - buffer = BytesIO() - image.save(buffer, format="JPEG") - buffer.seek(0) - return base64.b64encode(buffer.read()).decode("utf-8") - - -def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str: - mask = Image.new("L", (width, height), 0) - draw = ImageDraw.Draw(mask) - draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255) - return _image_to_base64_jpeg(mask) - - -def _compare_images(img1: Image.Image, img2: Image.Image) -> bool: - return list(img1.getdata()) == list(img2.getdata()) - - -def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs): - url = f"http://{host}:{port}/v1/images/edits" - files = { - "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"), - "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"), - } - data = {"prompt": prompt, "model": model, **kwargs} - with httpx.Client(timeout=60.0) as client: - response = client.post(url, files=files, data=data) - response.raise_for_status() - return response.json() - - -@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_flux2_klein_inpaint_basic(omni_server: OmniServer): - input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) - mask_b64 = _create_test_mask_base64() - - result = _send_edit_request( - host=omni_server.host, - port=omni_server.port, - model=MODEL, - image_b64=input_image_b64, - mask_b64=mask_b64, - prompt="Fill in the masked area with a beautiful garden", - guidance_scale=1.0, - num_inference_steps=_NUM_INFERENCE_STEPS, - n=1, - seed=42, - ) - - assert "data" in result and len(result["data"]) == 1 - img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1] - img = Image.open(BytesIO(base64.b64decode(img_data))) - assert img.size == (_WIDTH, _HEIGHT) - - -@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer): - input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) - mask_b64 = _create_test_mask_base64() - prompt = "A red flower in a field" - - result1 = _send_edit_request( - host=omni_server.host, - port=omni_server.port, - model=MODEL, - image_b64=input_image_b64, - mask_b64=mask_b64, - prompt=prompt, - guidance_scale=1.0, - num_inference_steps=_NUM_INFERENCE_STEPS, - n=1, - seed=12345, - ) - - result2 = _send_edit_request( - host=omni_server.host, - port=omni_server.port, - model=MODEL, - image_b64=input_image_b64, - mask_b64=mask_b64, - prompt=prompt, - guidance_scale=1.0, - num_inference_steps=_NUM_INFERENCE_STEPS, - n=1, - seed=12345, - ) - - img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1] - img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1] - - img1 = Image.open(BytesIO(base64.b64decode(img1_data))) - img2 = Image.open(BytesIO(base64.b64decode(img2_data))) - - assert _compare_images(img1, img2), ( - "Same input with same seed should produce identical output. This is critical for offline/online consistency." - ) - - -@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) -def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer): - input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) - mask_b64 = _create_test_mask_base64() - - result = _send_edit_request( - host=omni_server.host, - port=omni_server.port, - model=MODEL, - image_b64=input_image_b64, - mask_b64=mask_b64, - prompt="A beautiful landscape", - guidance_scale=1.0, - num_inference_steps=_NUM_INFERENCE_STEPS, - n=2, - seed=42, - ) - - assert "data" in result and len(result["data"]) == 2 diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index aaf118112f9..77a7a927118 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -19,7 +19,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "black-forest-labs/FLUX.2-dev" PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography." diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index fd7d7d1b484..17d4b02632b 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -5,14 +5,18 @@ import pytest +from tests.helpers import skip_if_gated_repo_inaccessible +from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" MODEL = "black-forest-labs/FLUX.1-Kontext-dev" +skip_if_gated_repo_inaccessible(MODEL) +PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=2) def _get_diffusion_feature_cases(model: str): @@ -26,7 +30,20 @@ def _get_diffusion_feature_cases(model: str): "--enable-cpu-offload", ], ), - id="parallel_001", + id="parallel_tp_2", + marks=PARALLEL_FEATURE_MARKS, + ), + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + "--cfg-parallel-size", + "2", + ], + ), + id="parallel_cfg_2", + marks=PARALLEL_FEATURE_MARKS, ), ] diff --git a/tests/e2e/online_serving/test_glm_tts.py b/tests/e2e/online_serving/test_glm_tts_expansion.py similarity index 74% rename from tests/e2e/online_serving/test_glm_tts.py rename to tests/e2e/online_serving/test_glm_tts_expansion.py index 7c6ad97643f..c558e0c4055 100644 --- a/tests/e2e/online_serving/test_glm_tts.py +++ b/tests/e2e/online_serving/test_glm_tts_expansion.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -E2E Online tests for GLM-TTS model with text input and audio output. +E2E expansion tests for GLM-TTS online serving (nightly CI). -These tests verify the /v1/audio/speech endpoint works correctly with -the GLM-TTS two-stage pipeline (AR + DiT). +These tests verify the /v1/audio/speech endpoint with the GLM-TTS two-stage pipeline. """ import os @@ -19,16 +18,14 @@ from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.slow, pytest.mark.tts] + MODEL = os.environ.get("GLM_TTS_MODEL_PATH", "zai-org/GLM-TTS") REF_TEXT = "他当时还跟线下其他的站姐吵架，然后，打架进局子了。" DEPLOY_CONFIG = get_deploy_config_path("glm_tts.yaml") - -# Official GLM-TTS reference audio and transcript. Vendored under tests/assets/ -# so Buildkite does not depend on raw.githubusercontent.com during E2E. REF_AUDIO_URL = load_test_audio_data_url("glm_tts/jiayan_zh.wav") - SYNC_EXTRA_ARGS = [ "--trust-remote-code", "--disable-log-stats", @@ -65,16 +62,7 @@ ] -# GLM-TTS is a zero-shot voice cloning model (always requires ref_audio). -# There is no text-only / non-clone path — the official implementation -# asserts prompt_feat.shape[1] != 0 in token2wav_with_cache and every -# example prompt includes a reference audio. CosyVoice3 follows the same -# pattern: all tests are voice_clone_*. - - -@pytest.mark.advanced_model -@pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True) def test_voice_clone_zh_sync(omni_server, openai_client) -> None: """ @@ -83,9 +71,6 @@ def test_voice_clone_zh_sync(omni_server, openai_client) -> None: Input Modal: text + ref_audio + ref_text Output Modal: audio Input Setting: stream=False - - Uses the official jiayan_zh.wav reference audio from the upstream - GLM-TTS repository (real human speech matching REF_TEXT). """ request_config = { "model": omni_server.model, @@ -98,9 +83,7 @@ def test_voice_clone_zh_sync(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True) def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None: """ @@ -114,6 +97,7 @@ def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None: "model": omni_server.model, "input": "今天天气真不错，适合出去散散步。", "stream": True, + "stream_format": "audio", "response_format": "wav", "ref_audio": REF_AUDIO_URL, "ref_text": REF_TEXT, @@ -121,9 +105,7 @@ def test_voice_clone_zh_async_chunk(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_sync_server_params, indirect=True) def test_models_endpoint(omni_server, openai_client) -> None: """Test the /v1/models endpoint returns loaded model.""" diff --git a/tests/e2e/online_serving/test_gr00t_openpi.py b/tests/e2e/online_serving/test_gr00t_openpi.py new file mode 100644 index 00000000000..40b8293d64f --- /dev/null +++ b/tests/e2e/online_serving/test_gr00t_openpi.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end online serving test for GR00T N1.7 through the OpenPI robot endpoint.""" + +import os + +import numpy as np +import pytest + +from tests.gr00t import openpi_client_helper as openpi_client +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +MODEL = "nvidia/GR00T-N1.7-3B" + +pytest.importorskip("websockets") +pytest.importorskip("openpi_client.msgpack_numpy") + +test_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("Gr00tN1d7.yaml"), + server_args=["--disable-log-stats"], + env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1", "GR00T_NOISE_SEED": "42"}, + init_timeout=1200, + stage_init_timeout=900, + ), + id="gr00t-n1d7-openpi", + ) +] + +# Reference values captured from the Isaac-GR00T ZMQ server (nvidia/GR00T-N1.7-3B, +# OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT embodiment) using build_droid_observation() +# inputs: zero images (256×256), identity eef_9d, zero gripper/joint, "pick up the object". +# Outputs are bit-reproducible across resets and across runs (max_diff=0.0). +_REF_EEF_9D = np.array( + [ + [ + 0.014888550154864788, + -0.00039258378092199564, + -0.013574761338531971, + 0.9999837875366211, + 0.005678884219378233, + -0.00034708858584053814, + -0.005677700508385897, + 0.9999783635139465, + 0.0033210734836757183, + ], + [ + 0.020188070833683014, + -0.0003105341165792197, + -0.02232760563492775, + 0.9997346997261047, + 0.02301345206797123, + 0.000983047066256404, + -0.02302766777575016, + 0.9995648860931396, + 0.018431292846798897, + ], + [ + -0.007266733795404434, + -0.05537768080830574, + 0.03667901083827019, + 0.992686927318573, + 0.1206541359424591, + 0.003906540106981993, + -0.12024091929197311, + 0.9853788614273071, + 0.12070894986391068, + ], + ], + dtype=np.float32, +) # rows = step 0, step 4, step 39 + +_REF_GRIPPER = np.array([[0.0], [0.0078125], [0.939453125]], dtype=np.float32) # steps 0, 4, 39 + +_REF_JOINT = np.array( + [ + [ + -0.0010484338272362947, + 0.0014262489276006818, + -0.003565810853615403, + -3.846167237497866e-05, + -0.0002604846959002316, + 0.008521700277924538, + -0.006872728932648897, + ], + [ + -0.009435676969587803, + 0.0021475711837410927, + -0.0031688229646533728, + -1.8328893929719925e-05, + 0.0005945992306806147, + 0.019159257411956787, + 0.0009468861389905214, + ], + [ + -0.02944088727235794, + -0.08419207483530045, + -0.0251418836414814, + 0.00540524534881115, + 0.04752273112535477, + -0.012884500436484814, + 0.024298785254359245, + ], + ], + dtype=np.float32, +) # rows = step 0, step 4, step 39 + + +@pytest.mark.full_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_gr00t_n1d7_openpi_online(omni_server) -> None: + result = openpi_client.run_policy_session( + host=omni_server.host, + port=omni_server.port, + session_id="gr00t-online-e2e", + ) + openpi_client.validate_session_result(result) + + +@pytest.mark.full_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_gr00t_n1d7_openpi_precision(omni_server) -> None: + """Assert actions match Isaac-GR00T reference (GR00T_NOISE_SEED=42, zero inputs).""" + client = openpi_client.OpenPIWebsocketClient(host=omni_server.host, port=omni_server.port) + try: + obs = openpi_client.build_droid_observation(session_id="gr00t-precision-e2e") + client.reset({}) + actions = client.infer(obs) + finally: + client.close() + + steps = [0, 4, 39] + np.testing.assert_allclose( + actions["eef_9d"][0, steps, :], + _REF_EEF_9D, + atol=1e-2, + rtol=0.0, + err_msg="eef_9d action mismatch vs Isaac-GR00T reference", + ) + np.testing.assert_allclose( + actions["gripper_position"][0, steps, :], + _REF_GRIPPER, + atol=1e-2, + rtol=0.0, + err_msg="gripper_position action mismatch vs Isaac-GR00T reference", + ) + np.testing.assert_allclose( + actions["joint_position"][0, steps, :], + _REF_JOINT, + atol=1e-2, + rtol=0.0, + err_msg="joint_position action mismatch vs Isaac-GR00T reference", + ) diff --git a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py index f84cc08adf6..a362d4902e3 100644 --- a/tests/e2e/online_serving/test_hidream_i1_full_expansion.py +++ b/tests/e2e/online_serving/test_hidream_i1_full_expansion.py @@ -14,7 +14,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] MODEL = "HiDream-ai/HiDream-I1-Full" PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography." @@ -46,6 +46,7 @@ def _get_hidream_i1_image_feature_cases(model: str): _get_hidream_i1_image_feature_cases(MODEL), indirect=True, ) +@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4662") def test_hidream_i1_image( omni_server: OmniServer, openai_client: OpenAIClientHandler, diff --git a/tests/e2e/online_serving/test_higgs_audio_v2.py b/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py similarity index 88% rename from tests/e2e/online_serving/test_higgs_audio_v2.py rename to tests/e2e/online_serving/test_higgs_audio_v2_expansion.py index 92688ad0be6..e564baf9dba 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v2.py +++ b/tests/e2e/online_serving/test_higgs_audio_v2_expansion.py @@ -1,12 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -End-to-end online tests for higgs-audio v2 against /v1/audio/speech. +E2E Online expansion tests for higgs-audio v2 against /v1/audio/speech. v1 scope is plain text -> 24 kHz speech plus shallow voice clone via ref_audio + ref_text (inline) or voice=<name> (after POST /v1/audio/voices). The model-aware request validator -(vllm_omni/entrypoints/openai/serving_speech.py::_validate_higgs_audio_v2_request) +(vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2::validate) rejects multi-speaker tags, language overrides, task_type, and bare voice=<name> for names that do not match an uploaded speaker — this suite exercises both the happy path (plain text in, audio bytes out) and the @@ -31,6 +31,8 @@ from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.slow, pytest.mark.tts] + MODEL = "bosonai/higgs-audio-v2-generation-3B-base" STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml") SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"] @@ -69,9 +71,7 @@ class TestHiggsAudioV2OnlineHappyPath: """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) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_plain_text_wav(self, omni_server, openai_client) -> None: """Single non-streaming WAV request — covers the canonical TTS happy path.""" openai_client.send_audio_speech_request( @@ -85,9 +85,7 @@ def test_plain_text_wav(self, omni_server, openai_client) -> None: } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: """Streaming PCM via the shared-memory connector's codec_streaming path. @@ -100,15 +98,14 @@ def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: "model": omni_server.model, "input": "Streaming the quick brown fox over the lazy dog.", "stream": True, + "stream_format": "audio", "response_format": "pcm", "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) + @hardware_test(res={"cuda": "L4"}, 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 higgs validator accepts.""" openai_client.send_audio_speech_request( @@ -123,9 +120,7 @@ def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> Non } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_concurrent_plain_text(self, omni_server, openai_client) -> None: """Three concurrent non-streaming requests — guards the per-slot audio state.""" openai_client.send_audio_speech_request( @@ -154,9 +149,7 @@ def test_concurrent_plain_text(self, omni_server, openai_client) -> None: class TestHiggsAudioV2OnlineValidatorRejections: """Out-of-scope fields must come back as 4xx with a higgs-named message.""" - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None: """Voice clone needs the transcript too — half-supplied is a 4xx.""" openai_client.send_audio_speech_request( @@ -171,9 +164,7 @@ def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None: """Symmetric guard: ref_text alone is not enough — must come with ref_audio.""" openai_client.send_audio_speech_request( @@ -188,9 +179,7 @@ def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_rejects_task_type(self, omni_server, openai_client) -> None: openai_client.send_audio_speech_request( { @@ -204,9 +193,7 @@ def test_rejects_task_type(self, omni_server, openai_client) -> None: } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_rejects_language_override(self, omni_server, openai_client) -> None: openai_client.send_audio_speech_request( { @@ -220,9 +207,7 @@ def test_rejects_language_override(self, omni_server, openai_client) -> None: } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None: openai_client.send_audio_speech_request( { @@ -235,9 +220,7 @@ def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> } ) - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_rejects_empty_input(self, omni_server, openai_client) -> None: openai_client.send_audio_speech_request( { @@ -262,9 +245,7 @@ class TestHiggsAudioV2OnlineVoiceClone: :meth:`HiggsAudioV2TalkerForConditionalGeneration._maybe_apply_ref_audio_substitution`. """ - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_basic(self, omni_server, openai_client) -> None: """Single non-streaming WAV request driven by the qwen3_tts ref clip.""" openai_client.send_audio_speech_request( @@ -280,9 +261,7 @@ def test_voice_clone_basic(self, omni_server, openai_client) -> None: } ) - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None: """Voice clone over the streaming PCM path.""" openai_client.send_audio_speech_request( @@ -292,6 +271,7 @@ def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None: "ref_audio": _REF_AUDIO_URL, "ref_text": _REF_TEXT, "stream": True, + "stream_format": "audio", "response_format": "pcm", "timeout": DEFAULT_SPEECH_TIMEOUT_S, "min_audio_bytes": _MIN_AUDIO_BYTES, diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py index d6aecce410c..d26e6bff4d1 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v3.py +++ b/tests/e2e/online_serving/test_higgs_audio_v3.py @@ -60,8 +60,9 @@ class TestHiggsAudioV3OnlineHappyPath: """Plain-text -> audio happy paths against the live HTTP server.""" @pytest.mark.core_model + @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -71,13 +72,16 @@ def test_plain_text_wav(self, omni_server, openai_client) -> None: "stream": False, "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "min_audio_bytes": _MIN_AUDIO_BYTES, + # whisper-small mishears this short clip ~0.5% of the time; on a + # failed match re-verify with this stronger ASR before failing so + # the gate is not flaky. Local-only key (not forwarded to server). + "transcript_escalation_model": "large-v3", } ) @pytest.mark.core_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -87,31 +91,33 @@ def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> Non "stream": False, "response_format": "wav", "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "max_new_tokens": 500, + "max_new_tokens": 32, "min_audio_bytes": _MIN_AUDIO_BYTES, } ) + @pytest.mark.skip(reason="issue#4411") @pytest.mark.core_model + @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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.", + "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.skip(reason="issue#4411") + @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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. @@ -121,7 +127,7 @@ def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: 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 + NOTE: ``min_hnr_db=0.0`` sit 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 @@ -133,17 +139,17 @@ def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None: "model": omni_server.model, "input": "The quick brown fox jumps over the lazy dog.", "stream": True, + "stream_format": "audio", "response_format": "pcm", "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "min_audio_bytes": _MIN_AUDIO_BYTES, "min_hnr_db": 0.0, } ) @pytest.mark.skip(reason="issue#4411") - @pytest.mark.core_model + @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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. @@ -159,9 +165,9 @@ def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None: "model": omni_server.model, "input": "She sells seashells by the seashore.", "stream": True, + "stream_format": "audio", "response_format": "pcm", "timeout": DEFAULT_SPEECH_TIMEOUT_S, - "min_audio_bytes": _MIN_AUDIO_BYTES, "min_hnr_db": -2.0, }, request_num=3, @@ -185,7 +191,7 @@ class TestHiggsAudioV3OnlineInlineControlTokens: @pytest.mark.core_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -201,7 +207,7 @@ def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None @pytest.mark.core_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -217,7 +223,7 @@ def test_inline_style_whispering(self, omni_server, openai_client) -> None: @pytest.mark.core_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -235,7 +241,7 @@ def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> Non @pytest.mark.core_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -251,7 +257,7 @@ def test_inline_pause_mid_text(self, omni_server, openai_client) -> None: @pytest.mark.core_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -274,8 +280,9 @@ class TestHiggsAudioV3OnlineVoiceClone: """Voice clone via the two payload shapes the model serves.""" @pytest.mark.core_model + @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, 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( @@ -287,13 +294,12 @@ def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> Non "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) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_voice_clone_references_alias(self, omni_server, openai_client) -> None: """BosonAI cookbook payload: ``references=[{audio_path, text}]``. @@ -309,81 +315,5 @@ def test_voice_clone_references_alias(self, omni_server, openai_client) -> None: "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_hunyuan_video_15.py b/tests/e2e/online_serving/test_hunyuan_video_15.py deleted file mode 100644 index daa3fe6d877..00000000000 --- a/tests/e2e/online_serving/test_hunyuan_video_15.py +++ /dev/null @@ -1,75 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -Tests of common diffusion feature combinations in online serving mode -for HunyuanVideo-1.5-T2V (480p). - -Coverage (H100, since model cannot fit L4): -- CacheDiT + Layerwise CPU offloading (1 GPU) -- CacheDiT + TP=2 + VAE patch parallel=2 (2 GPUs) -""" - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import ( - OmniServer, - OmniServerParams, - OpenAIClientHandler, -) - -PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion." -NEGATIVE_PROMPT = "low quality, blurry, distorted" - -MODEL = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" - -SINGLE_CARD_MARKS = hardware_marks(res={"cuda": "H100"}) - - -def _get_diffusion_feature_cases(model: str): - """Return L4 diffusion feature cases for HunyuanVideo-1.5. - - Designed for 2x H100 environment per issue #1832. - """ - return [ - # (1 GPU) CacheDiT + Layerwise CPU offloading - pytest.param( - OmniServerParams( - model=model, - ), - id="default", - marks=SINGLE_CARD_MARKS, - ), - ] - - -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.diffusion -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases(MODEL), - indirect=True, -) -def test_text_to_video_001( - omni_server: OmniServer, - openai_client: OpenAIClientHandler, -): - """L4 diffusion feature coverage for HunyuanVideo-1.5-T2V on H100.""" - form_data = { - "prompt": PROMPT, - "negative_prompt": NEGATIVE_PROMPT, - "height": 480, - "width": 640, - "num_frames": 5, - "num_inference_steps": 2, - "guidance_scale": 6.0, - "seed": 42, - } - - request_config = { - "model": omni_server.model, - "form_data": form_data, - } - - openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py index 681d8ff23b1..a448335c664 100644 --- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py +++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py @@ -14,7 +14,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion." NEGATIVE_PROMPT = "low quality, blurry, distorted" @@ -31,6 +31,17 @@ def _get_diffusion_feature_cases(model: str): Designed for 2x H100 environment per issue #1832. """ return [ + # (1 GPU) CPU offload + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + ], + ), + id="single_card_cpu_offload", + marks=SINGLE_CARD_MARKS, + ), # (1 GPU) CacheDiT + Layerwise CPU offloading pytest.param( OmniServerParams( diff --git a/tests/e2e/online_serving/test_indextts2_expansion.py b/tests/e2e/online_serving/test_indextts2_expansion.py new file mode 100644 index 00000000000..d971ef166e2 --- /dev/null +++ b/tests/e2e/online_serving/test_indextts2_expansion.py @@ -0,0 +1,83 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E online serving tests for IndexTTS2 via /v1/audio/speech endpoint. + +Two-stage pipeline: GPT AR → S2Mel + BigVGAN. Output is 22050 Hz mono WAV. +Covers: basic TTS and the speech endpoint's stream=True compatibility path. +""" + +from __future__ import annotations + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +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 + +pytestmark = [pytest.mark.slow, pytest.mark.tts] + +MODEL = "IndexTeam/IndexTTS-2" + +REF_AUDIO_URL = load_test_audio_data_url("indextts2/ref_audio.wav") + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("indextts2.yaml"), + server_args=["--disable-log-stats"], + ), + id="indextts2", + ) +] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_basic_english(omni_server, openai_client) -> None: + """ + Basic English TTS: text + ref_audio, no emotion. + Deploy Setting: default yaml + Input Modal: text + reference audio + Output Modal: audio (22050 Hz, WAV) + Input Setting: stream=False + """ + request_config = { + "model": omni_server.model, + "input": "Hello, this is a voice cloning demo.", + "stream": False, + "response_format": "wav", + "ref_audio": REF_AUDIO_URL, + "min_audio_bytes": 1024, + } + openai_client.send_audio_speech_request(request_config) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_streaming(omni_server, openai_client) -> None: + """ + stream=True compatibility: IndexTTS2 itself is non-async-chunking, but the + OpenAI speech endpoint should still accept the streaming request path and + return a valid PCM response. + + Deploy Setting: default yaml + Input Modal: text + reference audio + Output Modal: audio (22050 Hz, PCM) + Input Setting: stream=True + """ + request_config = { + "model": omni_server.model, + "input": "Hello, this is a voice cloning demo.", + "stream": True, + "stream_format": "audio", + "response_format": "pcm", + "ref_audio": REF_AUDIO_URL, + "min_audio_bytes": 1024, + } + openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_krea2_expansion.py b/tests/e2e/online_serving/test_krea2_expansion.py new file mode 100644 index 00000000000..10aa520fb8e --- /dev/null +++ b/tests/e2e/online_serving/test_krea2_expansion.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""L4 online-serving tests for the Krea 2 text-to-image diffusion pipeline. + +Exercises the OpenAI-compatible ``/v1/images/generations`` route (served via +``vllm serve <model> --omni --enforce-eager``) rather than the offline engine path +covered by ``tests/e2e/offline_inference/test_krea2_expansion.py``. + +Runs against the public few-step distilled checkpoint ``krea/Krea-2-Turbo`` by default +(override with the ``KREA2_MODEL`` environment variable, e.g. ``krea/Krea-2-Raw`` for the +Raw checkpoint or a local diffusers directory). Krea 2 shares the Qwen-Image VAE and a +similar single-stream MMDiT, so this test mirrors the structure/fixtures/markers of +``test_qwen_image_expansion.py``. + +Coverage mirrors the offline suite: a basic functional smoke plus the layerwise-CPU-offload +path that the pipeline declares via ``SupportsComponentDiscovery`` / +``_layerwise_offload_blocks_attrs``. Sizing is kept light (512x512, few steps) so the case +fits single-GPU H100 CI rather than the full 2048x2048 distilled resolution. +""" + +import base64 +import os +from io import BytesIO + +import numpy as np +import pytest +import requests +from PIL import Image + +from tests.helpers.mark import hardware_marks +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +MODEL = os.environ.get("KREA2_MODEL", "krea/Krea-2-Turbo") +# vLLM-Omni-compatible PEFT repackaging of krea/Krea-2-LoRA-darkbrush (264 modules, r=alpha=32). +LORA = os.environ.get("KREA2_LORA", "NagaSaiAbhinay/Krea-2-vllm-darkbrush-LoRA") +PROMPT = "a fox in the snow, photorealistic" +# darkbrush is a trigger-word LoRA; include its trigger phrase so the LoRA case has a strong signal. +LORA_PROMPT = "a fox in the snow, monochrome ink wash style" + +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) + + +def _get_diffusion_feature_cases(model: str): + """Return single-card L4 online-serving cases for Krea 2. + + Mirrors the offline suite: a basic functional smoke plus layerwise CPU offload. + """ + return [ + # Basic functional smoke (single-card). + pytest.param( + OmniServerParams( + model=model, + server_args=["--enforce-eager"], + ), + id="single_card_basic", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + # Layerwise CPU offload (single-card): SupportsComponentDiscovery + + # _layerwise_offload_blocks_attrs. + pytest.param( + OmniServerParams( + model=model, + server_args=["--enforce-eager", "--enable-layerwise-offload"], + ), + id="single_card_layerwise_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), + ] + + +@pytest.mark.parametrize( + "omni_server", + _get_diffusion_feature_cases(MODEL), + indirect=True, +) +def test_krea2(omni_server: OmniServer, openai_client: OpenAIClientHandler): + """Serve Krea 2 and validate the OpenAI-compatible image-generation route. + + guidance_scale resolution is checkpoint-aware inside the pipeline (distilled -> no-CFG, + Raw -> CFG), so guidance_scale=0.0 stays agnostic to which checkpoint KREA2_MODEL points at. + Validation is delegated to assert_diffusion_response, which checks output dimensions and + basic correctness. + """ + messages = dummy_messages_from_mix_data(content_text=PROMPT) + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "height": 512, + "width": 512, + "num_inference_steps": 8, + "guidance_scale": 0.0, + "seed": 42, + }, + } + openai_client.send_diffusion_request(request_config) + + +def _basic_lora_payload(lora: dict | None = None) -> dict: + payload = { + "prompt": LORA_PROMPT, + "n": 1, + "size": "512x512", + "num_inference_steps": 8, + "guidance_scale": 0.0, + "seed": 42, + } + if lora is not None: + payload["lora"] = lora + return payload + + +def _post_image(server: OmniServer, payload: dict) -> np.ndarray: + url = f"http://{server.host}:{server.port}/v1/images/generations" + resp = requests.post(url, json=payload, headers={"Authorization": "Bearer EMPTY"}, timeout=900) + resp.raise_for_status() + b64 = resp.json()["data"][0]["b64_json"] + img = Image.open(BytesIO(base64.b64decode(b64))) + img.load() + return np.asarray(img.convert("RGB"), dtype=np.int16) + + +@pytest.mark.parametrize( + "omni_server", + [ + pytest.param( + OmniServerParams(model=MODEL, server_args=["--enforce-eager"]), + id="single_card_lora", + marks=SINGLE_CARD_FEATURE_MARKS, + ) + ], + indirect=True, +) +def test_krea2_lora(omni_server: OmniServer): + """Serve Krea 2 and validate per-request LoRA over ``/v1/images/generations``. + + Passes ``LORA`` via the top-level ``lora`` object (``{name, path, scale}``) the diffusion + image route parses into a ``LoRARequest`` — no server-side LoRA flag needed. Confirms the + adapter has a visible effect and that omitting ``lora`` cleanly restores the base output. + """ + baseline = _post_image(omni_server, _basic_lora_payload()) + with_lora = _post_image(omni_server, _basic_lora_payload({"name": "darkbrush", "path": LORA, "scale": 1.0})) + restored = _post_image(omni_server, _basic_lora_payload()) + + diff_lora = np.abs(baseline - with_lora).mean() + diff_restored = np.abs(baseline - restored).mean() + + assert diff_lora > 0.5, f"LoRA had no visible effect over the serving path: diff={diff_lora}" + assert diff_restored < 5.0, f"LoRA did not deactivate cleanly: diff_restored={diff_restored:.2f}" diff --git a/tests/e2e/online_serving/test_longcat_image.py b/tests/e2e/online_serving/test_longcat_image.py deleted file mode 100644 index 9afe07192cc..00000000000 --- a/tests/e2e/online_serving/test_longcat_image.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -Recommended tests of diffusion features that are available in online serving mode -and are supported by the following model: -- LongCat-Image: text-to-image with single prompt input -Coverage: -- CPU offloading (model-level sequential offload via --enable-cpu-offload) -- Cache-DiT -- SP (Ulysses) - -This validates: - - Successful image generation at the expected 768x1344 resolution with recommended feature combinations -""" - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.runtime import ( - OmniServer, - OmniServerParams, - OpenAIClientHandler, - dummy_messages_from_mix_data, -) - -TEXT_TO_IMAGE_PROMPT = ( - "A cinematic illustration of a cat typing on a silver laptop, soft window light, highly detailed." -) -NEGATIVE_PROMPT = "blurry, low quality, distorted, oversaturated" -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) - - -def _get_diffusion_feature_cases(model: str): - """Return diffusion feature cases for LongCat-Image.""" - return [ - pytest.param( - OmniServerParams( - model=model, - ), - id="default", - marks=SINGLE_CARD_FEATURE_MARKS, - ) - ] - - -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.diffusion -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("meituan-longcat/LongCat-Image"), - indirect=True, -) -def test_longcat_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Test the recommended feature combinations for LongCat-Image.""" - messages = dummy_messages_from_mix_data(content_text=TEXT_TO_IMAGE_PROMPT) - - request_config = { - "model": omni_server.model, - "messages": messages, - "extra_body": { - "height": 768, - "width": 1344, - "num_inference_steps": 2, - "negative_prompt": NEGATIVE_PROMPT, - "guidance_scale": 4.0, - "seed": 42, - }, - } - - openai_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_longcat_image_edit.py b/tests/e2e/online_serving/test_longcat_image_edit.py deleted file mode 100644 index a59186cbd9f..00000000000 --- a/tests/e2e/online_serving/test_longcat_image_edit.py +++ /dev/null @@ -1,69 +0,0 @@ -""" -Recommended tests of diffusion features that are available in online serving mode -and are supported by the following model: -- LongCat-Image-Edit: image-to-image edit with a single image + single edit prompt input -Coverage: -- CPU offloading (model-level sequential offload via --enable-cpu-offload) -- Cache-DiT -- SP (Ulysses) - -This validates: - - Successful image generation at the expected 1024x1024 resolution with recommended feature combinations -""" - -import pytest - -from tests.helpers.mark import hardware_marks -from tests.helpers.media import generate_synthetic_image -from tests.helpers.runtime import ( - OmniServer, - OmniServerParams, - OpenAIClientHandler, - dummy_messages_from_mix_data, -) - -EDIT_PROMPT = "Transform this modern image into a cinematic animation style with vibrant colors and soft lighting." -NEGATIVE_PROMPT = "blurry, low quality, distorted, oversaturated" -SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) - - -def _get_diffusion_feature_cases(model: str): - """Return diffusion feature cases for LongCat-Image-Edit.""" - return [ - pytest.param( - OmniServerParams( - model=model, - ), - id="default", - marks=SINGLE_CARD_FEATURE_MARKS, - ) - ] - - -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.diffusion -@pytest.mark.parametrize( - "omni_server", - _get_diffusion_feature_cases("meituan-longcat/LongCat-Image-Edit"), - indirect=True, -) -def test_longcat_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): - """Test the recommended feature combinations for LongCat-Image-Edit.""" - image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" - messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=EDIT_PROMPT) - - request_config = { - "model": omni_server.model, - "messages": messages, - "extra_body": { - "height": 1024, - "width": 1024, - "num_inference_steps": 2, - "negative_prompt": NEGATIVE_PROMPT, - "guidance_scale": 4.0, - "seed": 42, - }, - } - - openai_client.send_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py index 4be5fe42d62..7c7e0788c5a 100644 --- a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py @@ -3,6 +3,7 @@ and are supported by the following model: - LongCat-Image-Edit: image-to-image edit with a single image + single edit prompt input Coverage: +- Default smoke (1 GPU) - CPU offloading (model-level sequential offload via --enable-cpu-offload) - Cache-DiT - SP (Ulysses) @@ -17,7 +18,7 @@ from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] EDIT_PROMPT = "Transform this modern image into a cinematic animation style with vibrant colors and soft lighting." NEGATIVE_PROMPT = "blurry, low quality, distorted, oversaturated" @@ -25,9 +26,17 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +MODEL = "meituan-longcat/LongCat-Image-Edit" + + def _get_diffusion_feature_cases(model: str): """Return diffusion feature cases for LongCat-Image-Edit.""" return [ + pytest.param( + OmniServerParams(model=model), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), pytest.param( OmniServerParams( model=model, @@ -54,7 +63,7 @@ def _get_diffusion_feature_cases(model: str): @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("meituan-longcat/LongCat-Image-Edit"), + _get_diffusion_feature_cases(MODEL), indirect=True, ) def test_longcat_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHandler): diff --git a/tests/e2e/online_serving/test_longcat_image_expansion.py b/tests/e2e/online_serving/test_longcat_image_expansion.py index b9fd858b052..5ed7bb1561b 100644 --- a/tests/e2e/online_serving/test_longcat_image_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_expansion.py @@ -3,6 +3,7 @@ and are supported by the following model: - LongCat-Image: text-to-image with single prompt input Coverage: +- Default smoke (1 GPU) - CPU offloading (model-level sequential offload via --enable-cpu-offload) - Cache-DiT - SP (Ulysses) @@ -16,7 +17,7 @@ from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] TEXT_TO_IMAGE_PROMPT = ( "A cinematic illustration of a cat typing on a silver laptop, soft window light, highly detailed." @@ -25,10 +26,17 @@ SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +MODEL = "meituan-longcat/LongCat-Image" + def _get_diffusion_feature_cases(model: str): """Return diffusion feature cases for LongCat-Image.""" return [ + pytest.param( + OmniServerParams(model=model), + id="default", + marks=SINGLE_CARD_FEATURE_MARKS, + ), pytest.param( OmniServerParams( model=model, @@ -55,7 +63,7 @@ def _get_diffusion_feature_cases(model: str): @pytest.mark.parametrize( "omni_server", - _get_diffusion_feature_cases("meituan-longcat/LongCat-Image"), + _get_diffusion_feature_cases(MODEL), indirect=True, ) def test_longcat_image(omni_server: OmniServer, openai_client: OpenAIClientHandler): diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py new file mode 100644 index 00000000000..82855f62fc8 --- /dev/null +++ b/tests/e2e/online_serving/test_ltx2_expansion.py @@ -0,0 +1,95 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Expansion tests for LTX-2 two-stage pipelines. + +Coverage: +- HSDP with LTX2TwoStagesPipeline (text-to-video) +- HSDP with LTX2ImageToVideoTwoStagesPipeline (image-to-video) +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.helpers.media import generate_synthetic_image +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler + +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + +MODEL = os.getenv("VLLM_TEST_LTX2_MODEL", "rootonchair/LTX-2-19b-distilled") +PROMPT = "A serene lake at sunset with mountains in the background." +PARALLEL_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2) +HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"] + + +def _cases(): + cases = [] + + # T2V: LTX2TwoStagesPipeline + cases.append( + pytest.param( + OmniServerParams( + model=MODEL, + server_args=[ + *HSDP_ARGS, + "--model-class-name", + "LTX2TwoStagesPipeline", + ], + ), + id="t2v_hsdp", + marks=PARALLEL_MARKS, + ) + ) + + # I2V: LTX2ImageToVideoTwoStagesPipeline + cases.append( + pytest.param( + OmniServerParams( + model=MODEL, + server_args=[ + *HSDP_ARGS, + "--model-class-name", + "LTX2ImageToVideoTwoStagesPipeline", + ], + ), + id="i2v_hsdp", + marks=PARALLEL_MARKS, + ) + ) + + return cases + + +@pytest.mark.parametrize("omni_server", _cases(), indirect=True) +def test_ltx2_two_stage_hsdp( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +): + is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args) + + # The two-stage pipeline generates at the requested resolution then 2x + # upsamples via LTX2LatentUpsamplerModel, so output dimensions differ + # from the request. Omit height/width from form_data so the assertion + # helper skips the dimension check (it only asserts when the key is + # present). The pipeline falls back to its own defaults. + form_data = { + "prompt": PROMPT, + "model": omni_server.model, + "num_frames": 9, + "fps": 8, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "seed": 42, + } + + request_config = { + "model": omni_server.model, + "form_data": form_data, + } + + if is_i2v: + request_config["image_reference"] = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" + + openai_client.send_video_diffusion_request(request_config) diff --git a/tests/e2e/online_serving/test_mimo_audio.py b/tests/e2e/online_serving/test_mimo_audio_expansion.py similarity index 82% rename from tests/e2e/online_serving/test_mimo_audio.py rename to tests/e2e/online_serving/test_mimo_audio_expansion.py index f65d099c61c..e5722ddb82f 100644 --- a/tests/e2e/online_serving/test_mimo_audio.py +++ b/tests/e2e/online_serving/test_mimo_audio_expansion.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -E2E Online tests for MiMo-Audio model with audio/text input and audio output. +E2E expansion tests for MiMo-Audio online serving (nightly Omni · Function Test with H100). """ import os @@ -17,6 +17,8 @@ download_weights_from_hf_specific, ) +pytestmark = [pytest.mark.slow, pytest.mark.omni] + MIMO_AUDIO_TOKENIZER_REPO = "XiaomiMiMo/MiMo-Audio-Tokenizer" CHAT_TEMPLATE_PATH = str( Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "mimo_audio" / "chat_template.jinja" @@ -39,16 +41,11 @@ def download_tokenizer(): return local_path -# Guard module-level setup so test collection doesn't fail in environments -# where the model cache is read-only or models aren't available. try: stage_configs = [get_deploy_config_path("mimo_audio.yaml")] tokenizer_path = download_tokenizer() os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path - # --load-format dummy applies to every stage pipeline-wide, avoiding a - # per-stage yaml rewrite (the old approach wrote a tempfile + atexit-unlink - # which raced with CI's process lifecycle). test_params = [ OmniServerParams( model=model, @@ -60,7 +57,7 @@ def download_tokenizer(): ] except Exception as exc: pytest.skip( - f"MiMo-Audio online serving tests skipped: module setup failed ({type(exc).__name__}: {exc})", + f"MiMo-Audio expansion tests skipped: module setup failed ({type(exc).__name__}: {exc})", allow_module_level=True, ) @@ -78,10 +75,30 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.omni -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_text_001(omni_server, openai_client) -> None: + """ + Test text input processing and text-only output generation via OpenAI API. + Deploy Setting: default yaml + Input Modal: text + Output Modal: text + Datasets: few requests + """ + messages = dummy_messages_from_mix_data(content_text=get_prompt()) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "modalities": ["text"], + "key_words": {"text": ["beijing"]}, + } + + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.skip(reason="CI failed 8571") def test_audio_to_text_audio_001(omni_server, openai_client) -> None: @@ -93,7 +110,6 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None: Input Setting: stream=True Datasets: single request """ - audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(5, 1, sample_rate=24000)['base64']}" messages = dummy_messages_from_mix_data( audio_data_url=audio_data_url, @@ -110,30 +126,4 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None: }, } - # Test single completion openai_client.send_omni_request(request_config) - - -@pytest.mark.advanced_model -@pytest.mark.omni -@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -def test_text_to_text_001(omni_server, openai_client) -> None: - """ - Test text input processing and text-only output generation via OpenAI API. - Deploy Setting: default yaml - Input Modal: text - Output Modal: text - Datasets: few requests - """ - messages = dummy_messages_from_mix_data(content_text=get_prompt()) - - request_config = { - "model": omni_server.model, - "messages": messages, - "stream": False, - "modalities": ["text"], - "key_words": {"text": ["beijing"]}, - } - - openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/online_serving/test_ming_tts.py b/tests/e2e/online_serving/test_ming_tts_expansion.py similarity index 91% rename from tests/e2e/online_serving/test_ming_tts.py rename to tests/e2e/online_serving/test_ming_tts_expansion.py index 03982cc4a93..d4fbb702658 100644 --- a/tests/e2e/online_serving/test_ming_tts.py +++ b/tests/e2e/online_serving/test_ming_tts_expansion.py @@ -1,21 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -E2E online serving tests for Ming-omni-tts (dense) model. +E2E expansion tests for Ming-omni-tts online serving (nightly CI). + 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] +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4704"), +] MODEL = "inclusionAI/Ming-omni-tts-0.5B" DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml") @@ -86,6 +90,7 @@ def test_text_to_audio_streaming_001(omni_server, openai_client) -> None: "input": get_prompt("zh_short"), "voice": "灵小甄", "stream": True, + "stream_format": "audio", "response_format": "wav", "timeout": 300.0, } diff --git a/tests/e2e/online_serving/test_moss_tts_expansion.py b/tests/e2e/online_serving/test_moss_tts_expansion.py index 80ee5f36026..2d048fe5274 100644 --- a/tests/e2e/online_serving/test_moss_tts_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_expansion.py @@ -3,8 +3,8 @@ """ E2E online tests for the full MOSS-TTS family via /v1/audio/speech. -The offline suite (``tests/e2e/offline_inference/test_moss_tts.py`` and -``test_moss_tts_v1_5.py``) covers the engine path; this file covers the +The offline suite (``tests/e2e/offline_inference/test_moss_tts_*_expansion.py``) +covers the engine path; this file covers the *serving* path (``serving_speech.py`` → ``_detect_moss_variant`` → the delay/realtime prompt builders), which the Nano online test does not exercise. @@ -30,11 +30,17 @@ from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -# TODO: Fix this test -# pytestmark = [pytest.mark.full_model, pytest.mark.tts] +pytestmark = [ + pytest.mark.slow, + pytest.mark.tts, + pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/4700"), +] MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime" REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav" +# Voice-clone output is not reliably transcribed by the Whisper check used at +# full_model run_level; assert non-trivial WAV payload size instead. +_MIN_AUDIO_BYTES = 10_000 @pytest.fixture(scope="session") @@ -94,6 +100,10 @@ def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> No Output Modal: audio (24 kHz, WAV) Input Setting: stream=False Datasets: single request + + NOTE: ``min_audio_bytes`` skips Whisper transcript similarity — cloned + voice output is often empty or mismatched under ASR without indicating a + real serving regression. """ request_config = { "model": omni_server.model, @@ -101,6 +111,7 @@ def test_text_to_audio_001(omni_server, openai_client, ref_audio_data_url) -> No "stream": False, "response_format": "wav", "ref_audio": ref_audio_data_url, + "min_audio_bytes": _MIN_AUDIO_BYTES, } openai_client.send_audio_speech_request(request_config) @@ -126,6 +137,7 @@ def test_text_to_audio_002_streaming(omni_server, openai_client, ref_audio_data_ "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "response_format": "pcm", "ref_audio": ref_audio_data_url, "min_hnr_db": -5.0, @@ -144,6 +156,8 @@ def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_ur Output Modal: audio (24 kHz, WAV) Input Setting: stream=False Datasets: single request + + NOTE: same ``min_audio_bytes`` rationale as test_text_to_audio_001. """ request_config = { "model": omni_server.model, @@ -151,6 +165,7 @@ def test_text_to_audio_003_chinese(omni_server, openai_client, ref_audio_data_ur "stream": False, "response_format": "wav", "ref_audio": ref_audio_data_url, + "min_audio_bytes": _MIN_AUDIO_BYTES, } openai_client.send_audio_speech_request(request_config) diff --git a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py index 66c0217a960..3468a9656c7 100644 --- a/tests/e2e/online_serving/test_moss_tts_nano_expansion.py +++ b/tests/e2e/online_serving/test_moss_tts_nano_expansion.py @@ -24,10 +24,9 @@ from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -# ``omni`` for all tests; 001/002 also get ``full_model`` via ``TestMossTtsNanoFull`` (003 is core+advanced only, no full_model). +# ``omni`` for all tests; 001/002 also get ``slow`` via ``TestMossTtsNanoFull`` (003 is core+advanced only, no slow). pytestmark = [ - pytest.mark.skip(reason="issue#4361"), - pytest.mark.full_model, + pytest.mark.slow, pytest.mark.tts, ] @@ -130,6 +129,7 @@ def test_text_to_audio_002(omni_server, openai_client, ref_audio_data_url) -> No "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "response_format": "pcm", "ref_audio": ref_audio_data_url, "min_hnr_db": -5.0, diff --git a/tests/e2e/online_serving/test_nextstep_expansion.py b/tests/e2e/online_serving/test_nextstep_expansion.py index 75290db5d29..3810793200f 100644 --- a/tests/e2e/online_serving/test_nextstep_expansion.py +++ b/tests/e2e/online_serving/test_nextstep_expansion.py @@ -14,7 +14,7 @@ dummy_messages_from_mix_data, ) -pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] +pytestmark = [pytest.mark.slow, pytest.mark.diffusion] # L4: 4 GPUs + TP=4; XPU B60: 2 cards (use num_cards={"cuda": 4, "xpu": 4} if needed) FOUR_CARD_MARKS = hardware_marks( @@ -31,6 +31,20 @@ def _get_diffusion_feature_cases(model: str): """Single online config: TP=4, explicit pipeline class.""" return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--tensor-parallel-size", + "2", + "--model-class-name", + "NextStep11Pipeline", + "--enable-cpu-offload", + ], + ), + id="nextstep_tp2_cpu_offload", + marks=FOUR_CARD_MARKS, + ), pytest.param( OmniServerParams( model=model, diff --git a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py index c49b42aa7c1..bbf9ca2b6fb 100644 --- a/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen2_5_omni_expansion.py @@ -55,7 +55,7 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @@ -93,7 +93,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model +@pytest.mark.slow @pytest.mark.omni @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_text_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index aa176c0770a..3ded285cee5 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -7,6 +7,7 @@ import os import pytest +from openai import BadRequestError from tests.helpers.mark import hardware_test from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video @@ -40,11 +41,21 @@ def get_batch_token_config(default_path): Uses the new flat-stage schema (``stages.<id>.<field>``); the legacy ``stage_args.<id>.engine_args.<field>`` path no longer applies because the deploy YAML doesn't nest engine fields under ``engine_args:``. + + Note: + 64 was the original value but caused a GPU hang (NVIDIA watchdog killed + the process after ~2 s) when 5 concurrent requests arrived, due to an + upstream scheduler change (``throttle_prefills`` API in vLLM commit + 0b131b16c). 512 was first tried but caused output quality degradation for + long audio (~120 s) inputs. 2048 gives the scheduler enough token budget + to handle 5 concurrent requests even with multi-second audio prefill, + while still being well below the default 32768. + If the upstream scheduler changes again, this value may need adjustment. """ return modify_stage_config( default_path, updates={ - "stages": {0: {"max_num_batched_tokens": 64}, 1: {"max_num_batched_tokens": 64}}, + "stages": {0: {"max_num_batched_tokens": 2048}, 1: {"max_num_batched_tokens": 2048}}, }, ) @@ -63,8 +74,6 @@ def get_batch_token_config(default_path): use_stage_cli=True, server_args=[ "--no-async-chunk", - "--stage-overrides", - '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}', ], ), id="default", @@ -402,6 +411,44 @@ def test_audio_in_video_002(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_audio_in_video_default_loader_sampling_regression(omni_server, openai_client) -> None: + """ + Regression: ``use_audio_in_video`` with default video loader sampling (no + ``media_io_kwargs.video.fps``) after a prior request that pinned fps. + + Before the fix the second request could partial-hit mm-cache (audio cached, + video re-processed without audio) and fail with ``StopIteration`` in + ``replace_multimodal_special_tokens``. + """ + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 128, embed_audio=True)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + video_data_url=video_data_url, + content_text="Reply with one word: OK.", + ) + + base_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "use_audio_in_video": True, + "modalities": ["text"], + } + + # Prime mm-cache with explicit loader fps (historical failure mode). + openai_client.send_omni_request( + { + **base_config, + "extra_body": {"media_io_kwargs": {"video": {"fps": 1, "num_frames": 128}}}, + } + ) + + # Default loader sampling: must succeed without explicit media_io fps. + openai_client.send_omni_request(dict(base_config)) + + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_one_word_prompt_001(omni_server, openai_client) -> None: @@ -564,3 +611,21 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: text = responses[0].text_content if responses else "" word_count = len(text.split()) assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}" + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params[:1], indirect=True) +def test_invalid_audio_format_rejected(omni_server, openai_client) -> None: + """Ensure invalid audio format is rejected with 400 before streaming starts.""" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + content_text=get_prompt(), + ) + with pytest.raises(BadRequestError, match="audio format"): + openai_client.client.chat.completions.create( + model=omni_server.model, + messages=messages, + modalities=["text", "audio"], + audio={"voice": "alloy", "format": "aac"}, + stream=True, + ) diff --git a/tests/e2e/online_serving/test_qwen3_tts_base.py b/tests/e2e/online_serving/test_qwen3_tts_base.py index 5bacd7ced45..60caf8f8285 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base.py @@ -100,6 +100,7 @@ def test_text_to_audio_002(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, "response_format": "wav", "task_type": "Base", diff --git a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py index 866d7483c06..0edaa4668e4 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py @@ -80,6 +80,7 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "response_format": "wav", "task_type": "Base", "voice": "clone", diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py index a03e7ad8bf4..2aab8484009 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice.py @@ -15,10 +15,15 @@ from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" +_STAGE_CONFIG = modify_stage_config( + get_deploy_config_path("qwen3_tts.yaml"), + updates={"stages": {0: {"default_sampling_params.max_tokens": 500}}}, +) + def get_prompt(prompt_type="text"): """Text prompt for text-to-audio tests (same as test_qwen3_omni - beijing test case).""" @@ -38,7 +43,7 @@ def get_max_batch_size(size_type="few"): pytest.param( OmniServerParams( model=MODEL, - stage_config_path=get_deploy_config_path("qwen3_tts.yaml"), + stage_config_path=_STAGE_CONFIG, server_args=["--trust-remote-code"], ), id="async_chunk", @@ -90,6 +95,7 @@ def test_text_to_audio_002(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "response_format": "wav", "task_type": "CustomVoice", "voice": "vivian", diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py index ce3bec928e7..9c941cad865 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py @@ -75,6 +75,7 @@ def test_voice_001(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "response_format": "wav", "task_type": "CustomVoice", "voice": "uncle_fu", diff --git a/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py b/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py new file mode 100644 index 00000000000..e6c3969323b --- /dev/null +++ b/tests/e2e/online_serving/test_qwen3_tts_ref_audio_artifact_mode_e2e.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E regression for Qwen3-TTS ref_audio artifact mode collision (#5049). + +Reproduces the issue's exact sequence on a live server: an +``x_vector_only_mode=true`` request followed by an ICL +(``x_vector_only_mode=false`` + ``ref_text``) request using the **same** +``ref_audio``. + +Before the fix, the ref_audio artifact readiness was mode-agnostic: the +x-vector-only request (which caches a speaker embedding but no ``ref_code``) +marked the artifact "ready", so the ICL request took the artifact-only serving +path (``ref_audio`` stripped before dispatch), found no ``ref_code`` and no +audio to recompute from, and raised in ``build_prompt_embeds`` -> +``EngineDeadError`` — taking the whole server down rather than failing the one +request. + +After the fix (readiness keyed by ``(artifact_key, x_vector_only)``), the ICL +request no longer matches the x-vector artifact, so it sends ``ref_audio`` and +the worker computes ``ref_code``. Both requests succeed and the server stays up. +""" + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +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 = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" +DEFAULT_AUDIO_SPEECH_TIMEOUT_S = 180.0 + +# Vendored under tests/assets/qwen3_tts/clone_2.wav (see test_qwen3_tts_base.py). +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." +INPUT_TEXT = "The weather is nice today, perfect for a walk in the park." + +tts_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("qwen3_tts.yaml"), + server_args=["--trust-remote-code"], + ), + id="async_chunk", + ) +] + + +@pytest.mark.advanced_model +@pytest.mark.core_model +@pytest.mark.tts +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) +def test_xvector_then_icl_same_ref_audio_keeps_engine_alive(omni_server, openai_client) -> None: + """Regression for #5049: x-vector-only then ICL on the same ref_audio. + + Both requests must return audio and the engine must stay alive (pre-fix the + second request killed EngineCore). + """ + base_request = { + "model": omni_server.model, + "input": INPUT_TEXT, + "stream": False, + "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, + "response_format": "wav", + "task_type": "Base", + "ref_audio": REF_AUDIO_URL, + # Assert non-empty audio came back (not just HTTP 200). + "min_audio_bytes": 1, + } + + # 1) x-vector-only request: caches a speaker embedding (no ref_code) and, on + # completion, marks the ref_audio artifact "ready". + openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True}) + + # 2) ICL request with the SAME ref_audio. Pre-fix this reused the x-vector + # artifact via the artifact-only path, hit the missing ref_code, and + # killed EngineCore (#5049). It must now succeed with audio. + openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": False, "ref_text": REF_TEXT}) + + # 3) The server must still be serving after the ICL request (pre-fix this + # would fail with EngineDeadError / connection refused). + openai_client.send_audio_speech_request({**base_request, "x_vector_only_mode": True}) diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index 1457ad63600..b3d6f8b1949 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py @@ -29,6 +29,16 @@ # If a similar test only involves one model, one can just define a global list variable. def _get_diffusion_feature_cases(model: str): return [ + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--enable-cpu-offload", + ], + ), + id="cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), pytest.param( OmniServerParams( model=model, diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index f1cf449f7c0..c900ca4ceb3 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -25,6 +25,14 @@ def _get_diffusion_feature_cases(model: str): return [ + pytest.param( + OmniServerParams( + model=model, + server_args=["--enable-cpu-offload"], + ), + id="cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), pytest.param( OmniServerParams(model=model, server_args=["--step-execution"]), id="step_execution", diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index 4e79beab7a0..315fba6ef3f 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -7,7 +7,7 @@ Kept cases (maximum feature coverage, no cpu-offload): sp_001 : cache_dit + Ulysses-SP 2 (2×H100) cfg_parallel_002 : cache_dit + CFG-Parallel 2 (2×H100) - layers_guard_001 : layerwise offload + layers=3 (1×H100, issue #1969 guard) + layers_guard_001 : layerwise offload + layers=2 (1×H100, issue #1969 guard) Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload. """ @@ -32,6 +32,14 @@ # sp_001 : cache_dit + Ulysses-SP 2 # cfg_parallel_002: cache_dit + CFG-Parallel 2 FEATURE_CASES = [ + pytest.param( + OmniServerParams( + model=MODEL, + server_args=["--enable-cpu-offload"], + ), + id="cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), pytest.param( OmniServerParams( model=MODEL, @@ -108,8 +116,8 @@ def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler): model=MODEL, server_args=["--enable-layerwise-offload"], ), - 3, - id="layers_guard_001_layers3", + 2, + id="layers_guard_001_layers2", marks=SINGLE_CARD_FEATURE_MARKS, ), ] diff --git a/tests/e2e/online_serving/test_sd3_expansion.py b/tests/e2e/online_serving/test_sd3_expansion.py index 6be0fe4e9bb..f881a212bcb 100644 --- a/tests/e2e/online_serving/test_sd3_expansion.py +++ b/tests/e2e/online_serving/test_sd3_expansion.py @@ -12,7 +12,7 @@ dummy_messages_from_mix_data, ) -pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] +pytestmark = [pytest.mark.diffusion, pytest.mark.slow] FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) POSITIVE_PROMPT = "A serene mountain landscape at sunset" @@ -24,6 +24,20 @@ # that provide good performance improvements. def _get_diffusion_feature_cases(model: str): return [ + # CFG Parallel + Tensor Parallel + CPU offload + pytest.param( + OmniServerParams( + model=model, + server_args=[ + "--cfg-parallel-size", + "2", + "--tensor-parallel-size", + "2", + "--enable-cpu-offload", + ], + ), + marks=FOUR_CARD_FEATURE_MARKS, + ), # Cache-DiT + CFG Parallel + Tensor Parallel pytest.param( OmniServerParams( diff --git a/tests/e2e/online_serving/test_sd3.py b/tests/e2e/online_serving/test_sdxl.py similarity index 64% rename from tests/e2e/online_serving/test_sd3.py rename to tests/e2e/online_serving/test_sdxl.py index 538e1bda8b0..4d7fc34d5fc 100644 --- a/tests/e2e/online_serving/test_sd3.py +++ b/tests/e2e/online_serving/test_sdxl.py @@ -1,5 +1,7 @@ """ -Tests for Stable Diffusion 3.5 medium model. +Tests for Stable Diffusion XL Base 1.0 model. + +Validates single-GPU text-to-image generation with CFG (guidance_scale > 1). """ import pytest @@ -12,23 +14,19 @@ dummy_messages_from_mix_data, ) -FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) -POSITIVE_PROMPT = "A serene mountain landscape at sunset" +SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}) +POSITIVE_PROMPT = "A serene mountain landscape at sunset, photorealistic" NEGATIVE_PROMPT = "blurry, low quality, distorted" -# For now, we only test one or two good configurations for this model to keep -# the CI light. The test cases focus on the most common feature combinations -# that provide good performance improvements. def _get_diffusion_feature_cases(model: str): return [ - # Cache-DiT + CFG Parallel + Tensor Parallel pytest.param( OmniServerParams( model=model, ), id="default", - marks=FOUR_CARD_FEATURE_MARKS, + marks=SINGLE_CARD_FEATURE_MARKS, ), ] @@ -39,7 +37,7 @@ def _get_diffusion_feature_cases(model: str): @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases( - model="stabilityai/stable-diffusion-3.5-medium", + model="stabilityai/stable-diffusion-xl-base-1.0", ), indirect=True, ) @@ -51,9 +49,9 @@ def test_text_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientH "extra_body": { "height": 1024, "width": 1024, - "num_inference_steps": 28, + "num_inference_steps": 30, "negative_prompt": NEGATIVE_PROMPT, - "guidance_scale": 4.5, + "guidance_scale": 7.5, "seed": 42, }, } diff --git a/tests/e2e/online_serving/test_step_audio2.py b/tests/e2e/online_serving/test_step_audio2.py deleted file mode 100644 index a23aaa0ffe4..00000000000 --- a/tests/e2e/online_serving/test_step_audio2.py +++ /dev/null @@ -1,368 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -E2E tests for Step-Audio2 online serving with concurrent requests. - -Tests: -1. Single request via API -2. Concurrent throughput (1/2/4/8 concurrent requests) -3. Latency percentiles (p50/p95/p99) -""" - -import asyncio -import os -import socket -import subprocess -import sys -import time -from pathlib import Path - -import aiohttp -import numpy as np -import pytest -from vllm.utils.network_utils import get_open_port - -MODEL = "stepfun-ai/Step-Audio-2-mini" -STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml") - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - - -class OmniServer: - """Omniserver for vLLM-Omni tests.""" - - def __init__( - self, - model: str, - serve_args: list[str], - *, - env_dict: dict[str, str] | None = None, - ) -> None: - self.model = model - self.serve_args = serve_args - self.env_dict = env_dict - self.proc: subprocess.Popen | None = None - self.host = "127.0.0.1" - self.port = get_open_port() - - def _start_server(self) -> None: - """Start the vLLM-Omni server subprocess.""" - env = os.environ.copy() - env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - if self.env_dict is not None: - env.update(self.env_dict) - - cmd = [ - sys.executable, - "-m", - "vllm_omni.entrypoints.cli.main", - "serve", - self.model, - "--omni", - "--host", - self.host, - "--port", - str(self.port), - ] + self.serve_args - - print(f"Launching OmniServer with: {' '.join(cmd)}") - self.proc = subprocess.Popen( - cmd, - env=env, - cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - ) - - max_wait = 600 # 10 minutes - start_time = time.time() - while time.time() - start_time < max_wait: - try: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.settimeout(1) - result = sock.connect_ex((self.host, self.port)) - if result == 0: - print(f"Server ready on {self.host}:{self.port}") - return - except Exception: - pass - time.sleep(2) - - raise RuntimeError(f"Server failed to start within {max_wait} seconds") - - def __enter__(self): - self._start_server() - return self - - def __exit__(self, exc_type, exc_val, exc_tb): - if self.proc: - self.proc.terminate() - try: - self.proc.wait(timeout=30) - except subprocess.TimeoutExpired: - self.proc.kill() - self.proc.wait() - - -def create_dummy_audio_base64(duration_sec: float = 5.0, sample_rate: int = 16000) -> str: - """Create dummy audio as base64.""" - import base64 - import io - import wave - - num_samples = int(duration_sec * sample_rate) - t = np.linspace(0, duration_sec, num_samples, dtype=np.float32) - audio = (0.5 * np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16) - - buffer = io.BytesIO() - with wave.open(buffer, "wb") as wav_file: - wav_file.setnchannels(1) - wav_file.setsampwidth(2) - wav_file.setframerate(sample_rate) - wav_file.writeframes(audio.tobytes()) - - return base64.b64encode(buffer.getvalue()).decode("utf-8") - - -def create_s2st_request(audio_base64: str) -> dict: - """Create S2ST request payload.""" - # Use audio_url format with base64 data URL - audio_url = f"data:audio/wav;base64,{audio_base64}" - return { - "model": MODEL, - "messages": [ - {"role": "system", "content": "请仔细聆听这段语音，然后复述其内容。"}, - { - "role": "user", - "content": [ - { - "type": "audio_url", - "audio_url": {"url": audio_url}, - }, - { - "type": "text", - "text": "<tts_start>", - }, - ], - }, - ], - "max_tokens": 512, - "temperature": 0.7, - "top_p": 0.9, - } - - -async def send_request(session: aiohttp.ClientSession, url: str, payload: dict) -> tuple[float, int, str]: - """Send single request and return (latency, num_tokens, response_text).""" - start = time.perf_counter() - async with session.post(url, json=payload) as response: - result = await response.json() - latency = time.perf_counter() - start - - if "choices" in result and len(result["choices"]) > 0: - text = result["choices"][0].get("message", {}).get("content", "") - tokens = result.get("usage", {}).get("completion_tokens", len(text.split())) - else: - text = str(result) - tokens = 0 - - return latency, tokens, text - - -async def benchmark_concurrent( - url: str, - payload: dict, - num_concurrent: int, - num_requests: int, - warmup: int = 2, -) -> dict: - """Run concurrent benchmark with true sustained concurrency. - - Uses semaphore to maintain exactly num_concurrent requests in flight at all times, - rather than batch-and-wait approach. Measures real wall clock time. - """ - semaphore = asyncio.Semaphore(num_concurrent) - - async def bounded_request(session: aiohttp.ClientSession): - async with semaphore: - return await send_request(session, url, payload) - - async with aiohttp.ClientSession() as session: - if warmup > 0: - print(f" Warming up ({warmup} requests)...", end=" ", flush=True) - for _ in range(warmup): - try: - await send_request(session, url, payload) - except Exception as e: - print(f"Warmup failed: {e}") - print("done") - - print(f" Running {num_requests} requests with concurrency={num_concurrent}...", end=" ", flush=True) - wall_start = time.perf_counter() - - tasks = [bounded_request(session) for _ in range(num_requests)] - results = await asyncio.gather(*tasks, return_exceptions=True) - - wall_time = time.perf_counter() - wall_start - print(f"done ({wall_time:.2f}s)") - - all_latencies = [] - all_tokens = [] - errors = 0 - - for result in results: - if isinstance(result, Exception): - print(f" Request failed: {result}") - errors += 1 - continue - latency, tokens, _ = result - all_latencies.append(latency) - all_tokens.append(tokens) - - if not all_latencies: - return {"error": "All requests failed"} - - latencies = np.array(all_latencies) - total_tokens = sum(all_tokens) - successful_requests = len(all_latencies) - - return { - "num_concurrent": num_concurrent, - "num_requests": num_requests, - "successful_requests": successful_requests, - "errors": errors, - "wall_time": wall_time, - "p50_latency": float(np.percentile(latencies, 50)), - "p95_latency": float(np.percentile(latencies, 95)), - "p99_latency": float(np.percentile(latencies, 99)), - "mean_latency": float(np.mean(latencies)), - "min_latency": float(np.min(latencies)), - "max_latency": float(np.max(latencies)), - "throughput_req_per_sec": successful_requests / wall_time, - "throughput_tokens_per_sec": total_tokens / wall_time, - } - - -class TestStepAudio2OnlineServing: - """Online serving tests for Step-Audio2.""" - - @pytest.fixture(scope="class") - def server(self): - """Start vLLM-Omni server.""" - with OmniServer(MODEL, ["--stage-configs-path", STAGE_CONFIG]) as omni_server: - yield f"http://{omni_server.host}:{omni_server.port}/v1/chat/completions" - - @pytest.mark.asyncio - async def test_single_request(self, server): - """Test single request works.""" - audio_base64 = create_dummy_audio_base64(duration_sec=2.0) - payload = create_s2st_request(audio_base64) - - async with aiohttp.ClientSession() as session: - latency, tokens, text = await send_request(session, server, payload) - - assert latency > 0 - assert tokens >= 0 - print(f"Single request: latency={latency:.2f}s, tokens={tokens}") - - @pytest.mark.asyncio - @pytest.mark.parametrize("num_concurrent", [1, 2, 4]) - async def test_concurrent_throughput(self, server, num_concurrent): - """Test concurrent throughput at different concurrency levels.""" - audio_base64 = create_dummy_audio_base64(duration_sec=5.0) - payload = create_s2st_request(audio_base64) - - num_requests = num_concurrent * 3 - - stats = await benchmark_concurrent(server, payload, num_concurrent, num_requests) - - assert "error" not in stats - print(f"\nConcurrency={num_concurrent}:") - print(f" p50={stats['p50_latency']:.2f}s, p95={stats['p95_latency']:.2f}s") - print(f" throughput={stats['throughput_req_per_sec']:.2f} req/s") - - -async def run_throughput_benchmark( - server_url: str, - audio_duration: float = 5.0, - concurrency_levels: list[int] = [1, 2, 4, 8], - requests_per_level: int = 10, - warmup: int = 2, -): - """Run full throughput benchmark.""" - audio_base64 = create_dummy_audio_base64(duration_sec=audio_duration) - payload = create_s2st_request(audio_base64) - - print(f"{'=' * 70}") - print("Step-Audio2 Online Serving Throughput Benchmark") - print(f"{'=' * 70}") - print(f"Server: {server_url}") - print(f"Audio duration: {audio_duration}s") - print(f"Requests per level: {requests_per_level}") - print(f"Warmup requests: {warmup}") - print(f"Concurrency levels: {concurrency_levels}") - print(f"{'=' * 70}") - - results = [] - for num_concurrent in concurrency_levels: - print(f"\n[Concurrency={num_concurrent}]") - stats = await benchmark_concurrent(server_url, payload, num_concurrent, requests_per_level, warmup) - results.append(stats) - - if "error" not in stats: - print( - f" Result: p50={stats['p50_latency']:.2f}s, " - f"p95={stats['p95_latency']:.2f}s, " - f"throughput={stats['throughput_req_per_sec']:.2f} req/s, " - f"{stats['throughput_tokens_per_sec']:.1f} tok/s" - ) - - print(f"\n{'=' * 80}") - print("SUMMARY") - print(f"{'=' * 80}") - print( - f"{'Conc':<6} {'Reqs':<6} {'Err':<5} {'Wall(s)':<9} {'p50':<8} {'p95':<8} {'p99':<8} {'Req/s':<8} {'Tok/s':<8}" - ) - print("-" * 80) - - for stats in results: - if "error" in stats: - print(f"{stats.get('num_concurrent', '?'):<6} ERROR: {stats['error']}") - continue - print( - f"{stats['num_concurrent']:<6} " - f"{stats['successful_requests']:<6} " - f"{stats['errors']:<5} " - f"{stats['wall_time']:<9.2f} " - f"{stats['p50_latency']:<8.2f} " - f"{stats['p95_latency']:<8.2f} " - f"{stats['p99_latency']:<8.2f} " - f"{stats['throughput_req_per_sec']:<8.2f} " - f"{stats['throughput_tokens_per_sec']:<8.1f}" - ) - - print("=" * 80) - return results - - -if __name__ == "__main__": - import argparse - - parser = argparse.ArgumentParser(description="Step-Audio2 Online Serving Benchmark") - parser.add_argument( - "--server-url", type=str, default="http://localhost:8000/v1/chat/completions", help="vLLM server URL" - ) - parser.add_argument("--audio-duration", type=float, default=5.0, help="Input audio duration in seconds") - parser.add_argument("--concurrency", type=str, default="1,2,4,8", help="Comma-separated concurrency levels to test") - parser.add_argument("--requests-per-level", type=int, default=10, help="Number of requests per concurrency level") - parser.add_argument("--warmup", type=int, default=2, help="Number of warmup requests before each benchmark") - - args = parser.parse_args() - concurrency_levels = [int(x) for x in args.concurrency.split(",")] - - asyncio.run( - run_throughput_benchmark( - args.server_url, - args.audio_duration, - concurrency_levels, - args.requests_per_level, - args.warmup, - ) - ) diff --git a/tests/e2e/online_serving/test_step_audio2_expansion.py b/tests/e2e/online_serving/test_step_audio2_expansion.py new file mode 100644 index 00000000000..05beb1131c7 --- /dev/null +++ b/tests/e2e/online_serving/test_step_audio2_expansion.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E expansion tests for Step-Audio2 online serving (weekly CI). + +Tests speech-to-speech translation via /v1/chat/completions. +Request shape follows examples/online_serving/step_audio2/openai_chat_completion_client.py. +""" + +import os +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import generate_synthetic_audio +from tests.helpers.runtime import OmniServerParams + +os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + +MODEL = "stepfun-ai/Step-Audio-2-mini" +STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml") +SAMPLE_RATE = 16000 +SEED = 42 +# pyttsx3 synthetic input phrase; Step-Audio2 text output includes codec tokens +# (<audio_NNNN>) so we validate audio only, not text↔audio similarity. +SYNTH_PHRASE = "how are you" + +TEST_PARAMS = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + env_dict={"VLLM_IMAGE_FETCH_TIMEOUT": "60"}, + ), + id="step_audio2", + ) +] + +pytestmark = [ + pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True), + pytest.mark.slow, + pytest.mark.tts, +] + + +def _synthetic_audio_base64(duration_sec: int = 2) -> str: + return generate_synthetic_audio( + duration_sec, + 1, + SAMPLE_RATE, + phrase_text=SYNTH_PHRASE, + )["base64"] + + +def _default_sampling_params_list() -> list[dict]: + return [ + { + "temperature": 0.7, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 256, + "seed": SEED, + "detokenize": True, + "repetition_penalty": 1.05, + }, + { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 1, + "seed": SEED, + "detokenize": False, + }, + ] + + +def _build_s2st_request_config(omni_server, audio_base64: str) -> dict: + """Build audio-to-audio (S2ST) chat completion request config.""" + audio_url = f"data:audio/wav;base64,{audio_base64}" + return { + "model": omni_server.model, + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "You are an audio processing assistant. Listen and repeat the audio content.", + } + ], + }, + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": audio_url}}, + {"type": "text", "text": "Please listen to this audio and repeat its content."}, + ], + }, + {"role": "assistant", "content": "<tts_start>"}, + ], + # Thinker text includes undetokenized audio codec placeholders; assert audio only. + "modalities": ["audio"], + "audio_ref_text": SYNTH_PHRASE, + "sampling_params_list": _default_sampling_params_list(), + "extra_body": { + "continue_final_message": True, + "add_generation_prompt": False, + }, + } + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_single_s2st_request(omni_server, openai_client) -> None: + """Test single speech-to-speech request via chat completions API.""" + request_config = _build_s2st_request_config( + omni_server, + _synthetic_audio_base64(duration_sec=2), + ) + openai_client.send_omni_request(request_config) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("request_num", [2, 4]) +def test_concurrent_s2st_requests(omni_server, openai_client, request_num: int) -> None: + """Test concurrent speech-to-speech requests.""" + request_config = _build_s2st_request_config( + omni_server, + _synthetic_audio_base64(duration_sec=2), + ) + responses = openai_client.send_omni_request(request_config, request_num=request_num) + assert len(responses) == request_num diff --git a/tests/e2e/online_serving/test_voxcpm2_tts.py b/tests/e2e/online_serving/test_voxcpm2_tts.py index cf493bc1493..19ddc0401a7 100644 --- a/tests/e2e/online_serving/test_voxcpm2_tts.py +++ b/tests/e2e/online_serving/test_voxcpm2_tts.py @@ -89,6 +89,7 @@ def test_text_to_audio_002(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, "response_format": "wav", "voice": "default", diff --git a/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py b/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py index d09570811fe..996585a1d83 100644 --- a/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py +++ b/tests/e2e/online_serving/test_voxcpm2_tts_expansion.py @@ -62,6 +62,7 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: "model": omni_server.model, "input": get_prompt(), "stream": True, + "stream_format": "audio", "timeout": DEFAULT_AUDIO_SPEECH_TIMEOUT_S, "response_format": "wav", "voice": "default", @@ -73,7 +74,6 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) -@pytest.mark.skip(reason="#4335") def test_response_format_001(omni_server, openai_client) -> None: """ Test voice-clone non-stream PCM output via OpenAI API. @@ -92,5 +92,6 @@ def test_response_format_001(omni_server, openai_client) -> None: "voice": "default", "ref_audio": REF_AUDIO_URL, "min_audio_bytes": _MIN_AUDIO_BYTES, + "min_hnr_db": -2.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 deleted file mode 100644 index a86de210074..00000000000 --- a/tests/e2e/online_serving/test_voxtral_tts.py +++ /dev/null @@ -1,180 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -""" -E2E Online tests for Voxtral TTS model with text input and audio output. - -These tests verify the /v1/audio/speech endpoint works correctly with -actual model inference, not mocks. -""" - -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 - -MODEL = "mistralai/Voxtral-4B-TTS-2603" - -STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml") -EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"] -TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)] - - -def verify_wav_audio(content: bytes) -> bool: - """Verify that content is valid WAV audio data.""" - if len(content) < 44: - return False - return content[:4] == b"RIFF" and content[8:12] == b"WAVE" - - -# Minimum expected audio size for a short sentence (~1 second of 24kHz 16-bit mono WAV) -_MIN_AUDIO_BYTES_BASIC = 10000 - - -@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -class TestVoxtralTTSFixedVoice: - """E2E tests for Voxtral TTS model.""" - - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) - def test_speech_english_basic(self, omni_server, openai_client) -> None: - """Test basic English TTS generation.""" - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Hello, how are you?", - "voice": "casual_female", - "language": "English", - "response_format": "wav", - "timeout": 120.0, - "min_audio_bytes": _MIN_AUDIO_BYTES_BASIC, - } - ) - - @pytest.mark.core_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) - def test_speech_english_streaming(self, omni_server, openai_client) -> None: - """Test basic streaming English TTS generation (PCM via streaming API).""" - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Hello, how are you?", - "voice": "casual_female", - "language": "English", - "stream": True, - "response_format": "pcm", - "timeout": 120.0, - "min_audio_bytes": _MIN_AUDIO_BYTES_BASIC, - } - ) - - @pytest.mark.advanced_model - @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) - def test_speech_different_voices(self, omni_server, openai_client) -> None: - """Test TTS with different voice presets.""" - voices = ["casual_female", "neutral_male"] - for voice in voices: - openai_client.send_audio_speech_request( - { - "model": omni_server.model, - "input": "Testing voice selection.", - "voice": voice, - "language": "English", - "response_format": "wav", - "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/e2e/online_serving/test_voxtral_tts_expansion.py b/tests/e2e/online_serving/test_voxtral_tts_expansion.py new file mode 100644 index 00000000000..1448567d31b --- /dev/null +++ b/tests/e2e/online_serving/test_voxtral_tts_expansion.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E tests for Voxtral TTS online serving. + +Smoke tests (core_model + advanced_model) and expansion scenarios (slow) +for the /v1/audio/speech endpoint. +""" + +import os + +os.environ.setdefault("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 + +MODEL = "mistralai/Voxtral-4B-TTS-2603" +STAGE_CONFIG = get_deploy_config_path("voxtral_tts.yaml") +EXTRA_ARGS = ["--trust-remote-code", "--enforce-eager", "--disable-log-stats"] +TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, server_args=EXTRA_ARGS)] +_MIN_AUDIO_BYTES_BASIC = 10000 + +pytestmark = [ + pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True), + pytest.mark.slow, + pytest.mark.tts, +] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_english_basic(omni_server, openai_client) -> None: + """Test basic English TTS generation.""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "how are you", + "voice": "casual_female", + "language": "English", + "response_format": "wav", + "timeout": 120.0, + "min_audio_bytes": _MIN_AUDIO_BYTES_BASIC, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_english_streaming(omni_server, openai_client) -> None: + """Test basic streaming English TTS generation (PCM via streaming API).""" + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Hello, how are you?", + "voice": "casual_female", + "language": "English", + "stream": True, + "stream_format": "audio", + "response_format": "pcm", + "timeout": 120.0, + "min_audio_bytes": _MIN_AUDIO_BYTES_BASIC, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_different_voices(omni_server, openai_client) -> None: + """Test TTS with different voice presets.""" + voices = ["casual_female", "neutral_male"] + for voice in voices: + openai_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "Testing voice selection.", + "voice": voice, + "language": "English", + "response_format": "wav", + "timeout": 120.0, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_speed(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, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_instructions(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, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_response_formats(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, + } + ) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_speech_batches(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/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py index fd0221ddc85..85eee128a9d 100644 --- a/tests/e2e/online_serving/test_wan22_expansion.py +++ b/tests/e2e/online_serving/test_wan22_expansion.py @@ -85,7 +85,22 @@ def _get_wan22_feature_cases(): """ cases = [] - # ---- CUDA cases (unchanged) ---- + # ---- CUDA cases ---- + + # Single-card: CPU offload (applies to all models) + for model_path, model_key in WAN22_MODELS: + cases.append( + pytest.param( + OmniServerParams( + model=model_path, + server_args=["--enable-cpu-offload"], + ), + id=f"{model_key}_cpu_offload", + marks=CUDA_SINGLE_CARD_MARKS, + ) + ) + + # Single-card: Cache-DiT (applies to all models) for model_path, model_key in WAN22_MODELS: cases.append( pytest.param( diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py index 4f0e9644721..d58bb9e3222 100644 --- a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py +++ b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py @@ -37,6 +37,18 @@ def _get_vace_feature_cases(): return [ + # Single GPU: CPU offload + pytest.param( + OmniServerParams( + model=MODEL, + server_args=[ + "--enable-cpu-offload", + "--vae-use-tiling", + ], + ), + id="single_card_cpu_offload", + marks=SINGLE_CARD_FEATURE_MARKS, + ), # Single GPU: Cache-DiT + layerwise CPU offload pytest.param( OmniServerParams( diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index f810bb80e31..7944753889d 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -9,13 +9,13 @@ from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path -from vllm_omni.config.stage_config import _PIPELINE_REGISTRY +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.entrypoints.async_omni import AsyncOmni os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" SEED = 42 -pipeline = _PIPELINE_REGISTRY["qwen2_5_omni_thinker_only"] +pipeline = OMNI_PIPELINES["qwen2_5_omni_thinker_only"] # Single-stage thinker-only deploy, materialized from tests.helpers.stage_config._CI_OVERLAYS. stage_config = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml") diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index bc163873b19..e98c58d749b 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,21 +1,20 @@ -import concurrent.futures +import contextlib import importlib import os -import threading import time import types -from contextlib import contextmanager import pytest from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.stage_init_utils import ( LogicalStageInitPlan, ReplicaInitPlan, build_stage0_input_processor, compute_replica_layout, ) +from vllm_omni.engine.stage_runtime import StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -63,7 +62,7 @@ def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output def _make_llm_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, vllm_config: object, num_replicas: int = 1, final_output: bool = False, @@ -73,7 +72,7 @@ def _make_llm_plan( replicas: list[ReplicaInitPlan] = [] for replica_id in range(num_replicas): stage_cfg = types.SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="llm", runtime=types.SimpleNamespace(devices=str(replica_id)), engine_args={}, @@ -85,7 +84,7 @@ def _make_llm_plan( launch_mode="local", stage_cfg=stage_cfg, metadata=_make_llm_metadata( - configured_stage_id, + stage_id, replica_id=replica_id, final_output=final_output, final_output_type=final_output_type, @@ -99,7 +98,7 @@ def _make_llm_plan( ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=replicas, ) @@ -107,13 +106,13 @@ def _make_llm_plan( def _make_diffusion_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, num_replicas: int = 1, ): replicas: list[ReplicaInitPlan] = [] for replica_id in range(num_replicas): stage_cfg = types.SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="diffusion", runtime=types.SimpleNamespace(devices=str(replica_id)), engine_args={}, @@ -124,14 +123,14 @@ def _make_diffusion_plan( num_replicas=num_replicas, launch_mode="local", stage_cfg=stage_cfg, - metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id), + metadata=_make_diffusion_metadata(stage_id, replica_id=replica_id), stage_connector_spec={}, omni_kv_connector=(None, None, None), ) ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=replicas, ) @@ -165,7 +164,7 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): shared = types.SimpleNamespace(name="shared") extra = types.SimpleNamespace(name="extra") - cleanup_clients = AsyncOmniEngine._collect_initialized_clients_for_cleanup( + cleanup_clients = StageRuntime._collect_initialized_clients_for_cleanup( stage_pools=[types.SimpleNamespace(clients=[shared, None])], initialized_clients_by_stage={0: [shared], 1: [extra]}, ) @@ -173,35 +172,36 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients(): assert cleanup_clients == [shared, extra] -def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.engine.stage_engine_startup import StageReplicaResources from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.num_stages = 1 - engine.diffusion_batch_size = 1 - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) - plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0] + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] env_var = current_omni_platform.device_control_env_var old_env = os.environ.get(env_var) os.environ[env_var] = "0,1" + runtime._init_visible_devices_baseline = "0,1" - @contextmanager - def _fake_stage_runtime_setup(_stage_id, _runtime_cfg): - current_omni_platform.set_device_control_env_var("1") - yield - - monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup) - monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) - monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace()) + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) + monkeypatch.setattr( + runtime_mod, + "launch_diffusion_stage_replica", + lambda **_: (types.SimpleNamespace(), StageReplicaResources()), + ) try: - engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock()) + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1) assert os.environ.get(env_var) == "0,1" finally: if old_env is None: @@ -210,68 +210,63 @@ def _fake_stage_runtime_setup(_stage_id, _runtime_cfg): os.environ[env_var] = old_env -def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod + from vllm_omni.engine.stage_engine_startup import StageReplicaResources - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.num_stages = 1 - engine.diffusion_batch_size = 4 - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=4, + async_chunk=False, + ) - plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0] + plan = _make_diffusion_plan(0, stage_id=0).replicas[0] captured: dict[str, object] = {} - @contextmanager - def _noop_stage_runtime_setup(*_): - yield - - monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup) - monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None) + monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None) - def _capture_initialize_diffusion_stage( - stage_id, _model, _stage_cfg, _metadata, *, stage_init_timeout, batch_size, use_inline - ): - captured["stage_id"] = stage_id - captured["stage_init_timeout"] = stage_init_timeout - captured["batch_size"] = batch_size - captured["use_inline"] = use_inline - return types.SimpleNamespace() + def _capture_launch_diffusion_stage_replica(**kwargs): + captured["stage_id"] = kwargs["metadata"].stage_id + captured["stage_init_timeout"] = kwargs["stage_init_timeout"] + captured["batch_size"] = kwargs["batch_size"] + captured["use_inline"] = kwargs["use_inline"] + captured["omni_master_server"] = kwargs["omni_master_server"] + return types.SimpleNamespace(), StageReplicaResources() - monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _capture_initialize_diffusion_stage) + monkeypatch.setattr(runtime_mod, "launch_diffusion_stage_replica", _capture_launch_diffusion_stage_replica) - engine._initialize_diffusion_replica(plan, stage_init_timeout=302, stage_launch_lock=threading.Lock()) + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=302) assert captured == { "stage_id": 0, "stage_init_timeout": 302, "batch_size": 4, "use_inline": True, + "omni_master_server": None, } -def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_input_processor(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_stage_runtime_initializes_stage_pools(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.config_path = "dummy-config" - engine.num_stages = 2 - engine.async_chunk = False - engine.diffusion_batch_size = 1 - engine.single_stage_mode = False - engine._single_stage_id_filter = None - engine._omni_master_server = None - engine.stage_configs = [types.SimpleNamespace(), types.SimpleNamespace()] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace(), types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) stage_plans = [ - _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True), - _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, final_output=True), + _make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True), + _make_llm_plan(1, stage_id=1, vllm_config=cfg1, final_output=True), ] stage0_client_r0 = types.SimpleNamespace( @@ -302,51 +297,43 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu stage0_output_processor = object() stage1_output_processor = object() - top_level_input_processor = object() - - monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) - monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2, 1], {})) - monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) - monkeypatch.setattr(engine, "_initialize_stage_replicas", lambda *_: initialized_clients) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) + monkeypatch.setattr(runtime, "_initialize_stage_replicas", lambda *_: initialized_clients) monkeypatch.setattr( - engine_mod, + runtime_mod, "build_llm_stage_output_processor", lambda plan, _cfg, **_kw: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor, ) - monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor) - engine._initialize_stages(stage_init_timeout=1) + runtime.initialize() - assert len(engine.stage_pools) == 2 - assert engine.input_processor is top_level_input_processor - assert engine.stage_clients == [stage0_client_r0, stage1_client_r0] - assert engine.stage_vllm_configs == [cfg0, cfg1] - assert engine.output_processors == [stage0_output_processor, stage1_output_processor] - assert engine.default_sampling_params_list == [ - stage0_client_r0.default_sampling_params, - stage1_client_r0.default_sampling_params, - ] - assert engine.stage_metadata == [ - StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm"), - StageRuntimeInfo(final_output=True, final_output_type=None, stage_type="llm"), - ] + assert len(runtime.stage_pools) == 2 + assert runtime.stage_pools[0].stage_client is stage0_client_r0 + assert runtime.stage_pools[1].stage_client is stage1_client_r0 + assert runtime.stage_pools[0].stage_vllm_config is cfg0 + assert runtime.stage_pools[1].stage_vllm_config is cfg1 + assert runtime.stage_pools[0].output_processor is stage0_output_processor + assert runtime.stage_pools[1].output_processor is stage1_output_processor def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.async_chunk = False - engine.single_stage_mode = False - engine._single_stage_id_filter = None - engine.stage_configs = [ - types.SimpleNamespace(stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")), - types.SimpleNamespace( - stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3") - ), - ] + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = StageRuntime( + stage_configs=[ + types.SimpleNamespace( + stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0") + ), + types.SimpleNamespace( + stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3") + ), + ], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) metadata_by_stage = { 0: _make_llm_metadata(0), @@ -354,78 +341,46 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc } monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) monkeypatch.setattr( - engine_mod, + runtime_mod, "build_vllm_config", lambda stage_cfg, *_args, **_kwargs: (types.SimpleNamespace(tag=f"cfg-{stage_cfg.stage_id}"), object), ) - stage_plans, prompt_expand_func = engine._build_logical_stage_init_plans( + stage_plans = runtime._build_logical_stage_init_plans( omni_transfer_config=None, replicas_per_stage=[1, 3], replica_devices_map={1: ["1", "2", "3"]}, ) - assert prompt_expand_func is None - assert [plan.configured_stage_id for plan in stage_plans] == [0, 1] + assert [plan.stage_id for plan in stage_plans] == [0, 1] assert [replica.stage_cfg.runtime.devices for replica in stage_plans[1].replicas] == ["1", "2", "3"] assert [replica.replica_id for replica in stage_plans[1].replicas] == [0, 1, 2] assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas) -def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch): - engine = object.__new__(AsyncOmniEngine) - engine._stage_init_executor = None - - cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) - stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)] - - created_executors = [] - - class FakeExecutor: - def __init__(self, *args, **kwargs): - self.shutdown_calls = [] - created_executors.append(self) - - def submit(self, fn, *args, **kwargs): - future = concurrent.futures.Future() - future.set_result(fn(*args, **kwargs)) - return future - - def shutdown(self, wait=False, cancel_futures=False): - self.shutdown_calls.append((wait, cancel_futures)) - - monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor) - monkeypatch.setattr( - engine, - "_initialize_replica", - lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}", - ) - - first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1) - second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1) - - assert first == {0: ["client-0"]} - assert second == {0: ["client-0"]} - assert len(created_executors) == 1 - assert engine._stage_init_executor is created_executors[0] - - def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch): - engine = object.__new__(AsyncOmniEngine) + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=123, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) stage_plans = [ - _make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2), - _make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, num_replicas=2), + _make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2), + _make_llm_plan(1, stage_id=1, vllm_config=cfg1, num_replicas=2), ] clients = { @@ -435,13 +390,13 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk (1, 1): types.SimpleNamespace(name="stage1-replica1"), } - def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + def _initialize_replica(plan, _stage_init_timeout): time.sleep(0.02 * (3 - plan.metadata.stage_id - plan.replica_id)) return clients[(plan.metadata.stage_id, plan.replica_id)] - monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) - initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123) + initialized_clients = runtime._initialize_stage_replicas(stage_plans, stage_init_timeout=123) assert initialized_clients == { 0: [clients[(0, 0)], clients[(0, 1)]], @@ -449,148 +404,120 @@ def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): } -def test_shutdown_closes_long_lived_stage_init_executor(): - engine = object.__new__(AsyncOmniEngine) - engine._shutdown_called = False - engine._weak_finalizer = None - engine._omni_master_server = None - engine._coordinator_runtime = None - engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False) - - class FakeQueue: - def __init__(self): - self.closed = False - self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None) - - def close(self): - self.closed = True - - class FakeExecutor: - def __init__(self): - self.shutdown_calls = [] - - def shutdown(self, wait=False, cancel_futures=False): - self.shutdown_calls.append((wait, cancel_futures)) - - engine.request_queue = FakeQueue() - engine.output_queue = FakeQueue() - engine.rpc_output_queue = FakeQueue() - engine._stage_init_executor = FakeExecutor() +def test_remote_replicas_use_distinct_init_group_keys(): + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=123, + diffusion_batch_size=1, + async_chunk=False, + ) + plan = _make_llm_plan( + 1, + stage_id=1, + vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)), + num_replicas=2, + ) - engine.shutdown() + for replica in plan.replicas: + replica.launch_mode = "remote" + replica.metadata.runtime_cfg = None - assert engine._stage_init_executor is None - assert engine.request_queue.closed is True - assert engine.output_queue.closed is True - assert engine.rpc_output_queue.closed is True + assert [runtime._replica_init_group_key(replica) for replica in plan.replicas] == [ + "remote:1:0", + "remote:1:1", + ] def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.config_path = "dummy-config" - engine.num_stages = 1 - engine.async_chunk = False - engine.diffusion_batch_size = 1 - engine.single_stage_mode = False - engine._single_stage_id_filter = None - engine._omni_master_server = None - engine._coordinator_runtime = None - engine.stage_configs = [types.SimpleNamespace()] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) - stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)] + stage_plans = [_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2)] initialized_client = types.SimpleNamespace(shutdown=lambda: None) - monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) - monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {})) - monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) - def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): + def _initialize_replica(plan, _stage_init_timeout): if plan.replica_id == 0: return initialized_client time.sleep(0.05) raise RuntimeError("replica launch failed") - monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica) captured_cleanup: list[list[object]] = [] def _capture_shutdown(clients): captured_cleanup.append(list(clients)) - monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown) + monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown) with pytest.raises(RuntimeError, match="replica launch failed"): - engine._initialize_stages(stage_init_timeout=1) + runtime.initialize() assert captured_cleanup == [[initialized_client]] def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_replica_failure(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.config_path = "dummy-config" - engine.num_stages = 1 - engine.async_chunk = False - engine.diffusion_batch_size = 1 - engine.single_stage_mode = False - engine._single_stage_id_filter = None - engine._omni_master_server = None - engine._coordinator_runtime = None - engine.stage_configs = [types.SimpleNamespace()] + runtime = StageRuntime( + stage_configs=[types.SimpleNamespace()], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=1, + diffusion_batch_size=1, + async_chunk=False, + ) cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)) - stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)] + stage_plans = [_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2)] initialized_client = types.SimpleNamespace(shutdown=lambda: None) - replica1_started = threading.Event() - monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None) - monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None) - monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {})) - monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None)) + monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans) - def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock): - if plan.replica_id == 0: - assert replica1_started.wait(timeout=1) - raise RuntimeError("replica launch failed") - replica1_started.set() - time.sleep(0.05) - return initialized_client + def _initialize_stage_replicas(_stage_plans, _stage_init_timeout): + exc = RuntimeError("replica launch failed") + exc._initialized_clients_by_stage = {0: [None, initialized_client]} + raise exc - monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica) + monkeypatch.setattr(runtime, "_initialize_stage_replicas", _initialize_stage_replicas) captured_cleanup: list[list[object]] = [] def _capture_shutdown(clients): captured_cleanup.append(list(clients)) - monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown) + monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown) with pytest.raises(RuntimeError, match="replica launch failed"): - engine._initialize_stages(stage_init_timeout=1) + runtime.initialize() assert captured_cleanup == [[initialized_client]] -def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): - import vllm_omni.engine.async_omni_engine as engine_mod +def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "dummy-model" - engine.single_stage_mode = False - engine._omni_master_server = None - engine.stage_configs = [] + runtime = StageRuntime( + stage_configs=[], + model="dummy-model", + config_path="dummy-config", + stage_init_timeout=302, + diffusion_batch_size=1, + async_chunk=False, + ) fake_vllm_config = types.SimpleNamespace() fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None) - fake_proc = types.SimpleNamespace() captured_timeout: int | None = None plan = ReplicaInitPlan( @@ -603,34 +530,38 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand omni_kv_connector=(None, None, None), stage_vllm_config=fake_vllm_config, executor_class=object, + engine_args_dict={}, ) device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" - @contextmanager - def _noop_stage_runtime_setup(*_): - yield + def _capture_acquire_device_locks(*_args): + nonlocal captured_timeout + captured_timeout = _args[2] + return [] - monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) - monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: []) - monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake")) + monkeypatch.setattr(runtime_mod, "acquire_device_locks", _capture_acquire_device_locks) - def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout): - nonlocal captured_timeout - captured_timeout = handshake_timeout + from vllm_omni.engine.stage_engine_startup import StageReplicaResources - monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout) + @contextlib.contextmanager + def _fake_launch_stage_replica(**_kwargs): + yield StageReplicaResources( + manager=types.SimpleNamespace(shutdown=lambda: None), + addresses=fake_addresses, + ) + + monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica) monkeypatch.setattr( - engine_mod.StageEngineCoreClientBase, + runtime_mod.StageEngineCoreClientBase, "make_async_mp_client", staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)), ) try: - engine._initialize_llm_replica(plan, 302, threading.Lock()) + runtime._initialize_local_llm_replica(plan, 302) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) @@ -660,6 +591,24 @@ def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_p assert engine_args["tokenizer"] == "/external/tokenizer" +def test_build_engine_args_stage_model_overrides_parent_model(): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model": "/stage/model"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict( + stage_cfg, + "/parent/model", + ) + + assert engine_args["model"] == "/stage/model" + + def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path): from vllm_omni.engine.stage_init_utils import build_engine_args_dict @@ -790,12 +739,13 @@ def test_resolve_stage_configs_injects_global_diffusion_attention_when_missing(m monkeypatch.setattr( engine_mod, "load_and_resolve_stage_configs", - lambda *args, **kwargs: ("dummy-config", [stage_cfg]), + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), ) _config_path, stage_configs = engine._resolve_stage_configs( model="dummy-model", kwargs={"diffusion_attention_backend": "FLASH_ATTN"}, + trust_remote_code=False, ) diffusion_attention_config = stage_configs[0].engine_args.diffusion_attention_config @@ -823,12 +773,13 @@ def test_resolve_stage_configs_preserves_stage_diffusion_attention(monkeypatch): monkeypatch.setattr( engine_mod, "load_and_resolve_stage_configs", - lambda *args, **kwargs: ("dummy-config", [stage_cfg]), + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), ) _config_path, stage_configs = engine._resolve_stage_configs( model="dummy-model", kwargs={"diffusion_attention_backend": "FLASH_ATTN"}, + trust_remote_code=False, ) assert stage_configs[0].engine_args.diffusion_attention_config is existing_attention @@ -849,12 +800,13 @@ def test_resolve_stage_configs_does_not_inject_diffusion_attention_into_llm_stag monkeypatch.setattr( engine_mod, "load_and_resolve_stage_configs", - lambda *args, **kwargs: ("dummy-config", [stage_cfg]), + lambda *args, **kwargs: ("dummy-config", [stage_cfg], None), ) _config_path, stage_configs = engine._resolve_stage_configs( model="dummy-model", kwargs={"diffusion_attention_backend": "TORCH_SDPA"}, + trust_remote_code=False, ) assert stage_configs[0].engine_args.attention_config == {"backend": "FLASH_ATTN"} @@ -925,3 +877,54 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key(): assert "attention_config" not in engine_args_dict assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN" + + +def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch): + """Regression: two stages must never draw the same ZMQ port. + + ``get_open_ports_list`` only dedups within a single call, so per-route + allocation used to let a later stage reuse an earlier stage's port. The + second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address + already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/ + code2wav). ``OmniMasterServer`` now dedups every port it hands out. + """ + import itertools + + from vllm_omni.engine import stage_engine_startup as ses + + # A colliding prefix (repeats + the master port 9000) followed by an endless + # fresh stream, so the only way to succeed is to redraw the collisions. + supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002)) + + def fake_get_open_ports_list(count): + return [next(supply) for _ in range(count)] + + monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list) + + server = ses.OmniMasterServer( + master_address="127.0.0.1", + master_port=9000, # seed: a route must not reuse the registration port + stage_ids=[0, 1, 2], + ) + + ports = [] + for sid in (0, 1, 2): + alloc = server.get_allocation(sid) + for addr in ( + alloc.handshake_bind_address, + alloc.input_bind_address, + alloc.output_bind_address, + ): + ports.append(ses._port_from_zmq_address(addr)) + + assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}" + assert 9000 not in ports, "route reused the master registration port" + + +def test_port_from_zmq_address_parsing(): + from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address + + assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277 + assert _port_from_zmq_address(None) is None + assert _port_from_zmq_address("ipc:///tmp/sock") is None + assert _port_from_zmq_address("tcp://host:not-a-port") is None diff --git a/tests/engine/test_membership_controller.py b/tests/engine/test_membership_controller.py new file mode 100644 index 00000000000..c00a7f64600 --- /dev/null +++ b/tests/engine/test_membership_controller.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import threading +from types import SimpleNamespace + +import pytest + +from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus +from vllm_omni.engine.membership_controller import MembershipController + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeHub: + def __init__(self, snapshots=None): + self.snapshots = list(snapshots or []) + self.closed = False + + def get_replica_list(self): + if self.snapshots: + return self.snapshots.pop(0) + return SimpleNamespace(replicas=[]) + + def close(self): + self.closed = True + + +class FakePool: + def __init__(self, stage_id: int): + self.stage_id = stage_id + self.clients = [] + self.added = [] + self.removed = [] + self.invalidated = [] + self.hub = None + self.lb = None + + def attach_hub(self, hub): + self.hub = hub + + def attach_load_balancer(self, lb): + self.lb = lb + + def add_client(self, input_addr, client): + self.added.append((input_addr, client)) + self.clients.append(client) + return len(self.clients) - 1 + + def invalidate_addr(self, input_addr): + self.invalidated.append(input_addr) + return ["req-1", "req-2"] + + def remove_client(self, input_addr): + self.removed.append(input_addr) + return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown")) + + +def _snapshot(*replicas): + return SimpleNamespace(replicas=list(replicas)) + + +def _replica(stage_id: int, input_addr: str, status=ReplicaStatus.UP): + return SimpleNamespace(stage_id=stage_id, input_addr=input_addr, status=status) + + +def _controller(monkeypatch, pool, hub, remote_replica_factory=None): + import vllm_omni.engine.membership_controller as membership_mod + + monkeypatch.setattr(membership_mod, "OmniCoordClientForHub", lambda _addr: hub) + if remote_replica_factory is None: + + def remote_replica_factory(stage_id, replica_id): + return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}) + + return MembershipController( + stage_pools=[pool], + coordinator_pub_address="tcp://127.0.0.1:12345", + load_balancer_factory=lambda: object(), + remote_replica_factory=remote_replica_factory, + ) + + +@pytest.mark.asyncio +async def test_watch_replica_list_unregisters_disappeared_replicas(monkeypatch): + pool = FakePool(stage_id=0) + hub = FakeHub( + snapshots=[ + _snapshot(_replica(0, "tcp://gone")), + _snapshot(), + ] + ) + controller = _controller(monkeypatch, pool, hub) + controller.WATCH_INTERVAL_S = 0 + unregistered = [] + + async def _handle_unregister(stage_id, input_addr, *args, **kwargs): + unregistered.append((stage_id, input_addr)) + controller._shutdown_event.set() + + controller.handle_unregister = _handle_unregister # type: ignore[method-assign] + + await asyncio.wait_for(controller._watch_replica_list(), timeout=1) + await controller.drain_tasks(timeout=1) + + assert unregistered == [(0, "tcp://gone")] + + +@pytest.mark.asyncio +async def test_shutdown_closes_hub_then_cancels_watcher(monkeypatch): + pool = FakePool(stage_id=0) + hub = FakeHub() + controller = _controller(monkeypatch, pool, hub) + watcher = asyncio.create_task(asyncio.sleep(10)) + controller._watcher_task = watcher + + controller.shutdown() + await asyncio.sleep(0) + + assert hub.closed is True + assert watcher.cancelled() + + +@pytest.mark.asyncio +async def test_drain_tasks_waits_for_membership_tasks(monkeypatch): + pool = FakePool(stage_id=0) + controller = _controller(monkeypatch, pool, FakeHub()) + completed = [] + + async def _task(): + await asyncio.sleep(0) + completed.append(True) + + controller._spawn_task(_task(), label="unit") + await controller.drain_tasks(timeout=1) + + assert completed == [True] + + +@pytest.mark.asyncio +async def test_do_register_offloads_remote_factory(monkeypatch): + pool = FakePool(stage_id=0) + factory_thread_ids = [] + + def _factory(stage_id, replica_id): + factory_thread_ids.append(threading.get_ident()) + return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"}) + + controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory) + + await controller._do_register(0, 1) + + assert factory_thread_ids + assert factory_thread_ids[0] != threading.get_ident() + assert pool.added[0][0] == "tcp://stage-0-replica-1" diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index b86cc0af985..23f8cfaa256 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -11,11 +11,9 @@ from typing import Any import janus -import psutil import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams -from vllm.v1.engine.core_client import AsyncMPClient from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -26,8 +24,12 @@ ShutdownRequestMessage, StageSubmissionMessage, ) -from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState -from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + _build_terminal_empty_output, + _infer_stage_audio_sample_rate, +) from vllm_omni.engine.stage_pool import StagePool from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -138,6 +140,20 @@ def push_diffusion_output(self, output) -> None: self._diffusion_outputs.put_nowait(output) +def test_terminal_empty_audio_output_uses_stage_sample_rate() -> None: + final_stage = FakeStageClient(final_output=True, final_output_type="audio") + final_stage.sample_rate = 44100 + final_pool = SimpleNamespace(stage_client=final_stage, _stage_vllm_config=None) + + terminal_output = _build_terminal_empty_output( + "req-1", + final_output_type="audio", + audio_sample_rate=_infer_stage_audio_sample_rate(final_pool), + ) + + assert terminal_output.outputs[0].multimodal_output["sr"] == 44100 + + class FakeCollectiveRpcStageClient(FakeStageClient): def __init__(self, *args, rpc_result: Any = None, **kwargs) -> None: super().__init__(*args, **kwargs) @@ -180,42 +196,6 @@ def update_scheduler_stats(self, _scheduler_stats) -> None: return None -class _FakeProc: - pid = 1234 - - def __init__(self): - self.terminated = False - self.killed = False - self.join_calls = [] - - def is_alive(self): - return not self.terminated and not self.killed - - def terminate(self): - self.terminated = True - - def kill(self): - self.killed = True - - def join(self, timeout=None): - self.join_calls.append(timeout) - - -class _FakeChildProc: - def __init__(self): - self.terminated = False - self.killed = False - - def is_running(self): - return not self.terminated and not self.killed - - def terminate(self): - self.terminated = True - - def kill(self): - self.killed = True - - def _sampling_params(max_tokens: int = 4) -> SamplingParams: return SamplingParams(max_tokens=max_tokens) @@ -437,61 +417,6 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids)) -def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch): - fake_proc = _FakeProc() - fake_child = _FakeChildProc() - - class FakePsutilProcess: - def __init__(self, pid): - assert pid == fake_proc.pid - - def children(self, recursive=True): - assert recursive - return [fake_child] - - def fail_base_shutdown(self, **kwargs): - raise RuntimeError("base shutdown failed") - - monkeypatch.setattr(psutil, "Process", FakePsutilProcess) - monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), [])) - monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown) - - client = object.__new__(StageEngineCoreClient) - client._proc = fake_proc - - with pytest.raises(RuntimeError, match="base shutdown failed"): - client.shutdown() - - assert fake_proc.terminated - assert fake_proc.join_calls == [5] - assert fake_child.terminated - - -def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch): - fake_proc = _FakeProc() - fake_child = _FakeChildProc() - - class FakePsutilProcess: - def __init__(self, pid): - assert pid == fake_proc.pid - - def children(self, recursive=True): - assert recursive - return [fake_child] - - monkeypatch.setattr(psutil, "Process", FakePsutilProcess) - monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs))) - monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None) - - client = object.__new__(StageEngineCoreClient) - client._proc = fake_proc - - client.shutdown() - - assert fake_child.terminated - assert fake_child.killed - - @pytest.fixture def orchestrator_factory(): fixtures: list[OrchestratorFixture] = [] @@ -598,6 +523,67 @@ async def test_run_single_stage_diffusion(orchestrator_factory) -> None: await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +async def test_run_single_stage_diffusion_streaming_forwards_intermediate_chunks(orchestrator_factory) -> None: + """Intermediate diffusion chunks (finished=False) reach the frontend before the final chunk.""" + 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-stream", + prompt={"prompt": "draw a cat"}, + original_prompt={"prompt": "draw a cat"}, + sampling_params_list=[params], + final_stage_id=0, + ) + + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + stage0.push_diffusion_output( + OmniRequestOutput.from_diffusion( + request_id="req-stream", + images=[], + final_output_type="image", + custom_output={"chunk": 0}, + finished=False, + ) + ) + stage0.push_diffusion_output( + OmniRequestOutput.from_diffusion( + request_id="req-stream", + images=[], + final_output_type="image", + custom_output={"chunk": 1}, + finished=True, + ) + ) + + output_msgs: list[OutputMessage] = [] + deadline = time.monotonic() + 2.0 + while not output_msgs or not output_msgs[-1].finished: + if time.monotonic() >= deadline: + raise AssertionError( + f"Timed out waiting for finished orchestrator output, got {len(output_msgs)} message(s)" + ) + try: + msg = orchestrator_fixture.output_sync_q.get_nowait() + except queue.Empty: + await asyncio.sleep(0.01) + continue + if isinstance(msg, OutputMessage): + output_msgs.append(msg) + + assert [msg.request_id for msg in output_msgs] == ["req-stream", "req-stream"] + assert [msg.finished for msg in output_msgs] == [False, True] + assert [msg.engine_outputs.finished for msg in output_msgs] == [False, True] + assert [msg.engine_outputs.custom_output["chunk"] for msg in output_msgs] == [0, 1] + await _wait_for(lambda: "req-stream" not in orchestrator_fixture.orchestrator.request_states) + finally: + await _shutdown_orchestrator(orchestrator_fixture) + + @pytest.mark.asyncio async def test_run_llm_to_diffusion(orchestrator_factory) -> None: stage0 = FakeStageClient(stage_type="llm", final_output=False) diff --git a/tests/engine/test_orchestrator_monitor.py b/tests/engine/test_orchestrator_monitor.py new file mode 100644 index 00000000000..ee9126cf470 --- /dev/null +++ b/tests/engine/test_orchestrator_monitor.py @@ -0,0 +1,101 @@ +import json +import time + +import pytest + +import vllm_omni.engine.orchestrator_monitor as monitor_mod +from vllm_omni.engine.orchestrator_monitor import ( + _NULL_ORCH_MONITOR, + OrchestratorMonitor, + create_orch_monitor, + replica_key, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_replica_key_format(): + assert replica_key(0, 0) == "stage=0,replica=0" + assert replica_key(2, 1) == "stage=2,replica=1" + + +def test_create_orch_monitor_disabled_returns_null_singleton(): + monitor = create_orch_monitor(enabled=False, replica_sampler=lambda: {}) + assert monitor is _NULL_ORCH_MONITOR + assert create_orch_monitor(enabled=False, replica_sampler=lambda: {}) is monitor + monitor.note_loop(idle=True) + monitor.flush() + + +def test_note_loop_and_flush_json(tmp_path, monkeypatch): + out_path = tmp_path / "orch_monitor.json" + monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path)) + + now = [0.0] + + def fake_monotonic() -> float: + return now[0] + + monkeypatch.setattr(time, "monotonic", fake_monotonic) + + def sampler() -> dict[str, tuple[int, int]]: + return {replica_key(0, 0): (5, 10)} + + monitor = create_orch_monitor(enabled=True, replica_sampler=sampler) + monitor.register_replica(0, 0) + + monitor.note_loop(idle=False) + monitor.note_loop(idle=True) + + now[0] = 1.5 + monitor.flush() + + payload = json.loads(out_path.read_text()) + assert payload["configured_window_s"] == monitor_mod._WINDOW_S + assert payload["windows"]["loop_idle"] == [1] + assert payload["windows"]["loop_active"] == [1] + assert payload["replicas"]["stage=0,replica=0"]["outputs_queue_size"] == [5] + assert payload["replicas"]["stage=0,replica=0"]["inflight"] == [10] + + +def test_register_replica_backfills_prior_windows(tmp_path, monkeypatch): + out_path = tmp_path / "orch_monitor.json" + monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path)) + + now = [0.0] + monkeypatch.setattr(time, "monotonic", lambda: now[0]) + + monitor = OrchestratorMonitor(replica_sampler=lambda: {}) + monitor.note_loop(idle=True) + now[0] = 1.5 + monitor.register_replica(1, 0) + monitor.note_loop(idle=False) + monitor.flush() + + series = json.loads(out_path.read_text())["replicas"]["stage=1,replica=0"] + assert series["outputs_queue_size"] == [0, 0] + assert series["inflight"] == [0, 0] + + +def test_note_loop_survives_sampler_failure(): + def bad_sampler() -> dict[str, tuple[int, int]]: + raise RuntimeError("sampler failed") + + monitor = OrchestratorMonitor(replica_sampler=bad_sampler) + monitor.register_replica(0, 0) + monitor.note_loop(idle=True) + monitor.flush() + + +def test_flush_is_idempotent(tmp_path, monkeypatch): + out_path = tmp_path / "orch_monitor.json" + monkeypatch.setenv("VLLM_OMNI_ORCH_MONITOR_PATH", str(out_path)) + + monitor = OrchestratorMonitor(replica_sampler=lambda: {}) + monitor.register_replica(0, 0) + monitor.note_loop(idle=False) + monitor.flush() + monitor.flush() + + payload = json.loads(out_path.read_text()) + assert len(payload["windows"]["duration_s"]) == 1 diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py new file mode 100644 index 00000000000..27ee11df4f9 --- /dev/null +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import queue +from types import SimpleNamespace +from typing import Any + +import janus +import pytest +from vllm.outputs import CompletionOutput, RequestOutput +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_pool import StagePool + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class FakeStageClient: + def __init__( + self, + *, + next_inputs: list[dict[str, Any]] | None = None, + final_output: bool = False, + ) -> None: + self.stage_id = 0 + self.replica_id = 0 + self.stage_type = "llm" + self.final_output = final_output + self.final_output_type = "text" + self.default_sampling_params = SamplingParams(max_tokens=1) + self.requires_multimodal_data = False + self.engine_input_source = [0] + self.is_comprehension = False + self.model_stage = None + self.custom_process_input_func = None + self.next_inputs = list(next_inputs or []) + self.add_request_calls: list[tuple[Any, ...]] = [] + self._engine_core_outputs = queue.Queue() + + async def add_request_async(self, *args, **_kwargs) -> None: + self.add_request_calls.append(args) + + async def get_output_async(self): + try: + return self._engine_core_outputs.get_nowait() + except queue.Empty: + return SimpleNamespace(outputs=[]) + + def process_engine_inputs(self, _source_outputs, prompt=None, streaming_context=None): + return list(self.next_inputs) + + async def abort_requests_async(self, _request_ids: list[str]) -> None: + return None + + def set_engine_outputs(self, _outputs) -> None: + return None + + def check_health(self) -> None: + return None + + def shutdown(self) -> None: + return None + + +class FakeOutputProcessor: + def add_request(self, *args, **kwargs) -> None: + return None + + +class FakeInputProcessor: + def __init__(self) -> None: + self.calls: list[dict[str, Any]] = [] + + def process_inputs(self, **kwargs): + self.calls.append(kwargs) + return SimpleNamespace( + request_id=kwargs["request_id"], + prompt_token_ids=[101, 102], + prompt_embeds=None, + external_req_id=None, + ) + + +def _request_output(request_id: str) -> RequestOutput: + completion = CompletionOutput( + index=0, + text="transcript", + token_ids=[11, 12], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + return RequestOutput( + request_id=request_id, + prompt="prompt", + prompt_token_ids=[1, 2], + prompt_logprobs=None, + outputs=[completion], + finished=True, + metrics=None, + lora_request=None, + ) + + +@pytest.mark.asyncio +async def test_forward_text_prompt_uses_target_stage_input_processor() -> None: + stage0 = FakeStageClient(final_output=True) + stage1 = FakeStageClient( + final_output=True, + next_inputs=[{"prompt": "hello", "multi_modal_data": {"video": ["frame"]}}], + ) + stage_pools = [ + StagePool( + 0, + [stage0], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ), + StagePool( + 1, + [stage1], + output_processor=FakeOutputProcessor(), + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ), + ] + request_q = janus.Queue() + output_q = janus.Queue() + rpc_q = janus.Queue() + orchestrator = Orchestrator( + request_async_queue=request_q.async_q, + output_async_queue=output_q.async_q, + rpc_async_queue=rpc_q.async_q, + stage_pools=stage_pools, + async_chunk=False, + ) + input_processor = FakeInputProcessor() + orchestrator._stage_input_processors[1] = input_processor + req_state = OrchestratorRequestState( + request_id="req-text", + prompt={"prompt": "original"}, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=1, + ) + + await orchestrator._forward_to_next_stage("req-text", 0, _request_output("req-text"), req_state) + + assert input_processor.calls + assert input_processor.calls[0]["prompt"] == {"prompt": "hello", "multi_modal_data": {"video": ["frame"]}} + assert stage1.add_request_calls + submitted_request = stage1.add_request_calls[0][0] + assert submitted_request.prompt_token_ids == [101, 102] + assert submitted_request.external_req_id == "req-text" diff --git a/tests/engine/test_output_modality.py b/tests/engine/test_output_modality.py index 0abfe6ea091..ff15e582479 100644 --- a/tests/engine/test_output_modality.py +++ b/tests/engine/test_output_modality.py @@ -14,7 +14,7 @@ # ── Load modules without triggering vllm_omni.__init__ ───────────── pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -_ENGINE_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "engine" +_OUTPUTS_DIR = Path(__file__).resolve().parents[2] / "vllm_omni" / "outputs" def _load_module(name: str, filepath: Path): @@ -26,12 +26,16 @@ def _load_module(name: str, filepath: Path): _om_mod = _load_module( - "vllm_omni.engine.output_modality", - _ENGINE_DIR / "output_modality.py", + "vllm_omni.outputs.output_modality", + _OUTPUTS_DIR / "output_modality.py", +) +_load_module( + "vllm_omni.outputs.utils", + _OUTPUTS_DIR / "utils.py", ) _mm_mod = _load_module( - "vllm_omni.engine.mm_outputs", - _ENGINE_DIR / "mm_outputs.py", + "vllm_omni.outputs.mm_outputs", + _OUTPUTS_DIR / "mm_outputs.py", ) OutputModality = _om_mod.OutputModality diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 56447264a3f..86e71fcbc88 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -4,14 +4,15 @@ from unittest.mock import MagicMock +import numpy as np import pytest import torch from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind from vllm.v1.engine import FinishReason -from vllm_omni.engine.output_modality import OutputModalityNames -from vllm_omni.engine.output_processor import OmniRequestState +from vllm_omni.outputs.output_modality import OutputModalityNames +from vllm_omni.outputs.output_processor import OmniRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -250,6 +251,25 @@ def test_no_detokenizer_make_request_output(): assert AUDIO in result.outputs[0].multimodal_output +def test_no_detokenizer_make_request_output_with_routed_experts(): + """make_request_output accepts the routed_experts arg that the multimodal + output channel (_process_mm_only_outputs) passes positionally, and attaches + it to the completion output. + + Regression: the call site passes 6 positional args + (..., kv_transfer_params, routed_experts) but the override previously took + only 5, raising TypeError on every generation-stage output. + """ + s = _make_no_detok_state(RequestOutputKind.CUMULATIVE) + s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO) + routed_experts = np.zeros((2, 3), dtype=np.int32) + # Mirror the exact call shape of _process_mm_only_outputs. + result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts) + assert result is not None + assert not isinstance(result, PoolingRequestOutput) + assert result.outputs[0].routed_experts is routed_experts + + def test_no_detokenizer_stream_interval_skipped(): """stream_interval > 1 does not crash when detokenizer is None.""" kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5} diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 2be2e9fb225..dc06c0c8d5f 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -3,8 +3,7 @@ from __future__ import annotations import os -import threading -from contextlib import contextmanager, nullcontext +from contextlib import contextmanager from types import SimpleNamespace from typing import Any @@ -18,10 +17,14 @@ OmniMasterServer, StageAllocation, StageCoordinatorAddresses, + StageRegistrationResponse, + StageReplicaResources, + _launch_omni_core_engines, + connect_remote_diffusion_proc, connect_remote_engine_cores, - launch_omni_core_engines, ) from vllm_omni.engine.stage_init_utils import LogicalStageInitPlan, ReplicaInitPlan +from vllm_omni.engine.stage_runtime import DistStageRuntime, StageRuntime pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -43,13 +46,13 @@ def _make_stage_cfg(stage_id: int, stage_type: str = "llm"): def _make_llm_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, launch_mode: str, vllm_config: Any | None = None, ) -> LogicalStageInitPlan: - stage_cfg = _make_stage_cfg(configured_stage_id) + stage_cfg = _make_stage_cfg(stage_id) metadata = SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="llm", runtime_cfg={"devices": "0"}, prompt_expand_func=None, @@ -63,7 +66,7 @@ def _make_llm_plan( ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=[ ReplicaInitPlan( replica_id=0, @@ -76,6 +79,7 @@ def _make_llm_plan( stage_vllm_config=vllm_config or SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)), executor_class=object, + engine_args_dict={}, ) ], ) @@ -84,12 +88,12 @@ def _make_llm_plan( def _make_diffusion_plan( stage_idx: int, *, - configured_stage_id: int, + stage_id: int, launch_mode: str, ) -> LogicalStageInitPlan: - stage_cfg = _make_stage_cfg(configured_stage_id, stage_type="diffusion") + stage_cfg = _make_stage_cfg(stage_id, stage_type="diffusion") metadata = SimpleNamespace( - stage_id=configured_stage_id, + stage_id=stage_id, stage_type="diffusion", runtime_cfg={"devices": "0"}, prompt_expand_func=None, @@ -103,7 +107,7 @@ def _make_diffusion_plan( ) return LogicalStageInitPlan( stage_idx=stage_idx, - configured_stage_id=configured_stage_id, + stage_id=stage_id, replicas=[ ReplicaInitPlan( replica_id=0, @@ -415,7 +419,7 @@ def test_master_address_and_port_stored(self, mocker: MockerFixture): def test_omni_master_server_starts_as_none(self, mocker: MockerFixture): engine = self._make_engine_no_thread(mocker) - assert engine._omni_master_server is None + assert not hasattr(engine, "_omni_master_server") # --------------------------------------------------------------------------- @@ -424,34 +428,28 @@ def test_omni_master_server_starts_as_none(self, mocker: MockerFixture): class TestSingleStageInitialization: - def _build_engine( - self, stage_cfgs: list[Any], *, single_stage_mode: bool, stage_id_filter: int | None - ) -> AsyncOmniEngine: - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.config_path = "/fake/stages.yaml" - engine.num_stages = len(stage_cfgs) - engine.stage_configs = stage_cfgs - engine.single_stage_mode = single_stage_mode - engine._single_stage_id_filter = stage_id_filter - engine._omni_master_address = "127.0.0.1" - engine._omni_master_port = 26000 - engine._omni_master_server = None - engine._omni_heartbeat_timeout = 30.0 - engine._coordinator_runtime = None - engine.async_chunk = False - engine.diffusion_batch_size = 2 - return engine + def _build_runtime(self, stage_cfgs: list[Any], *, stage_id_filter: int | None) -> DistStageRuntime: + return DistStageRuntime( + stage_configs=stage_cfgs, + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=stage_id_filter, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) def test_build_logical_stage_init_plans_marks_non_matching_stage_remote(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - stage_cfgs = [_make_stage_cfg(7), _make_stage_cfg(11)] - engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=7) + stage_cfgs = [_make_stage_cfg(0), _make_stage_cfg(1)] + runtime = self._build_runtime(stage_cfgs, stage_id_filter=0) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: SimpleNamespace( stage_id=cfg.stage_id, @@ -460,24 +458,46 @@ def test_build_logical_stage_init_plans_marks_non_matching_stage_remote(self, mo runtime_cfg={}, ), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) - monkeypatch.setattr(engine_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(), object)) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(), object)) try: - stage_plans, _ = engine._build_logical_stage_init_plans(None, [1, 1], {}) + stage_plans = runtime._build_logical_stage_init_plans(None, [1, 1], {}) finally: monkeypatch.undo() assert [plan.replicas[0].launch_mode for plan in stage_plans] == ["local", "remote"] - def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_build_logical_stage_init_plans_rejects_non_contiguous_stage_ids(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = self._build_runtime([_make_stage_cfg(7), _make_stage_cfg(11)], stage_id_filter=7) + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr( + runtime_mod, + "extract_stage_metadata", + lambda cfg: SimpleNamespace( + stage_id=cfg.stage_id, + stage_type=getattr(cfg, "stage_type", "llm"), + prompt_expand_func=None, + runtime_cfg={}, + ), + ) + try: + with pytest.raises(ValueError, match="stage_id must match its position"): + runtime._build_logical_stage_init_plans(None, [1, 1], {}) + finally: + monkeypatch.undo() + + def test_start_omni_master_server_uses_stage_ids(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.distributed import omni_coordinator as omni_coord_mod - engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([], stage_id_filter=0) mock_oms = mocker.Mock(spec=OmniMasterServer) - mocker.patch.object(engine_mod, "OmniMasterServer", return_value=mock_oms) + mocker.patch.object(runtime_mod, "OmniMasterServer", return_value=mock_oms) mocker.patch.object( omni_coord_mod, "OmniCoordinatorRuntime", @@ -485,51 +505,51 @@ def test_start_omni_master_server_uses_configured_stage_ids(self, mocker: Mocker ) stage_plans = [ - _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), - _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote"), + _make_llm_plan(0, stage_id=0, launch_mode="local"), + _make_diffusion_plan(1, stage_id=1, launch_mode="remote"), ] - engine._start_omni_master_server(stage_plans) + runtime._start_omni_master_server(stage_plans) - call_kwargs = engine_mod.OmniMasterServer.call_args.kwargs + call_kwargs = runtime_mod.OmniMasterServer.call_args.kwargs assert call_kwargs["master_address"] == "127.0.0.1" assert call_kwargs["master_port"] == 26000 - assert call_kwargs["stage_ids"] == [7, 11] - assert call_kwargs["stage_replica_counts"] == {7: 1, 11: 1} + assert call_kwargs["stage_ids"] == [0, 1] + assert call_kwargs["stage_replica_counts"] == {0: 1, 1: 1} # head_local_replicas reserves slots that the head will register - # itself (launch_mode == "local"). Stage 11 is remote, so it must + # itself (launch_mode == "local"). Stage 1 is remote, so it must # NOT appear in the head-owned set — that slot is for the headless # to fill via auto-assign. - assert call_kwargs["head_local_replicas"] == {7: [0]} + assert call_kwargs["head_local_replicas"] == {0: [0]} mock_oms.start.assert_called_once() def test_start_omni_master_server_duplicate_stage_ids_raise(self): - engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([], stage_id_filter=0) stage_plans = [ - _make_llm_plan(0, configured_stage_id=7, launch_mode="local"), - _make_llm_plan(1, configured_stage_id=7, launch_mode="remote"), + _make_llm_plan(0, stage_id=0, launch_mode="local"), + _make_llm_plan(1, stage_id=0, launch_mode="remote"), ] with pytest.raises(ValueError, match="Duplicate stage_id"): - engine._start_omni_master_server(stage_plans) + runtime._start_omni_master_server(stage_plans) def test_start_omni_master_server_missing_address_raises(self): - engine = self._build_engine([], single_stage_mode=True, stage_id_filter=7) - engine._omni_master_address = None + runtime = self._build_runtime([], stage_id_filter=0) + runtime._omni_master_address = None with pytest.raises(ValueError, match="requires both"): - engine._start_omni_master_server([_make_llm_plan(0, configured_stage_id=7, launch_mode="local")]) + runtime._start_omni_master_server([_make_llm_plan(0, stage_id=0, launch_mode="local")]) def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_single_stage_mode( self, mocker: MockerFixture ): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - engine = self._build_engine([_make_stage_cfg(7)], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([_make_stage_cfg(0)], stage_id_filter=0) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: SimpleNamespace( stage_id=cfg.stage_id, @@ -538,16 +558,16 @@ def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_s runtime_cfg={"devices": "0"}, ), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) - monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {}) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {}) monkeypatch.setattr( - engine_mod, + runtime_mod, "build_vllm_config", lambda *_, **__: (SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)), object), ) try: - stage_plans, _ = engine._build_logical_stage_init_plans(None, [1], {}) + stage_plans = runtime._build_logical_stage_init_plans(None, [1], {}) finally: monkeypatch.undo() @@ -556,22 +576,39 @@ def test_build_logical_stage_init_plans_preserves_runtime_cfg_for_local_llm_in_s def test_validate_single_stage_mode_allows_diffusion_replicas(self): stage_cfg = _make_stage_cfg(0, stage_type="diffusion") stage_cfg.runtime.num_replicas = 2 - engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=0) + runtime = self._build_runtime([stage_cfg], stage_id_filter=0) + + runtime._validate_single_stage_mode_replica_constraints() + + def test_validate_single_stage_mode_does_not_swallow_num_replica_validation_errors(self): + class RuntimeConfig: + @property + def num_replicas(self): + return 1 + + @num_replicas.setter + def num_replicas(self, value): + raise ValueError("invalid num_replicas") - engine._validate_single_stage_mode_replica_constraints() + stage_cfg = _make_stage_cfg(0, stage_type="llm") + stage_cfg.runtime = RuntimeConfig() + runtime = self._build_runtime([stage_cfg], stage_id_filter=0) + + with pytest.raises(ValueError, match="invalid num_replicas"): + runtime._validate_single_stage_mode_replica_constraints() def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_single_stage_mode( self, mocker: MockerFixture ): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod - stage_cfg = _make_stage_cfg(7, stage_type="diffusion") + stage_cfg = _make_stage_cfg(0, stage_type="diffusion") stage_cfg.runtime.devices = "0,1" - engine = self._build_engine([stage_cfg], single_stage_mode=True, stage_id_filter=7) + runtime = self._build_runtime([stage_cfg], stage_id_filter=0) monkeypatch = pytest.MonkeyPatch() monkeypatch.setattr( - engine_mod, + runtime_mod, "extract_stage_metadata", lambda cfg: SimpleNamespace( stage_id=cfg.stage_id, @@ -587,10 +624,10 @@ def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_singl replica_id=0, ), ) - monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {}) - monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) + monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {}) + monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None)) try: - stage_plans, _ = engine._build_logical_stage_init_plans(None, [2], {0: ["0", "1"]}) + stage_plans = runtime._build_logical_stage_init_plans(None, [2], {0: ["0", "1"]}) finally: monkeypatch.undo() @@ -600,11 +637,11 @@ def test_build_logical_stage_init_plans_preserves_diffusion_runtime_cfg_in_singl assert [replica.metadata.runtime_cfg for replica in replicas] == [{"devices": "0"}, {"devices": "1"}] def test_initialize_stages_calls_master_server_only_in_single_stage_mode(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - stage_plan = _make_llm_plan(0, configured_stage_id=0, launch_mode="local") + runtime = self._build_runtime(stage_cfgs, stage_id_filter=0) + stage_plan = _make_llm_plan(0, stage_id=0, launch_mode="local") client = SimpleNamespace( stage_type="llm", is_comprehension=False, @@ -613,59 +650,61 @@ def test_initialize_stages_calls_master_server_only_in_single_stage_mode(self, m default_sampling_params=SimpleNamespace(), ) - mocker.patch.object(engine_mod, "prepare_engine_environment") - mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) - mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) - mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) - mock_start = mocker.patch.object(engine, "_start_omni_master_server") - mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [client]}) - mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) - mocker.patch.object(engine_mod, "build_llm_stage_output_processor", return_value=object()) + mocker.patch.object(runtime_mod, "prepare_engine_environment") + mocker.patch.object(runtime_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(runtime, "_prepare_stage_plans", return_value=[stage_plan]) + mock_start = mocker.patch.object(runtime, "_start_omni_master_server") + mocker.patch.object(runtime, "_initialize_stage_replicas", return_value={0: [client]}) + mocker.patch.object(runtime_mod, "build_llm_stage_output_processor", return_value=object()) - engine._initialize_stages(stage_init_timeout=60) + runtime.initialize() mock_start.assert_called_once() - engine = self._build_engine(stage_cfgs, single_stage_mode=False, stage_id_filter=None) - mocker.patch.object(engine_mod, "prepare_engine_environment") - mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) - mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) - mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) - mock_start = mocker.patch.object(engine, "_start_omni_master_server") - mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [client]}) - mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) - mocker.patch.object(engine_mod, "build_llm_stage_output_processor", return_value=object()) + plain_runtime = StageRuntime( + stage_configs=stage_cfgs, + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + ) + mocker.patch.object(runtime_mod, "prepare_engine_environment") + mocker.patch.object(runtime_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(plain_runtime, "_prepare_stage_plans", return_value=[stage_plan]) + mocker.patch.object(plain_runtime, "_initialize_stage_replicas", return_value={0: [client]}) + mocker.patch.object(runtime_mod, "build_llm_stage_output_processor", return_value=object()) - engine._initialize_stages(stage_init_timeout=60) - mock_start.assert_not_called() + plain_runtime.initialize() def test_initialize_stages_stops_master_server_and_shuts_down_initialized_clients_on_failure( self, mocker: MockerFixture, ): - import vllm_omni.engine.async_omni_engine as engine_mod + import vllm_omni.engine.stage_runtime as runtime_mod stage_cfgs = [_make_stage_cfg(0)] - engine = self._build_engine(stage_cfgs, single_stage_mode=True, stage_id_filter=0) - stage_plan = _make_llm_plan(0, configured_stage_id=0, launch_mode="local") + runtime = self._build_runtime(stage_cfgs, stage_id_filter=0) + stage_plan = _make_llm_plan(0, stage_id=0, launch_mode="local") initialized_client = mocker.Mock() mock_master = mocker.Mock(spec=OmniMasterServer) - mocker.patch.object(engine_mod, "prepare_engine_environment") - mocker.patch.object(engine_mod, "load_omni_transfer_config_for_model", return_value=None) - mocker.patch.object(engine_mod, "compute_replica_layout", return_value=([1], {})) - mocker.patch.object(engine, "_build_logical_stage_init_plans", return_value=([stage_plan], None)) + mocker.patch.object(runtime_mod, "prepare_engine_environment") + mocker.patch.object(runtime_mod, "load_omni_transfer_config_for_model", return_value=None) + mocker.patch.object(runtime_mod, "compute_replica_layout", return_value=([1], {})) + mocker.patch.object(runtime, "_prepare_stage_plans", return_value=[stage_plan]) def _start_master(_plans): - engine._omni_master_server = mock_master + runtime._omni_master_server = mock_master - mocker.patch.object(engine, "_start_omni_master_server", side_effect=_start_master) - mocker.patch.object(engine, "_initialize_stage_replicas", return_value={0: [initialized_client]}) - mocker.patch.object(engine_mod, "build_stage0_input_processor", return_value=object()) - mocker.patch.object(engine, "_assemble_stage_pools", side_effect=RuntimeError("assemble failed")) - mock_shutdown = mocker.patch.object(engine, "_shutdown_initialized_clients") + mocker.patch.object(runtime, "_start_omni_master_server", side_effect=_start_master) + mocker.patch.object(runtime, "_initialize_stage_replicas", return_value={0: [initialized_client]}) + mocker.patch.object(runtime, "_finalize_initialized_stages", side_effect=RuntimeError("assemble failed")) + mock_shutdown = mocker.patch.object(runtime, "_shutdown_initialized_clients") with pytest.raises(RuntimeError, match="assemble failed"): - engine._initialize_stages(stage_init_timeout=60) + runtime.initialize() mock_shutdown.assert_called_once_with([initialized_client]) mock_master.stop.assert_called_once() @@ -673,13 +712,21 @@ def _start_master(_plans): class TestSingleStageReplicaInitialization: def test_initialize_llm_replica_remote_uses_connect_remote_engine_cores(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = {"stage_id": 7, "stage_type": "llm"} + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"} fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)) fake_addresses = SimpleNamespace( @@ -693,52 +740,74 @@ def test_initialize_llm_replica_remote_uses_connect_remote_engine_cores(self, mo def _fake_connect(**kwargs): events.append("enter") try: - yield fake_manager, fake_coordinator, fake_addresses, None + yield StageReplicaResources( + manager=fake_manager, + coordinator=fake_coordinator, + addresses=fake_addresses, + ) finally: events.append("exit") plan = _make_llm_plan( 0, - configured_stage_id=7, + stage_id=1, launch_mode="remote", vllm_config=fake_vllm_config, ).replicas[0] sentinel_client = SimpleNamespace() - mock_connect = mocker.patch.object(engine_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + mock_connect = mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect) mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", side_effect=lambda **_: (events.append("attach"), sentinel_client)[1], ) - result = engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + result = runtime._initialize_remote_replica(plan, stage_init_timeout=60) assert result is sentinel_client - engine._omni_master_server.get_stage_config.assert_called_once_with(7, timeout_s=60, replica_id=0) - assert fake_vllm_config.parallel_config.data_parallel_size_local == 0 - assert mock_connect.call_args.kwargs["stage_id"] == 7 + runtime._omni_master_server.get_stage_config.assert_called_once_with(1, timeout_s=60, replica_id=0) + assert mock_connect.call_args.kwargs["vllm_config"].parallel_config.data_parallel_size_local == 0 + assert mock_connect.call_args.kwargs["stage_id"] == 1 assert mock_connect.call_args.kwargs["replica_id"] == 0 assert events == ["enter", "exit", "attach"] def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture): - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = None + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = None - plan = _make_llm_plan(0, configured_stage_id=7, launch_mode="remote").replicas[0] + plan = _make_llm_plan(0, stage_id=1, launch_mode="remote").replicas[0] with pytest.raises(ValueError, match="registered without stage config"): - engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + runtime._initialize_remote_replica(plan, stage_init_timeout=60) def test_initialize_llm_replica_remote_attach_failure_cleans_up_started_resources(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = {"stage_id": 7, "stage_type": "llm"} + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"} fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(data_parallel_size_local=1)) fake_addresses = SimpleNamespace( @@ -749,15 +818,19 @@ def test_initialize_llm_replica_remote_attach_failure_cleans_up_started_resource @contextmanager def _fake_connect(**kwargs): - yield fake_manager, fake_coordinator, fake_addresses, None + yield StageReplicaResources( + manager=fake_manager, + coordinator=fake_coordinator, + addresses=fake_addresses, + ) plan = _make_llm_plan( 0, - configured_stage_id=7, + stage_id=1, launch_mode="remote", vllm_config=fake_vllm_config, ).replicas[0] - mocker.patch.object(engine_mod, "connect_remote_engine_cores", side_effect=_fake_connect) + mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect) mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", @@ -765,21 +838,29 @@ def _fake_connect(**kwargs): ) with pytest.raises(RuntimeError, match="attach failed"): - engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + runtime._initialize_remote_replica(plan, stage_init_timeout=60) fake_manager.shutdown.assert_called_once() fake_coordinator.shutdown.assert_called_once() - def test_initialize_llm_replica_single_stage_local_uses_launch_omni_core_engines(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_initialize_local_llm_replica_uses_shared_launcher(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._coordinator_runtime = None - engine.stage_configs = [] + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=2, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._coordinator_runtime = None + runtime._stage_configs = [] fake_vllm_config = SimpleNamespace(parallel_config=SimpleNamespace()) fake_addresses = SimpleNamespace( @@ -788,11 +869,14 @@ def test_initialize_llm_replica_single_stage_local_uses_launch_omni_core_engines @contextmanager def _fake_launch(**kwargs): - yield mocker.Mock(), None, fake_addresses + yield StageReplicaResources( + manager=mocker.Mock(), + addresses=fake_addresses, + ) plan = _make_llm_plan( 0, - configured_stage_id=3, + stage_id=0, launch_mode="local", vllm_config=fake_vllm_config, ).replicas[0] @@ -801,19 +885,19 @@ def _fake_launch(**kwargs): device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" - mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext()) - mocker.patch.object(engine_mod, "build_engine_args_dict", return_value={}) - mocker.patch.object(engine_mod, "acquire_device_locks", return_value=[]) - mocker.patch.object(engine_mod, "release_device_locks") - mock_launch = mocker.patch.object(engine_mod, "launch_omni_core_engines", side_effect=_fake_launch) + mocker.patch.object(runtime_mod, "build_engine_args_dict", return_value={}) + mocker.patch.object(runtime_mod, "acquire_device_locks", return_value=[]) + mocker.patch.object(runtime_mod, "release_device_locks") + mock_launch = mocker.patch.object(runtime_mod, "launch_stage_replica", side_effect=_fake_launch) mocker.patch.object( StageEngineCoreClientBase, "make_async_mp_client", side_effect=lambda **_: sentinel_client, ) try: - result = engine._initialize_llm_replica(plan, stage_init_timeout=60, llm_stage_launch_lock=threading.Lock()) + result = runtime._initialize_local_llm_replica(plan, stage_init_timeout=60) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) @@ -821,85 +905,121 @@ def _fake_launch(**kwargs): os.environ[device_env_var] = prev_device_env assert result is sentinel_client - assert mock_launch.call_args.kwargs["stage_id"] == 3 + assert mock_launch.call_args.kwargs["stage_id"] == 0 assert mock_launch.call_args.kwargs["stage_config"] is plan.stage_cfg assert mock_launch.call_args.kwargs["replica_id"] == 0 + assert mock_launch.call_args.kwargs["omni_master_server"] is runtime._omni_master_server def test_initialize_diffusion_replica_remote_uses_from_addresses(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod - - engine = object.__new__(AsyncOmniEngine) - engine.single_stage_mode = True - engine.diffusion_batch_size = 4 - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.get_stage_config.return_value = {"stage_id": 11, "stage_type": "diffusion"} - engine._omni_master_server.get_zmq_addresses.return_value = SimpleNamespace( - inputs=["tcp://in"], - outputs=["tcp://out"], + import vllm_omni.engine.stage_runtime as runtime_mod + + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=4, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "diffusion"} - remote_metadata = _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote").replicas[0].metadata - plan = _make_diffusion_plan(1, configured_stage_id=11, launch_mode="remote").replicas[0] + def _fake_connect(**kwargs): + @contextmanager + def _ctx(): + yield StageReplicaResources( + addresses=SimpleNamespace( + inputs=["tcp://in"], + outputs=["tcp://out"], + ) + ) + + return _ctx() + + remote_metadata = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0].metadata + plan = _make_diffusion_plan(1, stage_id=1, launch_mode="remote").replicas[0] sentinel_client = SimpleNamespace() - mocker.patch.object(engine_mod, "extract_stage_metadata", return_value=remote_metadata) - mock_from_addresses = mocker.patch.object( - engine_mod.StageDiffusionClient, "from_addresses", return_value=sentinel_client + mocker.patch.object(runtime_mod, "extract_stage_metadata", return_value=remote_metadata) + mock_connect = mocker.patch.object(runtime_mod, "connect_remote_diffusion_proc", side_effect=_fake_connect) + mock_from_addresses = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses", + return_value=sentinel_client, ) - result = engine._initialize_diffusion_replica(plan, stage_init_timeout=60, stage_launch_lock=threading.Lock()) + result = runtime._initialize_remote_replica(plan, stage_init_timeout=60) assert result is sentinel_client - engine._omni_master_server.get_stage_config.assert_called_once_with(11, timeout_s=60, replica_id=0) - engine._omni_master_server.get_zmq_addresses.assert_called_once_with(11, replica_id=0) + runtime._omni_master_server.get_stage_config.assert_called_once_with(1, timeout_s=60, replica_id=0) + mock_connect.assert_called_once_with( + omni_master_server=runtime._omni_master_server, + stage_id=1, + replica_id=0, + ) mock_from_addresses.assert_called_once() - def test_initialize_diffusion_replica_single_stage_local_registers_with_master(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_initialize_local_diffusion_replica_registers_with_master(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine.diffusion_batch_size = 4 - engine.stage_configs = [] - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.address = "127.0.0.1" - engine._omni_master_server.port = 25000 - engine._coordinator_runtime = None - - plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=4, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.address = "127.0.0.1" + runtime._omni_master_server.port = 25000 + runtime._coordinator_runtime = None + + plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0] sentinel_client = SimpleNamespace() proc = mocker.Mock() device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" - mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext()) - mocker.patch.object(engine_mod, "inject_kv_stage_info") - mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config") - mock_register = mocker.patch.object( - engine_mod, - "register_stage_with_omni_master", - return_value=("tcp://hs", "tcp://req", "tcp://resp"), + mocker.patch.object(runtime_mod, "inject_kv_stage_info") + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + mock_register = mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address=None, + ), ) - mock_spawn = mocker.patch.object( - engine_mod, - "spawn_diffusion_proc", - return_value=(proc, None, None, None), + fake_manager = SimpleNamespace( + proc=proc, + addresses=SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ), + ) + mock_manager = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager", + return_value=fake_manager, ) - mock_handshake = mocker.patch.object(engine_mod, "complete_diffusion_handshake") - mock_from_addresses = mocker.patch.object( - engine_mod.StageDiffusionClient, - "from_addresses", + mock_from_addresses = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_client.StageDiffusionClient.from_addresses", return_value=sentinel_client, ) try: - result = engine._initialize_diffusion_replica( - plan, stage_init_timeout=60, stage_launch_lock=threading.Lock() - ) + result = runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) @@ -910,78 +1030,84 @@ def test_initialize_diffusion_replica_single_stage_local_registers_with_master(s mock_register.assert_called_once_with( omni_master_address="127.0.0.1", omni_master_port=25000, - omni_stage_id=5, + omni_stage_id=0, omni_stage_config=plan.stage_cfg, - return_addresses=True, replica_id=0, ) - mock_spawn.assert_called_once_with( - "fake-model", - "diffusion-config", - handshake_address="tcp://hs", - request_address="tcp://req", - response_address="tcp://resp", + mock_manager.assert_called_once_with( + model="fake-model", + od_config="diffusion-config", + stage_init_timeout=60, + handshake_address="tcp://127.0.0.1:26001", + addresses=mocker.ANY, omni_coordinator_address=None, - omni_stage_id=5, + omni_stage_id=0, omni_replica_id=0, ) - mock_handshake.assert_called_once_with(proc, "tcp://hs", 60) + assert mock_manager.call_args.kwargs["addresses"].inputs == ["tcp://127.0.0.1:26002"] + assert mock_manager.call_args.kwargs["addresses"].outputs == ["tcp://127.0.0.1:26003"] mock_from_addresses.assert_called_once_with( plan.metadata, - request_address="tcp://req", - response_address="tcp://resp", - proc=proc, + request_address="tcp://127.0.0.1:26002", + response_address="tcp://127.0.0.1:26003", + proc_manager=mocker.ANY, batch_size=4, ) + assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc - def test_initialize_diffusion_replica_local_failure_terminates_proc(self, mocker: MockerFixture): - import vllm_omni.engine.async_omni_engine as engine_mod + def test_initialize_local_diffusion_replica_failure_terminates_proc(self, mocker: MockerFixture): + import vllm_omni.engine.stage_runtime as runtime_mod from vllm_omni.platforms import current_omni_platform - engine = object.__new__(AsyncOmniEngine) - engine.model = "fake-model" - engine.single_stage_mode = True - engine.diffusion_batch_size = 4 - engine.stage_configs = [] - engine._omni_master_server = mocker.Mock(spec=OmniMasterServer) - engine._omni_master_server.address = "127.0.0.1" - engine._omni_master_server.port = 25000 - engine._coordinator_runtime = None - - plan = _make_diffusion_plan(0, configured_stage_id=5, launch_mode="local").replicas[0] - proc = mocker.Mock() + runtime = DistStageRuntime( + stage_configs=[], + model="fake-model", + config_path="/fake/stages.yaml", + stage_init_timeout=60, + diffusion_batch_size=4, + async_chunk=False, + single_stage_id_filter=None, + omni_master_address="127.0.0.1", + omni_master_port=26000, + ) + runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer) + runtime._omni_master_server.address = "127.0.0.1" + runtime._omni_master_server.port = 25000 + runtime._coordinator_runtime = None + + plan = _make_diffusion_plan(0, stage_id=0, launch_mode="local").replicas[0] device_env_var = current_omni_platform.device_control_env_var prev_device_env = os.environ.get(device_env_var) os.environ[device_env_var] = "0" + runtime._init_visible_devices_baseline = "0" - mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext()) - mocker.patch.object(engine_mod, "inject_kv_stage_info") - mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config") - mocker.patch.object( - engine_mod, - "register_stage_with_omni_master", - return_value=("tcp://hs", "tcp://req", "tcp://resp"), + mocker.patch.object(runtime_mod, "inject_kv_stage_info") + mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value="diffusion-config") + mocker.patch( + "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://127.0.0.1:26002", + output_address="tcp://127.0.0.1:26003", + replica_id=0, + coordinator_router_address=None, + ), ) - mocker.patch.object( - engine_mod, - "spawn_diffusion_proc", - return_value=(proc, None, None, None), + mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager", + side_effect=RuntimeError("handshake failed"), ) - mocker.patch.object(engine_mod, "complete_diffusion_handshake", side_effect=RuntimeError("handshake failed")) - mock_terminate = mocker.patch.object(engine_mod, "terminate_alive_proc") try: with pytest.raises(RuntimeError, match="handshake failed"): - engine._initialize_diffusion_replica(plan, stage_init_timeout=60, stage_launch_lock=threading.Lock()) + runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=60) finally: if prev_device_env is None: os.environ.pop(device_env_var, None) else: os.environ[device_env_var] = prev_device_env - mock_terminate.assert_called_once_with(proc) - # --------------------------------------------------------------------------- # Stage engine startup helpers @@ -1025,14 +1151,16 @@ def fake_socket_ctx(*args, **kwargs): yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, stage_id=7, replica_id=2, - ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue): - assert yielded_coordinator is None + ) as resources: + assert resources.coordinator is None + yielded_addresses = resources.addresses + assert yielded_addresses is not None assert yielded_addresses.coordinator_input == "tcp://coord-in" assert yielded_addresses.coordinator_output == "tcp://coord-out" assert yielded_addresses.frontend_stats_publish_address == "tcp://stats" @@ -1058,17 +1186,48 @@ def fake_socket_ctx(*args, **kwargs): yield mocker.Mock() mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) - mocker.patch("vllm_omni.engine.stage_engine_startup._wait_for_omni_engine_startup") + mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") with connect_remote_engine_cores( vllm_config=vllm_config, omni_master_server=omni_master_server, stage_id=7, - ) as (_, yielded_coordinator, yielded_addresses, _tensor_queue): - assert yielded_coordinator is None + ) as resources: + assert resources.coordinator is None + yielded_addresses = resources.addresses + assert yielded_addresses is not None assert yielded_addresses.coordinator_input is None assert yielded_addresses.coordinator_output is None assert yielded_addresses.frontend_stats_publish_address is None + def test_connect_remote_diffusion_proc_waits_for_headless_remote(self, mocker: MockerFixture): + omni_master_server = mocker.Mock(spec=OmniMasterServer) + omni_master_server.get_zmq_addresses.return_value = EngineZmqAddresses( + inputs=["tcp://client-in"], + outputs=["tcp://client-out"], + ) + omni_master_server.get_allocation.return_value = mocker.Mock(handshake_bind_address="tcp://127.0.0.1:26001") + + @contextmanager + def fake_socket_ctx(*args, **kwargs): + yield mocker.Mock() + + mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) + mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") + + with connect_remote_diffusion_proc( + omni_master_server=omni_master_server, + stage_id=7, + replica_id=2, + ) as resources: + assert resources.addresses is omni_master_server.get_zmq_addresses.return_value + + omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2) + omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2) + mock_wait.assert_called_once() + _, _, core_engines, parallel_config, *_ = mock_wait.call_args.args + assert core_engines[0].local is False + assert parallel_config.data_parallel_size_local == 0 + class TestLaunchOmniCoreEngines: def test_registers_stage_once_and_reuses_handshake_for_all_local_engines(self, mocker: MockerFixture): @@ -1093,7 +1252,13 @@ def fake_socket_ctx(*args, **kwargs): mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://client-in", + output_address="tcp://client-out", + replica_id=2, + coordinator_router_address=None, + ), ) mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mock_manager_cls = mocker.patch( @@ -1101,7 +1266,7 @@ def fake_socket_ctx(*args, **kwargs): return_value=local_engine_manager, ) mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") - with launch_omni_core_engines( + with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=mocker.Mock(), log_stats=False, @@ -1164,7 +1329,13 @@ def fake_socket_ctx(*args, **kwargs): mocker.patch("vllm_omni.engine.stage_engine_startup.DPCoordinator", return_value=coordinator) mock_register = mocker.patch( "vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master", - return_value="tcp://127.0.0.1:26001", + return_value=StageRegistrationResponse( + handshake_address="tcp://127.0.0.1:26001", + input_address="tcp://client-in", + output_address="tcp://client-out", + replica_id=3, + coordinator_router_address=None, + ), ) mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx()) mocker.patch( @@ -1172,7 +1343,7 @@ def fake_socket_ctx(*args, **kwargs): return_value=mocker.Mock(), ) mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup") - with launch_omni_core_engines( + with _launch_omni_core_engines( vllm_config=vllm_config, executor_class=mocker.Mock(), log_stats=False, diff --git a/tests/engine/test_stage_engine_core_client.py b/tests/engine/test_stage_engine_core_client.py index dde0927af2d..36c95cde7d9 100644 --- a/tests/engine/test_stage_engine_core_client.py +++ b/tests/engine/test_stage_engine_core_client.py @@ -1,15 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Tests for StageEngineCoreClient.check_health(). - -Uses object.__new__ to construct a minimal client — check_health only -touches self.resources, self.stage_id, and self._proc. -""" +"""Tests for StageEngineCoreClient.check_health().""" from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock import pytest from vllm.v1.engine.exceptions import EngineDeadError @@ -19,28 +14,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _make_client(*, engine_dead=False, proc_alive=True): +def _make_client(*, engine_dead=False): client = object.__new__(StageEngineCoreClient) client.stage_id = 0 client.resources = SimpleNamespace(engine_dead=engine_dead) - client._proc = MagicMock(is_alive=MagicMock(return_value=proc_alive), exitcode=1) return client def test_check_health_passes_when_alive(): - client = _make_client(engine_dead=False, proc_alive=True) + client = _make_client(engine_dead=False) client.check_health() # no exception def test_check_health_raises_when_resources_engine_dead(): - client = _make_client(engine_dead=True, proc_alive=True) + client = _make_client(engine_dead=True) with pytest.raises(EngineDeadError, match="engine core is dead"): client.check_health() - - -def test_check_health_raises_when_proc_not_alive(): - client = _make_client(engine_dead=False, proc_alive=False) - with pytest.raises(EngineDeadError, match="not alive"): - client.check_health() - # Verify it set resources.engine_dead as a side effect - assert client.resources.engine_dead is True diff --git a/tests/engine/test_strategy_lb_policy.py b/tests/engine/test_strategy_lb_policy.py new file mode 100644 index 00000000000..8531c2a999b --- /dev/null +++ b/tests/engine/test_strategy_lb_policy.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for strategy-derived omni_lb_policy threading into AsyncOmniEngine. + +``AsyncOmniEngine._apply_strategy_lb_policy`` only touches ``self._omni_lb_policy`` +and module-level logging, so we exercise it with a lightweight stub ``self`` to +avoid standing up a full engine (which needs GPUs/model weights). +""" + +from __future__ import annotations + +import pytest + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _Stub: + def __init__(self, current: str = "random") -> None: + self._omni_lb_policy = current + + +_apply = AsyncOmniEngine._apply_strategy_lb_policy + + +def test_no_derived_policy_is_noop(): + stub = _Stub("random") + _apply(stub, None, {}) + assert stub._omni_lb_policy == "random" + + +def test_derived_overrides_default_random(): + # "random" is the engine default == "unset", so the strategy value wins. + stub = _Stub("random") + _apply(stub, "round-robin", {}) + assert stub._omni_lb_policy == "round-robin" + + +def test_derived_overrides_when_flag_absent(): + stub = _Stub("random") + _apply(stub, "least-queue-length", {"omni_lb_policy": None}) + assert stub._omni_lb_policy == "least-queue-length" + + +def test_explicit_matching_flag_is_fine(): + stub = _Stub("round-robin") + _apply(stub, "round-robin", {"omni_lb_policy": "round-robin"}) + assert stub._omni_lb_policy == "round-robin" + + +def test_explicit_conflicting_flag_raises(): + stub = _Stub("round-robin") + with pytest.raises(ValueError, match="Conflicting load-balancer policy"): + _apply(stub, "least-queue-length", {"omni_lb_policy": "round-robin"}) + + +def test_explicit_random_is_treated_as_unset(): + # An explicit --omni-lb-policy=random is indistinguishable from the default, + # so a strategy value still overrides it (documented Phase-1 semantics). + stub = _Stub("random") + _apply(stub, "round-robin", {"omni_lb_policy": "random"}) + assert stub._omni_lb_policy == "round-robin" diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index d108b932380..5a7ef7bf032 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -11,10 +11,11 @@ import io import json from argparse import Namespace +from http import HTTPStatus from types import SimpleNamespace import pytest -from fastapi import FastAPI +from fastapi import FastAPI, HTTPException from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture @@ -23,7 +24,7 @@ from vllm.sampling_params import RequestOutputKind from vllm_omni.entrypoints.async_omni import AsyncOmni -from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router +from vllm_omni.entrypoints.openai.api_server import _check_max_generated_image_size, _DiffusionServingModels, router from vllm_omni.entrypoints.openai.image_api_utils import ( encode_image_base64, parse_size, @@ -1243,7 +1244,7 @@ def test_parameter_validation(): # Invalid layers for layered models (must stay within the backend-supported range) with pytest.raises(ValueError): - ImageGenerationRequest(prompt="test", layers=2) + ImageGenerationRequest(prompt="test", layers=1) with pytest.raises(ValueError): ImageGenerationRequest(prompt="test", layers=11) @@ -1300,6 +1301,58 @@ def test_generate_images_rejects_model_mismatch(test_client): assert "model mismatch" in response.json()["detail"].lower() +def test_image_file_response_format_multiple(test_client): + """Test response_format=file with n>1 returns ZIP archive""" + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "a dog", + "n": 3, + "response_format": "file", + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"] == "application/zip" + assert "attachment" in response.headers.get("content-disposition", "") + assert ".zip" in response.headers.get("content-disposition", "") + + # Verify it's a valid ZIP with 3 PNG files + import zipfile + + zip_buffer = io.BytesIO(response.content) + with zipfile.ZipFile(zip_buffer, "r") as zf: + files = zf.namelist() + assert len(files) == 3 + assert all(f.endswith(".png") for f in files) + + # Verify each file is a valid PNG + for filename in files: + img_bytes = zf.read(filename) + img = Image.open(io.BytesIO(img_bytes)) + assert img.format == "PNG" + + +def test_image_file_response_format_single(test_client): + """Test response_format=file with n=1 returns a single image file.""" + response = test_client.post( + "/v1/images/generations", + json={ + "prompt": "a dog", + "n": 1, + "response_format": "file", + }, + ) + + assert response.status_code == 200 + assert response.headers["content-type"] == "image/png" + assert "attachment" in response.headers.get("content-disposition", "") + assert ".png" in response.headers.get("content-disposition", "") + + img = Image.open(io.BytesIO(response.content)) + assert img.format == "PNG" + + def make_test_image_bytes(size=(64, 64)) -> bytes: img = Image.new( "RGB", @@ -1610,13 +1663,13 @@ def test_image_edit_invalid_layers(async_omni_test_client): files=[("image", img_bytes)], data={ "prompt": "test", - "layers": 2, + "layers": 1, }, ) assert response.status_code == 400 detail = response.json()["detail"] assert "Invalid layers" in detail - assert "layers must be between 3 and 10 inclusive" in detail + assert "layers must be between 2 and 10 inclusive" in detail # Test layers above the supported range response = async_omni_test_client.post( @@ -1630,7 +1683,7 @@ def test_image_edit_invalid_layers(async_omni_test_client): assert response.status_code == 400 detail = response.json()["detail"] assert "Invalid layers" in detail - assert "layers must be between 3 and 10 inclusive" in detail + assert "layers must be between 2 and 10 inclusive" in detail def test_image_edit_resolution_and_size_conflict(async_omni_test_client): @@ -1979,6 +2032,56 @@ def __init__(self): assert images[0].size == (32, 32) +# --------------------------------------------------------------------------- +# _check_max_generated_image_size unit tests +# --------------------------------------------------------------------------- + + +def test_width_height_within_limit_passes(): + args = SimpleNamespace(max_generated_image_size=1024 * 1024) + # Exactly at limit is allowed (> not >=) + _check_max_generated_image_size(args, 1024, 1024) + # Below limit + _check_max_generated_image_size(args, 512, 512) + + +def test_width_height_exceeds_limit_raises_400(): + limit = 1024 * 1024 + args = SimpleNamespace(max_generated_image_size=limit) + with pytest.raises(HTTPException) as exc_info: + _check_max_generated_image_size(args, 1025, 1024) + assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST.value + assert "1025x1024" in exc_info.value.detail + assert str(limit) in exc_info.value.detail + + +def test_width_height_error_message_contains_size_hint(): + args = SimpleNamespace(max_generated_image_size=512 * 512) + with pytest.raises(HTTPException) as exc_info: + _check_max_generated_image_size(args, 1024, 512) + assert "--max-generated-image-size" in exc_info.value.detail + + +def test_resolution_within_limit_passes(): + args = SimpleNamespace(max_generated_image_size=1024 * 1024) + # Exactly at limit is allowed (> not >=): 1024*1024 == limit + _check_max_generated_image_size(args, None, None, resolution=1024) + # Below limit + _check_max_generated_image_size(args, None, None, resolution=512) + + +def test_resolution_exceeds_limit_raises_400(): + limit = 1024 * 1024 + args = SimpleNamespace(max_generated_image_size=limit) + with pytest.raises(HTTPException) as exc_info: + _check_max_generated_image_size(args, None, None, resolution=1025) + assert exc_info.value.status_code == HTTPStatus.BAD_REQUEST.value + detail = exc_info.value.detail + assert "1025" in detail + assert "1025x1025" in detail + assert str(limit) in detail + + def test_image_edits_size_auto_preserves_bridge_size(async_omni_stage_configs_only_client): """size=auto must NOT pin the diffusion stage sampling_params.height/width. diff --git a/tests/entrypoints/openai_api/test_openai_audio_list_endpoints.py b/tests/entrypoints/openai_api/test_openai_audio_list_endpoints.py index 32a2dedbd5d..68acb7c5918 100644 --- a/tests/entrypoints/openai_api/test_openai_audio_list_endpoints.py +++ b/tests/entrypoints/openai_api/test_openai_audio_list_endpoints.py @@ -29,7 +29,7 @@ class TestOpenAIAudioListEndpoints: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_list_voices_endpoint(self, omni_server) -> None: """``GET /v1/audio/voices`` returns a non-empty ``voices`` list.""" url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/voices" @@ -45,7 +45,7 @@ def test_list_voices_endpoint(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) def test_models_endpoint(self, omni_server) -> None: """``GET /v1/models`` returns at least one model.""" url = f"http://{omni_server.host}:{omni_server.port}/v1/models" diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 83d17a28cd3..398ae649415 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -88,6 +88,47 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts(): np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32)) +def test_unpack_accepts_openpi_client_ndarray_markers(): + image = np.arange(6, dtype=np.uint8).reshape(2, 3) + payload = { + "observation/exterior_image_0_left": { + b"__ndarray__": True, + b"dtype": image.dtype.str, + b"shape": image.shape, + b"data": image.tobytes(), + } + } + + decoded = openpi_connection._unpack_numpy(payload) + + np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image) + assert decoded["observation/exterior_image_0_left"].dtype == np.uint8 + + +def test_unpack_openpi_client_packed_observation(): + openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy") + image = np.zeros((180, 320, 3), dtype=np.uint8) + obs = { + "observation/exterior_image_0_left": image, + "observation/joint_position": np.zeros(7, dtype=np.float32), + } + + decoded = openpi_connection._unpack(openpi_msgpack.packb(obs)) + + np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image) + np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"]) + + +def test_pack_uses_openpi_client_ndarray_markers(): + openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy") + actions = np.asarray([[1.0, 2.0]], dtype=np.float32) + + packed = openpi_connection._pack(actions) + decoded = openpi_msgpack.unpackb(packed) + + np.testing.assert_allclose(decoded, actions) + + def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged(): payload = { "metadata": { diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index bbe50a9255d..50f8eef911c 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -198,12 +198,28 @@ def _assert_realtime_smoke(result: dict) -> None: assert result["output_sample_rate"] > 0 -def _assert_realtime_accuracy(result: dict) -> None: +def _assert_realtime_accuracy( + result: dict, + whisper_model_size: str = "large-v3", +) -> None: + """Assert that whisper transcription of audio output matches model text. + + Args: + result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``. + whisper_model_size: Whisper model used to transcribe the generated audio + for the accuracy check. Defaults to ``large-v3``: the default + ``small`` model mishears short Chinese TTS clips (observed: + 北京→韦京 and a dropped leading sentence, sim=0.443), which + caused spurious sim<0.8 failures under async_chunk codec + variability even though audio generation was correct. large-v3 + transcribes these clips reliably, so a failure here now points + at the model, not the ASR grader. + """ final_text = (result["transcription_text"] or "").strip() assert final_text, "Expected non-empty transcription (model text stream)" wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"]) - whisper_text = convert_audio_bytes_to_text(wav_out).strip() + whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip() assert whisper_text, "Whisper returned empty string for synthesized output audio" sim = cosine_similarity_text(whisper_text.lower(), final_text.lower()) @@ -214,27 +230,6 @@ def _assert_realtime_accuracy(result: dict) -> None: class TestQwen3OmniRealtimeWebSocket: - @pytest.mark.core_model - @pytest.mark.omni - @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) - @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True) - def test_streaming_audio_input_pcm_output_async_chunk_smoke(self, omni_server) -> None: - """Ready CI: async_chunk on, no send delay, smoke only.""" - pcm16 = _synthetic_pcm16_input() - - result = asyncio.run( - _run_realtime_audio_roundtrip( - omni_server.host, - omni_server.port, - omni_server.model, - pcm16, - chunk_ms=100, - send_delay_ms=0, - ) - ) - - _assert_realtime_smoke(result) - @pytest.mark.advanced_model @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py index d82e1a68b5d..d51ac653a7b 100644 --- a/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py +++ b/tests/entrypoints/openai_api/test_qwen3_tts_speaker_embedding.py @@ -141,6 +141,7 @@ def test_speaker_embedding_streaming(self, omni_server) -> None: "x_vector_only_mode": True, "response_format": "pcm", "stream": True, + "stream_format": "audio", "max_new_tokens": MAX_NEW_TOKENS, } with httpx.Client(timeout=120.0) as client: diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py b/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py index 296138564eb..ecae2115840 100644 --- a/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py @@ -102,7 +102,7 @@ async def _run_ws_session(host: str, port: int, model: str) -> dict: class TestQwen3TTSWebSocket: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "L4"}, num_cards=4) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_ws_server_params, indirect=True) def test_streaming_pcm_output(self, omni_server) -> None: result = asyncio.run(_run_ws_session(omni_server.host, omni_server.port, omni_server.model)) @@ -113,9 +113,9 @@ def test_streaming_pcm_output(self, omni_server) -> None: session_done = result["session_done"] assert session_done is not None - assert session_done["total_sentences"] == 2 - assert len(starts) == 2 - assert len(dones) == 2 + assert session_done["total_sentences"] == 1 + assert len(starts) == 1 + assert len(dones) == 1 for idx, start in enumerate(starts): assert start["type"] == "audio.start" 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 507018ef717..0dea9717957 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py +++ b/tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py @@ -1,8 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 -"""Regression tests for multistage diffusion generation input construction.""" +"""Regression tests for multistage generation input construction.""" from __future__ import annotations +import asyncio from types import SimpleNamespace import pytest @@ -93,6 +94,79 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser assert engine.default_sampling_params_list[2].lora_request is None +def test_prepare_multistage_multimodal_inputs_defers_downstream_modalities(serving_chat): + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + image_payload = object() + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="asr", requires_multimodal_data=True), + SimpleNamespace(model_stage="aura", requires_multimodal_data=True), + ] + ) + serving_chat.model_config = SimpleNamespace( + allowed_local_media_path="", + allowed_media_domains=None, + ) + request = SimpleNamespace(media_io_kwargs=None) + messages = [ + { + "role": "user", + "content": [ + {"type": "audio_url", "audio_url": {"url": "data:audio/wav;base64,abc"}}, + {"type": "image_pil", "image_pil": image_payload}, + {"type": "text", "text": "What changed?"}, + ], + } + ] + + stripped_messages, deferred_data = asyncio.run( + OmniOpenAIServingChat._prepare_multistage_multimodal_inputs( + serving_chat, + messages, + request, + ) + ) + + assert deferred_data == {"image": [image_payload]} + stripped_content = stripped_messages[0]["content"] + assert {"type": "image_pil", "image_pil": image_payload} not in stripped_content + assert {"type": "text", "text": "What changed?"} in stripped_content + + +def test_needs_multistage_multimodal_split_noops_for_single_stage(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="thinker", requires_multimodal_data=True), + ] + ) + + assert serving_chat._needs_multistage_multimodal_split() is False + + +def test_needs_multistage_multimodal_split_noops_for_qwen3_omni(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="thinker", requires_multimodal_data=True), + SimpleNamespace(model_stage="talker", requires_multimodal_data=False), + SimpleNamespace(model_stage="code2wav", requires_multimodal_data=False), + ] + ) + + assert serving_chat._needs_multistage_multimodal_split() is False + + +def test_needs_multistage_multimodal_split_detects_aura_handoff(serving_chat): + serving_chat.engine_client = SimpleNamespace( + stage_configs=[ + SimpleNamespace(model_stage="asr", requires_multimodal_data=True), + SimpleNamespace(model_stage="aura", requires_multimodal_data=True), + ] + ) + + assert serving_chat._needs_multistage_multimodal_split() is True + + def test_build_multistage_generation_inputs_multi_image_emits_n_img_placeholders(serving_chat): """N reference images with bot_task set must emit N <img> placeholders. diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index c885d907ca4..a6c7baffe71 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -6,6 +6,8 @@ are correctly applied to the comprehension stage while preserving YAML defaults. """ +from types import SimpleNamespace + import pytest from pytest_mock import MockerFixture from vllm.sampling_params import SamplingParams @@ -487,6 +489,35 @@ def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock assert comprehension_params.stop_token_ids == [] +def test_to_sampling_params_list_pads_missing_tail_stage_with_defaults(): + """AURA callers may pass 3 semantic model params for a 4-stage engine pipeline.""" + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + instance = object.__new__(OmniOpenAIServingChat) + default_params = [ + SamplingParams(max_tokens=10), + SamplingParams(max_tokens=20), + SamplingParams(max_tokens=30), + SamplingParams(max_tokens=40), + ] + instance.engine_client = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="llm") for _ in range(4)], + default_sampling_params_list=default_params, + ) + + result = instance._to_sampling_params_list( + [ + {"max_tokens": 1}, + {"max_tokens": 2}, + {"max_tokens": 3}, + ] + ) + + assert len(result) == 4 + assert [params.max_tokens for params in result] == [1, 2, 3, 40] + assert result[3] is not default_params[3] + + # ============================================================================= # Tests for _get_comprehension_stage_index # ============================================================================= diff --git a/tests/entrypoints/openai_api/test_serving_chat_speaker.py b/tests/entrypoints/openai_api/test_serving_chat_speaker.py index 7d6119b6f26..b0977c0c23e 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_speaker.py +++ b/tests/entrypoints/openai_api/test_serving_chat_speaker.py @@ -79,6 +79,7 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M serving_chat.renderer.get_tokenizer.return_value = mocker.MagicMock() serving_chat.reasoning_parser_cls = None serving_chat.tool_parser = None + serving_chat.parser_cls = None serving_chat.use_harmony = False serving_chat.enable_auto_tools = False serving_chat.exclude_tools_when_tool_choice_none = False @@ -86,8 +87,9 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M serving_chat.chat_template = None serving_chat.chat_template_content_format = "string" serving_chat.default_chat_template_kwargs = {} - serving_chat._validate_chat_template = mocker.MagicMock(return_value=None) - serving_chat._prepare_extra_chat_template_kwargs = mocker.MagicMock(return_value={}) + serving_chat.online_renderer = mocker.MagicMock() + serving_chat.online_renderer.validate_chat_template.return_value = None + serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={}) serving_chat._preprocess_chat = mocker.AsyncMock( side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian") ) diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 9958d24e1c5..8faca1b78b4 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -39,6 +39,8 @@ OmniOpenAIServingSpeech, _create_wav_header, ) +from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest, SpeechServingContext +from vllm_omni.entrypoints.openai.tts_adapters.ming_tts import MingTTSAdapter from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, build_fish_voice_clone_prompt_ids, @@ -972,6 +974,29 @@ def test_upload_voice_embedding_wrong_dims_rejected(self, speech_server): ) ) + def test_upload_ming_voice_embedding_wrong_dims_rejected_without_replacing_existing(self, speech_server): + """Ming embedding uploads must be 192-dim before replacing stored voices.""" + speech_server._tts_model_type = "ming_tts" + existing = { + "name": "bad_emb_voice", + "file_path": "/tmp/voice_samples/bad_emb_voice.safetensors", + "mime_type": "application/x-safetensors", + "embedding_source": "direct", + "embedding_dim": 192, + } + speech_server.uploaded_speakers = {"bad_emb_voice": existing.copy()} + + with pytest.raises(ValueError, match="Ming speaker embedding must have 192 dims, got 191"): + asyncio.run( + speech_server.upload_voice_embedding( + embedding_json=json.dumps([0.0] * 191), + consent="consent", + name="bad_emb_voice", + ) + ) + + assert speech_server.uploaded_speakers["bad_emb_voice"] == existing + def test_base_task_requires_ref_audio_or_speaker_embedding(self, speech_server): """Base task without ref_audio or speaker_embedding is rejected.""" req = OpenAICreateSpeechRequest(input="Hello", task_type="Base") @@ -1524,7 +1549,7 @@ def test_speaker_alias_in_base_task_with_uploaded_voice(self, speech_server, moc req = OpenAICreateSpeechRequest.model_validate({"input": "Hello", "speaker": "UTESF", "task_type": "Base"}) assert req.voice == "UTESF" mocker.patch("pathlib.Path.exists", return_value=True) - result = speech_server._validate_qwen_tts_request(req) + result = speech_server._validate_tts_request(req) assert result is None # ── uploaded voice with embedding ── @@ -1556,6 +1581,86 @@ def test_build_tts_params_with_uploaded_voice_embedding(self, speech_server, moc assert params["x_vector_only_mode"] == [True] assert "ref_audio" not in params + def test_ming_adapter_uploaded_direct_embedding_uses_embedding_not_audio( + self, speech_server, mocker: MockerFixture + ): + """Ming uploaded direct embeddings are loaded into prompt construction.""" + speech_server.uploaded_speakers = { + "emb_voice": { + "name": "emb_voice", + "file_path": "/tmp/voice_samples/emb_voice.safetensors", + "mime_type": "application/x-safetensors", + "embedding_source": "direct", + "embedding_dim": 192, + } + } + fake_embedding = [0.1] * 192 + mock_get_emb = mocker.patch.object( + speech_server, "_get_uploaded_speaker_embedding", return_value=fake_embedding + ) + mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data") + mock_prompt = mocker.patch.object( + speech_server, + "_build_ming_dense_prompt", + return_value={"additional_information": {"speaker_count": 1}}, + ) + + adapter = MingTTSAdapter(SpeechServingContext(server=speech_server)) + req = OpenAICreateSpeechRequest(input="Hello", voice="emb_voice") + prepared = asyncio.run(adapter.build(req, [], False)) + + mock_get_emb.assert_called_once_with("emb_voice") + mock_get_audio.assert_not_called() + mock_prompt.assert_called_once_with(req, ref_audio_data=None) + prompt_request = mock_prompt.call_args.args[0] + assert prompt_request.speaker_embedding == fake_embedding + assert len(prompt_request.speaker_embedding) == 192 + assert prepared.prompt["additional_information"]["speaker_count"] == 1 + assert prepared.model_type == "ming_tts" + + def test_ming_adapter_uploaded_audio_path_preserves_ref_text(self, speech_server, mocker: MockerFixture): + """Ming uploaded audio voices still resolve audio and stored ref_text.""" + speech_server.uploaded_speakers = { + "audio_voice": { + "name": "audio_voice", + "file_path": "/tmp/voice_samples/audio_voice.wav", + "mime_type": "audio/wav", + "ref_text": "Reference transcript.", + } + } + ref_audio_source = "data:audio/wav;base64,ZmFrZWF1ZGlv" + ref_audio_data = ([0.0, 0.1], 16000) + fake_embedding = [0.2] * 192 + mock_get_audio = mocker.patch.object(speech_server, "_get_uploaded_audio_data", return_value=ref_audio_source) + mock_get_emb = mocker.patch.object(speech_server, "_get_uploaded_speaker_embedding") + mocker.patch.object( + speech_server, + "_resolve_ref_audio", + new=mocker.AsyncMock(return_value=ref_audio_data), + ) + mock_extract = mocker.patch.object( + speech_server, + "_extract_ming_speaker_embeddings_from_ref_audio", + return_value=[fake_embedding], + ) + mock_prompt = mocker.patch.object( + speech_server, + "_build_ming_dense_prompt", + return_value={"additional_information": {"speaker_count": 1}}, + ) + + adapter = MingTTSAdapter(SpeechServingContext(server=speech_server)) + req = OpenAICreateSpeechRequest(input="Hello", voice="audio_voice") + prepared = asyncio.run(adapter.build(req, [], False)) + + mock_get_audio.assert_called_once_with("audio_voice") + mock_get_emb.assert_not_called() + mock_extract.assert_called_once_with([ref_audio_data]) + mock_prompt.assert_called_once_with(req, ref_audio_data=ref_audio_data) + assert req.ref_text == "Reference transcript." + assert req.speaker_embedding == fake_embedding + assert prepared.model_type == "ming_tts" + # ── regression: full flow from issue #1603 ── def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_server, mocker: MockerFixture): @@ -1579,7 +1684,7 @@ def test_regression_1603_speaker_key_with_uploaded_audio_voice(self, speech_serv # Validation should pass (file exists) mocker.patch("pathlib.Path.exists", return_value=True) - err = speech_server._validate_qwen_tts_request(req) + err = speech_server._validate_tts_request(req) assert err is None, f"Validation failed: {err}" # Build params should auto-set ref_audio from stored file @@ -1617,7 +1722,7 @@ def test_regression_1603_speaker_key_with_uploaded_embedding_voice(self, speech_ # Validation should pass mocker.patch("pathlib.Path.exists", return_value=True) - err = speech_server._validate_qwen_tts_request(req) + err = speech_server._validate_tts_request(req) assert err is None, f"Validation failed: {err}" # Build params should use embedding, NOT audio @@ -1761,6 +1866,35 @@ def test_validate_instructions_length_uses_cached_value(self, mocker: MockerFixt assert error is not None assert "max 10 characters" in error + def test_batch_voice_not_treated_as_inline_ref_audio( + self, + speech_server, + mocker: MockerFixture, + ): + speech_server._check_model = mocker.AsyncMock(return_value=None) + speech_server._batch_max_items = 32 + + def validate_and_resolve_uploaded_voice(request): + assert request.voice == "alice" + assert request.ref_audio is None + + request.ref_audio = "data:audio/wav;base64,dGVzdA==" + request.ref_text = "reference text" + return None + + speech_server._validate_tts_request = mocker.Mock(side_effect=validate_and_resolve_uploaded_voice) + speech_server._generate_audio_bytes = mocker.AsyncMock(return_value=("YWJj", "audio/wav")) + + batch = BatchSpeechRequest(voice="alice", items=[SpeechBatchItem(input="hello")]) + + response = asyncio.run(speech_server.create_speech_batch(batch)) + call = speech_server._generate_audio_bytes.await_args + request = call.args[0] + + assert response.results[0].status == "success" + assert request.ref_audio is not None + assert call.kwargs["has_inline_ref_audio"] is False + class TestFileValidationFunctions: """Unit tests for file validation helper functions.""" @@ -1827,27 +1961,52 @@ def test_validate_path_within_directory(self, tmp_path): class TestStreamingProtocolValidation: - """Unit tests for the stream field validators in OpenAICreateSpeechRequest.""" + """Unit tests for streaming validators in OpenAICreateSpeechRequest.""" + + def test_default_is_non_streaming(self): + req = OpenAICreateSpeechRequest(input="Hello") + assert req.stream is False + assert req.stream_format is None + assert req.is_streaming() is False def test_stream_validation_errors(self): - """stream=True requires response_format not in ('pcm', 'wav') and speed=1.0.""" + """stream=True requires response_format in ('pcm', 'wav') and speed=1.0.""" with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"): OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="mp3") with pytest.raises(ValidationError, match="Speed adjustment is not supported"): OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0) + def test_stream_format_audio_validation_errors(self): + with pytest.raises(ValidationError, match="requires response_format='pcm' or 'wav'"): + OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="mp3") + with pytest.raises(ValidationError, match="Speed adjustment is not supported"): + OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="pcm", speed=2.0) + def test_stream_valid(self): - """stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted.""" + """stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted as SSE.""" req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm") assert req.stream is True + assert req.is_sse_stream() is True + assert req.is_raw_audio_stream() is False req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav") assert req.stream is True + assert req.is_sse_stream() is True + assert req.is_raw_audio_stream() is False - def test_sse_stream_format_is_blocked(self): - """stream_format='sse' is blocked.""" - with pytest.raises(ValidationError, match="sse"): - OpenAICreateSpeechRequest(input="Hello", stream_format="sse") + def test_stream_format_audio_is_valid(self): + req = OpenAICreateSpeechRequest(input="Hello", stream_format="audio", response_format="pcm") + assert req.stream_format == "audio" + assert req.is_raw_audio_stream() is True + assert req.is_sse_stream() is False + + def test_sse_stream_format_is_valid(self): + """stream_format='sse' is accepted for /audio/speech.""" + req = OpenAICreateSpeechRequest(input="Hello", stream_format="sse") + + assert req.stream_format == "sse" + assert req.is_sse_stream() is True + assert req.is_raw_audio_stream() is False class TestStreamingResponse: @@ -1920,14 +2079,166 @@ async def awaitable_create_speech(*args, **kwargs): app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) return app + @staticmethod + def _assert_sse_audio_response(response, response_format: str = "pcm"): + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + body = response.text + assert "event: speech.audio.delta" in body + assert "event: speech.audio.done" in body + data_line = next(line for line in body.splitlines() if line.startswith("data: ")) + payload = json.loads(data_line.removeprefix("data: ")) + assert payload["type"] == "speech.audio.delta" + assert payload["response_format"] == response_format + assert base64.b64decode(payload["audio"]) + def test_streaming(self, streaming_app): - """stream=True must return audio/pcm with non-empty body.""" + """stream=True defaults to OpenAI speech.audio.* SSE events.""" client = TestClient(streaming_app) response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "pcm"}) + self._assert_sse_audio_response(response) + + def test_stream_format_audio_streaming(self, streaming_app): + """stream_format=audio without stream=True returns raw audio/pcm chunks.""" + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "audio", "response_format": "pcm"}, + ) + assert response.status_code == 200 + assert "audio/pcm" in response.headers["content-type"] + assert "text/event-stream" not in response.headers["content-type"] + assert len(response.content) > 0 + + def test_sse_streaming(self, streaming_app): + """stream_format=sse without stream=True returns audio deltas as SSE.""" + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "sse", "response_format": "pcm"}, + ) + + self._assert_sse_audio_response(response) + + def test_stream_true_with_stream_format_sse_uses_sse(self, streaming_app): + """stream=True and stream_format=sse both select SSE streaming.""" + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream": True, "stream_format": "sse", "response_format": "pcm"}, + ) + + self._assert_sse_audio_response(response) + + def test_stream_format_audio_with_stream_true_opts_into_raw_audio(self, streaming_app): + """stream_format=audio remains an explicit raw audio opt-in.""" + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream": True, "stream_format": "audio", "response_format": "pcm"}, + ) + assert response.status_code == 200 assert "audio/pcm" in response.headers["content-type"] + assert "text/event-stream" not in response.headers["content-type"] assert len(response.content) > 0 + def test_sse_rejects_unsupported_response_format(self, streaming_app): + """stream_format=sse with a non-pcm/wav format must fail before streaming starts.""" + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "sse", "response_format": "mp3"}, + ) + + assert response.status_code in (400, 422) + assert "text/event-stream" not in response.headers.get("content-type", "") + + def test_sse_rejects_speed_adjustment(self, streaming_app): + """stream_format=sse with speed != 1.0 must fail before streaming starts.""" + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "sse", "response_format": "pcm", "speed": 2.0}, + ) + + assert response.status_code in (400, 422) + assert "text/event-stream" not in response.headers.get("content-type", "") + + def test_stream_format_audio_rejects_unsupported_response_format(self, streaming_app): + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "audio", "response_format": "mp3"}, + ) + + assert response.status_code in (400, 422) + assert "audio/" not in response.headers.get("content-type", "") + + def test_stream_format_audio_rejects_speed_adjustment(self, streaming_app): + client = TestClient(streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "audio", "response_format": "pcm", "speed": 2.0}, + ) + + assert response.status_code in (400, 422) + assert "audio/" not in response.headers.get("content-type", "") + + @pytest.fixture + def erroring_streaming_app(self, mocker: MockerFixture): + """Test app whose mock engine raises mid-stream, to exercise the SSE error event.""" + + async def mock_generate_streaming(*args, **kwargs): + raise RuntimeError("boom: simulated engine failure") + yield # pragma: no cover - generator marker, unreachable + + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming) + mock_engine_client.default_sampling_params_list = [{}] + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + + speech_server = OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + + original_create_speech = speech_server.create_speech + sig = signature(original_create_speech) + new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"] + new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation) + + async def awaitable_create_speech(*args, **kwargs): + return await original_create_speech(*args, **kwargs) + + awaitable_create_speech.__signature__ = new_sig + speech_server.create_speech = awaitable_create_speech + + app = FastAPI() + app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None) + return app + + def test_sse_emits_error_event_on_generator_failure(self, erroring_streaming_app): + """An exception inside the SSE generator must surface as a speech.audio.error event.""" + client = TestClient(erroring_streaming_app) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream_format": "sse", "response_format": "pcm"}, + ) + + assert response.status_code == 200 + assert "text/event-stream" in response.headers["content-type"] + body = response.text + assert "event: speech.audio.error" in body + error_line = next(line for line in body.splitlines() if line.startswith("data: ")) + payload = json.loads(error_line.removeprefix("data: ")) + assert payload["type"] == "speech.audio.error" + assert "error" in payload + assert payload["error"]["message"] + def test_non_streaming_unchanged(self, streaming_app): """Non-streaming path must still return audio/wav.""" client = TestClient(streaming_app) @@ -2353,7 +2664,7 @@ def test_api_server_delete_voice_exception_returns_500(mocker: MockerFixture): def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixture): fake_base = _patch_api_server_base(mocker) raw_request = _make_api_server_request(None, path="/v1/audio/speech") - raw_request.app.state.openai_serving_tokenization = fake_base + raw_request.app.state.serving_tokenization = fake_base request = OpenAICreateSpeechRequest(input="Hello") response = asyncio.run(api_server_module.create_speech(request, raw_request)) @@ -2366,7 +2677,7 @@ def test_api_server_create_speech_without_handler_returns_404(mocker: MockerFixt def test_api_server_create_speech_batch_without_handler_returns_404(mocker: MockerFixture): fake_base = _patch_api_server_base(mocker) raw_request = _make_api_server_request(None, path="/v1/audio/speech/batch") - raw_request.app.state.openai_serving_tokenization = fake_base + raw_request.app.state.serving_tokenization = fake_base request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi")]) response = asyncio.run(api_server_module.create_speech_batch(request, raw_request)) @@ -2376,11 +2687,66 @@ def test_api_server_create_speech_batch_without_handler_returns_404(mocker: Mock ) +def test_api_server_create_speech_batch_omits_null_fields(mocker: MockerFixture): + # The batch response must omit optional null fields rather than serialize them + # as null (issue #4646 follow-up): errored items drop usage/audio_data/media_type, + # successful items drop error. This is the shape documented in speech_api.md. + from vllm_omni.entrypoints.openai.protocol.audio import ( + BatchSpeechResponse, + SpeechBatchItemResult, + SpeechInputTokenDetails, + SpeechTokenUsage, + ) + + handler = mocker.MagicMock() + handler.create_speech_batch = mocker.AsyncMock( + return_value=BatchSpeechResponse( + id="speech-batch-test", + results=[ + SpeechBatchItemResult( + index=0, + status="success", + audio_data="YWJj", + media_type="audio/wav", + usage=SpeechTokenUsage( + input_tokens=119, + output_tokens=77, + total_tokens=196, + input_token_details=SpeechInputTokenDetails(text_tokens=18, audio_tokens=101), + ), + ), + SpeechBatchItemResult(index=1, status="error", error="Input text cannot be empty"), + ], + total=2, + succeeded=1, + failed=1, + ) + ) + raw_request = _make_api_server_request(handler, path="/v1/audio/speech/batch") + request = BatchSpeechRequest(items=[SpeechBatchItem(input="hi"), SpeechBatchItem(input="")]) + + response = asyncio.run(api_server_module.create_speech_batch(request, raw_request)) + + assert isinstance(response, JSONResponse) + assert response.status_code == 200 + body = json.loads(response.body) + success, errored = body["results"][0], body["results"][1] + # Successful item carries usage and drops the null `error`. + assert success["usage"]["total_tokens"] == 196 + assert success["usage"]["input_token_details"] == {"text_tokens": 18, "audio_tokens": 101} + assert "error" not in success + # Errored item drops usage/audio_data/media_type instead of serializing null. + assert "usage" not in errored + assert "audio_data" not in errored + assert "media_type" not in errored + assert errored["error"] == "Input text cannot be empty" + + def test_api_server_create_audio_generate_without_handler_returns_404(mocker: MockerFixture): fake_base = _patch_api_server_base(mocker) raw_request = _make_api_server_request(None, path="/v1/audio/generate") raw_request.app.state.openai_serving_audio_generate = None - raw_request.app.state.openai_serving_tokenization = fake_base + raw_request.app.state.serving_tokenization = fake_base request = OpenAICreateAudioGenerateRequest(input="a bird singing") response = asyncio.run(api_server_module.create_audio_generate(request, raw_request)) @@ -2807,7 +3173,10 @@ async def awaitable_create_speech(*args, **kwargs): def test_wav_streaming_success(self, wav_streaming_app): """Test WAV format streaming returns correct content type and includes WAV header.""" client = TestClient(wav_streaming_app) - response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "wav"}) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream": True, "stream_format": "audio", "response_format": "wav"}, + ) assert response.status_code == 200 assert "audio/wav" in response.headers["content-type"] @@ -2826,7 +3195,10 @@ def test_streaming_unsupported_format_rejected(self, wav_streaming_app): unsupported_formats = ["mp3"] for fmt in unsupported_formats: - response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": fmt}) + response = client.post( + "/v1/audio/speech", + json={"input": "Hello", "stream": True, "stream_format": "audio", "response_format": fmt}, + ) assert response.status_code == 422 @@ -2878,25 +3250,25 @@ def test_cosyvoice3_stage_registered(self): def test_validate_cosyvoice3_empty_input(self, cosyvoice3_server): request = OpenAICreateSpeechRequest(input="", ref_audio="data:audio/wav;base64,abc", ref_text="hello") - error = cosyvoice3_server._validate_cosyvoice3_request(request) + error = cosyvoice3_server._validate_tts_request(request) assert error is not None assert "empty" in error.lower() def test_validate_cosyvoice3_missing_ref_audio(self, cosyvoice3_server): request = OpenAICreateSpeechRequest(input="Hello", ref_text="hello") - error = cosyvoice3_server._validate_cosyvoice3_request(request) + error = cosyvoice3_server._validate_tts_request(request) assert error is not None assert "ref_audio" in error.lower() def test_validate_cosyvoice3_missing_ref_text(self, cosyvoice3_server): request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc") - error = cosyvoice3_server._validate_cosyvoice3_request(request) + error = cosyvoice3_server._validate_tts_request(request) assert error is not None assert "ref_text" in error.lower() def test_validate_cosyvoice3_invalid_ref_audio_format(self, cosyvoice3_server): request = OpenAICreateSpeechRequest(input="Hello", ref_audio="/local/path.wav", ref_text="hello") - error = cosyvoice3_server._validate_cosyvoice3_request(request) + error = cosyvoice3_server._validate_tts_request(request) assert error is not None assert "url" in error.lower() or "format" in error.lower() @@ -2906,7 +3278,7 @@ def test_validate_cosyvoice3_valid_request(self, cosyvoice3_server): ref_audio="data:audio/wav;base64,abc123", ref_text="Reference transcript", ) - error = cosyvoice3_server._validate_cosyvoice3_request(request) + error = cosyvoice3_server._validate_tts_request(request) assert error is None def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server): @@ -2916,7 +3288,7 @@ def test_validate_cosyvoice3_max_new_tokens_range(self, cosyvoice3_server): ref_text="hello", max_new_tokens=0, ) - error = cosyvoice3_server._validate_cosyvoice3_request(request) + error = cosyvoice3_server._validate_tts_request(request) assert error is not None assert "max_new_tokens" in error @@ -2982,13 +3354,13 @@ def glm_tts_server(mocker: MockerFixture): class TestGLMTTSServing: def test_validate_glm_tts_requires_ref_audio(self, glm_tts_server): request = OpenAICreateSpeechRequest(input="Hello", ref_text="Reference transcript") - error = glm_tts_server._validate_glm_tts_request(request) + error = glm_tts_server._validate_tts_request(request) assert error is not None assert "ref_audio" in error def test_validate_glm_tts_requires_ref_text(self, glm_tts_server): request = OpenAICreateSpeechRequest(input="Hello", ref_audio="data:audio/wav;base64,abc") - error = glm_tts_server._validate_glm_tts_request(request) + error = glm_tts_server._validate_tts_request(request) assert error is not None assert "ref_text" in error @@ -3095,7 +3467,7 @@ def test_prepare_speech_generation_awaits_voxtral_async(self, voxtral_server, mo def test_prepare_speech_generation_awaits_qwen3_tts_async(self, qwen3_tts_server, mocker: MockerFixture): """Qwen3 TTS path should call _estimate_prompt_len_async.""" - qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) qwen3_tts_server._build_tts_params = mocker.MagicMock( return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} ) @@ -3105,6 +3477,154 @@ def test_prepare_speech_generation_awaits_qwen3_tts_async(self, qwen3_tts_server 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_default_seed_sets_tts_local_seed( + self, qwen3_tts_server, mocker: MockerFixture + ): + """Deploy default seed should seed Qwen3 TTS residual MTP sampling.""" + qwen3_tts_server.engine_client.default_sampling_params_list = [ + SimpleNamespace(max_tokens=2048, seed=42, extra_args=None) + ] + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + request = OpenAICreateSpeechRequest(input="hello") + + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + stage0_params = qwen3_tts_server.engine_client.generate.call_args.kwargs["sampling_params_list"][0] + assert stage0_params.seed == 42 + assert stage0_params.extra_args["tts_local_seed"] == 42 + assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None + + def test_prepare_speech_generation_uses_adapter_model_type_label( + self, + voxtral_server, + mocker: MockerFixture, + ): + """Adapter model_type should replace the legacy _tts_model_type label ladder.""" + legacy_tts_model_type = "dummy_tts" + adapter_model_type = "adapter_dummy_tts" + + class FakeAdapter: + def validate(self, request): + return None + + async def build(self, request, sampling_params_list, has_inline_ref_audio): + return PreparedRequest( + prompt={"prompt": request.input}, + tts_params={}, + model_type=adapter_model_type, + ) + + voxtral_server._tts_model_type = legacy_tts_model_type + mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter()) + log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info") + + asyncio.run(voxtral_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="hello"))) + + assert adapter_model_type != legacy_tts_model_type + assert any( + call.args + and call.args[0] == "TTS speech request %s: text=%r, model=%s" + and call.args[3] == adapter_model_type + for call in log_info.call_args_list + ) + + def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture): + """stream_format=sse should request delta-style multimodal outputs.""" + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=qwen3_tts_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", stream_format="sse") + + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True) + + def test_prepare_speech_generation_treats_stream_true_as_streaming(self, qwen3_tts_server, mocker: MockerFixture): + """stream=True should request delta-style multimodal outputs for SSE streaming.""" + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=qwen3_tts_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", stream=True, response_format="pcm") + + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True) + + def test_prepare_speech_generation_treats_audio_as_streaming(self, qwen3_tts_server, mocker: MockerFixture): + """stream_format=audio should request delta-style multimodal outputs.""" + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=qwen3_tts_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm") + + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, True) + + def test_prepare_speech_generation_no_async_chunk_stream_uses_final_only( + self, qwen3_tts_server, mocker: MockerFixture + ): + """Full-payload TTS streaming should not request delta multimodal outputs.""" + qwen3_tts_server.engine_client.model_config.async_chunk = False + qwen3_tts_server._adapter.validate = mocker.MagicMock(return_value=None) + qwen3_tts_server._build_tts_params = mocker.MagicMock( + return_value={"text": ["hello"], "task_type": ["CustomVoice"], "speaker": ["Vivian"]} + ) + qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=qwen3_tts_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", stream_format="audio", response_format="pcm") + + asyncio.run(qwen3_tts_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(qwen3_tts_server.engine_client.default_sampling_params_list, False) + + def test_prepare_speech_generation_no_async_chunk_stream_keeps_delta_for_non_qwen3( + self, voxtral_server, mocker: MockerFixture + ): + """FINAL_ONLY streaming for async_chunk=False is scoped to qwen3_tts only.""" + voxtral_server.engine_client.model_config.async_chunk = False + mocker.patch.object(voxtral_server._get_tts_adapter(), "validate", return_value=None) + voxtral_server._build_voxtral_prompt_async = mocker.AsyncMock( + return_value={ + "prompt_token_ids": [1, 2, 3], + "additional_information": {"voice": ["test"]}, + } + ) + mock_coerce = mocker.patch( + "vllm_omni.entrypoints.openai.serving_speech.coerce_param_message_types", + return_value=voxtral_server.engine_client.default_sampling_params_list, + ) + request = OpenAICreateSpeechRequest(input="hello", voice="test", stream_format="audio", response_format="pcm") + + asyncio.run(voxtral_server._prepare_speech_generation(request)) + + mock_coerce.assert_called_once_with(voxtral_server.engine_client.default_sampling_params_list, True) + def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false( self, qwen3_tts_server, mocker: MockerFixture ): @@ -3162,7 +3682,7 @@ def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self 24000, artifact_key, ) - qwen3_tts_server._ref_audio_model_artifact_ready.add(artifact_key) + qwen3_tts_server._ref_audio_model_artifact_ready.add((artifact_key, False)) qwen3_tts_server._codec_frame_rate = 25.0 qwen3_tts_server._tts_tokenizer = lambda _text, padding=False: {"input_ids": list(range(10))} qwen3_tts_server.engine_client.model_config.hf_config.talker_config = SimpleNamespace( @@ -3192,10 +3712,10 @@ def test_qwen3_ref_audio_artifact_ready_is_evicted_with_resolve_cache(self, qwen qwen3_tts_server._ref_audio_resolve_cache_max_bytes = 1_000_000 qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") - qwen3_tts_server._ref_audio_model_artifact_ready.add("artifact-a") + qwen3_tts_server._ref_audio_model_artifact_ready.add(("artifact-a", False)) qwen3_tts_server._put_resolved_ref_audio("ref-b", [0.0] * 8, 24000, "artifact-b") - assert "artifact-a" not in qwen3_tts_server._ref_audio_model_artifact_ready + assert ("artifact-a", False) not in qwen3_tts_server._ref_audio_model_artifact_ready assert "artifact-b" in {entry[3] for entry in qwen3_tts_server._ref_audio_resolve_cache.values()} @pytest.mark.asyncio @@ -3204,13 +3724,13 @@ async def failing_generator(): raise ValueError("boom") yield # pragma: no cover - qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = "artifact-fail" + qwen3_tts_server._request_ref_audio_artifact_keys["req-fail"] = ("artifact-fail", False) with pytest.raises(ValueError, match="boom"): await anext(qwen3_tts_server._generate_audio_chunks(failing_generator(), "req-fail")) assert "req-fail" not in qwen3_tts_server._request_ref_audio_artifact_keys - assert "artifact-fail" not in qwen3_tts_server._ref_audio_model_artifact_ready + assert ("artifact-fail", False) not in qwen3_tts_server._ref_audio_model_artifact_ready @pytest.mark.asyncio async def test_generate_audio_chunks_discards_ref_audio_artifact_warmup_on_close(self, qwen3_tts_server): @@ -3223,22 +3743,48 @@ async def pcm_generator(): ) await asyncio.sleep(0) - qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = "artifact-close" + qwen3_tts_server._request_ref_audio_artifact_keys["req-close"] = ("artifact-close", False) stream = qwen3_tts_server._generate_audio_chunks(pcm_generator(), "req-close") assert await anext(stream) await stream.aclose() assert "req-close" not in qwen3_tts_server._request_ref_audio_artifact_keys - assert "artifact-close" not in qwen3_tts_server._ref_audio_model_artifact_ready + assert ("artifact-close", False) not in qwen3_tts_server._ref_audio_model_artifact_ready def test_qwen3_ref_audio_artifact_ready_requires_live_resolve_cache_entry(self, qwen3_tts_server): - qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = "artifact-evicted" + qwen3_tts_server._request_ref_audio_artifact_keys["req-evicted"] = ("artifact-evicted", False) qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-evicted") assert "req-evicted" not in qwen3_tts_server._request_ref_audio_artifact_keys - assert "artifact-evicted" not in qwen3_tts_server._ref_audio_model_artifact_ready + assert ("artifact-evicted", False) not in qwen3_tts_server._ref_audio_model_artifact_ready + + def test_qwen3_xvector_ready_artifact_does_not_enable_icl_artifact_only(self, qwen3_tts_server): + # An x-vector-only artifact (speaker embedding, no ref_code) must not enable + # the artifact-only path for a later ICL request with the same ref_audio (#5049). + qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") + qwen3_tts_server._track_ref_audio_artifact_warmup("req-xvec", "artifact-a", x_vector_only=True) + qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-xvec") + + icl_params = {"task_type": ["Base"], "x_vector_only_mode": [False]} + assert qwen3_tts_server._qwen3_tts_can_use_ref_audio_artifact_only(icl_params, "artifact-a") is False + + def test_qwen3_xvector_ready_artifact_still_reusable_by_xvector_request(self, qwen3_tts_server): + qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") + qwen3_tts_server._track_ref_audio_artifact_warmup("req-xvec", "artifact-a", x_vector_only=True) + qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-xvec") + + xvec_params = {"task_type": ["Base"], "x_vector_only_mode": [True]} + assert qwen3_tts_server._qwen3_tts_can_use_ref_audio_artifact_only(xvec_params, "artifact-a") is True + + def test_qwen3_icl_ready_artifact_enables_icl_artifact_only(self, qwen3_tts_server): + qwen3_tts_server._put_resolved_ref_audio("ref-a", [0.0] * 8, 24000, "artifact-a") + qwen3_tts_server._track_ref_audio_artifact_warmup("req-icl", "artifact-a", x_vector_only=False) + qwen3_tts_server._mark_ref_audio_artifact_ready_for_request("req-icl") + + icl_params = {"task_type": ["Base"], "x_vector_only_mode": [False]} + assert qwen3_tts_server._qwen3_tts_can_use_ref_audio_artifact_only(icl_params, "artifact-a") is True def test_shutdown_is_idempotent(self, mocker: MockerFixture): """Calling shutdown() twice should not raise.""" diff --git a/tests/entrypoints/openai_api/test_serving_speech_batch.py b/tests/entrypoints/openai_api/test_serving_speech_batch.py index 054100ce843..686dba7d00d 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_batch.py +++ b/tests/entrypoints/openai_api/test_serving_speech_batch.py @@ -88,7 +88,7 @@ class TestSpeechBatchE2E: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_basic_two_items(self, omni_server) -> None: """Batch with two items returns two successful base64-encoded results.""" @@ -115,14 +115,17 @@ def test_batch_basic_two_items(self, omni_server) -> None: assert result["index"] == i assert result["status"] == "success" assert result["audio_data"] is not None - assert result["error"] is None + # Successful items omit `error` entirely (the batch response is + # serialized with exclude_none), so probe with .get rather than + # indexing a key that is only present on errored items. + assert result.get("error") is None audio_bytes = base64.b64decode(result["audio_data"]) assert verify_wav_bytes(audio_bytes), f"Item {i}: invalid WAV" assert len(audio_bytes) > MIN_AUDIO_BYTES, f"Item {i}: audio too small ({len(audio_bytes)} bytes)" @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_single_item(self, omni_server) -> None: """Batch with a single item works correctly.""" @@ -144,7 +147,7 @@ def test_batch_single_item(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_per_item_voice_override(self, omni_server) -> None: """Per-item voice overrides the batch-level default.""" @@ -172,7 +175,7 @@ def test_batch_per_item_voice_override(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_multiple_languages(self, omni_server) -> None: """Batch items with different languages per item.""" @@ -197,7 +200,7 @@ def test_batch_multiple_languages(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_whisper_transcription(self, omni_server) -> None: """Whisper transcription of batch output matches input text.""" @@ -232,7 +235,7 @@ def test_batch_whisper_transcription(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_result_indices_ordered(self, omni_server) -> None: """Result indices match the input item order.""" @@ -259,7 +262,7 @@ class TestSpeechBatchValidation: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_empty_items_rejected(self, omni_server) -> None: """Empty items list returns a 4xx error.""" @@ -273,7 +276,7 @@ def test_batch_empty_items_rejected(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", default_server_params, indirect=True) def test_batch_exceeds_max_items(self, omni_server) -> None: """Batch exceeding 32 items returns 400 error.""" @@ -331,7 +334,7 @@ class TestSpeechBatchSize2: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", batch2_server_params, indirect=True) def test_batch2_produces_valid_audio(self, omni_server) -> None: """Batch of 2 items with batched engine produces valid audio.""" @@ -368,7 +371,7 @@ def test_batch2_produces_valid_audio(self, omni_server) -> None: @pytest.mark.advanced_model @pytest.mark.tts - @hardware_test(res={"cuda": "H100"}, num_cards=1) + @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", batch2_server_params, indirect=True) def test_batch2_vs_sequential_timing(self, omni_server) -> None: """Compare batch-of-2 vs 2 sequential single requests. diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 1b93ef58e24..8817173633e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -1,4 +1,6 @@ import asyncio +import base64 +from types import SimpleNamespace import pytest from fastapi import FastAPI, WebSocket @@ -9,6 +11,7 @@ from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler +from vllm_omni.utils.forced_aligner import WordTimestamp pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -25,10 +28,11 @@ def _build_test_app( speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav")) speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {})) + speech_service.forced_aligner_config = None - async def mock_generate_pcm_chunks(_generator, _request_id): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): for chunk in (b"\x01\x02", b"\x03\x04\x05"): - yield chunk + yield (chunk, 24000) if include_sample_rate else chunk speech_service._generate_pcm_chunks = mock_generate_pcm_chunks speech_service.engine_client = mocker.MagicMock() @@ -56,6 +60,7 @@ def test_non_streaming_single_frame(self, mocker: MockerFixture): with client.websocket_connect("/v1/audio/speech/stream") as ws: ws.send_json({"type": "session.config", "voice": "Vivian"}) ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) start = ws.receive_json() assert start["type"] == "audio.start" @@ -69,7 +74,6 @@ def test_non_streaming_single_frame(self, mocker: MockerFixture): done = ws.receive_json() assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False} - ws.send_json({"type": "input.done"}) session_done = ws.receive_json() assert session_done == {"type": "session.done", "total_sentences": 1} @@ -82,6 +86,7 @@ def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_config = None async def mock_prepare_speech_generation(request): captured_requests.append(request) @@ -89,9 +94,9 @@ async def mock_prepare_speech_generation(request): speech_service._prepare_speech_generation = mock_prepare_speech_generation - async def mock_generate_pcm_chunks(_generator, _request_id): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"): - yield chunk + yield (chunk, 24000) if include_sample_rate else chunk speech_service._generate_pcm_chunks = mock_generate_pcm_chunks app, _ = _build_test_app(speech_service) @@ -108,6 +113,7 @@ async def mock_generate_pcm_chunks(_generator, _request_id): } ) ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) start = ws.receive_json() assert start["type"] == "audio.start" @@ -121,7 +127,6 @@ async def mock_generate_pcm_chunks(_generator, _request_id): done = ws.receive_json() assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False} - ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} assert len(captured_requests) == 1 @@ -130,6 +135,167 @@ async def mock_generate_pcm_chunks(_generator, _request_id): assert captured_requests[0].initial_codec_chunk_frames == 12 assert speech_service._generate_audio_bytes.await_count == 0 + def test_word_timestamps_requires_configured_aligner(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "stream_audio": True, + "response_format": "pcm", + "word_timestamps": True, + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + error = ws.receive_json() + assert error["type"] == "error" + assert "without --forced-aligner" in error["message"] + + def test_word_timestamps_emit_sidecar_json_frame(self, mocker: MockerFixture): + captured_requests = [] + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_config = SimpleNamespace(model="aligner") + + async def mock_prepare_speech_generation(request): + captured_requests.append(request) + return "req-stream", object(), {} + + speech_service._prepare_speech_generation = mock_prepare_speech_generation + + first_chunk = b"\x01" * 1000 + second_chunk = b"\x02" * 1000 + + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + for chunk in (first_chunk, second_chunk): + yield (chunk, 1000) if include_sample_rate else chunk + + speech_service._generate_pcm_chunks = mock_generate_pcm_chunks + # Sentence-level: aligner runs once over the whole sentence audio. + mock_align = mocker.AsyncMock( + return_value=[ + WordTimestamp("Hello", 0, 200), + WordTimestamp("world", 200, 900), + ] + ) + mocker.patch.object(streaming_speech_module, "forced_align", mock_align) + app, _ = _build_test_app(speech_service) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "stream_audio": True, + "response_format": "pcm", + "word_timestamps": True, + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + start = ws.receive_json() + assert start["type"] == "audio.start" + assert start["word_timestamps"] is True + + # Audio streams first; timestamps are null until the sentence + # is fully aligned. + chunk = ws.receive_json() + assert chunk["type"] == "audio.chunk" + assert chunk["chunk_id"] == 0 + assert chunk["chunk_start_ms"] == 0 + assert chunk["chunk_end_ms"] == 500 + assert chunk["sample_rate"] == 1000 + assert base64.b64decode(chunk["audio_b64"]) == first_chunk + assert chunk["timestamps"] is None + + chunk = ws.receive_json() + assert chunk["type"] == "audio.chunk" + assert chunk["chunk_id"] == 1 + assert chunk["chunk_start_ms"] == 500 + assert chunk["chunk_end_ms"] == 1000 + assert chunk["sample_rate"] == 1000 + assert base64.b64decode(chunk["audio_b64"]) == second_chunk + assert chunk["timestamps"] is None + + # Final frame: empty audio carrying the whole-sentence timestamps. + chunk = ws.receive_json() + assert chunk["type"] == "audio.chunk" + assert chunk["chunk_id"] == 2 + assert chunk["chunk_start_ms"] == 0 + assert chunk["chunk_end_ms"] == 1000 + assert chunk["sample_rate"] == 1000 + assert base64.b64decode(chunk["audio_b64"]) == b"" + assert chunk["timestamps"] == [ + {"word": "Hello", "start_ms": 0, "end_ms": 200}, + {"word": "world", "start_ms": 200, "end_ms": 900}, + ] + + done = ws.receive_json() + assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False} + + assert captured_requests[0].word_timestamps is True + assert mock_align.await_count == 1 + + def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture): + # The streaming layer forwards the aligner's (already monotonic, + # non-overlapping) words as JSON dicts in the trailing frame. + speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) + speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav")) + speech_service.engine_client = mocker.MagicMock() + speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_config = SimpleNamespace(model="aligner") + speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {})) + + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): + chunk = b"\x01" * 1000 + yield (chunk, 1000) if include_sample_rate else chunk + + speech_service._generate_pcm_chunks = mock_generate_pcm_chunks + mock_align = mocker.AsyncMock( + return_value=[ + WordTimestamp("Hello", 0, 1000), + WordTimestamp("world", 1000, 1200), + ] + ) + mocker.patch.object(streaming_speech_module, "forced_align", mock_align) + app, _ = _build_test_app(speech_service) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json( + { + "type": "session.config", + "voice": "Vivian", + "stream_audio": True, + "response_format": "pcm", + "word_timestamps": True, + } + ) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + # Real-time audio chunk (timestamps null), then the timestamp frame. + assert ws.receive_json()["timestamps"] is None + final = ws.receive_json() + assert final["type"] == "audio.chunk" + timestamps = final["timestamps"] + assert timestamps == [ + {"word": "Hello", "start_ms": 0, "end_ms": 1000}, + {"word": "world", "start_ms": 1000, "end_ms": 1200}, + ] + for ts in timestamps: + assert ts["end_ms"] >= ts["start_ms"] + def test_flush_on_input_done(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -179,16 +345,19 @@ def test_empty_input_text_emits_no_audio(self, mocker: MockerFixture): assert speech_service._generate_audio_bytes.await_count == 0 - def test_multiple_sentences_increment_indices(self, mocker: MockerFixture): + def test_multiple_sentences_are_buffered_into_one_request(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: ws.send_json({"type": "session.config", "voice": "Vivian"}) - ws.send_json({"type": "input.text", "text": "First sentence. Second sentence. "}) + ws.send_json({"type": "input.text", "text": "First sentence. "}) + ws.send_json({"type": "input.text", "text": "Second sentence. "}) + ws.send_json({"type": "input.done"}) - first_start = ws.receive_json() - assert first_start["sentence_index"] == 0 + start = ws.receive_json() + assert start["sentence_index"] == 0 + assert start["sentence_text"] == "First sentence. Second sentence." ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", @@ -196,19 +365,7 @@ def test_multiple_sentences_increment_indices(self, mocker: MockerFixture): "total_bytes": 36, "error": False, } - - second_start = ws.receive_json() - assert second_start["sentence_index"] == 1 - ws.receive_bytes() - assert ws.receive_json() == { - "type": "audio.done", - "sentence_index": 1, - "total_bytes": 36, - "error": False, - } - - ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 2} + assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -222,6 +379,7 @@ def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture): assert error == {"type": "error", "message": "Unknown message type: unknown"} ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) assert ws.receive_json()["type"] == "audio.start" ws.receive_bytes() assert ws.receive_json() == { @@ -231,7 +389,6 @@ def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture): "error": False, } - ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} def test_config_timeout_closes_session(self, mocker: MockerFixture): @@ -249,18 +406,19 @@ def test_generation_error_marks_audio_done(self, mocker: MockerFixture): speech_service._generate_pcm_chunks = mocker.AsyncMock() speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_config = None app, _ = _build_test_app(speech_service) with TestClient(app) as client: with client.websocket_connect("/v1/audio/speech/stream") as ws: ws.send_json({"type": "session.config", "voice": "Vivian"}) ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) assert ws.receive_json()["type"] == "audio.start" assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"} assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True} - ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture): @@ -269,8 +427,9 @@ def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {})) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_config = None - async def mock_generate_pcm_chunks(_generator, _request_id): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): yield b"\x01\x02" raise RuntimeError("stream boom") @@ -288,6 +447,7 @@ async def mock_generate_pcm_chunks(_generator, _request_id): } ) ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) assert ws.receive_json()["type"] == "audio.start" assert ws.receive_bytes() == b"\x01\x02" @@ -302,7 +462,6 @@ async def mock_generate_pcm_chunks(_generator, _request_id): "error": True, } - ws.send_json({"type": "input.done"}) assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture): @@ -361,8 +520,9 @@ def test_disconnect_aborts_streaming_request(self, mocker: MockerFixture): speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {})) speech_service.engine_client = mocker.MagicMock() speech_service.engine_client.abort = mocker.AsyncMock() + speech_service.forced_aligner_config = None - async def mock_generate_pcm_chunks(_generator, _request_id): + async def mock_generate_pcm_chunks(_generator, _request_id, *, include_sample_rate=False): yield b"\x01\x02" speech_service._generate_pcm_chunks = mock_generate_pcm_chunks @@ -385,7 +545,9 @@ async def mock_generate_pcm_chunks(_generator, _request_id): config.ref_audio = None config.ref_text = None config.x_vector_only_mode = None + config.speaker_embedding = None config.stream_audio = True + config.word_timestamps = False with pytest.raises(WebSocketDisconnect): asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0)) diff --git a/tests/entrypoints/openai_api/test_serving_video_output_stream.py b/tests/entrypoints/openai_api/test_serving_video_output_stream.py new file mode 100644 index 00000000000..72d293853cb --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_video_output_stream.py @@ -0,0 +1,446 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for `/v1/realtime/video` WebSocket video output streaming.""" + +from __future__ import annotations + +import asyncio +import json +from collections.abc import AsyncGenerator, Callable +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np +import pytest +from fastapi import FastAPI, WebSocket +from pytest_mock import MockerFixture +from starlette.testclient import TestClient + +from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _fake_video_frames(num_frames: int = 2) -> list[np.ndarray]: + return [np.full((4, 4, 3), fill_value=i, dtype=np.uint8) for i in range(num_frames)] + + +def _build_test_app( + *, + mocker: MockerFixture, + streaming_chunks: list[tuple[bytes, bool]], + mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]], + final_chunk: bytes = b"", + stall_timeout: float = 5.0, +) -> tuple[FastAPI, OmniStreamingVideoOutputHandler, MagicMock]: + engine_client = mocker.MagicMock() + engine_client.abort = mocker.AsyncMock() + engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()] + engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] + engine_client.generate = mock_generate + + class FakeStreamingVideoEncoder: + encode_calls: list[int] = [] + + def encode(self, video): + del video + idx = len(self.encode_calls) + self.encode_calls.append(idx) + return streaming_chunks[idx][0] + + def close(self): + return final_chunk + + def _make_encoder(*, output_format, fps, video_codec_options=None): + del output_format, fps, video_codec_options + return FakeStreamingVideoEncoder() + + encoder_factory = mocker.MagicMock(side_effect=_make_encoder) + engine_client.streaming_encoder_factory = encoder_factory + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", + encoder_factory, + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", + return_value="diffusion", + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + + handler = OmniStreamingVideoOutputHandler( + engine_client=engine_client, + model_name="test-model", + stage_configs=engine_client.stage_configs, + stall_timeout=stall_timeout, + start_timeout=5.0, + ) + app = FastAPI() + + @app.websocket("/v1/realtime/video") + async def ws_endpoint(websocket: WebSocket): + await handler.handle_session(websocket) + + return app, handler, engine_client + + +class TestStreamingVideoOutputWebSocket: + """WebSocket protocol tests for streaming generated video output.""" + + def test_streaming_session_emits_video_start_binary_chunks_and_done(self, mocker: MockerFixture): + """A full session delivers video.start, binary chunks, then session.done.""" + + async def mock_generate(*_args, **_kwargs): + for frames, finished in [ + (_fake_video_frames(2), False), + (_fake_video_frames(3), True), + ]: + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[frames], + final_output_type="image", + finished=finished, + ) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)], + mock_generate=mock_generate, + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"}) + + start = ws.receive_json() + assert start["type"] == "video.start" + assert start["format"] == "m4s" + assert "format" not in start["config"] + assert "request_id" in start + + assert ws.receive_bytes() == b"mp4-chunk-0" + assert ws.receive_bytes() == b"mp4-chunk-1" + + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["chunks"] == 2 + assert done["stopped"] is False + + engine_client.abort.assert_not_called() + engine_client.streaming_encoder_factory.assert_called_once() + assert engine_client.streaming_encoder_factory.call_args.kwargs["output_format"] == "m4s" + + def test_streaming_session_rejects_unknown_format(self, mocker: MockerFixture): + async def mock_generate(*_args, **_kwargs): + if False: + yield None + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[], + mock_generate=mock_generate, + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "hello", "format": "webm"}) + err = ws.receive_json() + assert err["type"] == "error" + assert "format" in err["message"] + + engine_client.abort.assert_not_called() + + def test_streaming_session_emits_final_encoder_delta_before_done(self, mocker: MockerFixture): + """A non-empty encoder close delta is delivered as the final binary chunk.""" + + async def mock_generate(*_args, **_kwargs): + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=True, + ) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", True)], + mock_generate=mock_generate, + final_chunk=b"mp4-trailer", + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "A cat walking in the rain"}) + assert ws.receive_json()["type"] == "video.start" + assert ws.receive_bytes() == b"mp4-chunk-0" + assert ws.receive_bytes() == b"mp4-trailer" + + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["chunks"] == 2 + assert done["stopped"] is False + + engine_client.abort.assert_not_called() + + def test_generation_error_output_returns_websocket_error(self, mocker: MockerFixture): + """An erroneous OmniRequestOutput from generate is sent to the client as a WebSocket error.""" + error_message = "diffusion engine exploded" + + async def mock_generate(*_args, **_kwargs): + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=False, + ) + yield OmniRequestOutput.from_error("req-ws", error_message) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", False)], + mock_generate=mock_generate, + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "hello"}) + assert ws.receive_json()["type"] == "video.start" + assert ws.receive_bytes() == b"mp4-chunk-0" + err = ws.receive_json() + assert err["type"] == "error" + assert error_message in err["message"] + + engine_client.abort.assert_not_called() + + +class _MockVideoOutputWebSocket: + """Feeds scripted client messages and records server sends.""" + + def __init__(self, messages: list[str]) -> None: + self._messages: asyncio.Queue[str] = asyncio.Queue() + for message in messages: + self._messages.put_nowait(message) + self.sent: list[dict | bytes] = [] + + def enqueue(self, message: str) -> None: + self._messages.put_nowait(message) + + async def accept(self) -> None: + return None + + async def receive_text(self) -> str: + return await self._messages.get() + + async def send_json(self, data: dict) -> None: + self.sent.append(data) + + async def send_bytes(self, data: bytes) -> None: + self.sent.append(data) + + @property + def sent_json(self) -> list[dict]: + return [m for m in self.sent if isinstance(m, dict)] + + +def _build_handler_for_async_tests( + mocker: MockerFixture, + *, + mock_generate: Callable[..., AsyncGenerator[OmniRequestOutput, None]], + streaming_chunks: list[tuple[bytes, bool]] | None = None, + stall_timeout: float = 5.0, +) -> tuple[OmniStreamingVideoOutputHandler, MagicMock, _MockVideoOutputWebSocket]: + engine_client = mocker.MagicMock() + engine_client.abort = mocker.AsyncMock() + engine_client.default_sampling_params_list = [OmniDiffusionSamplingParams()] + engine_client.stage_configs = [SimpleNamespace(stage_type="diffusion")] + engine_client.generate = mock_generate + + chunks = streaming_chunks or [(b"mp4-chunk-0", True)] + + class FakeStreamingVideoEncoder: + def __init__(self) -> None: + self.encode_calls = 0 + + def encode(self, video): + del video + idx = self.encode_calls + self.encode_calls += 1 + return chunks[idx][0] + + def close(self): + return b"" + + encoder_factory = mocker.MagicMock(side_effect=lambda **_: FakeStreamingVideoEncoder()) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.create_streaming_video_encoder", + encoder_factory, + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_stage_type", + return_value="diffusion", + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.build_stage_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video_output_stream.get_default_sampling_params_list", + return_value=[OmniDiffusionSamplingParams()], + ) + + handler = OmniStreamingVideoOutputHandler( + engine_client=engine_client, + model_name="test-model", + stage_configs=engine_client.stage_configs, + stall_timeout=stall_timeout, + start_timeout=5.0, + ) + start_msg = json.dumps({"type": "session.start", "prompt": "async test"}) + ws = _MockVideoOutputWebSocket([start_msg]) + return handler, engine_client, ws + + +class TestStreamingVideoOutputStallTimeout: + """Stall-timeout and session.ping behavior.""" + + def test_long_generation_without_client_messages_succeeds(self, mocker: MockerFixture): + """Engine progress between yields keeps the session alive without client pings.""" + + async def mock_generate(*_args, **_kwargs): + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=False, + ) + await asyncio.sleep(0.3) + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=True, + ) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", False), (b"mp4-chunk-1", True)], + mock_generate=mock_generate, + stall_timeout=0.5, + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "slow generation"}) + assert ws.receive_json()["type"] == "video.start" + assert ws.receive_bytes() == b"mp4-chunk-0" + assert ws.receive_bytes() == b"mp4-chunk-1" + done = ws.receive_json() + assert done["type"] == "session.done" + assert done["stopped"] is False + + engine_client.abort.assert_not_called() + + def test_stall_timeout_aborts_when_engine_silent(self, mocker: MockerFixture): + """No engine output for longer than stall_timeout aborts the request.""" + + async def mock_generate(*_args, **_kwargs): + await asyncio.sleep(1.0) + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=True, + ) + + app, _handler, engine_client = _build_test_app( + mocker=mocker, + streaming_chunks=[(b"mp4-chunk-0", True)], + mock_generate=mock_generate, + stall_timeout=0.2, + ) + + with TestClient(app) as client: + with client.websocket_connect("/v1/realtime/video") as ws: + ws.send_json({"type": "session.start", "prompt": "stalled"}) + assert ws.receive_json()["type"] == "video.start" + err = ws.receive_json() + assert err["type"] == "error" + assert "Stall timeout" in err["message"] + + engine_client.abort.assert_called() + + @pytest.mark.asyncio + async def test_session_ping_returns_pong(self, mocker: MockerFixture): + """session.ping during generation receives session.pong.""" + + async def mock_generate(*_args, **_kwargs): + await asyncio.sleep(0.5) + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=True, + ) + + handler, _engine_client, ws = _build_handler_for_async_tests( + mocker, + mock_generate=mock_generate, + stall_timeout=5.0, + ) + + async def _send_ping_after_delay() -> None: + await asyncio.sleep(0.05) + ws.enqueue(json.dumps({"type": "session.ping"})) + + ping_task = asyncio.create_task(_send_ping_after_delay()) + await handler.handle_session(ws) # type: ignore[arg-type] + await ping_task + + assert any(m.get("type") == "session.pong" for m in ws.sent_json) + + @pytest.mark.asyncio + async def test_ping_prevents_stall_during_silent_engine(self, mocker: MockerFixture): + """Client pings refresh the stall clock while the engine produces no output.""" + + release = asyncio.Event() + + async def mock_generate(*_args, **_kwargs): + await release.wait() + yield OmniRequestOutput.from_diffusion( + request_id="req-ws", + images=[_fake_video_frames(2)], + final_output_type="image", + finished=True, + ) + + handler, engine_client, ws = _build_handler_for_async_tests( + mocker, + mock_generate=mock_generate, + stall_timeout=0.3, + ) + + async def _pinger() -> None: + while not release.is_set(): + ws.enqueue(json.dumps({"type": "session.ping"})) + await asyncio.sleep(0.1) + + pinger_task = asyncio.create_task(_pinger()) + session_task = asyncio.create_task(handler.handle_session(ws)) # type: ignore[arg-type] + await asyncio.sleep(0.8) + release.set() + await session_task + await pinger_task + + errors = [m for m in ws.sent_json if m.get("type") == "error"] + assert not any("Stall timeout" in m.get("message", "") for m in errors) + assert any(m.get("type") == "session.pong" for m in ws.sent_json) + assert any(m.get("type") == "session.done" for m in ws.sent_json) + engine_client.abort.assert_not_called() diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py index 0787eb8562c..de2fcc24dd7 100644 --- a/tests/entrypoints/openai_api/test_serving_video_stream.py +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -14,11 +14,12 @@ import pytest from PIL import Image -from vllm_omni.entrypoints.openai import serving_video_stream, video_stream_envs +from vllm_omni.entrypoints.openai import video_stream_base, video_stream_envs from vllm_omni.entrypoints.openai.serving_video_stream import ( - OmniStreamingVideoHandler, + QwenOmniStreamingVideoHandler, StreamingVideoSessionConfig, ) +from vllm_omni.entrypoints.openai.video_stream_base import OmniStreamingVideoHandler from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -149,7 +150,7 @@ async def _gen(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): captured_requests.append(request) return {"prompt": "x"} @@ -186,7 +187,7 @@ async def _gen(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): captured_requests.append(request) return {"prompt": "x"} @@ -223,7 +224,7 @@ async def _gen(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): captured_requests.append(request) return {"prompt": "x"} @@ -309,7 +310,7 @@ async def _gen(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _preprocess_to_engine_prompt(self, request): return {"prompt": "x"} @@ -374,7 +375,7 @@ async def test_new_query_cancels_in_flight_query(): query_cancelled = asyncio.Event() calls = 0 - class BlockingHandler(OmniStreamingVideoHandler): + class BlockingHandler(QwenOmniStreamingVideoHandler): async def _process_query(self, *args, **kwargs): nonlocal calls calls += 1 @@ -412,7 +413,7 @@ async def test_video_done_waits_for_in_flight_query(): allow_finish = asyncio.Event() query_finished = asyncio.Event() - class BlockingHandler(OmniStreamingVideoHandler): + class BlockingHandler(QwenOmniStreamingVideoHandler): async def _process_query(self, *args, **kwargs): query_started.set() await allow_finish.wait() @@ -452,11 +453,11 @@ def blocked_decode(raw_bytes: bytes): release_decode.wait(timeout=2.0) return Image.open(io.BytesIO(raw_bytes)).convert("RGB") - class BlockingHandler(OmniStreamingVideoHandler): + class BlockingHandler(QwenOmniStreamingVideoHandler): async def _process_query(self, *args, **kwargs): query_started.set() - monkeypatch.setattr(serving_video_stream, "_decode_frame_bytes", blocked_decode) + monkeypatch.setattr(video_stream_base, "_decode_frame_bytes", blocked_decode) ws = TimedWebSocket() handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) @@ -486,7 +487,7 @@ async def test_client_cannot_send_internal_frame_decode_failed_message(): captured_frames: list[list[str]] = [] frame = _b64(_make_jpeg()) - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _process_query( self, websocket, @@ -550,7 +551,7 @@ async def test_frame_filter_error_sends_invalid_image(monkeypatch): def fail_should_retain(self, frame_jpeg): raise ValueError("decode failed") - monkeypatch.setattr(serving_video_stream.FrameSimilarityFilter, "should_retain", fail_should_retain) + monkeypatch.setattr(video_stream_base.FrameSimilarityFilter, "should_retain", fail_should_retain) ws = TimedWebSocket() handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) @@ -579,7 +580,7 @@ async def _gen(): return _gen() - class CapturingHandler(OmniStreamingVideoHandler): + class CapturingHandler(QwenOmniStreamingVideoHandler): async def _process_query_engine( self, websocket, @@ -594,7 +595,7 @@ async def _process_query_engine( ): captured_audio_lengths.append(len(audio_buffer)) - monkeypatch.setattr(serving_video_stream, "_MAX_AUDIO_BUFFER_BYTES", 4) + monkeypatch.setattr(video_stream_base, "_MAX_AUDIO_BUFFER_BYTES", 4) ws = TimedWebSocket() handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=5.0) @@ -618,7 +619,7 @@ async def _process_query_engine( def test_build_messages_keeps_recent_history_text_only(): - handler = OmniStreamingVideoHandler(chat_service=object()) + handler = QwenOmniStreamingVideoHandler(chat_service=object()) old_frame = _b64(_make_jpeg(1, 2, 3)) current_frame = _b64(_make_jpeg(4, 5, 6)) history = [ diff --git a/tests/entrypoints/openai_api/test_speech_usage.py b/tests/entrypoints/openai_api/test_speech_usage.py new file mode 100644 index 00000000000..8211dab88b7 --- /dev/null +++ b/tests/entrypoints/openai_api/test_speech_usage.py @@ -0,0 +1,220 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for speech token-usage accounting (issue #4646). + +These exercise the pure usage logic without a model/server: input-token +breakdown (text + reference-audio gating) and output-token accumulation. +""" + +from dataclasses import dataclass + +import pytest + +from vllm_omni.entrypoints.openai.speech_usage import ( + SpeechOutputTokenCounter, + build_speech_usage, + gate_audio_tokens, + qwen3_tts_input_token_details, +) + + +# A deterministic stand-in for a real tokenizer: 1 token per whitespace word. +def word_count(text: str) -> int: + return len(text.split()) + + +# --- input breakdown: gating ------------------------------------------------ + + +def test_gate_audio_tokens_base_in_context_counts_ref_frames(): + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120) + == 120 + ) + + +def test_gate_audio_tokens_customvoice_is_zero(): + assert ( + gate_audio_tokens( + task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=None, ref_code_length=120 + ) + == 0 + ) + + +def test_gate_audio_tokens_x_vector_only_is_zero_even_with_ref(): + # x-vector cloning inserts NO codec frames, so audio_tokens must be 0 + # even though ref_code_length is populated. + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=None, ref_code_length=120) == 0 + ) + + +def test_gate_audio_tokens_icl_override_wins(): + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=False, ref_code_length=120) == 0 + ) + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=True, icl_mode_override=True, ref_code_length=80) == 80 + ) + + +def test_gate_audio_tokens_icl_override_ignored_for_non_base(): + # icl_mode_override is honored only for the Base task; CustomVoice stays 0 + # even if such a flag leaks into its params (mirrors the estimator). + assert ( + gate_audio_tokens( + task_type="CustomVoice", x_vector_only_mode=False, icl_mode_override=True, ref_code_length=120 + ) + == 0 + ) + + +def test_gate_audio_tokens_bad_ref_length_is_zero(): + assert ( + gate_audio_tokens(task_type="Base", x_vector_only_mode=False, icl_mode_override=None, ref_code_length="oops") + == 0 + ) + + +# --- input breakdown: Qwen3-TTS reproduces issue #4646 observation ---------- + + +def test_customvoice_text_scales_with_input_no_audio(): + short = qwen3_tts_input_token_details( + input_text="hi there", + instructions=None, + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + long = qwen3_tts_input_token_details( + input_text="hi there this is a much longer sentence", + instructions=None, + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + assert short.audio_tokens == 0 and long.audio_tokens == 0 + assert long.text_tokens > short.text_tokens # scales with input + + +def test_base_icl_text_tracks_input_audio_tracks_ref(): + # The core #4646 fix: text_tokens reflect the input (not dropped), and + # audio_tokens track the reference audio independently. + params = {"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]} + short = qwen3_tts_input_token_details( + input_text="hi", instructions=None, tts_params=params, count_text_tokens=word_count + ) + long = qwen3_tts_input_token_details( + input_text="hi there friend", instructions=None, tts_params=params, count_text_tokens=word_count + ) + # text now varies with input (the bug was that it was dropped for Base) + assert long.text_tokens > short.text_tokens + # audio fixed by the reference, independent of input text + assert short.audio_tokens == long.audio_tokens == 100 + + +def test_instructions_count_toward_text(): + without = qwen3_tts_input_token_details( + input_text="hello world", + instructions=None, + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + with_instr = qwen3_tts_input_token_details( + input_text="hello world", + instructions="speak slowly and warmly", + tts_params={"task_type": ["CustomVoice"]}, + count_text_tokens=word_count, + ) + assert with_instr.text_tokens == without.text_tokens + word_count("speak slowly and warmly") + + +# --- build usage ------------------------------------------------------------ + + +def test_build_speech_usage_aggregates_and_totals(): + details = qwen3_tts_input_token_details( + input_text="hi there", + instructions=None, + tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [100]}, + count_text_tokens=word_count, + ) + usage = build_speech_usage(details, output_tokens=250) + assert usage.input_tokens == details.text_tokens + 100 + assert usage.output_tokens == 250 + assert usage.total_tokens == usage.input_tokens + 250 + assert usage.input_token_details.text_tokens == 2 + assert usage.input_token_details.audio_tokens == 100 + + +def test_batch_item_result_serializes_usage(): + # Guards the batch per-item usage carrier (and that it is optional/None on error). + from vllm_omni.entrypoints.openai.protocol.audio import SpeechBatchItemResult + + details = qwen3_tts_input_token_details( + input_text="hello world", + instructions=None, + tts_params={"task_type": ["Base"], "x_vector_only_mode": [False], "ref_code_length": [50]}, + count_text_tokens=word_count, + ) + item = SpeechBatchItemResult( + index=0, + status="success", + audio_data="x", + media_type="audio/wav", + usage=build_speech_usage(details, output_tokens=30), + ) + dumped = item.model_dump() + assert dumped["usage"]["input_tokens"] == 2 + 50 + assert dumped["usage"]["output_tokens"] == 30 + assert dumped["usage"]["input_token_details"] == {"text_tokens": 2, "audio_tokens": 50} + # Errored items carry no usage. + assert SpeechBatchItemResult(index=1, status="error", error="boom").model_dump()["usage"] is None + + +# --- output token accumulation ---------------------------------------------- +# The generated codec-token count is read from per-stage pipeline metrics on the +# final output: res.metrics["stage_metrics"][<stage>]["num_tokens_out"]. Only the +# AR stage (0) is non-zero; code2wav (stage 1) reports 0. + + +@dataclass +class _FakeRes: + metrics: dict + + +def _final_res(stage0_out, stage1_out=0): + return _FakeRes( + metrics={ + "stage_metrics": { + "0": {"num_tokens_in": 110, "num_tokens_out": stage0_out}, + "1": {"num_tokens_in": 0, "num_tokens_out": stage1_out}, + } + } + ) + + +def test_output_tokens_from_stage_metrics(): + acc = SpeechOutputTokenCounter() + # Intermediate outputs carry empty metrics -> contribute nothing. + acc.observe(_FakeRes(metrics={})) + acc.observe(_FakeRes(metrics={})) + # Final output carries stage metrics; stage-0 num_tokens_out is the count. + acc.observe(_final_res(stage0_out=77)) + assert acc.total() == 77 + + +def test_output_tokens_takes_max_across_stages(): + acc = SpeechOutputTokenCounter() + acc.observe(_final_res(stage0_out=42, stage1_out=0)) + assert acc.total() == 42 + + +def test_output_tokens_zero_without_metrics(): + acc = SpeechOutputTokenCounter() + acc.observe(_FakeRes(metrics={})) + acc.observe(_FakeRes(metrics={"stage_metrics": {}})) + assert acc.total() == 0 + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/entrypoints/openai_api/test_storage_manager.py b/tests/entrypoints/openai_api/test_storage_manager.py new file mode 100644 index 00000000000..30e50576672 --- /dev/null +++ b/tests/entrypoints/openai_api/test_storage_manager.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit tests for the OpenAI video storage managers. +""" + +import asyncio +import os +import time + +import pytest + +from vllm_omni.entrypoints.openai import storage as storage_module +from vllm_omni.entrypoints.openai.storage import FileStorageHandle, LocalStorageManager, LocalStorageTTLManager + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_local_storage_open_returns_file_handle_for_saved_key(tmp_path): + storage = LocalStorageManager(storage_path=str(tmp_path / "storage")) + save_context = asyncio.run(storage.save(b"video-bytes", "video-123")) + + handle = asyncio.run(storage.open(save_context.key)) + + assert isinstance(handle, FileStorageHandle) + assert handle.path == os.path.join(storage.storage_path, "video-123") + with open(handle.path, "rb") as file: + assert file.read() == b"video-bytes" + + +def test_local_storage_open_returns_none_for_missing_key(tmp_path): + storage = LocalStorageManager(storage_path=str(tmp_path / "storage")) + + handle = asyncio.run(storage.open("missing-video")) + + assert handle is None + + +def test_local_storage_ttl_save_sets_expiration_metadata(tmp_path, monkeypatch): + storage = LocalStorageTTLManager( + storage_path=str(tmp_path / "storage"), + max_concurrency=1, + ttl_seconds=60, + sweep_interval_seconds=300, + ) + monkeypatch.setattr(storage_module.time, "time", lambda: 1_700_000_000) + + save_context = asyncio.run(storage.save(b"video-bytes", "video-ttl")) + + assert save_context.key == "video-ttl" + assert save_context.created_at == 1_700_000_000 + assert save_context.expires_at == 1_700_000_060 + + +def test_local_storage_ttl_sweeper_removes_expired_file(tmp_path): + storage = LocalStorageTTLManager( + storage_path=str(tmp_path / "storage"), + max_concurrency=1, + ttl_seconds=1, + sweep_interval_seconds=60, + ) + + async def setup_file() -> tuple[str, str]: + save_context = await storage.save(b"video-bytes", "video-expired") + file_path = storage.get_full_file_path(save_context.key) + return save_context.key, file_path + + storage_key, file_path = asyncio.run(setup_file()) + expired_mtime = time.time() - 10 + os.utime(file_path, (expired_mtime, expired_mtime)) + assert os.path.exists(file_path) + + deleted = asyncio.run(storage._sweep_once(time.time() - 1)) + + assert deleted == 1 + assert not os.path.exists(file_path) + assert asyncio.run(storage.open(storage_key)) is None + + +def test_local_storage_ttl_sweeper_keeps_files_when_path_missing(tmp_path): + storage = LocalStorageTTLManager( + storage_path=str(tmp_path / "storage"), + max_concurrency=1, + ttl_seconds=1, + sweep_interval_seconds=60, + ) + + missing_path = storage.get_full_file_path("video-missing") + assert not os.path.exists(missing_path) + + deleted = asyncio.run(storage._sweep_once(time.time() - 1)) + + assert deleted == 0 + assert not os.path.exists(missing_path) diff --git a/tests/entrypoints/openai_api/test_text_splitter.py b/tests/entrypoints/openai_api/test_text_splitter.py deleted file mode 100644 index b9022e015dd..00000000000 --- a/tests/entrypoints/openai_api/test_text_splitter.py +++ /dev/null @@ -1,253 +0,0 @@ -"""Tests for SentenceSplitter used in streaming TTS input.""" - -import pytest - -from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class TestSentenceSplitterEnglish: - """Tests for English sentence splitting.""" - - def test_single_sentence_no_boundary(self): - splitter = SentenceSplitter() - result = splitter.add_text("Hello world") - assert result == [] - assert splitter.buffer == "Hello world" - - def test_single_sentence_with_boundary(self): - splitter = SentenceSplitter() - result = splitter.add_text("Hello world. How are you?") - assert len(result) == 1 - assert result[0] == "Hello world." - - def test_multiple_sentences(self): - splitter = SentenceSplitter() - result = splitter.add_text("Hello. How are you? I am fine! ") - assert len(result) == 3 - assert result[0] == "Hello." - assert result[1] == "How are you?" - assert result[2] == "I am fine!" - - def test_exclamation_mark(self): - splitter = SentenceSplitter() - result = splitter.add_text("Wow, that is great! Tell me more.") - assert len(result) == 1 - assert result[0] == "Wow, that is great!" - - def test_question_mark(self): - splitter = SentenceSplitter() - result = splitter.add_text("Can you hear me? I hope so.") - assert len(result) == 1 - assert result[0] == "Can you hear me?" - - -class TestSentenceSplitterChinese: - """Tests for CJK sentence splitting.""" - - def test_chinese_period(self): - splitter = SentenceSplitter() - result = splitter.add_text("你好世界。你好吗") - assert len(result) == 1 - assert result[0] == "你好世界。" - - def test_chinese_exclamation(self): - splitter = SentenceSplitter() - result = splitter.add_text("太好了！谢谢你") - assert len(result) == 1 - assert result[0] == "太好了！" - - def test_chinese_question(self): - splitter = SentenceSplitter() - result = splitter.add_text("你是谁？我是小明") - assert len(result) == 1 - assert result[0] == "你是谁？" - - def test_chinese_comma_no_split(self): - """Chinese commas are clause-level and should not trigger a split.""" - splitter = SentenceSplitter() - result = splitter.add_text("你好，世界") - assert result == [] - assert splitter.buffer == "你好，世界" - - def test_chinese_semicolon_no_split(self): - """Chinese semicolons are clause-level and should not trigger a split.""" - splitter = SentenceSplitter() - result = splitter.add_text("第一点；第二点") - assert result == [] - assert splitter.buffer == "第一点；第二点" - - def test_chinese_multiple(self): - splitter = SentenceSplitter() - result = splitter.add_text("你好！你好吗？我很好。") - assert len(result) == 3 - assert result[0] == "你好！" - assert result[1] == "你好吗？" - assert result[2] == "我很好。" - - -class TestSentenceSplitterMixed: - """Tests for mixed-language sentence splitting.""" - - def test_mixed_english_chinese(self): - splitter = SentenceSplitter() - result = splitter.add_text("Hello世界。How are you? ") - assert len(result) == 2 - assert result[0] == "Hello世界。" - assert result[1] == "How are you?" - - -class TestSentenceSplitterIncremental: - """Tests for incremental (multi-chunk) text input.""" - - def test_accumulation_across_chunks(self): - splitter = SentenceSplitter() - # First chunk: no boundary - result1 = splitter.add_text("Hello ") - assert result1 == [] - - # Second chunk: completes a sentence - result2 = splitter.add_text("world. How") - assert len(result2) == 1 - assert result2[0] == "Hello world." - assert splitter.buffer == "How" - - def test_word_by_word(self): - splitter = SentenceSplitter() - words = ["Hello, ", "how ", "are ", "you? ", "I ", "am ", "fine."] - all_sentences = [] - for word in words: - all_sentences.extend(splitter.add_text(word)) - - assert len(all_sentences) == 1 - assert all_sentences[0] == "Hello, how are you?" - # "I am fine." stays in buffer (no trailing whitespace after period) - - def test_three_chunks(self): - splitter = SentenceSplitter() - splitter.add_text("The quick brown ") - splitter.add_text("fox jumps. ") - result = splitter.add_text("Over the lazy dog. ") - # "The quick brown fox jumps." should have been returned on second chunk - # "Over the lazy dog." on third chunk - assert len(result) == 1 - assert result[0] == "Over the lazy dog." - - -class TestSentenceSplitterFlush: - """Tests for flush behavior.""" - - def test_flush_returns_remaining(self): - splitter = SentenceSplitter() - splitter.add_text("Hello world") - result = splitter.flush() - assert result == "Hello world" - assert splitter.buffer == "" - - def test_flush_empty_buffer(self): - splitter = SentenceSplitter() - result = splitter.flush() - assert result is None - - def test_flush_after_sentence(self): - splitter = SentenceSplitter() - splitter.add_text("Hello world. Remaining text") - result = splitter.flush() - assert result == "Remaining text" - - def test_flush_whitespace_only(self): - splitter = SentenceSplitter() - splitter.add_text("Hello. ") - # "Hello." extracted, buffer is " " - result = splitter.flush() - # Whitespace-only should return None - assert result is None - - def test_flush_clears_buffer(self): - splitter = SentenceSplitter() - splitter.add_text("some text") - splitter.flush() - assert splitter.buffer == "" - # Second flush should return None - assert splitter.flush() is None - - -class TestSentenceSplitterEdgeCases: - """Edge case tests.""" - - def test_empty_input(self): - splitter = SentenceSplitter() - result = splitter.add_text("") - assert result == [] - assert splitter.buffer == "" - - def test_none_like_empty(self): - """Empty string should not affect buffer.""" - splitter = SentenceSplitter() - splitter.add_text("Hello") - splitter.add_text("") - assert splitter.buffer == "Hello" - - def test_only_punctuation(self): - splitter = SentenceSplitter() - result = splitter.add_text(". ") - # "." is 1 char, below default min_sentence_length of 2 - # It will be carried forward - assert result == [] - - def test_min_sentence_length(self): - splitter = SentenceSplitter(min_sentence_length=10) - result = splitter.add_text("Hi. Hello world. ") - # "Hi." is 3 chars (< 10), so it gets carried to "Hello world." - assert len(result) == 1 - assert "Hi." in result[0] - assert "Hello world." in result[0] - - def test_short_segments_are_carried_until_long_enough(self): - splitter = SentenceSplitter(min_sentence_length=10) - result = splitter.add_text("Hi. Ok. Hello there. ") - assert result == ["Hi.Ok.Hello there."] - assert splitter.buffer == "" - - def test_min_sentence_length_zero(self): - splitter = SentenceSplitter(min_sentence_length=0) - result = splitter.add_text("A. B. ") - assert len(result) == 2 - - def test_no_boundary_then_flush(self): - splitter = SentenceSplitter() - result = splitter.add_text("Hello world how are you") - assert result == [] - flushed = splitter.flush() - assert flushed == "Hello world how are you" - - def test_consecutive_punctuation(self): - splitter = SentenceSplitter() - result = splitter.add_text("Really?! Yes, really. ") - assert len(result) >= 1 - - def test_reuse_after_flush(self): - """Splitter can be reused after flush.""" - splitter = SentenceSplitter() - splitter.add_text("First session.") - splitter.flush() - - result = splitter.add_text("Second session. More text") - assert len(result) == 1 - assert result[0] == "Second session." - assert splitter.buffer == "More text" - - -class TestSentenceSplitterBufferLimit: - """Tests for buffer overflow protection.""" - - def test_buffer_overflow_raises(self): - from vllm_omni.entrypoints.openai.text_splitter import _MAX_BUFFER_SIZE - - splitter = SentenceSplitter() - # Fill buffer just under the limit - splitter.add_text("x" * (_MAX_BUFFER_SIZE - 1)) - # One more char should exceed the limit - with pytest.raises(ValueError, match="exceeded maximum size"): - splitter.add_text("xx") diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py index 9e732403fbb..7228196827e 100644 --- a/tests/entrypoints/openai_api/test_video_api_utils.py +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -48,6 +48,107 @@ def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): assert mux_calls[0]["audio"] is None +def test_fragmented_mp4_video_encoder_reuses_single_muxer(monkeypatch): + muxers = [] + + class FakeFragmentedMP4Muxer: + def __init__(self, *, width, height, fps, video_codec_options=None): + self.calls = [] + muxers.append( + { + "width": width, + "height": height, + "fps": fps, + "video_codec_options": video_codec_options, + "instance": self, + } + ) + + def mux_video_frames(self, frames): + self.calls.append(frames) + return f"fragment-{len(self.calls)}".encode() + + def close(self): + return b"trailer" + + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.FragmentedMP4Muxer", + FakeFragmentedMP4Muxer, + ) + + encoder = video_api_utils.FragmentedMP4VideoEncoder( + fps=12, + video_codec_options={"preset": "ultrafast"}, + ) + assert encoder.encode([np.zeros((4, 6, 3), dtype=np.float32)]) == b"fragment-1" + assert encoder.encode([np.ones((4, 6, 3), dtype=np.float32)]) == b"fragment-2" + assert encoder.close() == b"trailer" + + assert len(muxers) == 1 + assert muxers[0]["width"] == 6 + assert muxers[0]["height"] == 4 + assert muxers[0]["fps"] == 12.0 + assert muxers[0]["video_codec_options"] == {"preset": "ultrafast"} + assert len(muxers[0]["instance"].calls) == 2 + + +def test_create_streaming_video_encoder_selects_requested_format(): + assert isinstance( + video_api_utils.create_streaming_video_encoder(output_format="m4s", fps=12), + video_api_utils.FragmentedMP4VideoEncoder, + ) + + +def test_finalize_streaming_mp4_bytes_produces_progressive_mp4(): + """Fragment MP4 chunks are remuxed into decodable progressive MP4 bytes.""" + import numpy as np + + from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer, finalize_streaming_video_bytes + + def _read_mp4_video_info(mp4_bytes: bytes) -> tuple[int, float, float | None, int, int]: + import io + + import av + + with av.open(io.BytesIO(mp4_bytes), mode="r", format="mp4") as container: + stream = container.streams.video[0] + frame_count = sum(1 for _ in container.decode(stream)) + assert stream.average_rate is not None + fps = float(stream.average_rate) + duration_sec = None + if stream.duration is not None: + assert stream.time_base is not None + duration_sec = float(stream.duration * stream.time_base) + return frame_count, fps, duration_sec, stream.width, stream.height + + width = 32 + height = 32 + fps = 16.0 + input_frame_count = 2 + + muxer = FragmentedMP4Muxer(width=width, height=height, fps=fps) + frames = np.zeros((input_frame_count, height, width, 3), dtype=np.uint8) + streamed = muxer.mux_video_frames(frames) + muxer.close() + + finalized = finalize_streaming_video_bytes(streamed, input_format="m4s", fps=fps) + assert finalized + assert finalized != streamed + + streamed_info = _read_mp4_video_info(streamed) + final_info = _read_mp4_video_info(finalized) + expected_duration = input_frame_count / fps + + assert streamed_info[0] == input_frame_count + assert final_info[0] == input_frame_count + assert streamed_info[1] == pytest.approx(fps) + assert final_info[1] == pytest.approx(fps) + assert streamed_info[3:] == (width, height) + assert final_info[3:] == (width, height) + + assert final_info[2] == pytest.approx(expected_duration, rel=0.05) + assert streamed_info[2] == pytest.approx(expected_duration, rel=0.05) + + def test_rife_model_inference_runs_on_dummy_tensors(): model = rife_interpolator.Model().eval() img0 = torch.rand(1, 3, 32, 32) diff --git a/tests/entrypoints/openai_api/test_video_frame_filter.py b/tests/entrypoints/openai_api/test_video_frame_filter.py index a734aec75c0..70df6e34cab 100644 --- a/tests/entrypoints/openai_api/test_video_frame_filter.py +++ b/tests/entrypoints/openai_api/test_video_frame_filter.py @@ -29,8 +29,8 @@ def test_very_different_frames_retained(self): assert f.should_retain(make_jpeg(255, 255, 255)) is True assert f.should_retain(make_jpeg(0, 0, 0)) is True - def test_low_threshold_keeps_slightly_different(self): - f = FrameSimilarityFilter(threshold=0.50) + def test_high_threshold_keeps_slightly_different(self): + f = FrameSimilarityFilter(threshold=0.999) assert f.should_retain(make_jpeg(100, 100, 100)) is True assert f.should_retain(make_jpeg(110, 110, 110)) is True diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index b54afd54a88..ff69a2bf260 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -11,6 +11,7 @@ import os import threading import time +from contextlib import asynccontextmanager from types import SimpleNamespace import numpy as np @@ -43,23 +44,18 @@ def __init__( videos, audios=None, sample_rate=None, - custom_output=None, + multimodal_output=None, stage_durations=None, peak_memory_mb=0.0, ): - self.multimodal_output = {"video": videos} + self.multimodal_output = dict(multimodal_output or {"video": videos}) if audios is not None: self.multimodal_output["audio"] = audios if sample_rate is not None: self.multimodal_output["audio_sample_rate"] = sample_rate - self._custom_output = custom_output or {} self.stage_durations = stage_durations or {} self.peak_memory_mb = peak_memory_mb - @property - def custom_output(self): - return self._custom_output - class FakeAsyncOmni: def __init__(self): @@ -99,6 +95,14 @@ async def generate_video_bytes( raise +class FakeServerSocket: + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + @pytest.fixture(autouse=True) def isolated_video_backends(tmp_path, monkeypatch): """Use isolated in-memory metadata and local storage for each test.""" @@ -111,6 +115,91 @@ def isolated_video_backends(tmp_path, monkeypatch): return store, tasks, storage +@pytest.mark.asyncio +async def test_server_worker_keeps_engine_alive_until_http_shutdown(monkeypatch): + events: list[str] = [] + serve_started = asyncio.Event() + http_shutdown = asyncio.Event() + engine_context_exited = asyncio.Event() + sock = FakeServerSocket() + + class FakeEngine: + stage_configs = [] + + async def get_supported_tasks(self): + return ("generate",) + + @asynccontextmanager + async def fake_build_async_omni(*args, **kwargs): + del args, kwargs + events.append("engine_enter") + try: + yield FakeEngine() + finally: + events.append("engine_exit") + engine_context_exited.set() + + async def fake_serve_http(*args, **kwargs): + del args, kwargs + events.append("serve_http") + serve_started.set() + + async def wait_for_shutdown(): + await http_shutdown.wait() + events.append("http_shutdown") + + return asyncio.create_task(wait_for_shutdown()) + + async def fake_storage_start(): + events.append("storage_start") + + async def fake_get_vllm_config(engine_client): + del engine_client + return None + + async def fake_init_app_state(engine_client, state, args): + del engine_client, state, args + events.append("init_app_state") + + monkeypatch.setattr(api_server, "build_async_omni", fake_build_async_omni) + monkeypatch.setattr(api_server, "build_openai_app", lambda args, supported_tasks: FastAPI()) + monkeypatch.setattr(api_server, "serve_http", fake_serve_http) + monkeypatch.setattr(api_server.STORAGE_MANAGER, "start", fake_storage_start) + monkeypatch.setattr(api_server, "_get_vllm_config", fake_get_vllm_config) + monkeypatch.setattr(api_server, "omni_init_app_state", fake_init_app_state) + monkeypatch.setattr(api_server, "get_uvicorn_log_config", lambda args: None) + + args = SimpleNamespace( + tool_parser_plugin="", + reasoning_parser_plugin="", + reasoning_parser=None, + structured_outputs_config=SimpleNamespace(reasoning_parser=None), + enable_ssl_refresh=False, + host="127.0.0.1", + port=0, + uvicorn_log_level="info", + disable_uvicorn_access_log=True, + ssl_keyfile=None, + ssl_certfile=None, + ssl_ca_certs=None, + ssl_cert_reqs=None, + ssl_ciphers=None, + h11_max_incomplete_event_size=None, + h11_max_header_count=None, + ) + + worker_task = asyncio.create_task(api_server.omni_run_server_worker("127.0.0.1:0", sock, args)) + await asyncio.wait_for(serve_started.wait(), timeout=2) + + assert not engine_context_exited.is_set() + + http_shutdown.set() + await asyncio.wait_for(worker_task, timeout=2) + + assert sock.closed + assert events.index("http_shutdown") < events.index("engine_exit") + + @pytest.fixture def test_client(): app = FastAPI() @@ -423,6 +512,7 @@ def test_decode_video_bytes_can_keep_last_frames(): ) assert len(frames) == 2 + assert frames.fps == pytest.approx(8.0) red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames] assert red_means[0] > 100 assert red_means[1] > red_means[0] @@ -520,11 +610,49 @@ def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs): engine = test_client.app.state.openai_serving_video._engine_client captured = engine.captured_sampling_params_list[0] assert captured.num_frames == 72 - assert captured.fps == 24 - assert captured.frame_rate == 24.0 + # fps omitted -> sampling params carry None (the "not provided" signal); the 24 + # default is applied only at output encoding. + assert captured.fps is None + assert captured.frame_rate is None assert fps_values == [24] +def test_model_reported_fps_wins_when_request_fps_omitted(test_client, mocker: MockerFixture): + fps_values = [] + + def _fake_encode(video, fps, audio=None, audio_sample_rate=None, **kwargs): + del video, audio, audio_sample_rate, kwargs + fps_values.append(fps) + return b"fake-video" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + side_effect=_fake_encode, + ) + + engine = test_client.app.state.openai_serving_video._engine_client + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + result = MockVideoResult([object()]) + result.multimodal_output["fps"] = 8 + yield result + + engine.generate = _generate + + response = test_client.post("/v1/videos", data={"prompt": "source fps"}) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + captured = engine.captured_sampling_params_list[0] + # fps omitted -> None on the sampling params; the model-reported fps (8) wins for output. + assert captured.fps is None + assert captured.frame_rate is None + assert fps_values == [8] + + def test_size_param_sets_width_height(test_client, mocker: MockerFixture): mocker.patch( "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", @@ -702,7 +830,13 @@ async def _generate(prompt, request_id, sampling_params_list): engine.captured_sampling_params_list = sampling_params_list import numpy as np - yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2}) + yield MockVideoResult( + [np.zeros((1, 64, 64, 3), dtype=np.uint8)], + multimodal_output={ + "video": [np.zeros((1, 64, 64, 3), dtype=np.uint8)], + "metadata": {"video": {"video_fps_multiplier": 2}}, + }, + ) engine.generate = _generate @@ -806,11 +940,16 @@ async def _generate(prompt, request_id, sampling_params_list): yield MockVideoResult( [object()], - custom_output={ - "action": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32), - "raw_action_dim": 2, - "action_mode": "policy", - "domain_id": 7, + multimodal_output={ + "video": [object()], + "actions": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, }, ) @@ -848,11 +987,16 @@ async def _generate(prompt, request_id, sampling_params_list): engine.captured_sampling_params_list = sampling_params_list yield MockVideoResult( [object()], - custom_output={ - "action": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32), - "raw_action_dim": 2, - "action_mode": "policy", - "domain_id": 7, + multimodal_output={ + "video": [object()], + "actions": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, }, ) @@ -886,19 +1030,51 @@ def test_action_extraction_accepts_unbatched_action(): result = MockVideoResult( [object()], - custom_output={ - "action": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), - "raw_action_dim": 2, - "action_mode": "policy", - "domain_id": 7, + multimodal_output={ + "video": [object()], + "actions": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, }, ) actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1) + assert actions and actions[0] is not None + assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]] + assert actions[0].shape == [2, 2] + + +def test_action_extraction_accepts_multimodal_actions_payload(): + import numpy as np + + result = MockVideoResult([object()]) + result.multimodal_output.update( + { + "actions": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32), + "metadata": { + "actions": { + "raw_action_dim": 2, + "action_mode": "policy", + "domain_id": 7, + }, + }, + } + ) + + actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1) + assert actions[0] is not None assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]] assert actions[0].shape == [2, 2] + assert actions[0].raw_action_dim == 2 + assert actions[0].action_mode == "policy" + assert actions[0].domain_id == 7 def test_missing_handler_returns_503(): @@ -1020,6 +1196,31 @@ def test_invalid_lora_returns_400(test_client): assert "lora object" in failed["error"]["message"].lower() +def test_failed_generation_awaits_storage_cleanup(test_client, isolated_video_backends, mocker: MockerFixture): + """Regression (merge seam): when async generation raises, the failure handler + must ``await _cleanup_video(video_id)`` (single-arg, async) and still record + FAILED. Upstream carried a sync ``_cleanup_video(video_id, output_path)`` whose + stale call in the generic handler sat outside the conflict markers; against the + PR's async storage manager that raised NameError before the FAILED update, + wedging the job in IN_PROGRESS and orphaning the artifact.""" + _store, _tasks, storage = isolated_video_backends + delete_spy = mocker.spy(storage, "delete") + mocker.patch.object( + OmniOpenAIServingVideo, + "generate_video_bytes", + side_effect=RuntimeError("GPU exploded"), + ) + + response = test_client.post("/v1/videos", data={"prompt": "will fail"}) + 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"] == 500 + assert "GPU exploded" in failed["error"]["message"] + delete_spy.assert_called_once_with(video_id) + + def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture): mocker.patch.object( OmniOpenAIServingVideo, @@ -1178,7 +1379,7 @@ def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: Moc final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) file_name = final["file_name"] assert file_name is not None - file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, file_name) + file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id) assert os.path.exists(file_path) delete_resp = test_client.delete(f"/v1/videos/{video_id}") @@ -1189,6 +1390,30 @@ def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: Moc assert not os.path.exists(file_path) +def test_download_completed_job_uses_storage_open_and_download_name(test_client, mocker: MockerFixture): + video_bytes = b"stored-video-data" + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=video_bytes, + ) + create_resp = test_client.post("/v1/videos", data={"prompt": "Download this video"}) + assert create_resp.status_code == 200 + video_id = create_resp.json()["id"] + + final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + file_name = final["file_name"] + assert file_name == f"{video_id}.mp4" + + storage_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, video_id) + assert os.path.exists(storage_path) + + response = test_client.get(f"/v1/videos/{video_id}/content") + assert response.status_code == 200 + assert response.content == video_bytes + assert response.headers["content-type"] == "video/mp4" + assert file_name in response.headers["content-disposition"] + + def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client): handler = BlockingVideoHandler() test_client.app.state.openai_serving_video = handler @@ -1246,7 +1471,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture) extra_params = { "is_enable_stage2": True, "pyramid_num_stages": 3, - "pyramid_num_inference_steps_list": [20, 20, 20], + "pyramid_num_inference_steps_list": [1, 1, 1], "use_cfg_zero_star": True, } response = test_client.post( @@ -1264,7 +1489,7 @@ def test_extra_params_merged_into_extra_args(test_client, mocker: MockerFixture) captured = engine.captured_sampling_params_list[0] assert captured.extra_args["is_enable_stage2"] is True assert captured.extra_args["pyramid_num_stages"] == 3 - assert captured.extra_args["pyramid_num_inference_steps_list"] == [20, 20, 20] + assert captured.extra_args["pyramid_num_inference_steps_list"] == [1, 1, 1] assert captured.extra_args["use_cfg_zero_star"] is True @@ -1658,7 +1883,13 @@ def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker: async def _generate(prompt, request_id, sampling_params_list): engine.captured_prompt = prompt engine.captured_sampling_params_list = sampling_params_list - yield MockVideoResult([object()], custom_output={"video_fps_multiplier": 2}) + yield MockVideoResult( + [object()], + multimodal_output={ + "video": [object()], + "metadata": {"video": {"video_fps_multiplier": 2}}, + }, + ) engine.generate = _generate diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index b8637c1e990..7448cc5c010 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -8,6 +8,7 @@ from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.outputs import OmniRequestOutput pytestmark = [pytest.mark.core_model] @@ -198,6 +199,48 @@ async def collect_outputs(request_id): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_yields_streaming_diffusion_chunks_before_final(): + """AsyncOmni.generate yields every intermediate diffusion chunk before the final one.""" + + async def streaming_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts): + del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts + yield OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + final_output_type="image", + custom_output={"chunk": 0}, + finished=False, + ) + yield OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + final_output_type="image", + custom_output={"chunk": 1}, + finished=True, + ) + + async def run_test(): + omni = get_async_omni_instance() + omni._process_orchestrator_results = streaming_process_results + + outputs = [] + async for output in AsyncOmni.generate( + omni, + prompt={"prompt": "a cat"}, + request_id="req-stream", + sampling_params_list=[SimpleNamespace()], + output_modalities=["image"], + ): + outputs.append(output) + + assert len(outputs) == 2 + assert [output.finished for output in outputs] == [False, True] + assert [output.custom_output["chunk"] for output in outputs] == [0, 1] + + asyncio.run(run_test()) + + @pytest.mark.cpu @pytest.mark.parametrize( "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE] @@ -229,69 +272,72 @@ async def run(): # End to end tests for ensuring internal manipulation of request ID # in diffusion / Omni models don't leak back to the user. # -# NOTE: It seems like we currently need the shutdowns here, otherwise -# running the tests sequentially seems to leave a zombie process in diffusion -# that can OOM the Omni tests. -@hardware_test(res={"cuda": "L4"}, num_cards=1) -@pytest.mark.omni -@pytest.mark.asyncio -@pytest.mark.parametrize( - "req_id", - ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"], - ids=["plain", "prefixed-img", "prefixed-chat"], -) -async def test_diffusion_generate_preserves_request_id(req_id): - """Ensure diffusion model requests don't leak internal UUID-suffixed req id back to user.""" - engine = AsyncOmni(model=DIFFUSION_MODEL) - try: - async for output in engine.generate("a white cat", request_id=req_id): - assert output.request_id == req_id - finally: - engine.shutdown() +# One AsyncOmni per test function (all cases in a single asyncio loop) to avoid +# repeated cold starts. Do not use class/module-scoped engine fixtures here: +# pytest-asyncio uses a function-scoped event loop by default, so reusing an +# engine across tests can hang on the second generate() call. -@hardware_test(res={"cuda": "H100"}, num_cards=1) -@pytest.mark.omni -@pytest.mark.asyncio -@pytest.mark.parametrize( - "req_id", - ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"], - ids=["plain", "prefixed-img", "prefixed-chat"], -) -async def test_omni_generate_preserves_request_id(req_id): - """Ensure omni model requests don't leak internal UUID-suffixed req id back to user.""" - engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG) - try: - async for output in engine.generate( - "Say hello in one word.", - request_id=req_id, - output_modalities=["text"], - ): - assert output.request_id == req_id - finally: - engine.shutdown() +# Covers: +# * plain client ids (``my-req-1``) +# * OpenAI-style prefixed ids (``img_gen-*``, ``chatcmpl-*``) that AsyncOmni +# suffixes internally for engine routing — streamed outputs must still echo +# the caller-visible id, not the internal UUID-suffixed id +# * empty ``request_id`` — server assigns a non-empty id for the caller +_DIFFUSION_REQ_IDS = ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"] +_OMNI_REQ_IDS = ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"] @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.omni @pytest.mark.asyncio -async def test_diffusion_generate_empty_request_id(): - """Empty request_id should get a generated internal ID, not stay empty (diffusion)""" +async def test_diffusion_generate_request_id(): + """Diffusion E2E request-id contract (``riverclouds/qwen_image_random``). + + Scenarios (one engine, sequential ``generate`` calls): + - plain id ``my-req-1`` + - image-style prefix ``img_gen-abc123`` + - chat-style prefix ``chatcmpl-xyz`` + - empty ``request_id`` → output id is non-empty (auto-assigned) + + Each streaming output must expose the user-supplied id unchanged; internal + UUID suffixing must not leak into ``output.request_id``. + """ engine = AsyncOmni(model=DIFFUSION_MODEL) try: + for req_id in _DIFFUSION_REQ_IDS: + async for output in engine.generate("a white cat", request_id=req_id): + assert output.request_id == req_id async for output in engine.generate("a white cat", request_id=""): assert output.request_id != "" finally: engine.shutdown() -@hardware_test(res={"cuda": "H100"}, num_cards=1) +@hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.omni @pytest.mark.asyncio -async def test_omni_generate_empty_request_id(): - """Empty request_id should get a generated internal ID, not stay empty (omni)""" +async def test_omni_generate_request_id(): + """Omni E2E request-id contract (``Qwen/Qwen2.5-Omni-7B``, thinker-only stage). + + Same scenarios as ``test_diffusion_generate_request_id``: + - plain id ``my-req-1`` + - image-style prefix ``img_gen-abc123`` + - chat-style prefix ``chatcmpl-xyz`` + - empty ``request_id`` → output id is non-empty (auto-assigned) + + Text modality only; asserts caller-visible ids are preserved across the + multi-stage orchestrator path on H100. + """ engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG) try: + for req_id in _OMNI_REQ_IDS: + async for output in engine.generate( + "Say hello in one word.", + request_id=req_id, + output_modalities=["text"], + ): + assert output.request_id == req_id async for output in engine.generate( "Say hello in one word.", request_id="", diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index eba1878ec85..d6b33dc89d6 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -284,6 +284,29 @@ def test_serve_cli_accepts_diffusion_attention_backend(): assert diffusion_attention_config.default.backend == "FLASH_ATTN" +def test_serve_cli_accepts_request_batch_max_wait_ms(): + """Ensure diffusion serve CLI forwards request-batch admission wait to stage config.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "Qwen/Qwen-Image", + "--omni", + "--request-batch-max-wait-ms", + "250", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + + assert args.request_batch_max_wait_ms == 250.0 + assert stage_cfg["engine_args"]["request_batch_max_wait_ms"] == 250.0 + + def test_serve_cli_accepts_additional_config(): """Ensure diffusion serve CLI exposes additional_config and forwards it to stage config.""" parser = TrackingArgumentParser() @@ -320,7 +343,7 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo ) mocker.patch( "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", - return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage]), + return_value=("dummy.yaml", [fake_llm_stage, fake_diffusion_stage], None), ) engine = AsyncOmniEngine.__new__(AsyncOmniEngine) @@ -332,7 +355,45 @@ def test_resolve_stage_configs_injects_additional_config_into_diffusion_stage(mo "stage_configs_path": "dummy.yaml", "additional_config": {"torchair_graph_config": {"enabled": True}}, }, + trust_remote_code=False, ) assert not hasattr(stage_configs[0].engine_args, "additional_config") assert stage_configs[1].engine_args.additional_config == {"torchair_graph_config": {"enabled": True}} + + +def test_default_stage_config_includes_quantization_config(): + """Ensure structured quantization_config survives default diffusion-stage creation.""" + quantization_config = { + "method": "example_quant", + "weights": "weights.bin", + } + + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"quantization_config": quantization_config})[0] + + assert stage_cfg["engine_args"]["quantization_config"] == quantization_config + + +def test_resolve_stage_configs_injects_quantization_config_into_diffusion_stage(mocker): + fake_diffusion_stage = SimpleNamespace( + stage_type="diffusion", + engine_args=SimpleNamespace(quantization_config=None), + ) + mocker.patch( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + return_value=("dummy.yaml", [fake_diffusion_stage], None), + ) + + engine = AsyncOmniEngine.__new__(AsyncOmniEngine) + engine._strip_single_engine_args = lambda kwargs: kwargs + + _, stage_configs = engine._resolve_stage_configs( + "dummy-model", + { + "stage_configs_path": "dummy.yaml", + "quantization_config": {"method": "bitsandbytes"}, + }, + trust_remote_code=False, + ) + + assert stage_configs[0].engine_args.quantization_config == {"method": "bitsandbytes"} diff --git a/tests/entrypoints/test_omni_sleep_mode.py b/tests/entrypoints/test_omni_sleep_mode.py index c2a5ff0a103..efbb19fff2a 100644 --- a/tests/entrypoints/test_omni_sleep_mode.py +++ b/tests/entrypoints/test_omni_sleep_mode.py @@ -179,7 +179,7 @@ async def test_llm_sleep_ack(self, llm_engine: AsyncOmni): try: device_id = 0 used_before = get_device_global_memory_used_gib(device_id) - acks = await llm_engine.sleep(stage_ids=[0], level=2) + acks = await llm_engine.sleep(stage_ids=[0], level=1) await asyncio.sleep(1.5) used_after = get_device_global_memory_used_gib(device_id) drop_gib = used_before - used_after @@ -209,7 +209,7 @@ async def test_diffusion_sleep_handshake(self, diffusion_engine: AsyncOmni): """Diffusion Worker stage signal loop""" try: logger.info("Starting Diffusion Worker Handshake Test") - acks = await diffusion_engine.sleep(stage_ids=[0], level=2) + acks = await diffusion_engine.sleep(stage_ids=[0], level=1) def _get_status(ack): return ack.status if hasattr(ack, "status") else ack.get("status") @@ -264,9 +264,9 @@ async def test_diffusion_integrity_bit_level(self, diffusion_engine: AsyncOmni): base_output = output assert base_output is not None and len(base_output.images) > 0 logger.info("Baseline Generation successful.") - # Sleep Level 2 + # Sleep Level 1 logger.info("Entering Deep Sleep (VRAM Scavenging)...") - await diffusion_engine.sleep(stage_ids=[0], level=2) + await diffusion_engine.sleep(stage_ids=[0], level=1) # Wake-up logger.info("Waking up (Reloading Weights)...") await diffusion_engine.wake_up(stage_ids=[0]) @@ -353,8 +353,8 @@ async def test_diffusion_vram_lifecycle_audit(self, diffusion_engine: AsyncOmni) logger.info(f"Diffusion Initial VRAM: {vram_initial:.2f} GiB") # Sleep - logger.info("Triggering Level 2 Deep Sleep (Full Weight Offloading)...") - acks = await diffusion_engine.sleep(stage_ids=[0], level=2) + logger.info("Triggering Level 1 Deep Sleep (Partial Weight Offloading)...") + acks = await diffusion_engine.sleep(stage_ids=[0], level=1) reported_freed_bytes = sum(getattr(ack, "freed_bytes", 0) for ack in acks) reported_freed_gib = reported_freed_bytes / 1024**3 @@ -407,6 +407,47 @@ async def test_diffusion_vram_lifecycle_audit(self, diffusion_engine: AsyncOmni) diffusion_engine.shutdown() await asyncio.sleep(1) + @pytest.mark.asyncio + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + async def test_level2_sleep_wake_raises(self, llm_engine: AsyncOmni): + """Regression for #4473 Repro A: wake_up() after sleep(level=2) must raise + NotImplementedError instead of silently producing corrupted output.""" + try: + await llm_engine.sleep(stage_ids=[0], level=2) + with pytest.raises(NotImplementedError, match="sleep\\(level=2\\)"): + await llm_engine.wake_up(stage_ids=[0]) + finally: + llm_engine.shutdown() + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + async def test_partial_wake_blocks_generate(self, llm_engine: AsyncOmni): + """Regression for #4473 Repro B: generate() must be rejected if kv_cache + is still asleep after wake_up(tags=["weights"]), instead of crashing with + CUDA illegal memory access.""" + try: + await llm_engine.sleep(stage_ids=[0], level=1) + await llm_engine.wake_up(stage_ids=[0], tags=["weights"]) + with pytest.raises(RuntimeError, match="partially or fully asleep"): + async for _ in llm_engine.generate("test", sampling_params=SamplingParams(max_tokens=4)): + pass + finally: + llm_engine.shutdown() + + @pytest.mark.asyncio + @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) + async def test_duplicate_wake_is_idempotent(self, llm_engine: AsyncOmni): + """Regression for #4473 Repro C: duplicate wake_up(tags=None) must be a + safe no-op instead of raising a cumem CUDA invalid argument error.""" + try: + await llm_engine.sleep(stage_ids=[0], level=1) + first_acks = await llm_engine.wake_up(stage_ids=[0]) + assert len(first_acks) > 0, "First wake_up() should return ACKs" + second_acks = await llm_engine.wake_up(stage_ids=[0]) + assert second_acks == [], f"Duplicate wake_up() should return [] but got {second_acks}" + finally: + llm_engine.shutdown() + # --------------------------------------------------------------------------- # H100 (or MI325) e2e: BAGEL / pure diffusion, multi-TP @@ -444,7 +485,7 @@ async def test_diffusion_model_sleep_tp(tp_size: int): async for _ in engine.generate("test", sampling_params_list=[llm_sp, diff_sp]): pass - acks = await engine.sleep(level=2) + acks = await engine.sleep(level=1) statuses = [get_ack_info(ack, "status") for ack in acks] assert all(s == "SUCCESS" for s in statuses), f"Sleep failed. Statuses: {statuses}" @@ -500,7 +541,7 @@ async def test_multistage_sleep_h100(tp_size: int): async for _ in engine.generate("warmup", sampling_params_list=[SamplingParams(), sp]): pass - acks = await engine.sleep(stage_ids=[0, 1], level=2) + acks = await engine.sleep(stage_ids=[0, 1], level=1) assert len(acks) == 2 await engine.wake_up(stage_ids=[0, 1]) diff --git a/tests/entrypoints/test_resolve_dreamzero_config.py b/tests/entrypoints/test_resolve_dreamzero_config.py index a962c5262c1..cc459c269ee 100644 --- a/tests/entrypoints/test_resolve_dreamzero_config.py +++ b/tests/entrypoints/test_resolve_dreamzero_config.py @@ -24,15 +24,22 @@ def test_dreamzero_vla_resolves_to_dreamzero_config(monkeypatch): def test_dreamzero_config_sets_model_class_and_policy_config(monkeypatch): monkeypatch.setattr( - "vllm_omni.config.stage_config.StageConfigFactory._auto_detect_model_type", - classmethod(lambda _cls, _model, trust_remote_code=True: ("vla", None)), + "vllm_omni.config.config_factory.StageConfigFactory._try_infer_model_type", + classmethod(lambda _cls, model, trust_remote_code=True: "vla"), ) monkeypatch.setattr( - "vllm_omni.diffusion.utils.hf_utils._looks_like_dreamzero", + "vllm_omni.config.config_factory.StageConfigFactory.get_hf_config", + classmethod(lambda _cls, model, trust_remote_code=True: None), + ) + monkeypatch.setattr( + "vllm_omni.config.config_factory._looks_like_dreamzero", lambda _model: True, ) - stage_configs = load_stage_configs_from_model("GEAR-Dreams/DreamZero-DROID") + stage_configs, _ = load_stage_configs_from_model( + "GEAR-Dreams/DreamZero-DROID", + trust_remote_code=False, + ) engine_args = stage_configs[0].engine_args assert engine_args.model_class_name == "DreamZeroPipeline" diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index e05983d45b6..61ffed3d5ab 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -9,6 +9,7 @@ from pytest_mock import MockerFixture from vllm_omni.entrypoints.cli.serve import OmniServeCommand, run_headless +from vllm_omni.entrypoints.utils import parse_stage_overrides from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -31,6 +32,20 @@ def test_serve_parser_accepts_no_async_chunk_and_marks_it_explicit() -> None: assert not explicit["async_chunk"] +def test_serve_parser_accepts_strategy_config() -> None: + """``--strategy-config`` must parse onto the ``strategy_config`` dest and be + forwarded as an explicit kwarg so the engine can overlay the strategy.""" + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="subcommand") + cmd = OmniServeCommand() + cmd.subparser_init(subparsers) + + argv = ["serve", "fake-model", "--omni", "--strategy-config", "/tmp/strategy.yaml"] + args = parser.parse_args(argv) + assert args.strategy_config == "/tmp/strategy.yaml" + assert args.get_explicit_kwargs_dict()["strategy_config"] == "/tmp/strategy.yaml" + + def _make_headless_args(**kwargs) -> TrackingNamespace: defaults = { "model": "fake-model", @@ -80,13 +95,97 @@ def test_run_headless_rejects_non_multiprocess_worker_backend() -> None: run_headless(args) +# --------------------------------------------------------------------------- +# --stage-overrides parsing parity (headless vs standard path) +# --------------------------------------------------------------------------- + + +def test_parse_stage_overrides_valid_json() -> None: + """A valid JSON string is parsed into the nested per-stage dict.""" + parsed = parse_stage_overrides('{"0": {"devices": "0,1"}, "1": {"devices": "2"}}') + assert parsed == {"0": {"devices": "0,1"}, "1": {"devices": "2"}} + + +def test_parse_stage_overrides_none_and_empty_return_none() -> None: + """No overrides (None / empty string) resolve to ``None``.""" + assert parse_stage_overrides(None) is None + assert parse_stage_overrides("") is None + + +def test_parse_stage_overrides_empty_dict_returns_none() -> None: + """An empty dict is falsy and must resolve to ``None``, locking in parity + with the original standard-path ``if stage_overrides_json:`` falsy check.""" + assert parse_stage_overrides({}) is None + + +def test_parse_stage_overrides_passes_through_non_str() -> None: + """An already-parsed mapping is returned unchanged (identity).""" + overrides = {"0": {"devices": "0"}} + assert parse_stage_overrides(overrides) is overrides + + +def test_parse_stage_overrides_invalid_json_raises() -> None: + """Invalid JSON raises ValueError whose message matches the standard path + verbatim: the ``--stage-overrides is not valid JSON:`` prefix AND the + ``Got: <repr>`` suffix echoing the raw input.""" + bad = "{not valid json}" + with pytest.raises(ValueError) as excinfo: + parse_stage_overrides(bad) + message = str(excinfo.value) + assert message.startswith("--stage-overrides is not valid JSON:") + assert f"Got: {bad!r}" in message + + +def test_run_headless_parses_and_forwards_stage_overrides(mocker: MockerFixture) -> None: + """Regression: the headless path must parse ``--stage-overrides`` (a JSON + string) and forward the parsed dict to ``load_and_resolve_stage_configs``, + mirroring the standard engine path. Previously it was dropped entirely, + silently producing a different per-stage device layout.""" + captured: dict = {} + + def _fake_resolve(*args, **kwargs): + captured.update(kwargs) + # Return a stage that does NOT match stage_id=0 so run_headless stops + # right after the resolver call (we only care about how it was called). + return ("/fake/stages.yaml", [SimpleNamespace(stage_id=99)], None) + + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + side_effect=_fake_resolve, + ) + + args = _make_headless_args( + stage_id=0, + strategy_config="/tmp/strategy.yaml", + stage_overrides='{"0": {"devices": "0,1"}, "1": {"devices": "2"}}', + ) + with pytest.raises(ValueError, match="No stage config found for stage_id=0"): + run_headless(args) + + assert captured["stage_overrides"] == {"0": {"devices": "0,1"}, "1": {"devices": "2"}} + assert captured["strategy_config_path"] == "/tmp/strategy.yaml" + + +def test_run_headless_invalid_stage_overrides_raises(mocker: MockerFixture) -> None: + """Invalid ``--stage-overrides`` JSON in headless mode fails fast with the + shared ValueError instead of being silently ignored.""" + mocker.patch( + "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", + return_value=("/fake/stages.yaml", [SimpleNamespace(stage_id=0)], None), + ) + + args = _make_headless_args(stage_id=0, stage_overrides="{not valid json}") + with pytest.raises(ValueError, match="--stage-overrides is not valid JSON"): + run_headless(args) + + def test_run_headless_raises_when_stage_id_not_in_configs(mocker: MockerFixture) -> None: """Headless looks up its assigned stage_id in the loaded deploy YAML and fails fast when the launcher's --stage-id doesn't match any entry.""" other_stage = SimpleNamespace(stage_id=99) mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [other_stage]), + return_value=("/fake/stages.yaml", [other_stage], None), ) args = _make_headless_args(stage_id=0) @@ -119,7 +218,7 @@ def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace: def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None: """LLM headless: each loop iteration registers with auto-assigned replica_id (master picks a free slot) and spawns one - ``OmniCoreEngineProcManager`` per local replica.""" + ``StageEngineCoreProcManager`` per local replica.""" from vllm_omni.engine.stage_engine_startup import StageRegistrationResponse stage_cfg = _make_stage_cfg(0, stage_type="llm") @@ -134,7 +233,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -159,7 +258,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker ), ) mock_manager_cls = mocker.patch( - "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager", + "vllm_omni.engine.stage_engine_core_proc_manager.StageEngineCoreProcManager", return_value=engine_manager, ) mocker.patch("signal.signal") @@ -168,8 +267,8 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker # The launcher must request auto-assignment (replica_id=None) and the # full response so it can wire the master-allocated coordinator into the - # spawned subprocess. ``replica_binds_sockets=False`` is required for LLM - # because the head binds all three sockets (handshake, input, output). + # spawned subprocess. LLM uses head-owned sockets: the head binds all + # three sockets (handshake, input, output) and the worker connects. assert mock_register.call_count == 1 kwargs = mock_register.call_args.kwargs assert kwargs["omni_master_address"] == "127.0.0.1" @@ -177,8 +276,7 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker assert kwargs["omni_stage_id"] == 0 assert kwargs["omni_stage_config"] is stage_cfg assert kwargs["replica_id"] is None - assert kwargs["return_full_response"] is True - assert kwargs["replica_binds_sockets"] is False + assert "socket_ownership" not in kwargs assert mock_manager_cls.call_count == 1 mgr_kwargs = mock_manager_cls.call_args.kwargs @@ -211,7 +309,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -239,7 +337,7 @@ def test_run_headless_llm_launches_one_manager_per_omni_dp_size_local(mocker: Mo ], ) mock_manager_cls = mocker.patch( - "vllm_omni.engine.omni_core_engine_proc_manager.OmniCoreEngineProcManager", + "vllm_omni.engine.stage_engine_core_proc_manager.StageEngineCoreProcManager", side_effect=[manager_a, manager_b], ) mocker.patch("signal.signal") @@ -270,7 +368,7 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -294,18 +392,24 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) coordinator_router_address="tcp://127.0.0.1:26100", ), ) - mock_spawn = mocker.patch( - "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", - return_value=(proc, None, None, None), + fake_manager = SimpleNamespace( + proc=proc, + addresses=SimpleNamespace( + inputs=["tcp://127.0.0.1:26002"], + outputs=["tcp://127.0.0.1:26003"], + ), + shutdown=mocker.Mock(), + ) + mock_manager = mocker.patch( + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager.launch_headless", + return_value=fake_manager, ) - mock_handshake = mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") # Replace the blocking wait with one that returns the only proc's sentinel # immediately so the test does not hang. mocker.patch( "multiprocessing.connection.wait", side_effect=lambda sentinels: [sentinels[0]], ) - mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc") mocker.patch("signal.signal") run_headless(_make_headless_args(stage_id=1)) @@ -321,17 +425,15 @@ def test_run_headless_diffusion_registers_and_spawns_proc(mocker: MockerFixture) assert reg_kwargs["omni_stage_id"] == 1 assert reg_kwargs["omni_stage_config"] is stage_cfg assert reg_kwargs["replica_id"] is None - assert reg_kwargs["return_full_response"] is True - - spawn_kwargs = mock_spawn.call_args.kwargs - assert spawn_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" - assert spawn_kwargs["request_address"] == "tcp://127.0.0.1:26002" - assert spawn_kwargs["response_address"] == "tcp://127.0.0.1:26003" - assert spawn_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100" - assert spawn_kwargs["omni_stage_id"] == 1 - assert spawn_kwargs["omni_replica_id"] == 0 + assert "socket_ownership" not in reg_kwargs - mock_handshake.assert_called_once_with(proc, "tcp://127.0.0.1:26001", 600) + manager_kwargs = mock_manager.call_args.kwargs + assert manager_kwargs["handshake_address"] == "tcp://127.0.0.1:26001" + assert manager_kwargs["addresses"].inputs == ["tcp://127.0.0.1:26002"] + assert manager_kwargs["addresses"].outputs == ["tcp://127.0.0.1:26003"] + assert manager_kwargs["omni_coordinator_address"] == "tcp://127.0.0.1:26100" + assert manager_kwargs["omni_stage_id"] == 1 + assert manager_kwargs["omni_replica_id"] == 0 def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixture) -> None: @@ -346,7 +448,7 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur mocker.patch( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", - return_value=("/fake/stages.yaml", [stage_cfg]), + return_value=("/fake/stages.yaml", [stage_cfg], None), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -371,15 +473,13 @@ def test_run_headless_diffusion_raises_on_nonzero_proc_exit(mocker: MockerFixtur ), ) mocker.patch( - "vllm_omni.diffusion.stage_diffusion_proc.spawn_diffusion_proc", - return_value=(proc, None, None, None), + "vllm_omni.diffusion.stage_diffusion_proc.StageDiffusionProcManager.launch_headless", + return_value=SimpleNamespace(proc=proc, shutdown=mocker.Mock()), ) - mocker.patch("vllm_omni.diffusion.stage_diffusion_proc.complete_diffusion_handshake") mocker.patch( "multiprocessing.connection.wait", side_effect=lambda sentinels: [sentinels[0]], ) - mocker.patch("vllm_omni.engine.stage_init_utils.terminate_alive_proc") mocker.patch("signal.signal") with pytest.raises(RuntimeError, match=r"exited with code 137"): diff --git a/tests/entrypoints/test_stage_utils.py b/tests/entrypoints/test_stage_utils.py index 05593d647b3..c4c0667df84 100644 --- a/tests/entrypoints/test_stage_utils.py +++ b/tests/entrypoints/test_stage_utils.py @@ -4,7 +4,7 @@ import pytest from pytest_mock import MockerFixture -from vllm_omni.entrypoints.stage_utils import _map_device_list, set_stage_devices +from vllm_omni.entrypoints.stage_utils import _map_device_list, resolve_stage_physical_devices, set_stage_devices def _make_dummy_torch(call_log): @@ -207,6 +207,13 @@ def test_map_device_list_multi_replica_offset_cvd(): assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"] +@pytest.mark.core_model +@pytest.mark.cpu +def test_resolve_stage_physical_devices_uses_visible_baseline(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3") + assert resolve_stage_physical_devices(1, "1", visible_baseline="3,4,5") == "4" + + @pytest.mark.core_model @pytest.mark.cpu def test_map_device_list_raises_on_non_numeric(): diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 479ada96396..4c8522ae0e4 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -315,13 +315,17 @@ def mock_exists(path): class TestLoadAndResolveStageConfigs: - def test_load_and_resolve_with_kwargs(self): + def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture): """Ensure that dtype survives default stage creation.""" kwargs = {"dtype": torch.float32} - config_path, stage_configs = load_and_resolve_stage_configs( + mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None) + mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None)) + + config_path, stage_configs, _ = load_and_resolve_stage_configs( model="black-forest-labs/FLUX.2-klein-4B", stage_configs_path=None, kwargs=kwargs, + trust_remote_code=False, default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs), ) assert config_path is None @@ -349,20 +353,23 @@ def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas( ] load_stage_configs = mocker.patch( "vllm_omni.entrypoints.utils.load_stage_configs_from_model", - return_value=returned_stage_configs, + return_value=(returned_stage_configs, None), ) - config_path, stage_configs = load_and_resolve_stage_configs( + config_path, stage_configs, _ = load_and_resolve_stage_configs( model="dummy-model", stage_configs_path=str(deploy_path), kwargs={}, + trust_remote_code=True, ) load_stage_configs.assert_called_once_with( "dummy-model", + trust_remote_code=True, base_engine_args={}, deploy_config_path=str(deploy_path), stage_overrides=None, + strategy_config_path=None, ) assert config_path == str(deploy_path) assert len(stage_configs) == 2 diff --git a/tests/examples/offline_inference/test_image_to_video.py b/tests/examples/offline_inference/test_image_to_video.py new file mode 100644 index 00000000000..b53648cb334 --- /dev/null +++ b/tests/examples/offline_inference/test_image_to_video.py @@ -0,0 +1,81 @@ +""" +Offline inference tests: image-to-video. +See examples/offline_inference/image_to_video/README.md +""" + +import shutil +from pathlib import Path + +import pytest +import requests + +from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet +from tests.helpers.assertions import assert_video_valid +from tests.helpers.mark import hardware_marks + +pytestmark = [ + pytest.mark.usefixtures("clean_gpu_memory_between_tests"), + pytest.mark.full_model, + pytest.mark.example, + *hardware_marks(res={"cuda": "H100"}), +] + +I2V_SCRIPT = EXAMPLES / "offline_inference" / "image_to_video" / "image_to_video.py" +README_PATH = I2V_SCRIPT.with_name("README.md") +EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_i2v" + +_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg" +_IMAGE_NAME = "cherry_blossom.jpg" + +_SKIP_SECTIONS = { + "Prerequisites", + "Advanced Features", + "FAQ", + # VACE conditional tasks are covered by test_vace_video_generation.py, which + # runs the same README snippets with synthetic assets and smoke settings. The + # prep block here downloads inputs via wget (not portable, not collected), so + # exercising this section from the shared runner would only duplicate and fail. + "Wan2.1 VACE Conditional Tasks", +} + + +def _ensure_test_image(run_dir: Path) -> None: + """Download or copy the example image into run_dir so CLI snippets can find it.""" + dest = run_dir / _IMAGE_NAME + if dest.exists(): + return + src = I2V_SCRIPT.parent / _IMAGE_NAME + if src.exists(): + shutil.copy2(src, dest) + return + response = requests.get(_IMAGE_URL, timeout=60) + response.raise_for_status() + dest.write_bytes(response.content) + + +def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool, str]: + if h2_title in _SKIP_SECTIONS: + return True, f"README section '{h2_title}' is intentionally excluded for examples tests" + if language == "python": + return True, "Python API snippets produce video files that ExampleRunner does not auto-collect" + if "/path/to/" in code: + return True, "Snippet references a placeholder local model path" + return False, "" + + +README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet) + + +@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id) +def test_image_to_video(snippet: ReadmeSnippet, example_runner: ExampleRunner): + should_skip, reason = snippet.skip + if should_skip: + pytest.skip(reason) + + run_dir = example_runner.output_root / EXAMPLE_OUTPUT_SUBFOLDER / snippet.test_id + run_dir.mkdir(parents=True, exist_ok=True) + _ensure_test_image(run_dir) + + result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER)) + for asset in result.assets: + assert_video_valid(asset) diff --git a/tests/examples/offline_inference/test_vace_video_generation.py b/tests/examples/offline_inference/test_vace_video_generation.py index 11746e58a20..e4f6ca93d09 100644 --- a/tests/examples/offline_inference/test_vace_video_generation.py +++ b/tests/examples/offline_inference/test_vace_video_generation.py @@ -1,11 +1,9 @@ -""" -Offline inference tests: vace video generation. -See examples/offline_inference/vace/vace_video_generation.md -""" +"""VACE coverage through the shared text-to-video and image-to-video examples.""" from pathlib import Path import pytest +from PIL import Image from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet from tests.helpers.assertions import assert_video_valid @@ -14,22 +12,83 @@ pytestmark = [ pytest.mark.usefixtures("clean_gpu_memory_between_tests"), pytest.mark.full_model, + pytest.mark.diffusion, pytest.mark.example, *hardware_marks(res={"cuda": "H100"}), ] -VACE_SCRIPT = EXAMPLES / "offline_inference" / "vace" / "vace_video_generation.py" -README_PATH = VACE_SCRIPT.with_name("vace_video_generation.md") +TEXT_TO_VIDEO_README = EXAMPLES / "offline_inference" / "text_to_video" / "text_to_video.md" +IMAGE_TO_VIDEO_README = EXAMPLES / "offline_inference" / "image_to_video" / "README.md" EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_vace" -README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH) +TEXT_TO_VIDEO_SNIPPET = next( + snippet + for snippet in ReadmeSnippet.extract_readme_snippets(TEXT_TO_VIDEO_README) + if "vace_t2v_output.mp4" in snippet.code +) +CONDITIONAL_OUTPUTS = ( + "vace_i2v_output.mp4", + "vace_v2lf_output.mp4", + "vace_flf2v_output.mp4", + "vace_inpaint_output.mp4", + "vace_r2v_output.mp4", +) +IMAGE_TO_VIDEO_SNIPPETS = [ + snippet + for snippet in ReadmeSnippet.extract_readme_snippets(IMAGE_TO_VIDEO_README) + if snippet.h2_title == "Wan2.1 VACE Conditional Tasks" + and any(output in snippet.code for output in CONDITIONAL_OUTPUTS) +] + + +def _smoke_snippet(snippet: ReadmeSnippet, asset_dir: Path) -> ReadmeSnippet: + # Six fresh model launches share a 90-minute X2V job. Keep the documented + # quality settings intact while reducing only the CI executions. + replacements = { + "astronaut.jpg": str(asset_dir / "astronaut.jpg"), + "vace_first_frame.png": str(asset_dir / "vace_first_frame.png"), + "vace_last_frame.png": str(asset_dir / "vace_last_frame.png"), + "vace_center_mask.png": str(asset_dir / "vace_center_mask.png"), + "--num-frames 81": "--num-frames 5", + "--num-inference-steps 30": "--num-inference-steps 2", + } + code = snippet.code + for source, target in replacements.items(): + code = code.replace(source, target) + return snippet._replace(code=code) + + +@pytest.fixture +def vace_asset_dir(tmp_path: Path) -> Path: + asset_dir = tmp_path / "vace_assets" + asset_dir.mkdir(parents=True, exist_ok=True) + Image.new("RGB", (832, 480), (180, 180, 180)).save(asset_dir / "astronaut.jpg") + Image.new("RGB", (512, 512), (80, 140, 220)).save(asset_dir / "vace_first_frame.png") + Image.new("RGB", (512, 512), (220, 140, 80)).save(asset_dir / "vace_last_frame.png") + mask = Image.new("L", (832, 480), 0) + mask.paste(255, (336, 0, 496, 480)) + mask.save(asset_dir / "vace_center_mask.png") + return asset_dir + + +def test_vace_text_to_video_shared_example(example_runner: ExampleRunner) -> None: + snippet = TEXT_TO_VIDEO_SNIPPET._replace( + code=TEXT_TO_VIDEO_SNIPPET.code.replace("--num-frames 81", "--num-frames 5").replace( + "--num-inference-steps 30", "--num-inference-steps 2" + ) + ) + result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER)) + for asset in result.assets: + assert_video_valid(asset) -@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id) -def test_vace_video_generation(snippet: ReadmeSnippet, example_runner: ExampleRunner): - should_skip, reason = snippet.skip - if should_skip: - pytest.skip(reason) +@pytest.mark.parametrize("snippet", IMAGE_TO_VIDEO_SNIPPETS, ids=lambda snippet: snippet.test_id) +def test_vace_conditional_shared_examples( + snippet: ReadmeSnippet, + example_runner: ExampleRunner, + vace_asset_dir: Path, +) -> None: + snippet = _smoke_snippet(snippet, vace_asset_dir) result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER)) for asset in result.assets: assert_video_valid(asset) diff --git a/tests/fullduplex/test_joyvl_brain.py b/tests/fullduplex/test_joyvl_brain.py new file mode 100644 index 00000000000..833ce802301 --- /dev/null +++ b/tests/fullduplex/test_joyvl_brain.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import parse_action +from vllm_omni.experimental.fullduplex.joyvl.memory.brain import InteractionBrain + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeSummarizer: + def __init__(self): + self.chunks = 0 + + async def summarize_chunk(self, chunk_index, frame_range, frames): + self.chunks += 1 + return f"summary-of-chunk-{chunk_index}" + + async def compress_to_long_term(self, mids): + return f"compressed({len(mids)})" + + +class _FakeDelegation: + def __init__(self): + self._ready = {} + + async def submit(self, question, note, frames): + self._ready["t1"] = question + return "t1" + + async def poll(self, task_id): + from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationResult + + return DelegationResult(task_id, "ready", digest=f"answer to {self._ready[task_id]}") + + +def test_query_freshness_and_archive_on_change(): + b = InteractionBrain() + assert b.update_query("count bottles") is True + assert b.update_query("count bottles") is False + b.record_response("1 bottle") + assert b.update_query("now describe") is True + assert b.memory.qa_history[-1].query == "count bottles" + assert b.memory.qa_history[-1].responses == [("0.0 seconds", "1 bottle")] + + +def test_query_in_frame_on_arrival_then_prefix_on_later_chunks(): + b = InteractionBrain(frame_seconds=1.0) + b.tick(3) + b.update_query("alert me if a fire breaks out") + assert b.query_in_current_chunk is True + assert "alert me if a fire breaks out" not in b.build_prefix() + b.record_response("a fire is breaking out") + b.close_chunk() + assert b.query_in_current_chunk is False + prefix = b.build_prefix() + assert "alert me if a fire breaks out" in prefix + assert "a fire is breaking out" in prefix + + +@pytest.mark.asyncio +async def test_flush_archives_qa_and_summarizes(): + summ = _FakeSummarizer() + b = InteractionBrain(summarizer=summ, chunk_frames=3, long_term_every_n_chunks=2, frame_seconds=1.0) + b.update_query("watch the room") + b.tick(3) + assert b.should_flush() is True + b.record_response("a person enters") + await b.flush([("0.0s", "u0"), ("2.0s", "u2")]) + assert summ.chunks == 1 + assert b.chunk_index == 2 + assert b.should_flush() is False + assert b.memory.qa_history[-1].responses == [("2.0 seconds", "a person enters")] + assert len(b.memory.mid_term_summaries) == 1 + + +@pytest.mark.asyncio +async def test_long_term_compression_after_n_chunks(): + summ = _FakeSummarizer() + b = InteractionBrain(summarizer=summ, chunk_frames=1, long_term_every_n_chunks=2, frame_seconds=1.0) + for _ in range(2): + await b.flush([("0.0s", "u")]) + assert b.memory.long_term_memory.startswith("compressed(") + assert b.memory.mid_term_summaries == [] + + +@pytest.mark.asyncio +async def test_long_term_window_caps_blocks(): + # Long-term memory must not grow unbounded over hours: oldest blocks are dropped + # once the sliding window is full (reference window = 15; here 2 for the test). + summ = _FakeSummarizer() + b = InteractionBrain( + summarizer=summ, chunk_frames=1, long_term_every_n_chunks=1, long_term_window=2, frame_seconds=1.0 + ) + for _ in range(4): + await b.flush([("0.0s", "u")]) + assert len(b.memory.long_term_blocks) == 2 # capped, not 4 + assert b.memory.long_term_memory.count("compressed(") == 2 + + +@pytest.mark.asyncio +async def test_delegation_submit_and_fold(): + b = InteractionBrain(delegation=_FakeDelegation()) + action = parse_action("</response> hold on <delegation> solve this") + info = await b.submit_delegation(action, []) + assert info["status"] == "submitted" + folded = await b.fold_delegations() + assert folded["status"] == "ready" + assert b.memory.qa_history[-1].responses[0][1] == "answer to solve this" diff --git a/tests/fullduplex/test_joyvl_policy.py b/tests/fullduplex/test_joyvl_policy.py new file mode 100644 index 00000000000..5fecdc50290 --- /dev/null +++ b/tests/fullduplex/test_joyvl_policy.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio + +import pytest + +from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import ( + DelegationResult, + ImageEditDelegationBridge, + ImageGenDelegationBridge, + OpenAIDelegationBridge, +) +from vllm_omni.experimental.fullduplex.joyvl.decision.policy import JoyVLPolicy, sample_frames + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_sample_frames_keeps_recent(): + assert sample_frames(["a", "b"], 4) == ["a", "b"] + out = sample_frames([str(i) for i in range(10)], 4) + assert out[-1] == "9" and len(out) == 4 + + +def test_build_messages_stable_head_and_query(): + p = JoyVLPolicy(num_frames=4) + p.tick(2) + p.set_query("count the bottles") + msgs, user = p.build_messages([{"type": "image_url", "image_url": {"url": "x"}}]) + assert msgs[0]["role"] == "system" and "</silence>" in msgs[0]["content"] + + assert any("count the bottles" in part.get("text", "") for part in user["content"]) + + +def test_commit_records_only_spoken(): + p = JoyVLPolicy() + p.set_query("q") + assert p.commit("</silence>").action.value == "silence" + assert p.brain.response_records == [] + assert p.commit("</response> hello").action.value == "response" + assert p.brain.response_records == [("0.0 seconds", "hello")] + + +def test_commit_dedup_default_is_exact_match_only(): + p = JoyVLPolicy() # threshold 1.0 == reference behavior + assert p.commit("</response> a red ball appears").action.value == "response" + assert p.commit("</response> a red ball appears").action.value == "silence" # exact repeat dropped + assert p.commit("</response> a red ball is appearing").action.value == "response" # near-dup kept + + +def test_commit_dedup_threshold_drops_near_duplicates(): + p = JoyVLPolicy(response_dedup_threshold=0.8) + assert p.commit("</response> a red ball appears").action.value == "response" + assert p.commit("</response> a red ball appears!").action.value == "silence" # near-dup dropped + assert p.commit("</response> the sky is now blue").action.value == "response" # distinct kept + + +def test_commit_dedup_rearms_after_silence(): + p = JoyVLPolicy(response_dedup_threshold=0.85) + assert p.commit("</response> someone raises a hand").action.value == "response" + # no gap: an immediate repeat is still suppressed + assert p.commit("</response> someone raises a hand").action.value == "silence" + # a genuine silent tick (hand lowered) re-arms; the recurrence speaks again + assert p.commit("</silence>").action.value == "silence" + assert p.commit("</response> someone raises a hand").action.value == "response" + + +class _FakeBrain(OpenAIDelegationBridge): + """Override only the network call so submit/poll plumbing is exercised for real.""" + + def __init__(self, answer, gate=None): + super().__init__("http://unused/v1", "background-model") + self._answer = answer + self._gate = gate + + async def _complete(self, messages): + if self._gate is not None: + await self._gate.wait() + return self._answer + + +@pytest.mark.asyncio +async def test_openai_delegation_bridge_is_nonblocking_and_resolves(): + gate = asyncio.Event() + bridge = _FakeBrain("the bottle is on the left shelf", gate=gate) + frames = [("0.0s", "data:image/jpeg;base64,AAA")] + task_id = await bridge.submit("where is the bottle?", "let me check", frames) + + # poll must not block on the slow background brain + assert (await bridge.poll(task_id)).status == "pending" + + gate.set() + await asyncio.sleep(0) + for _ in range(10): + result = await bridge.poll(task_id) + if result.status != "pending": + break + await asyncio.sleep(0) + assert result.is_ready + assert "bottle is on the left shelf" in result.digest + assert (await bridge.poll("missing")).status == "error" + + +class _FakeImageBrain(ImageGenDelegationBridge): + async def _call_images(self, payload): + # echo the prompt back so the test can assert it was passed through + return {"data": [{"b64_json": "QUJD"}]} # "ABC" + + +@pytest.mark.asyncio +async def test_image_delegation_bridge_returns_media_not_text(): + bridge = _FakeImageBrain("http://unused/v1") + task_id = await bridge.submit("draw a red dragon", "", []) + for _ in range(10): + result = await bridge.poll(task_id) + if result.status != "pending": + break + await asyncio.sleep(0) + assert result.is_ready + # the image rides in media as a data URL; digest is only a short text placeholder + assert result.media == "data:image/png;base64,QUJD" + assert "QUJD" not in result.digest and "red dragon" in result.digest + + +class _FakeEditBrain(ImageEditDelegationBridge): + def __init__(self, base_url, model): + super().__init__(base_url, model) + self.seen_frame = None + + async def _call_chat(self, frame, instruction): + self.seen_frame = frame + return {"choices": [{"message": {"content": [{"image_url": {"url": "data:image/png;base64,Q0FSVA=="}}]}}]} + + +@pytest.mark.asyncio +async def test_image_edit_delegation_conditions_on_frame_and_returns_media(): + bridge = _FakeEditBrain("http://unused/v1", "Qwen-Image-Edit") + frames = [("0.0s", "data:image/jpeg;base64,OLD"), ("1.0s", "data:image/jpeg;base64,NOW")] + task_id = await bridge.submit("把画面变成卡通风格", "", frames) + for _ in range(10): + result = await bridge.poll(task_id) + if result.status != "pending": + break + await asyncio.sleep(0) + assert result.is_ready + assert bridge.seen_frame == "data:image/jpeg;base64,NOW" # latest frame, not the old one + assert result.media == "data:image/png;base64,Q0FSVA==" + assert result.media not in result.digest # image rides in media, not context + + +@pytest.mark.asyncio +async def test_image_edit_delegation_errors_without_frame(): + bridge = _FakeEditBrain("http://unused/v1", "Qwen-Image-Edit") + task_id = await bridge.submit("cartoonify", "", []) + for _ in range(10): + result = await bridge.poll(task_id) + if result.status != "pending": + break + await asyncio.sleep(0) + assert result.status == "error" + + +class _NamedBridge: + def __init__(self, name): + self.name = name + + async def submit(self, question, note, frames): + return f"{self.name}-1" + + async def poll(self, task_id): + return DelegationResult(task_id, "ready", digest=self.name) + + +@pytest.mark.asyncio +async def test_routing_delegation_dispatches_by_request(): + from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import RoutingDelegationBridge + + chat, image, edit = _NamedBridge("chat"), _NamedBridge("image"), _NamedBridge("edit") + r = RoutingDelegationBridge(chat=chat, image=image, edit=edit) + + async def route(q): + tid = await r.submit(q, "", []) + return (await r.poll(tid)).digest + + assert await route("把你看到的画面变成卡通风格") == "edit" # restyle current view + assert await route("画一只红色的龙") == "image" # generate a new image + assert await route("这道数学题的答案是什么") == "chat" # hard question -> chat brain + + +class _Delegation: + async def submit(self, question, note, frames): + return "t1" + + async def poll(self, task_id): + return DelegationResult(task_id, "ready", digest="background answer") + + +def test_build_delegation_off_unless_backend_configured(): + from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig + from vllm_omni.experimental.fullduplex.joyvl.serving.server import SessionManager + + # No backend URL -> delegation disabled (None), NOT a silent stub that folds fake answers. + assert SessionManager._build_delegation(InteractionConfig()) is None + assert SessionManager._build_delegation(InteractionConfig(delegation_kind="router")) is None + # Stub only on explicit opt-in. + assert SessionManager._build_delegation(InteractionConfig(delegation_kind="stub")) is not None + # A real chat backend -> a live bridge. + cfg = InteractionConfig(delegation_backend_url="http://x/v1", delegation_model="m") + assert SessionManager._build_delegation(cfg) is not None + + +@pytest.mark.asyncio +async def test_delegation_cancel_drops_task(): + from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import StubDelegationBridge + + b = StubDelegationBridge(ready_after_ticks=1) + task_id = await b.submit("q", "", []) + b.cancel(task_id) + # after cancel the task is gone -> poll reports it as unknown rather than leaking + assert (await b.poll(task_id)).status == "error" + + +@pytest.mark.asyncio +async def test_delegation_aclose_clears_pending(): + from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import StubDelegationBridge + + b = StubDelegationBridge(ready_after_ticks=5) + task_id = await b.submit("q", "", []) + await b.aclose() + assert (await b.poll(task_id)).status == "error" + + +@pytest.mark.asyncio +async def test_session_reset_cancels_and_awaits_consolidation(): + from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig + from vllm_omni.experimental.fullduplex.joyvl.serving.session import InteractionSession + + class _Backend: + async def generate(self, *a, **k): + return "</silence>", None + + async def aclose(self): + pass + + class _Summarizer: + async def summarize_chunk(self, *a, **k): + await asyncio.sleep(100) # hang so the consolidation task is in-flight + return "x" + + async def compress_to_long_term(self, *a, **k): + return "" + + async def aclose(self): + pass + + sess = InteractionSession("s", InteractionConfig(), _Backend(), summarizer=_Summarizer()) + sess._spawn_consolidation(1, [("0.0s", "u")]) + assert len(sess._consolidating) == 1 + await sess.reset() # must cancel AND await the in-flight task, not just clear the set + assert len(sess._consolidating) == 0 + + +@pytest.mark.asyncio +async def test_manager_aclose_closes_without_error(): + from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig + from vllm_omni.experimental.fullduplex.joyvl.serving.server import SessionManager + + mgr = SessionManager(InteractionConfig(enable_memory=False, enable_delegation=False)) + await mgr.aclose() + await mgr.aclose() # idempotent + + +@pytest.mark.asyncio +async def test_delegation_submit_and_fold(): + p = JoyVLPolicy(delegation=_Delegation()) + action = p.commit("</response> hold on <delegation> hard question") + assert action.action.value == "delegate" + info = await p.submit_if_delegate(action) + assert info["status"] == "submitted" + folded = await p.fold_delegations() + assert folded["status"] == "ready" + assert p.brain.memory.qa_history[-1].responses[0][1] == "background answer" diff --git a/tests/fullduplex/test_runtime.py b/tests/fullduplex/test_runtime.py new file mode 100644 index 00000000000..fe25ce42f6b --- /dev/null +++ b/tests/fullduplex/test_runtime.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +from collections.abc import AsyncIterator + +import pytest + +from vllm_omni.experimental.fullduplex.core import protocol as ev +from vllm_omni.experimental.fullduplex.core.adapter import DuplexCapability, OutputChunk +from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime +from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig +from vllm_omni.experimental.fullduplex.joyvl.adapter import JoyVLDuplexAdapter + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +async def _feed(events): + for e in events: + yield e + + +def _collector(): + out: list[dict] = [] + + async def emit(event: dict) -> None: + out.append(event) + + return out, emit + + +class _FakeAdapter: + def __init__(self, chunks, barge_after=None): + self._chunks = chunks + self._barge_after = barge_after + + def capabilities(self): + return DuplexCapability(frozenset({"text"}), frozenset({"text"}), proactive=False) + + async def on_input(self, session, modality, data): + pass + + def should_respond(self, session): + return True + + async def respond(self, session) -> AsyncIterator[OutputChunk]: + for i, c in enumerate(self._chunks): + if self._barge_after is not None and i == self._barge_after: + session.barge_in() + yield OutputChunk("text", c) + + async def on_barge_in(self, session): + pass + + async def on_playback_ack(self, session, cursor): + pass + + +@pytest.mark.asyncio +async def test_runtime_basic_response(): + session = DuplexSession("s", DuplexSessionConfig(output_modalities=("text",))) + rt = DuplexRuntime(session, _FakeAdapter(["a", "b"])) + out, emit = _collector() + await rt.run(_feed([{"type": ev.INPUT_COMMIT}, {"type": ev.CLOSE}]), emit) + types = [e["type"] for e in out] + assert types == [ev.RESPONSE_CREATED, ev.RESPONSE_DELTA, ev.RESPONSE_DELTA, ev.RESPONSE_DONE] + assert [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA] == ["a", "b"] + + +@pytest.mark.asyncio +async def test_runtime_barge_in_drops_stale_output(): + session = DuplexSession("s") + rt = DuplexRuntime(session, _FakeAdapter(["a", "b", "c"], barge_after=1)) + out, emit = _collector() + await rt.run(_feed([{"type": ev.INPUT_COMMIT}, {"type": ev.CLOSE}]), emit) + data = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA] + assert data == ["a"] + assert ev.RESPONSE_DONE not in [e["type"] for e in out] + + +class _SlowAdapter: + def capabilities(self): + return DuplexCapability(frozenset({"text"}), frozenset({"text"}), proactive=False) + + async def on_input(self, session, modality, data): + pass + + def should_respond(self, session): + return True + + async def respond(self, session) -> AsyncIterator[OutputChunk]: + for c in ["a", "b", "c", "d", "e"]: + await asyncio.sleep(0.02) + yield OutputChunk("text", c) + + async def on_barge_in(self, session): + pass + + async def on_playback_ack(self, session, cursor): + pass + + +@pytest.mark.asyncio +async def test_runtime_cancel_event_interrupts_active_response(): + session = DuplexSession("s") + rt = DuplexRuntime(session, _SlowAdapter()) + out, emit = _collector() + + async def feed(): + yield {"type": ev.INPUT_COMMIT} + await asyncio.sleep(0.03) + yield {"type": ev.RESPONSE_CANCEL} + yield {"type": ev.CLOSE} + + await rt.run(feed(), emit) + deltas = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA] + types = [e["type"] for e in out] + assert len(deltas) < 5 + assert ev.RESPONSE_CANCELLED in types + assert ev.RESPONSE_DONE not in types + + +@pytest.mark.asyncio +async def test_runtime_new_response_supersedes_inflight_without_blocking(): + # A second response trigger arriving while the first is still streaming must + # supersede it (barge-in) rather than block the input loop until it finishes. + session = DuplexSession("s") + rt = DuplexRuntime(session, _SlowAdapter()) + out, emit = _collector() + + async def feed(): + yield {"type": ev.INPUT_COMMIT} + await asyncio.sleep(0.03) # let the first response emit ~one chunk + yield {"type": ev.INPUT_COMMIT} # supersedes the in-flight first response + await asyncio.sleep(0.2) # let the second response finish + yield {"type": ev.CLOSE} + + await rt.run(feed(), emit) + created = [e for e in out if e["type"] == ev.RESPONSE_CREATED] + done = [e for e in out if e["type"] == ev.RESPONSE_DONE] + # two responses created; only the latest one completes (the first was cancelled) + assert len(created) == 2 + assert len(done) == 1 + assert done[0]["response_index"] == created[-1]["response_index"] + + +@pytest.mark.asyncio +async def test_joyvl_adapter_speaks_then_silences(): + replies = iter(["</response> a fire is breaking out", "</silence>"]) + + async def fake_generate(messages): + return next(replies) + + adapter = JoyVLDuplexAdapter(fake_generate, num_frames=4) + cfg = DuplexSessionConfig(input_modalities=("video", "text"), output_modalities=("text",), proactive=True) + rt = DuplexRuntime(DuplexSession("s", cfg), adapter) + out, emit = _collector() + await rt.run( + _feed( + [ + {"type": ev.INPUT_APPEND, "modality": "text", "data": "alert me if a fire breaks out"}, + {"type": ev.INPUT_APPEND, "modality": "video", "data": "data:image/jpeg;base64,AAA"}, + {"type": ev.INPUT_APPEND, "modality": "video", "data": "data:image/jpeg;base64,BBB"}, + {"type": ev.CLOSE}, + ] + ), + emit, + ) + deltas = [e["data"] for e in out if e["type"] == ev.RESPONSE_DELTA] + assert "a fire is breaking out" in deltas + + assert deltas.count("a fire is breaking out") == 1 diff --git a/tests/gr00t/__init__.py b/tests/gr00t/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/tests/gr00t/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/gr00t/openpi_client_helper.py b/tests/gr00t/openpi_client_helper.py new file mode 100644 index 00000000000..e5943ca3307 --- /dev/null +++ b/tests/gr00t/openpi_client_helper.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass +from typing import Any + +import numpy as np + +try: + import websockets.sync.client as websockets_client +except ImportError: # pragma: no cover - optional e2e dependency + websockets_client = None + +try: + from openpi_client import msgpack_numpy +except ImportError: # pragma: no cover - optional e2e dependency + msgpack_numpy = None + +PING_INTERVAL_SECS = 300 +PING_TIMEOUT_SECS = 3600 +DEFAULT_HOST = "127.0.0.1" +DEFAULT_PORT = 8000 +DEFAULT_PATH = "/v1/realtime/robot/openpi" +DEFAULT_SESSION_ID = "gr00t-smoke" +ACTION_KEYS = {"eef_9d", "gripper_position", "joint_position"} +LANGUAGE_KEY = "annotation.language.language_instruction" + + +def _identity_eef_9d_state() -> np.ndarray: + state = np.zeros((1, 1, 9), dtype=np.float32) + state[..., 3:] = np.array([1, 0, 0, 0, 1, 0], dtype=np.float32) + return state + + +def require_dependencies() -> None: + missing = [] + if websockets_client is None: + missing.append("websockets") + if msgpack_numpy is None: + missing.append("openpi-client") + if missing: + raise ModuleNotFoundError(f"GR00T OpenPI test dependencies are missing: {', '.join(missing)}") + + +@dataclass(frozen=True) +class Gr00tServerMetadata: + action_horizon: int + action_keys: set[str] + embodiment_tag: str + needs_session_id: bool + + @classmethod + def from_dict(cls, payload: dict[str, Any]): + required_keys = ("action_horizon", "action_keys", "embodiment_tag", "needs_session_id") + missing_keys = [key for key in required_keys if key not in payload] + if missing_keys: + raise ValueError(f"Missing GR00T metadata keys: {missing_keys}") + + return cls( + action_horizon=int(payload["action_horizon"]), + action_keys={str(key) for key in payload["action_keys"]}, + embodiment_tag=str(payload["embodiment_tag"]), + needs_session_id=bool(payload["needs_session_id"]), + ) + + +class OpenPIWebsocketClient: + def __init__( + self, + *, + host: str = DEFAULT_HOST, + port: int = DEFAULT_PORT, + path: str = DEFAULT_PATH, + ) -> None: + require_dependencies() + self._uri = f"ws://{host}:{port}{path}" + self._packer = msgpack_numpy.Packer() + self._ws, self._server_metadata = self._connect() + + def _connect(self): + conn = websockets_client.connect( + self._uri, + compression=None, + max_size=None, + ping_interval=PING_INTERVAL_SECS, + ping_timeout=PING_TIMEOUT_SECS, + ) + metadata = msgpack_numpy.unpackb(conn.recv()) + if not isinstance(metadata, dict): + raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}") + return conn, metadata + + def get_server_metadata(self) -> dict[str, Any]: + return dict(self._server_metadata) + + def infer(self, obs: dict[str, Any]) -> dict[str, np.ndarray]: + payload = dict(obs) + payload["endpoint"] = "infer" + self._ws.send(self._packer.pack(payload)) + response = msgpack_numpy.unpackb(self._ws.recv()) + if isinstance(response, dict) and response.get("type") == "error": + raise RuntimeError(f"Inference failed: {response['message']}") + if not isinstance(response, dict): + raise RuntimeError(f"Expected dict actions from GR00T OpenPI endpoint, got {type(response)!r}") + return {str(key): np.asarray(value, dtype=np.float32) for key, value in response.items()} + + def reset(self, reset_info: dict[str, Any] | None = None) -> str: + payload = dict(reset_info or {}) + payload["endpoint"] = "reset" + self._ws.send(self._packer.pack(payload)) + response = msgpack_numpy.unpackb(self._ws.recv()) + if not isinstance(response, dict) or response.get("status") != "reset successful": + raise RuntimeError(f"Unexpected reset response: {response!r}") + return str(response["status"]) + + def close(self) -> None: + self._ws.close() + + +def build_droid_observation(*, session_id: str = DEFAULT_SESSION_ID) -> dict[str, Any]: + return { + "session_id": session_id, + "video": { + "exterior_image_1_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), + "wrist_image_left": np.zeros((1, 2, 256, 256, 3), dtype=np.uint8), + }, + "state": { + "eef_9d": _identity_eef_9d_state(), + "gripper_position": np.zeros((1, 1, 1), dtype=np.float32), + "joint_position": np.zeros((1, 1, 7), dtype=np.float32), + }, + "language": {LANGUAGE_KEY: [["pick up the object"]]}, + } + + +def run_policy_session( + *, + host: str = DEFAULT_HOST, + port: int = DEFAULT_PORT, + path: str = DEFAULT_PATH, + session_id: str = DEFAULT_SESSION_ID, +) -> dict[str, Any]: + client = OpenPIWebsocketClient(host=host, port=port, path=path) + try: + metadata = client.get_server_metadata() + actions = client.infer(build_droid_observation(session_id=session_id)) + reset_status = client.reset({}) + return { + "metadata": metadata, + "actions": actions, + "reset_status": reset_status, + "session_id": session_id, + } + finally: + client.close() + + +def validate_session_result(result: dict[str, Any]) -> None: + metadata = Gr00tServerMetadata.from_dict(result["metadata"]) + if metadata.embodiment_tag != "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT": + raise AssertionError(f"Unexpected embodiment_tag: {metadata.embodiment_tag}") + if not metadata.needs_session_id: + raise AssertionError("GR00T test expects needs_session_id metadata") + if metadata.action_keys != ACTION_KEYS: + raise AssertionError(f"Unexpected action keys: {metadata.action_keys}") + if result["reset_status"] != "reset successful": + raise AssertionError(f"Unexpected reset status: {result['reset_status']!r}") + + actions = result["actions"] + if set(actions) != ACTION_KEYS: + raise AssertionError(f"Unexpected action keys: {set(actions)}") + expected_shapes = { + "eef_9d": (1, metadata.action_horizon, 9), + "gripper_position": (1, metadata.action_horizon, 1), + "joint_position": (1, metadata.action_horizon, 7), + } + for key, expected_shape in expected_shapes.items(): + if actions[key].shape != expected_shape: + raise AssertionError(f"Action {key} shape mismatch: expected {expected_shape}, got {actions[key].shape}") + if actions[key].dtype != np.float32: + raise AssertionError(f"Action {key} dtype mismatch: expected float32, got {actions[key].dtype}") + if not np.isfinite(actions[key]).all(): + raise AssertionError(f"Action {key} contains non-finite values") diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index da62ceab246..11e86ed8cfb 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -17,11 +17,22 @@ from PIL import Image from tests.helpers.media import ( + convert_audio_bytes_to_text, cosine_similarity_text, + preprocess_text, ) _GENDER_PIPELINE = None _GENDER_PIPELINE_LOCK = threading.Lock() +# Transcript gates default to whisper ``small`` for speed. ``small`` mishears a +# short TTS clip ~0.5% of the time (e.g. "Hello"->"fellow", or hallucinating a +# leading SFX token), which flakes the deterministic similarity gate. Short +# clips first get a conservative containment fallback for minor ASR repeats/noise. +# A test can also opt in to ASR escalation by setting +# ``transcript_escalation_model`` to a whisper model name (e.g. ``"large-v3"``) +# in its request_config: on a failed fast pass the clip is re-transcribed with +# that stronger ASR before the test fails, so a weak-ASR mishear is rescued while +# a genuine model artifact still fails (the strong ASR mismatches too). _PCM_SPEECH_SAMPLE_RATE_HZ = 24_000 _MIN_PCM_SPEECH_HNR_DB = 1.0 _PRESET_VOICE_GENDER_MAP: dict[str, str] = { @@ -33,6 +44,23 @@ } +def _short_transcript_contains_expected(transcript: str, expected: str) -> bool: + """Allow minor ASR repeats/noise for very short speech clips.""" + transcript_clean = preprocess_text(transcript) + expected_clean = preprocess_text(expected) + if not transcript_clean or not expected_clean: + return False + + transcript_words = transcript_clean.split() + expected_words = expected_clean.split() + if not transcript_words or not expected_words: + return False + + short_text = min(len(transcript_clean), len(expected_clean)) <= 15 + small_word_delta = len(transcript_words) <= len(expected_words) + 2 + return short_text and small_word_delta and expected_clean in transcript_clean + + def assert_image_diffusion_response( response: "DiffusionResponse", request_config: dict[str, Any], @@ -115,6 +143,11 @@ def assert_video_diffusion_response( expected_height = _maybe_int(form_data.get("height")) expected_fps = _maybe_int(form_data.get("fps")) + # Skip num_frames assertion for Helios models because they round up frames + model = request_config.get("model", "") + if "Helios" in model: + expected_frames = None + for vid_bytes in response.videos: assert_video_valid( vid_bytes, @@ -369,8 +402,15 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str: return "unknown" -def _assert_preset_voice_gender_from_audio(audio_bytes: bytes | None, voice_name: str | None) -> None: +def _assert_preset_voice_gender_from_audio( + audio_bytes: bytes | None, + voice_name: str | None, + *, + response_format: str | None = None, +) -> None: """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown).""" + if response_format == "pcm": + return if not voice_name or not audio_bytes: return key = str(voice_name).lower() @@ -427,6 +467,69 @@ def _assert_pcm_int16_speech_hnr(audio_bytes: bytes, min_hnr_db: float = _MIN_PC ) +def _response_has_audio_output(response: Any) -> bool: + if response.audio_bytes: + return len(response.audio_bytes) > 0 + if isinstance(getattr(response, "audio_content", None), str) and response.audio_content.strip(): + return True + audio_data = getattr(response, "audio_data", None) + return bool(audio_data) + + +def _omni_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool: + if run_level not in {"advanced_model", "full_model"}: + return False + modalities = request_config.get("modalities", ["text", "audio"]) + if "audio" not in modalities: + return False + keywords_dict = request_config.get("key_words", {}) or {} + # When text is not an output modality, the keyword loop validates keywords + # against the audio transcript -- for keywords under ANY word_type + # (text/image/audio/video), not just "audio". Mirror that here so the + # transcript is actually computed; otherwise the loop hits + # `assert transcript is not None` with transcript=None (e.g. an audio-only + # request carrying key_words={"text": [...]}). + if "text" not in modalities and any( + keywords_dict.get(word_type) for word_type in ("text", "image", "audio", "video") + ): + return True + if request_config.get("audio_ref_text"): + return True + return "text" in modalities + + +def _speech_assertion_needs_audio_transcript(request_config: dict[str, Any], run_level: str) -> bool: + if run_level not in {"advanced_model", "full_model"}: + return False + if request_config.get("response_format") == "pcm": + return False + return bool(request_config.get("input")) + + +def _resolve_audio_transcript( + response: Any, + request_config: dict[str, Any], + run_level: str, + *, + speech_api: bool, +) -> str | None: + """Run Whisper only when this run_level / request_config needs a transcript for assertions.""" + needs = ( + _speech_assertion_needs_audio_transcript(request_config, run_level) + if speech_api + else _omni_assertion_needs_audio_transcript(request_config, run_level) + ) + if not needs: + return None + existing = getattr(response, "audio_content", None) + if isinstance(existing, str) and existing.strip(): + return existing + audio_bytes = getattr(response, "audio_bytes", None) + if not audio_bytes: + return None + return convert_audio_bytes_to_text(audio_bytes) + + def assert_omni_response(response: Any, request_config: dict[str, Any], run_level): """ Validate response results. @@ -442,15 +545,18 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve modalities = request_config.get("modalities", ["text", "audio"]) if run_level in {"advanced_model", "full_model"}: + transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=False) # Verify output success if "audio" in modalities: - assert response.audio_content is not None, "No audio output is generated" - print(f"audio content is: {response.audio_content}") + assert _response_has_audio_output(response), "No audio output is generated" + if transcript is not None: + print(f"audio content is: {transcript}") speaker = request_config.get("speaker") if speaker: _assert_preset_voice_gender_from_audio( response.audio_bytes, speaker, + response_format=request_config.get("response_format"), ) if "text" in modalities: assert response.text_content is not None, "No text output is generated" @@ -469,7 +575,8 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve ) else: if keywords: - audio_lower = response.audio_content.lower() + assert transcript is not None, "No audio transcript for keyword validation" + audio_lower = transcript.lower() assert any(str(kw).lower() in audio_lower for kw in keywords), ( "The output does not contain any of the keywords." ) @@ -479,7 +586,7 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve audio_ref_text = request_config.get("audio_ref_text") similarity_threshold = request_config.get("similarity_threshold", 0.8) if "text" in modalities: - transcript = (response.audio_content or "").strip() + assert transcript is not None, "No audio transcript for similarity validation" text_output = (response.text_content or "").strip() # For very short outputs (e.g. one-word answers), n-gram cosine # similarity with length penalty is unreliable because Whisper @@ -508,8 +615,9 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve print(f"similarity is: {similarity}") assert similarity > similarity_threshold, "The audio content is not same as the text" if audio_ref_text: + assert transcript is not None, "No audio transcript for reference-text validation" audio_similarity = cosine_similarity_text( - response.audio_content.lower(), + transcript.strip().lower(), str(audio_ref_text).lower(), ) assert audio_similarity > similarity_threshold, ( @@ -517,6 +625,63 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve ) +def _assert_transcript_matches( + transcript: str, + audio_bytes: bytes | None, + expected_text: Any, + *, + threshold: float, + escalation_model: str | None = None, +) -> None: + """Assert spoken audio matches ``expected_text``. + + ``transcript`` is the fast whisper-``small`` result. If it clears + ``threshold`` the check passes immediately. + + If the cosine check fails, very short clips get a conservative containment + fallback that accepts minor ASR repeats/noise only when the expected text is + still present and the transcript has few extra words. + + When ``escalation_model`` is set (opt-in via the ``transcript_escalation_model`` + request_config key) and the fast check plus containment fallback fail, the + clip is re-transcribed with that stronger ASR and the assertion is decided on + its verdict -- so a weak whisper-``small`` mishear on a short clip does not + flake the gate, while a genuine model artifact still fails (the strong ASR + mismatches too). + """ + expected = str(expected_text).strip().lower() + similarity = cosine_similarity_text(transcript.strip().lower(), expected) + print(f"Cosine similarity: {similarity:.3f}") + if similarity > threshold: + return + + if _short_transcript_contains_expected(transcript, expected): + print("short speech containment check passed") + return + + if escalation_model and audio_bytes: + print( + f"whisper-small below threshold ({similarity:.2f} <= {threshold}); " + f"escalating to whisper-{escalation_model} to rule out an ASR mishear" + ) + strong_transcript = convert_audio_bytes_to_text(audio_bytes, model_size=escalation_model) + strong_similarity = cosine_similarity_text(strong_transcript.strip().lower(), expected) + print( + f"audio content (whisper-{escalation_model}): {strong_transcript}\n" + f"Cosine similarity (whisper-{escalation_model}): {strong_similarity:.3f}" + ) + assert strong_similarity > threshold, ( + f"Transcript doesn't match input after ASR escalation: " + f"input={expected_text!r}; whisper-small='{transcript}' (sim={similarity:.2f}); " + f"whisper-{escalation_model}='{strong_transcript}' (sim={strong_similarity:.2f})" + ) + return + + assert similarity > threshold, ( + f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'" + ) + + def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None: """Validate speech API results from :class:`~tests.helpers.runtime.OmniResponse`. @@ -552,8 +717,6 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], req_fmt = request_config.get("response_format") if req_fmt == "pcm" and response.audio_bytes: - min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB)) - _assert_pcm_int16_speech_hnr(response.audio_bytes, min_hnr_db=min_hnr_db) if response.audio_format: assert "pcm" in response.audio_format.lower(), ( f"Expected audio/pcm content-type, got {response.audio_format!r}" @@ -561,18 +724,29 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], elif req_fmt == "wav" and response.audio_format: assert req_fmt in response.audio_format - if run_level in {"advanced_model", "full_model"} and req_fmt != "pcm": - expected_text = request_config.get("input") - if expected_text: - transcript = (response.audio_content or "").strip() - print(f"audio content is: {transcript}") - print(f"input text is: {expected_text}") - similarity = cosine_similarity_text(transcript.lower(), expected_text.lower()) - print(f"Cosine similarity: {similarity:.3f}") - assert similarity > 0.9, ( - f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'" - ) - _assert_preset_voice_gender_from_audio(response.audio_bytes, request_config.get("voice")) + if run_level in {"advanced_model", "full_model"}: + if req_fmt == "pcm" and response.audio_bytes: + min_hnr_db = float(request_config.get("min_hnr_db", _MIN_PCM_SPEECH_HNR_DB)) + _assert_pcm_int16_speech_hnr(response.audio_bytes, min_hnr_db=min_hnr_db) + + transcript = _resolve_audio_transcript(response, request_config, run_level, speech_api=True) + if transcript is not None: + expected_text = request_config.get("input") + if expected_text: + print(f"audio content is: {transcript}") + print(f"input text is: {expected_text}") + _assert_transcript_matches( + transcript, + getattr(response, "audio_bytes", None), + expected_text, + threshold=0.9, + escalation_model=request_config.get("transcript_escalation_model"), + ) + _assert_preset_voice_gender_from_audio( + response.audio_bytes, + request_config.get("voice"), + response_format=request_config.get("response_format"), + ) def assert_diffusion_response(response: "DiffusionResponse", request_config: dict[str, Any], run_level: str = None): diff --git a/tests/helpers/env.py b/tests/helpers/env.py index 6dee5e6109c..dcf2484cdd0 100644 --- a/tests/helpers/env.py +++ b/tests/helpers/env.py @@ -201,6 +201,7 @@ def run_pre_test_cleanup() -> None: wait_for_gpu_memory_to_clear( devices=list(range(num_gpus)), threshold_ratio=0.05, + timeout_s=60, ) except Exception as e: print(f"Pre-test cleanup note: {e}") diff --git a/tests/helpers/fixtures/config.py b/tests/helpers/fixtures/config.py new file mode 100644 index 00000000000..2e309b672f2 --- /dev/null +++ b/tests/helpers/fixtures/config.py @@ -0,0 +1,14 @@ +"""Fixtures config creation & out of tree registry management.""" + +import pytest + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES + + +@pytest.fixture +def clean_pipeline_registry(): + """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it.""" + snapshot = dict(OMNI_PIPELINES) + yield + OMNI_PIPELINES.clear() + OMNI_PIPELINES.update(snapshot) diff --git a/tests/helpers/media.py b/tests/helpers/media.py index e1d9955464c..e50a9f00132 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -617,11 +617,11 @@ def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer: @contextmanager -def _serialize_whisper_small_model_download(): - """Serialize Whisper ``small`` cache writes across processes (Linux/Unix).""" +def _serialize_whisper_model_download(model_size: str = "small"): + """Serialize Whisper cache writes across processes (Linux/Unix), per model.""" import fcntl - lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock" + lock_path = Path.home() / ".cache" / "whisper" / f".{model_size}_model_download.lock" lock_path.parent.mkdir(parents=True, exist_ok=True) f = open(lock_path, "a+b") try: @@ -632,7 +632,7 @@ def _serialize_whisper_small_model_download(): f.close() -def _whisper_transcribe_in_current_process(output_path: str) -> str: +def _whisper_transcribe_in_current_process(output_path: str, model_size: str = "small") -> str: import whisper device_index = None @@ -656,8 +656,8 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str: use_accelerator = False device = "cpu" - with _serialize_whisper_small_model_download(): - model = whisper.load_model("small", device=device) + with _serialize_whisper_model_download(model_size): + model = whisper.load_model(model_size, device=device) try: text = model.transcribe( output_path, @@ -674,15 +674,15 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str: return text or "" -def convert_audio_file_to_text(output_path: str) -> str: +def convert_audio_file_to_text(output_path: str, model_size: str = "small") -> str: """Convert an audio file to text in an isolated subprocess.""" ctx = multiprocessing.get_context("spawn") with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor: - future = executor.submit(_whisper_transcribe_in_current_process, output_path) + future = executor.submit(_whisper_transcribe_in_current_process, output_path, model_size) return future.result() -def convert_audio_bytes_to_text(raw_bytes: bytes) -> str: +def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") -> str: output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav") os.close(output_fd) if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"): @@ -690,7 +690,7 @@ def convert_audio_bytes_to_text(raw_bytes: bytes) -> str: 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}") - return convert_audio_file_to_text(output_path) + return convert_audio_file_to_text(output_path, model_size) __all__ = [ diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index b82742b40c7..263f34da256 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -19,7 +19,7 @@ from dataclasses import asdict, dataclass from io import BytesIO from pathlib import Path -from typing import Any, NamedTuple +from typing import Any, NamedTuple, cast from urllib.parse import quote import psutil @@ -45,9 +45,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.media import ( _merge_base64_audio_to_segment, - convert_audio_bytes_to_text, decode_b64_image, ) +from tests.model_tests.diffusion.utils import resolve_tiny_model_path from vllm_omni.config.stage_config import resolve_deploy_yaml from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.outputs import OmniRequestOutput @@ -898,16 +898,13 @@ def _process_stream_omni_response(self, chat_completion, *, wall_start: float) - if details := getattr(chunk.usage, "prompt_tokens_details", None): result.cached_tokens = details.cached_tokens - audio_content = None if audio_data: merged_seg = _merge_base64_audio_to_segment(audio_data) wav_buf = BytesIO() merged_seg.export(wav_buf, format="wav") result.audio_bytes = wav_buf.getvalue() - audio_content = convert_audio_bytes_to_text(result.audio_bytes) result.text_content = text_content result.audio_data = audio_data - result.audio_content = audio_content result.e2e_latency = time.perf_counter() - wall_start result.success = True except Exception as e: @@ -932,12 +929,9 @@ def _process_non_stream_omni_response(self, chat_completion, *, wall_start: floa result.prompt_tokens = usage.prompt_tokens if details := getattr(usage, "prompt_tokens_details", None): result.cached_tokens = details.cached_tokens - audio_content = None if audio_data: result.audio_bytes = base64.b64decode(audio_data) - audio_content = convert_audio_bytes_to_text(result.audio_bytes) result.text_content = text_content - result.audio_content = audio_content result.e2e_latency = time.perf_counter() - wall_start if chat_completion.choices and chat_completion.choices[0].logprobs is not None: result.logprobs = chat_completion.choices[0].logprobs.content @@ -1610,6 +1604,8 @@ def send_omni_request(self, request_config: dict[str, Any], request_num: int = 1 extra_body["mm_processor_kwargs"] = mm if "sampling_params_list" in request_config: extra_body["sampling_params_list"] = request_config["sampling_params_list"] + if request_config.get("extra_body"): + extra_body.update(request_config["extra_body"]) create_kwargs: dict[str, Any] = { "model": request_config.get("model"), @@ -1671,8 +1667,8 @@ def _process_stream_audio_speech_response( Process streaming /v1/audio/speech responses into an OmniResponse. This mirrors _process_stream_omni_response but operates on low-level - audio bytes and produces an OmniResponse with audio_content filled - from Whisper transcription. + audio bytes. Whisper transcription runs in assert_audio_speech_response + when the run_level requires it. """ result = OmniResponse() @@ -1708,14 +1704,9 @@ def _process_stream_audio_speech_response( raise TypeError(f"Unsupported audio speech streaming response type: {type(response)}") raw_bytes = bytes(data) - if response_format == "pcm": - transcript = None - else: - transcript = convert_audio_bytes_to_text(raw_bytes) # Populate OmniResponse. result.audio_bytes = raw_bytes - result.audio_content = transcript result.e2e_latency = time.perf_counter() - wall_start result.success = True result.audio_format = getattr(response, "response", None) @@ -1748,13 +1739,7 @@ def _process_non_stream_audio_speech_response( else: raise TypeError(f"Unsupported audio speech response type: {type(response)}") - if response_format == "pcm": - transcript = None - else: - transcript = convert_audio_bytes_to_text(raw_bytes) - result.audio_bytes = raw_bytes - result.audio_content = transcript result.e2e_latency = time.perf_counter() - wall_start result.success = True result.audio_format = getattr(response, "response", None) @@ -1809,6 +1794,8 @@ def send_audio_speech_request(self, request_config: dict[str, Any], request_num: "seed", "instructions", "speed", + "stream_format", + "x_vector_only_mode", ): if key in request_config: extra_body[key] = request_config[key] @@ -2107,6 +2094,9 @@ def send_video_diffusion_request( normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None} files: dict[str, tuple[str, BytesIO, str]] = {} image_reference = request_config.get("image_reference") + video_reference = request_config.get("video_reference") + if image_reference and video_reference: + raise ValueError("Only one of image_reference or video_reference can be provided") if image_reference: if image_reference.startswith("data:image"): header, encoded = image_reference.split(",", 1) @@ -2116,6 +2106,15 @@ def send_video_diffusion_request( files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) else: normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference}) + if video_reference: + if video_reference.startswith("data:video"): + header, encoded = video_reference.split(",", 1) + content_type = header.split(";")[0].removeprefix("data:") + extension = content_type.split("/")[-1] + file_data = base64.b64decode(encoded) + files["input_reference"] = (f"reference.{extension}", BytesIO(file_data), content_type) + else: + normalized_form_data["video_reference"] = json.dumps({"video_url": video_reference}) result = DiffusionResponse() create_url = self._build_url("/v1/videos") @@ -2191,6 +2190,102 @@ def _post_form_endpoint( timeout=timeout, ) + def send_streaming_video_diffusion_request( + self, + request_config: dict[str, Any], + request_num: int = 1, + *, + timeout_seconds: float = 600.0, + ) -> list[DiffusionResponse]: + """ + Send a native ``/v1/realtime/video`` WebSocket request and return one + finalized MP4 artifact assembled from the streamed binary fragments. + """ + if request_num != 1: + raise NotImplementedError("Concurrent streaming video diffusion requests are not currently implemented") + + response = asyncio.run( + self._send_streaming_video_diffusion_request_once( + request_config, + timeout_seconds=timeout_seconds, + ) + ) + assert_diffusion_response(response, request_config, run_level=self.run_level) + if response.e2e_latency is not None: + self._print_client_stat(f"[diffusion.stream] request#1 success in {response.e2e_latency:.3f}s") + else: + self._print_client_stat("[diffusion.stream] request#1 completed") + return [response] + + async def _send_streaming_video_diffusion_request_once( + self, + request_config: dict[str, Any], + *, + timeout_seconds: float, + ) -> DiffusionResponse: + form_data = request_config.get("form_data") + if not isinstance(form_data, dict): + raise ValueError("Video request_config must contain 'form_data'") + payload: dict[str, Any] = { + "type": "session.start", + **{key: value for key, value in form_data.items() if value is not None}, + } + model = request_config.get("model") + if model is not None: + payload["model"] = model + payload.setdefault("format", "m4s") + + fps = float(payload.get("fps") or 16) + stream_format = payload["format"] + url = self._build_ws_url("/v1/realtime/video") + + result = DiffusionResponse() + chunks: list[bytes] = [] + start_time = time.perf_counter() + deadline = start_time + timeout_seconds + + import websockets + + async with websockets.connect(url, max_size=None) as websocket: + await websocket.send(json.dumps(payload)) + + while True: + remaining = deadline - time.perf_counter() + if remaining <= 0: + raise TimeoutError(f"Streaming video request did not complete within {timeout_seconds}s") + + message = await asyncio.wait_for(websocket.recv(), timeout=remaining) + if isinstance(message, bytes): + chunks.append(message) + continue + + msg = json.loads(message) + msg_type = msg.get("type") + if msg_type == "video.start": + stream_format = msg.get("format") or stream_format + continue + if msg_type == "session.done": + break + if msg_type == "error": + raise RuntimeError(str(msg.get("message", msg))) + + from vllm_omni.diffusion.utils.media_utils import finalize_streaming_video_bytes + from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat + + streamed_bytes = b"".join(chunks) + if not streamed_bytes: + raise RuntimeError("Streaming video request completed without binary video chunks") + result.videos = [ + finalize_streaming_video_bytes( + streamed_bytes, + input_format=cast(StreamingVideoFormat, stream_format), + fps=fps, + ) + ] + result.e2e_latency = time.perf_counter() - start_time + result.success = True + return result + def _wait_until_video_completed( self, video_id: str, poll_interval_seconds: int = 2, timeout_seconds: int = 300 ) -> None: @@ -2711,25 +2806,6 @@ def stop_profile(self, stages: list[int] | None = None) -> list[Any]: # --------------------------------------------------------------------------- -def _core_model_stage_config_path_with_dummy_load_format(stage_config_path: str | None, run_level: str) -> str | None: - """For ``core_model`` runs, patch every stage in the deploy YAML to ``load_format: dummy``.""" - if run_level != "core_model" or stage_config_path is None: - return stage_config_path - from tests.helpers.stage_config import modify_stage_config - - with open(stage_config_path, encoding="utf-8") as f: - cfg = yaml.safe_load(f) or {} - new_schema_stages = cfg.get("stages") - stage_key = "stages" if new_schema_stages is not None else "stage_args" - update_path = "load_format" if new_schema_stages is not None else "engine_args.load_format" - stage_entries = cfg.get(stage_key, []) - stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage] - return modify_stage_config( - stage_config_path, - updates={stage_key: {stage_id: {update_path: "dummy"} for stage_id in stage_ids}}, - ) - - def iter_omni_server( request: Any, run_level: str, @@ -2737,27 +2813,28 @@ def iter_omni_server( omni_fixture_lock: threading.Lock, ) -> Generator[Any, Any, None]: """Start/stop an Omni HTTP server; used by ``omni_server`` / ``omni_server_function`` fixtures.""" - from tests.helpers.stage_config import modify_stage_config + from tests.helpers.stage_config import stage_config_path_for_run_level with omni_fixture_lock: params: OmniServerParams = request.param - model = model_prefix + params.model + # For now, when a tiny model is substituted, we preserve the original model + # name via --served-model-name (so that the server still accepts requests with + # the original name). We also do the same for server.model so that tests reading + # server.model send the correct name in requests. + # + # TODO: core models on this path currently do not clean up tiny models, although + # tiny model paths are deterministic, so it's not a huge footprint. Still, it would + # be ideal to cleanup consistently everywhere. + original_model = model_prefix + params.model + model = original_model + if run_level == "core_model" and request.node.get_closest_marker("diffusion"): + model = resolve_tiny_model_path(model) port = params.port - stage_config_path = params.stage_config_path - if run_level in {"advanced_model", "full_model"} and stage_config_path is not None: - with open(stage_config_path, encoding="utf-8") as f: - cfg = yaml.safe_load(f) or {} - new_schema_stages = cfg.get("stages") - stage_key = "stages" if new_schema_stages is not None else "stage_args" - delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format" - stage_entries = cfg.get(stage_key, []) - stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage] - stage_config_path = modify_stage_config( - stage_config_path, - deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}}, - ) + stage_config_path = stage_config_path_for_run_level(params.stage_config_path, run_level) server_args = params.server_args or [] + if model != original_model: + server_args = [*server_args, "--served-model-name", original_model] if params.use_omni and params.stage_init_timeout is not None: server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)] else: @@ -2783,6 +2860,8 @@ def iter_omni_server( port=port, env_dict=params.env_dict, ) as server: + if model != original_model: + server.model = original_model print("OmniServer started successfully") yield server print("OmniServer stopping...") @@ -2806,6 +2885,8 @@ def iter_omni_server( use_omni=params.use_omni, ) ) as server: + if model != original_model: + server.model = original_model print("OmniServer started successfully") yield server print("OmniServer stopping...") @@ -2820,6 +2901,8 @@ def iter_omni_runner( omni_fixture_lock: threading.Lock, ) -> Generator[Any, None, None]: """Yield an :class:`OmniRunner`; used by ``omni_runner`` / ``omni_runner_function`` fixtures.""" + from tests.helpers.stage_config import stage_config_path_for_run_level + with omni_fixture_lock: param = request.param if not isinstance(param, (tuple, list)) or len(param) not in (2, 3): @@ -2833,8 +2916,10 @@ def iter_omni_runner( else: model, stage_config_path, extra = param[0], param[1], param[2] extra_omni_kwargs = dict(extra) if extra is not None else {} - stage_config_path = _core_model_stage_config_path_with_dummy_load_format(stage_config_path, run_level) + stage_config_path = stage_config_path_for_run_level(stage_config_path, run_level) model = model_prefix + model + if run_level == "core_model" and request.node.get_closest_marker("diffusion"): + model = resolve_tiny_model_path(model) with OmniRunner(model, seed=42, stage_configs_path=stage_config_path, **extra_omni_kwargs) as runner: print("OmniRunner started successfully") yield runner diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index e34680be079..8706bb1bc2a 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -482,6 +482,7 @@ def delete_by_path(config_dict: dict, path: str) -> None: { "stage_id": 1, "max_num_seqs": 1, + "gpu_memory_utilization": 0.5, }, ], }, @@ -577,10 +578,15 @@ def delete_by_path(config_dict: dict, path: str) -> None: { "stage_id": 0, "max_num_seqs": 1, - "gpu_memory_utilization": 0.9, + # Tuned for the heavier rebased stack (vLLM v0.23.1rc1 + torch 2.11 + CUDA 13 + # + flashinfer) on a 24 GiB L4: the 16.78 GiB model weights + non-KV overhead + # + a 16384-token activation peak left 0 KV cache at util 0.85 (build 2354 + # OOM). The abort test only needs the engine to init, so a smaller batched- + # token budget is fine; outputs are unaffected (chunked prefill is on). + "gpu_memory_utilization": 0.90, "enforce_eager": True, "enable_prefix_caching": False, - "max_num_batched_tokens": 16384, + "max_num_batched_tokens": 2048, "max_model_len": 16384, "skip_mm_profiling": True, "mm_processor_cache_gb": 0, @@ -627,7 +633,93 @@ def get_deploy_config_path(rel_path: str) -> str: return str(_DEPLOY_DIR / rel_path) +def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]: + """Return one stage entry from a deploy yaml by ``stage_id``.""" + with open(get_deploy_config_path(rel_path), encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + + stage_key = "stages" if "stages" in cfg else "stage_args" + for stage in cfg.get(stage_key, []): + if stage.get("stage_id") == stage_id: + return stage + raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}") + + +def _get_config_value_by_path(config_dict: dict, path: str) -> Any: + """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``).""" + current: Any = config_dict + for key in path.split("."): + if not isinstance(current, dict) or key not in current: + return None + current = current[key] + return current + + +def _stage_load_format_paths(stage_config_path: str) -> tuple[str, str, list[int]]: + """Return ``(stage_key, load_format_field_path, stage_ids)`` for a deploy YAML.""" + with open(stage_config_path, encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + new_schema_stages = cfg.get("stages") + stage_key = "stages" if new_schema_stages is not None else "stage_args" + load_format_path = "load_format" if new_schema_stages is not None else "engine_args.load_format" + stage_entries = cfg.get(stage_key, []) + stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage] + return stage_key, load_format_path, stage_ids + + +def _add_dummy_load_format( + stage_config_path: str | None, + run_level: str, +) -> str | None: + """For ``core_model`` runs, patch every stage in the deploy YAML to ``load_format: dummy``.""" + if run_level != "core_model" or stage_config_path is None: + return stage_config_path + stage_key, load_format_path, stage_ids = _stage_load_format_paths(stage_config_path) + return modify_stage_config( + stage_config_path, + updates={stage_key: {stage_id: {load_format_path: "dummy"} for stage_id in stage_ids}}, + ) + + +def _delete_dummy_load_format( + stage_config_path: str | None, + run_level: str, +) -> str | None: + """For ``advanced_model`` / ``full_model``, strip ``load_format: dummy`` so real weights load.""" + if run_level not in {"advanced_model", "full_model"} or stage_config_path is None: + return stage_config_path + stage_key, load_format_path, _stage_ids = _stage_load_format_paths(stage_config_path) + with open(stage_config_path, encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + stage_entries = cfg.get(stage_key, []) + + deletes: dict[int, list[str]] = {} + for stage in stage_entries: + stage_id = stage.get("stage_id") + if stage_id is None: + continue + if _get_config_value_by_path(stage, load_format_path) == "dummy": + deletes[stage_id] = [load_format_path] + + if not deletes: + return stage_config_path + + return modify_stage_config( + stage_config_path, + deletes={stage_key: deletes}, + ) + + +def stage_config_path_for_run_level(stage_config_path: str | None, run_level: str) -> str | None: + """Apply run-level deploy YAML tweaks for weight loading (dummy vs real).""" + if run_level in {"advanced_model", "full_model"}: + return _delete_dummy_load_format(stage_config_path, run_level) + return _add_dummy_load_format(stage_config_path, run_level) + + __all__ = [ - "modify_stage_config", "get_deploy_config_path", + "get_deploy_config_stage", + "modify_stage_config", + "stage_config_path_for_run_level", ] diff --git a/tests/helpers/test_assertions.py b/tests/helpers/test_assertions.py new file mode 100644 index 00000000000..0bfa7f593aa --- /dev/null +++ b/tests/helpers/test_assertions.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from tests.helpers.assertions import _assert_transcript_matches + + +def test_short_transcript_repeat_passes_containment_fallback(): + _assert_transcript_matches( + " How... how are you?", + audio_bytes=None, + expected_text="how are you", + threshold=0.9, + ) + + +def test_short_transcript_unrelated_text_still_fails(): + with pytest.raises(AssertionError, match="Transcript doesn't match input"): + _assert_transcript_matches( + " I don't know, sorry.", + audio_bytes=None, + expected_text="how are you", + threshold=0.9, + ) diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index 1f79b186b49..f01933dbc82 100644 --- a/tests/metrics/test_prometheus.py +++ b/tests/metrics/test_prometheus.py @@ -12,10 +12,10 @@ _MODEL = "test-model" _PIPELINE_METRICS = [ - "vllm:omni_num_requests_running", - "vllm:omni_num_requests_waiting", - "vllm:omni_requests_success", - "vllm:omni_e2e_request_latency_s", + "vllm_omni:num_requests_running", + "vllm_omni:num_requests_waiting", + "vllm_omni:requests_success", + "vllm_omni:e2e_request_latency_s", ] @@ -58,32 +58,32 @@ def test_counter_values(self, scrape_output: str) -> None: # Per-reason buckets sourced from the merged completion Counter. stop = _sample_value( scrape_output, - f'vllm:omni_requests_success_total{{finished_reason="stop",model_name="{_MODEL}"}}', + f'vllm_omni:requests_success_total{{finished_reason="stop",model_name="{_MODEL}"}}', ) assert stop == 2.0 length = _sample_value( scrape_output, - f'vllm:omni_requests_success_total{{finished_reason="length",model_name="{_MODEL}"}}', + f'vllm_omni:requests_success_total{{finished_reason="length",model_name="{_MODEL}"}}', ) assert length == 1.0 abort = _sample_value( scrape_output, - f'vllm:omni_requests_success_total{{finished_reason="abort",model_name="{_MODEL}"}}', + f'vllm_omni:requests_success_total{{finished_reason="abort",model_name="{_MODEL}"}}', ) assert abort == 1.0 def test_gauge_values(self, scrape_output: str) -> None: running = _sample_value( scrape_output, - f'vllm:omni_num_requests_running{{model_name="{_MODEL}"}}', + f'vllm_omni:num_requests_running{{model_name="{_MODEL}"}}', ) assert running == 5.0 waiting = _sample_value( scrape_output, - f'vllm:omni_num_requests_waiting{{model_name="{_MODEL}"}}', + f'vllm_omni:num_requests_waiting{{model_name="{_MODEL}"}}', ) assert waiting == 2.0 @@ -93,7 +93,7 @@ def test_histogram_counts(self, scrape_output: str) -> None: # observing the latency histogram, so the count stays at 3. e2e_count = _sample_value( scrape_output, - f'vllm:omni_e2e_request_latency_s_count{{model_name="{_MODEL}"}}', + f'vllm_omni:e2e_request_latency_s_count{{model_name="{_MODEL}"}}', ) assert e2e_count == 3.0 @@ -110,8 +110,8 @@ def test_no_legacy_engine_label(self, scrape_output: str) -> None: def test_no_legacy_seconds_or_ms_families(self, scrape_output: str) -> None: # Renamed: *_time_ms → *_s; *_time_seconds dropped. for legacy in ( - "vllm:omni_request_queue_time_seconds", - "vllm:omni_e2e_request_latency_seconds", + "vllm_omni:request_queue_time_seconds", + "vllm_omni:e2e_request_latency_seconds", ): assert legacy not in scrape_output, f"legacy family {legacy} still registered" @@ -170,8 +170,8 @@ def test_running_and_waiting_zero_after_request_completes(self, registry: Collec assert obj.engine._running_counter.value == 0 assert len(obj.request_states) == 0 out = generate_latest(registry).decode() - assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test"}') == 0.0 - assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test"}') == 0.0 + assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test"}') == 0.0 + assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test"}') == 0.0 def test_gauges_reflect_remaining_requests_after_one_completes(self, registry: CollectorRegistry) -> None: from types import SimpleNamespace @@ -196,5 +196,5 @@ def test_gauges_reflect_remaining_requests_after_one_completes(self, registry: C obj._log_summary_and_cleanup("req-1") out = generate_latest(registry).decode() - assert _sample_value(out, 'vllm:omni_num_requests_running{model_name="lifecycle-test-2"}') == 1.0 - assert _sample_value(out, 'vllm:omni_num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0 + assert _sample_value(out, 'vllm_omni:num_requests_running{model_name="lifecycle-test-2"}') == 1.0 + assert _sample_value(out, 'vllm_omni:num_requests_waiting{model_name="lifecycle-test-2"}') == 0.0 diff --git a/tests/model_executor/models/glm_image/test_glm_image_init.py b/tests/model_executor/models/glm_image/test_glm_image_init.py deleted file mode 100644 index eb42e9649fe..00000000000 --- a/tests/model_executor/models/glm_image/test_glm_image_init.py +++ /dev/null @@ -1,76 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for GLM-Image model __init__.py lazy import pattern. - -The ``__init__.py`` uses ``__getattr__`` for lazy loading to avoid importing -``transformers.models.glm_image`` at module init, which may not be available -in all environments. -""" - -import pytest - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -class TestLazyImport: - """Test the __getattr__ lazy import pattern in glm_image __init__.py.""" - - def test_getattr_exists_and_is_callable(self): - """Verify __getattr__ exists and is callable.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - assert hasattr(glm_image_pkg, "__getattr__") - assert callable(glm_image_pkg.__getattr__) - - def test_getattr_returns_class_for_known_attribute(self): - """Verify __getattr__ returns GlmImageForConditionalGeneration for known attribute.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # Call __getattr__ directly to test the lazy import logic - result = glm_image_pkg.__getattr__("GlmImageForConditionalGeneration") - - # Verify we get a class (the actual GlmImageForConditionalGeneration) - assert result is not None - assert isinstance(result, type) - assert result.__name__ == "GlmImageForConditionalGeneration" - - def test_getattr_raises_for_unknown_attribute(self): - """Verify __getattr__ raises AttributeError for unknown attributes.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # Test unknown attribute via __getattr__ directly - with pytest.raises(AttributeError, match="has no attribute"): - glm_image_pkg.__getattr__("UnknownClass") - - def test___all___exports_correct_symbols(self): - """Verify __all__ contains the expected exported symbols.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - assert hasattr(glm_image_pkg, "__all__") - assert "GlmImageForConditionalGeneration" in glm_image_pkg.__all__ - - -class TestLazyImportDoesNotImportTransformersAtInit: - """Verify that importing the package does not eagerly load transformers.""" - - def test_glm_image_module_has_getattr(self): - """Test that the module has __getattr__ for lazy loading.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # The module should have __getattr__ - assert hasattr(glm_image_pkg, "__getattr__") - assert callable(glm_image_pkg.__getattr__) - - # And __all__ - assert hasattr(glm_image_pkg, "__all__") - - def test_module_has_proper_structure(self): - """Test that the module has proper Python module structure.""" - import vllm_omni.model_executor.models.glm_image as glm_image_pkg - - # Should be a module - assert hasattr(glm_image_pkg, "__name__") - assert glm_image_pkg.__name__ == "vllm_omni.model_executor.models.glm_image" - - # Should have __file__ (even if it's a package) - assert hasattr(glm_image_pkg, "__file__") 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 index 7aab6133b58..c5966dd60b8 100644 --- a/tests/model_executor/models/ming_tts/test_ming_shared_modules.py +++ b/tests/model_executor/models/ming_tts/test_ming_shared_modules.py @@ -137,8 +137,12 @@ def test_ming_compute_logits_falls_back_to_dummy_token_id(): def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids(): class FakeBackbone: - def __call__(self, **kwargs): - return kwargs["inputs_embeds"] + # MingLLMModel.forward calls the backbone positionally + # (input_ids, positions, intermediate_tensors, inputs_embeds) so the + # same signature works for both Qwen2Model (``positions``) and + # BailingMoeModel (``position_ids``). + def __call__(self, input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs): + return inputs_embeds model = _make_ming_logits_model() model.model = FakeBackbone() diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index c51db9fccb9..48db0d49768 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -19,9 +19,8 @@ import pytest -from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, PipelineConfig, StageExecutionType, ) @@ -34,13 +33,13 @@ class TestRegistryDeclaration: def test_declared_in_omni_pipelines(self) -> None: - assert _PIPELINE_KEY in _OMNI_PIPELINES + assert _PIPELINE_KEY in OMNI_PIPELINES def test_visible_in_central_registry(self) -> None: - assert _PIPELINE_KEY in _PIPELINE_REGISTRY + assert _PIPELINE_KEY in OMNI_PIPELINES def test_lazy_load_returns_pipeline_config(self) -> None: - pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY] + pipeline = OMNI_PIPELINES[_PIPELINE_KEY] assert isinstance(pipeline, PipelineConfig) assert pipeline.model_type == _PIPELINE_KEY assert pipeline.model_arch == "MiniCPMO45OmniForConditionalGeneration" @@ -49,7 +48,7 @@ def test_lazy_load_returns_pipeline_config(self) -> None: class TestPipelineTopology: @pytest.fixture(scope="class") def pipeline(self) -> PipelineConfig: - return _PIPELINE_REGISTRY[_PIPELINE_KEY] + return OMNI_PIPELINES[_PIPELINE_KEY] def test_two_stages(self, pipeline: PipelineConfig) -> None: assert len(pipeline.stages) == 2 @@ -99,7 +98,7 @@ class TestArchAliases: @pytest.fixture(scope="class") def pipeline(self) -> PipelineConfig: - return _PIPELINE_REGISTRY[_PIPELINE_KEY] + return OMNI_PIPELINES[_PIPELINE_KEY] def test_shared_minicpmo_alias_present(self, pipeline: PipelineConfig) -> None: # MiniCPM-o 4.5 ships ``architectures=["MiniCPMO"]`` in its HF config. @@ -123,7 +122,7 @@ class TestHfConfigPredicate: @pytest.fixture(scope="class") def predicate(self): - pipeline = _PIPELINE_REGISTRY[_PIPELINE_KEY] + pipeline = OMNI_PIPELINES[_PIPELINE_KEY] assert pipeline.hf_config_predicate is not None, ( "MiniCPM-o 4.5 pipeline must declare hf_config_predicate to " "avoid misrouting MiniCPM-o 2.6 checkpoints into the 4.5 path." diff --git a/tests/model_executor/models/moss_tts/test_per_request_codes.py b/tests/model_executor/models/moss_tts/test_per_request_codes.py new file mode 100644 index 00000000000..b35346c2f99 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_per_request_codes.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MOSS-TTS per-request code routing in the talker. + +Guards the cross-request audio-corruption class (RFC #4316 Issue 1, bug report +#4355): under continuous batching every Stage-1 request must decode its *own* +codes, not request 0's. The talker emits ``codes["audio"]`` as a batch-ordered +**list** so the generic splitter (``to_payload_element``) routes +``element[idx]`` to request ``idx``; the list must stay index-aligned with the +batch (placeholder for skipped requests) or a shorter list silently falls back +to entry 0 — the original corruption. The hidden-row mapping must also use the +runner's real per-request token spans, not an equal split, or mixed +prefill+decode steps sample each request's codes from the wrong rows. + +CPU-only and weight-free: ``make_omni_output`` is driven directly on a bare +talker with ``_sample_audio_codes`` mocked to echo the hidden row it was given, +so an assertion on the output value proves which row each request sampled from. +""" + +from __future__ import annotations + +import functools + +import pytest + +torch = pytest.importorskip("torch") + +N_VQ = 4 + + +@functools.lru_cache(maxsize=1) +def _delay_talker_cls(): + """Defer talker import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSDelayTalkerForGeneration, + ) + + return MossTTSDelayTalkerForGeneration + + +def _make_bare_delay_talker(): + """A talker whose only behaviour is the routing/span logic under test. + + ``_sample_audio_codes`` is replaced with a stub returning an (N_VQ,) tensor + filled with the value of the hidden row it received. With a ``hidden`` where + row r holds the scalar r, the emitted codes reveal which row each request + sampled. + """ + cls = _delay_talker_cls() + talker = cls.__new__(cls) + talker._batch_state = None + talker.n_vq = N_VQ + + def _stub_sample(last_h, state): + row_value = int(last_h.reshape(-1)[0].item()) + return torch.full((N_VQ,), row_value, dtype=torch.long) + + talker._sample_audio_codes = _stub_sample + return talker + + +def _row_indexed_hidden(num_rows: int) -> torch.Tensor: + """(num_rows, 1) tensor where row r holds the scalar r.""" + return torch.arange(num_rows, dtype=torch.float32).reshape(num_rows, 1) + + +def _audio_payload(result): + return result.multimodal_outputs["codes"]["audio"] + + +class TestPerRequestRouting: + def test_mixed_prefill_decode_routes_by_span(self): + """req 0 prefills 3 rows, req 1 decodes 1 row in the same step. + + With real spans, req 0 samples its last row (2) and req 1 its last row + (3); the payload is a batch-aligned list, element[idx] per request. + """ + talker = _make_bare_delay_talker() + hidden = _row_indexed_hidden(4) + info_dicts = [{}, {}] + spans = [(0, 3), (3, 4)] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert isinstance(audio, list) + assert len(audio) == 2 # one entry per request, batch-aligned + assert audio[0].shape == (1, N_VQ) + assert torch.all(audio[0] == 2) # req 0: last row of span [0,3) + assert torch.all(audio[1] == 3) # req 1: last row of span [3,4) + # No cross-request leakage. + assert not torch.equal(audio[0], audio[1]) + + def test_no_request_samples_outside_its_span(self): + talker = _make_bare_delay_talker() + spans = [(0, 2), (2, 5), (5, 6)] + hidden = _row_indexed_hidden(6) + info_dicts = [{}, {}, {}] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert len(audio) == 3 + sampled = [] + for i, (start, end) in enumerate(spans): + value = int(audio[i].reshape(-1)[0].item()) + assert start <= value < end, f"req {i} sampled row {value} outside [{start},{end})" + sampled.append(value) + assert len(set(sampled)) == len(sampled) # distinct rows, no sharing + + def test_skipped_request_keeps_list_aligned(self): + """An empty-span request gets an empty placeholder at its own index, so + later requests are not shifted onto request 0's codes.""" + talker = _make_bare_delay_talker() + # req 1 has an empty span (start == end) and must not displace req 2. + spans = [(0, 1), (1, 1), (1, 2)] + hidden = _row_indexed_hidden(2) + info_dicts = [{}, {}, {}] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert len(audio) == 3 + assert torch.all(audio[0] == 0) # req 0: row 0 + assert audio[1].numel() == 0 # req 1: empty placeholder, aligned + assert torch.all(audio[2] == 1) # req 2: row 1, not shifted + + def test_accumulates_into_per_request_history(self): + talker = _make_bare_delay_talker() + hidden = _row_indexed_hidden(2) + prior = torch.full((1, N_VQ), 7, dtype=torch.long) + info_dicts = [{"audio_codes": {"accumulated": prior}}, {}] + spans = [(0, 1), (1, 2)] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + + assert audio[0].shape == (2, N_VQ) + assert torch.all(audio[0][0] == 7) # prior history kept + assert torch.all(audio[0][1] == 0) # this step's row 0 appended + assert torch.all(audio[1] == 1) + + def test_runs_with_spans(self): + """request_token_spans are required for correct per-request routing; + verify that passing them produces a batch-aligned list.""" + talker = _make_bare_delay_talker() + hidden = _row_indexed_hidden(2) + info_dicts = [{}, {}] + spans = [(0, 1), (1, 2)] + + result = talker.make_omni_output( + hidden, + runtime_additional_information=info_dicts, + request_token_spans=spans, + ) + audio = _audio_payload(result) + assert isinstance(audio, list) + assert len(audio) == 2 + + def test_no_codes_returns_empty_payload(self): + """Zero hidden rows → no codes → empty multimodal payload (no list).""" + talker = _make_bare_delay_talker() + hidden = torch.zeros((0, 1), dtype=torch.float32) + result = talker.make_omni_output(hidden, runtime_additional_information=[{}], request_token_spans=[(0, 0)]) + assert result.multimodal_outputs == {} diff --git a/tests/model_executor/models/moss_tts/test_talker_logits_processor.py b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py new file mode 100644 index 00000000000..35c30997959 --- /dev/null +++ b/tests/model_executor/models/moss_tts/test_talker_logits_processor.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for MOSS-TTS talker logits processing.""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +pytest.importorskip("vllm") +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts] + + +class _RecordingLogitsProcessor: + def __init__(self) -> None: + self.args: tuple[object, ...] | None = None + self.kwargs: dict[str, object] | None = None + + def __call__(self, *args: object, **kwargs: object) -> torch.Tensor: + self.args = args + self.kwargs = kwargs + return torch.zeros((1, 4)) + + +def test_moss_tts_delay_compute_logits_does_not_forward_sampling_metadata() -> None: + from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( + MossTTSDelayTalkerForGeneration, + ) + + model = MossTTSDelayTalkerForGeneration.__new__(MossTTSDelayTalkerForGeneration) + nn.Module.__init__(model) + model.text_lm_head = object() + model.logits_processor = _RecordingLogitsProcessor() + model._batch_state = None + hidden_states = torch.randn(1, 4) + sampling_metadata = object() + + logits = model.compute_logits(hidden_states, sampling_metadata=sampling_metadata) + + assert logits.shape == (1, 4) + assert model.logits_processor.args == (model.text_lm_head, hidden_states) + assert model.logits_processor.kwargs == {} diff --git a/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py new file mode 100644 index 00000000000..e17cdd6f25c --- /dev/null +++ b/tests/model_executor/models/qwen2_5_omni/test_per_video_audio_mask.py @@ -0,0 +1,708 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from vllm.multimodal.inputs import MultiModalBatchedField, MultiModalSharedField +from vllm.multimodal.processing.processor import PlaceholderFeaturesInfo + +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY, + Qwen2_5OmniThinkerMultiModalProcessor, + _coerce_use_audio_in_video_for_hf_processor, + _filter_video_use_audio_in_video_for_uncached_items, + _normalize_use_audio_in_video, +) +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, + Qwen3OmniMoeThinkerMultiModalProcessor, +) +from vllm_omni.outputs import OmniModelRunnerOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +AUDIO_TOKEN_ID = 10 +VIDEO_TOKEN_ID = 20 +AUDIO_BOS_TOKEN_ID = 30 +AUDIO_EOS_TOKEN_ID = 31 +IMAGE_TOKEN_ID = 40 + + +class _Feature: + def __init__(self, data: torch.Tensor) -> None: + self.data = data + + +def _fake_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace(get_vocab=lambda: {"<|audio_pad|>": AUDIO_TOKEN_ID}) + processor = SimpleNamespace(audio_token="<|audio_pad|>") + config = SimpleNamespace(audio_token_id=AUDIO_TOKEN_ID) + info = SimpleNamespace( + get_hf_config=lambda: config, + get_hf_processor=lambda: processor, + get_tokenizer=lambda: tokenizer, + ) + return SimpleNamespace(info=info) + + +def _fake_qwen2_prompt_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace( + get_vocab=lambda: { + "<|audio_pad|>": AUDIO_TOKEN_ID, + "<|video_pad|>": VIDEO_TOKEN_ID, + "<|image_pad|>": IMAGE_TOKEN_ID, + } + ) + processor = SimpleNamespace( + audio_token="<|audio_pad|>", + video_token="<|video_pad|>", + image_token="<|image_pad|>", + ) + image_processor = SimpleNamespace(merge_size=1) + config = SimpleNamespace(audio_token_id=AUDIO_TOKEN_ID) + info = SimpleNamespace( + get_hf_config=lambda: config, + get_hf_processor=lambda **_kwargs: processor, + get_image_processor=lambda **_kwargs: image_processor, + get_tokenizer=lambda: tokenizer, + ) + + def omni_get_updates_use_audio_in_video(**_kwargs): + return [ + AUDIO_BOS_TOKEN_ID, + VIDEO_TOKEN_ID, + AUDIO_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + ] + + return SimpleNamespace( + info=info, + omni_get_updates_use_audio_in_video=omni_get_updates_use_audio_in_video, + ) + + +def _fake_processor_with_spatial_merge_size() -> SimpleNamespace: + hf_config = SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)) + return SimpleNamespace(info=SimpleNamespace(get_hf_config=lambda: hf_config)) + + +def _fake_qwen3_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace( + get_vocab=lambda: { + "<|audio_pad|>": AUDIO_TOKEN_ID, + "<|video_pad|>": VIDEO_TOKEN_ID, + } + ) + processor = SimpleNamespace( + audio_token="<|audio_pad|>", + video_token="<|video_pad|>", + ) + info = SimpleNamespace( + get_hf_processor=lambda: processor, + get_tokenizer=lambda: tokenizer, + ) + return SimpleNamespace(info=info) + + +def _fake_qwen3_prompt_processor() -> SimpleNamespace: + tokenizer = SimpleNamespace( + get_vocab=lambda: { + "<|audio_pad|>": AUDIO_TOKEN_ID, + "<|video_pad|>": VIDEO_TOKEN_ID, + "<|image_pad|>": IMAGE_TOKEN_ID, + } + ) + processor = SimpleNamespace( + audio_token="<|audio_pad|>", + video_token="<|video_pad|>", + image_token="<|image_pad|>", + ) + image_processor = SimpleNamespace(merge_size=1) + config = SimpleNamespace( + audio_token_id=AUDIO_TOKEN_ID, + video_token_id=VIDEO_TOKEN_ID, + audio_start_token_id=AUDIO_BOS_TOKEN_ID, + audio_end_token_id=AUDIO_EOS_TOKEN_ID, + vision_config=SimpleNamespace(spatial_merge_size=1), + position_id_per_seconds=1, + ) + info = SimpleNamespace( + get_hf_config=lambda: config, + get_hf_processor=lambda **_kwargs: processor, + get_image_processor=lambda **_kwargs: image_processor, + get_tokenizer=lambda: tokenizer, + ) + return SimpleNamespace(info=info) + + +def test_normalize_use_audio_in_video_accepts_bool_and_per_video_mask(): + assert _normalize_use_audio_in_video(True, 2) == [True, True] + assert _normalize_use_audio_in_video(False, 2) == [False, False] + assert _normalize_use_audio_in_video([True, False], 2) == [True, False] + assert _normalize_use_audio_in_video(torch.tensor([1, 0]), 2) == [ + True, + False, + ] + + +def test_normalize_use_audio_in_video_rejects_wrong_length(): + with pytest.raises(ValueError, match="one boolean per video"): + _normalize_use_audio_in_video([True], 2) + + +def test_coerce_use_audio_in_video_for_hf_processor_keeps_global_bool(): + mm_kwargs = {"use_audio_in_video": True} + + result = _coerce_use_audio_in_video_for_hf_processor({"videos": [object(), object()]}, mm_kwargs) + + assert result is mm_kwargs + assert result["use_audio_in_video"] is True + + +def test_coerce_use_audio_in_video_for_hf_processor_hides_per_video_mask_from_hf(): + mm_kwargs = {"use_audio_in_video": [True, False], "other": "kept"} + + result = _coerce_use_audio_in_video_for_hf_processor({"videos": [object(), object()]}, mm_kwargs) + + assert result is not mm_kwargs + assert result["use_audio_in_video"] is False + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [True, False] + assert result["other"] == "kept" + assert mm_kwargs["use_audio_in_video"] == [True, False] + + +def test_coerce_use_audio_in_video_for_hf_processor_does_not_validate_against_empty_mm_data(): + result = _coerce_use_audio_in_video_for_hf_processor( + {}, + {"use_audio_in_video": [True, False]}, + ) + + assert result["use_audio_in_video"] is False + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [True, False] + + +def test_coerce_use_audio_in_video_for_hf_processor_does_not_validate_against_partial_mm_data(): + result = _coerce_use_audio_in_video_for_hf_processor( + {"videos": [object()]}, + {"use_audio_in_video": [True, False]}, + ) + + assert result["use_audio_in_video"] is False + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [True, False] + + +def test_filter_video_use_audio_in_video_for_uncached_items_aligns_partial_cache_mask(): + result = _filter_video_use_audio_in_video_for_uncached_items( + { + "use_audio_in_video": [True, False, True], + "second_per_grid_ts": [0.5, 1.0, 1.5], + }, + [True, False, False], + ) + + assert result["use_audio_in_video"] == [False, True] + assert result[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] == [False, True] + assert result["second_per_grid_ts"] == [1.0, 1.5] + + +def test_filter_video_use_audio_in_video_for_uncached_items_validates_request_mask_length(): + with pytest.raises(ValueError, match="one boolean per video"): + _filter_video_use_audio_in_video_for_uncached_items( + {"use_audio_in_video": [True]}, + [False, False], + ) + + +def test_mm_fields_config_keeps_global_use_audio_in_video_shared(): + fake_self = _fake_processor_with_spatial_merge_size() + hf_inputs = { + "video_grid_thw": torch.tensor([[1, 4, 4], [1, 4, 4]]), + "use_audio_in_video": torch.tensor(False), + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_mm_fields_config( + fake_self, + hf_inputs, + {}, + ) + + assert isinstance(result["use_audio_in_video"].field, MultiModalSharedField) + + +def test_mm_fields_config_uses_batched_field_for_per_video_audio_mask(): + fake_self = _fake_processor_with_spatial_merge_size() + hf_inputs = { + "video_grid_thw": torch.tensor([[1, 4, 4], [1, 4, 4]]), + "use_audio_in_video": torch.tensor([True, False]), + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_mm_fields_config( + fake_self, + hf_inputs, + {_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY: [True, False]}, + ) + + assert isinstance(result["use_audio_in_video"].field, MultiModalBatchedField) + + +def test_mm_fields_config_uses_batched_field_for_vector_hf_audio_mask(): + fake_self = _fake_processor_with_spatial_merge_size() + hf_inputs = { + "video_grid_thw": torch.tensor([[1, 4, 4], [1, 4, 4]]), + "use_audio_in_video": torch.tensor([True, False]), + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_mm_fields_config( + fake_self, + hf_inputs, + {}, + ) + + assert isinstance(result["use_audio_in_video"].field, MultiModalBatchedField) + + +def test_get_video_use_audio_in_video_preserves_explicit_false(): + fake_self = _fake_processor() + mm_kwargs = { + "video": [ + {"use_audio_in_video": _Feature(torch.tensor(True))}, + {"use_audio_in_video": _Feature(torch.tensor(False))}, + ] + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + mm_kwargs, + {}, + ) + + assert result == [True, False] + + +def test_get_video_use_audio_in_video_treats_missing_item_as_false_with_explicit_mask(): + fake_self = _fake_processor() + mm_kwargs = { + "video": [ + {"use_audio_in_video": _Feature(torch.tensor(True))}, + None, + ] + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + mm_kwargs, + {}, + ) + + assert result == [True, False] + + +def test_get_video_use_audio_in_video_falls_back_for_partial_cache_hit_with_explicit_mask(): + fake_self = _fake_processor() + update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID])) + mm_kwargs = { + "video": [ + {"use_audio_in_video": _Feature(torch.tensor(True))}, + None, + ] + } + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + mm_kwargs, + {"video": [[object()], [update_with_audio]]}, + ) + + assert result == [True, True] + + +def test_qwen2_5_prompt_updates_apply_audio_mask_per_video(): + fake_self = _fake_qwen2_prompt_processor() + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 2]]), + } + ) + mm_items = SimpleNamespace( + get_all_counts=lambda: {"video": 2, "audio": 1}, + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True, False]}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + + assert video_update.resolve(0).content.full == [ + AUDIO_BOS_TOKEN_ID, + VIDEO_TOKEN_ID, + AUDIO_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + ] + assert video_update.resolve(1).content.full == [VIDEO_TOKEN_ID, VIDEO_TOKEN_ID] + + +def test_qwen2_5_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): + captured = {} + fake_self = _fake_qwen2_prompt_processor() + + def omni_get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.omni_get_updates_use_audio_in_video = omni_get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + "second_per_grid_ts": torch.tensor([0.25]), + } + ) + mm_items = SimpleNamespace( + get_all_counts=lambda: {"video": 1, "audio": 1}, + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.5])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.25 + + +def test_qwen2_5_prompt_updates_accept_tensor_second_per_grid_ts_from_kwargs(): + captured = {} + fake_self = _fake_qwen2_prompt_processor() + + def omni_get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.omni_get_updates_use_audio_in_video = omni_get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + } + ) + mm_items = SimpleNamespace( + get_all_counts=lambda: {"video": 1, "audio": 1}, + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.75, 1.0])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.75 + + +def test_qwen2_5_prompt_updates_validate_mask_against_request_video_count(): + fake_self = _fake_qwen2_prompt_processor() + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "video_grid_thw": torch.tensor([[1, 1, 2], [1, 1, 2]]), + } + ) + mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2}) + + with pytest.raises(ValueError, match="one boolean per video"): + Qwen2_5OmniThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True]}, + out_mm_kwargs, + ) + + +def test_qwen2_5_mm_only_dummy_counts_subtract_only_videos_using_audio(): + captured_counts = {} + fake_self = SimpleNamespace( + dummy_inputs=SimpleNamespace(get_dummy_text=lambda counts: captured_counts.update(counts) or "dummy"), + _apply_hf_processor_text_mm=lambda **_kwargs: ([], {}, False), + ) + mm_items = SimpleNamespace(get_all_counts=lambda: {"video": 2, "audio": 1}) + + Qwen2_5OmniThinkerMultiModalProcessor._apply_hf_processor_mm_only( + fake_self, + mm_items, + {"use_audio_in_video": [True, False]}, + {}, + ) + + assert captured_counts == {"video": 2, "audio": 0} + + +def test_get_video_use_audio_in_video_falls_back_to_prompt_updates_for_cache_hits(): + fake_self = _fake_processor() + update_with_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID])) + update_without_audio = SimpleNamespace(content=SimpleNamespace(full=[VIDEO_TOKEN_ID])) + + result = Qwen2_5OmniThinkerMultiModalProcessor._get_video_use_audio_in_video( + fake_self, + {"video": [None, None]}, + {"video": [[update_with_audio], [update_without_audio]]}, + ) + + assert result == [True, False] + + +def test_derive_audio_from_video_placeholders_only_pairs_true_videos(): + fake_self = _fake_processor() + video_placeholders = [ + PlaceholderFeaturesInfo( + modality="video", + item_idx=0, + start_idx=5, + tokens=[VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID], + is_embed=torch.tensor([True, True, True]), + ), + PlaceholderFeaturesInfo( + modality="video", + item_idx=1, + start_idx=12, + tokens=[VIDEO_TOKEN_ID, VIDEO_TOKEN_ID], + is_embed=torch.tensor([True, True]), + ), + ] + + result = Qwen2_5OmniThinkerMultiModalProcessor._derive_audio_from_video_placeholders( + fake_self, + {"video": video_placeholders}, + {"audio": [[object()]]}, + [True, False], + ) + + assert len(result["audio"]) == 1 + assert result["audio"][0].item_idx == 0 + assert result["audio"][0].start_idx == 5 + assert result["audio"][0].is_embed.tolist() == [False, True, False] + + assert len(result["video"]) == 2 + assert result["video"][0].item_idx == 0 + assert result["video"][0].is_embed.tolist() == [True, False, True] + assert result["video"][1].item_idx == 1 + assert result["video"][1].is_embed.tolist() == [True, True] + + +def test_qwen3_processor_inherits_vllm_omni_per_video_helpers(): + assert issubclass( + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen2_5OmniThinkerMultiModalProcessor, + ) + + +def test_qwen3_prompt_updates_prefer_hf_second_per_grid_ts_from_outputs(): + captured = {} + fake_self = _fake_qwen3_prompt_processor() + + def get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.get_updates_use_audio_in_video = get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + "second_per_grid_ts": torch.tensor([0.25]), + } + ) + mm_items = SimpleNamespace( + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen3OmniMoeThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.5])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.25 + + +def test_qwen3_prompt_updates_accept_tensor_second_per_grid_ts_from_kwargs(): + captured = {} + fake_self = _fake_qwen3_prompt_processor() + + def get_updates_use_audio_in_video(**kwargs): + captured["video_second_per_grid_t"] = kwargs["video_second_per_grid_t"] + return [VIDEO_TOKEN_ID] + + fake_self.get_updates_use_audio_in_video = get_updates_use_audio_in_video + out_mm_kwargs = SimpleNamespace( + get_data=lambda **_kwargs: { + "audio_feature_lengths": torch.tensor([100]), + "video_grid_thw": torch.tensor([[1, 1, 1]]), + } + ) + mm_items = SimpleNamespace( + get_items=lambda *_args, **_kwargs: SimpleNamespace(get=lambda _idx: [0] * 100), + ) + + updates = Qwen3OmniMoeThinkerMultiModalProcessor._get_prompt_updates( + fake_self, + mm_items, + {"use_audio_in_video": [True], "second_per_grid_ts": torch.tensor([0.75, 1.0])}, + out_mm_kwargs, + ) + video_update = next(update for update in updates if update.modality == "video") + video_update.resolve(0) + + assert captured["video_second_per_grid_t"] == 0.75 + + +def test_qwen3_derive_audio_from_video_placeholders_excludes_audio_wrappers_from_embeds(): + fake_self = _fake_qwen3_processor() + video_placeholders = [ + PlaceholderFeaturesInfo( + modality="video", + item_idx=0, + start_idx=5, + tokens=[ + AUDIO_BOS_TOKEN_ID, + VIDEO_TOKEN_ID, + AUDIO_TOKEN_ID, + AUDIO_EOS_TOKEN_ID, + VIDEO_TOKEN_ID, + ], + is_embed=torch.ones(5, dtype=torch.bool), + ) + ] + + result = Qwen3OmniMoeThinkerMultiModalProcessor._derive_audio_from_video_placeholders( + fake_self, + {"video": video_placeholders}, + {"audio": [[object()]]}, + [True], + ) + + assert result["audio"][0].is_embed.tolist() == [ + False, + False, + True, + False, + False, + ] + assert result["video"][0].is_embed.tolist() == [ + False, + True, + False, + False, + True, + ] + + +def test_qwen3_mrope_positions_match_plain_video_placeholder_length(): + fake_self = SimpleNamespace( + iter_mm_features=lambda _: iter( + [ + ( + 2, + "video", + { + "grid_t": 1, + "grid_h": 2, + "grid_w": 2, + "t_factor": 1.0, + "use_audio_in_video": False, + "placeholder_len": 4, + }, + ) + ] + ) + ) + + positions, _ = Qwen3OmniMoeThinkerForConditionalGeneration.get_mrope_input_positions( + fake_self, + list(range(8)), + [], + ) + + assert positions.shape == (3, 8) + + +def test_qwen3_mrope_positions_count_single_interleaved_audio_wrappers(): + fake_self = SimpleNamespace( + iter_mm_features=lambda _: iter( + [ + ( + 2, + "video", + { + "grid_t": 1, + "grid_h": 2, + "grid_w": 2, + "t_factor": 1.0, + "use_audio_in_video": True, + "placeholder_len": 8, + }, + ) + ] + ), + _compute_interleaved_positions=lambda _start_idx, _data: ( + torch.zeros((3, 6), dtype=torch.long).numpy(), + 6, + ), + ) + + positions, _ = Qwen3OmniMoeThinkerForConditionalGeneration.get_mrope_input_positions( + fake_self, + list(range(10)), + [], + ) + + assert positions.shape == (3, 10) + + +def test_qwen3_interleaved_mrope_prefers_placeholder_length_for_audio_count(): + fake_self = SimpleNamespace( + _compute_audio_token_count=lambda _audio_feature_length: 100, + ) + data = { + "grid_t": 1, + "grid_h": 2, + "grid_w": 2, + "t_factor": 1.0, + "audio_feature_length": 999, + "placeholder_len": 8, + } + + positions, token_count = Qwen3OmniMoeThinkerForConditionalGeneration._compute_interleaved_positions( + fake_self, + 0, + data, + ) + + assert token_count == 6 + assert positions.shape == (3, 6) + + +def test_omni_model_runner_output_can_carry_kv_connector_output_only(): + marker = object() + + output = OmniModelRunnerOutput.with_kv_conn_output_only(marker) + + assert output.req_ids == [] + assert output.req_id_to_index == {} + assert output.sampled_token_ids == [] + assert output.kv_connector_output is marker diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py new file mode 100644 index 00000000000..7be4f6afad7 --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for ``Qwen3OmniMoeCode2Wav.chunked_decode_streaming``. + +Background +---------- +The Code2Wav decoder is a causal conv / transposed-conv stack whose forward +emits a fixed number ``C`` *fewer* samples than ``frames * total_upsample`` (the +right-edge trim, ``C == 555`` for Qwen3-Omni-30B-A3B). ``chunked_decode_streaming`` +used to slice ``[left_context * up : code_seq_len * up]``, assuming the full +length. Because the real output is ``C`` short, the upper bound clamped and every +streaming chunk silently dropped its last ``C`` samples -> a gap/click at every +chunk boundary plus ~1.2% cumulative time compression. + +The fix measures ``tail`` from the actual decoder output and shifts each chunk's +start back by ``tail`` so the chunk refills the gap the previous chunk left. + +Why this is path-agnostic (eager *and* CUDA graph) +-------------------------------------------------- +Originally the CUDA-graph wrapper returned a *surplus* (nominal ``actual * up``, +with stale padded-tail samples), so ``tail`` evaluated to 0 and the graph path +looked immune. Since #4466 (``CUDAGraphDecoderWrapper._trim_replay_output``) the +graph wrapper returns the eager-equivalent ``actual * up - C`` instead, so the +same per-chunk tail drop happens on the graph branch too and the same fix +repairs it. These tests cover both branches; ``_FakeWrapper`` models the +post-#4466 contract (``decode`` returns ``actual * up - C``). + +These are pure length/slice-arithmetic tests: CPU-only, no weights, no GPU. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import ( + Qwen3OmniMoeCode2Wav, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_UP = 8 # total_upsample +_TRIM = 5 # causal right-edge trim C, 0 < C < up (mirrors 555 < 1920) +_Q = 4 # num_quantizers + + +class _ShortTrimDecoder: + """Minimal stand-in exposing only what ``chunked_decode_streaming`` touches. + + Models a *shift-equivariant* causal decoder: the waveform value at an + absolute output sample ``q`` is ``floor(q / up)`` == the index of the frame + that sample belongs to. With codes set so frame ``i`` carries value ``i``, + any window covering a frame (with enough right context) reproduces the same + global waveform -- which is exactly the property that lets the next chunk + refill the tail the previous chunk dropped. The forward emits ``trim`` fewer + samples than ``frames * up`` (the right-edge trim that lacks right context). + """ + + def __init__(self, *, up=_UP, trim=_TRIM, num_quantizers=_Q, cudagraph=False): + self.total_upsample = up + self.trim = trim + self.config = SimpleNamespace(num_quantizers=num_quantizers) + self._cudagraph_enabled = cudagraph + self._cudagraph_wrapper = _FakeWrapper(self) if cudagraph else None + + def _decode(self, codes: torch.Tensor) -> torch.Tensor: + # codes: [B, Q, W]; frame value (row 0) == absolute frame index. + batch, _, width = codes.shape + out_len = max(0, width * self.total_upsample - self.trim) + frame_of = torch.arange(out_len) // self.total_upsample # [out_len] + vals = codes[:, 0, :].to(torch.float32) # [B, W] + out = vals.gather(1, frame_of.unsqueeze(0).expand(batch, -1)) # [B, out_len] + return out.unsqueeze(1) # [B, 1, out_len] + + # ``chunked_decode_streaming`` calls ``self(codes)`` on the eager branch. + def __call__(self, codes: torch.Tensor) -> torch.Tensor: + return self._decode(codes) + + +class _FakeWrapper: + """Models ``CUDAGraphDecoderWrapper.decode`` after #4466: returns the + eager-equivalent short length (``actual * up - C``), values correct.""" + + def __init__(self, decoder: _ShortTrimDecoder): + self.decoder = decoder + + def decode(self, codes: torch.Tensor) -> torch.Tensor: + return self.decoder._decode(codes) + + +def _make_codes(frames: int, q: int = _Q) -> torch.Tensor: + """codes[0, :, i] == i so each frame carries its own absolute index.""" + idx = torch.arange(frames, dtype=torch.long) + return idx.view(1, 1, frames).expand(1, q, frames).contiguous() + + +def _decode_chunk(model, window, lc): + """Call the real production method for one streaming window.""" + win_frames = window.shape[-1] + out = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( + model, + window, + left_context_size=[lc], + seq_token_counts=[win_frames * model.config.num_quantizers], + ) + return out[0].reshape(-1) + + +def _drive_async(model, codes, *, initial=4, chunk=25, left_context=25): + """Replicate qwen3_omni's ``talker2code2wav_async_chunk`` 4/25/25 schedule.""" + total = codes.shape[-1] + chunks = [] + emitted = 0 + first = True + while emitted < total: + target = min(initial, total) if first else min(emitted + chunk, total) + lc = min(emitted, left_context) + window = codes[:, :, emitted - lc : target] + chunks.append(_decode_chunk(model, window, lc)) + emitted = target + first = False + return torch.cat(chunks, dim=-1) + + +def _single_shot(model, codes): + """Ground-truth: one full-utterance forward (length T*up - trim).""" + return model(codes)[0].reshape(-1) + + +def test_streaming_reconstructs_single_shot_eager(): + """Eager: concatenated streaming chunks == the single-shot full decode.""" + model = _ShortTrimDecoder() + codes = _make_codes(116) + + gt = _single_shot(model, codes) + streamed = _drive_async(model, codes) + + assert streamed.shape[-1] == gt.shape[-1] + torch.testing.assert_close(streamed, gt) + + +def test_streaming_reconstructs_single_shot_cudagraph(): + """CUDA-graph branch (post-#4466 contract): identical exact reconstruction. + + Answers the first-principles question: the non-eager path is NOT immune on + current main -- it drops the same tail without the fix, and the fix repairs + it the same way. + """ + model = _ShortTrimDecoder(cudagraph=True) + codes = _make_codes(116) + + gt = _single_shot(model, codes) # forward() is eager regardless of the branch + streamed = _drive_async(model, codes) + + assert streamed.shape[-1] == gt.shape[-1] + torch.testing.assert_close(streamed, gt) + + +def test_eager_and_cudagraph_streaming_match(): + """Both decode branches must yield bit-identical streamed audio.""" + codes = _make_codes(116) + eager = _drive_async(_ShortTrimDecoder(cudagraph=False), codes) + graph = _drive_async(_ShortTrimDecoder(cudagraph=True), codes) + torch.testing.assert_close(eager, graph) + + +def test_per_chunk_tail_drop_is_repaired_not_present(): + """The defect is exactly ``C`` lost per internal boundary; the fix removes it. + + Without the fix the streamed length would be short by + ``(num_internal_boundaries) * trim``; assert the boundary count and that the + repaired stream has the full single-shot length (no deficit). + """ + model = _ShortTrimDecoder() + codes = _make_codes(116) + gt = _single_shot(model, codes) + streamed = _drive_async(model, codes) + + # 116 frames @ initial=4, chunk=25 -> targets 4,29,54,79,104,116 => 5 boundaries + num_internal_boundaries = 5 + buggy_len = gt.shape[-1] - num_internal_boundaries * _TRIM + assert streamed.shape[-1] == gt.shape[-1] + assert streamed.shape[-1] != buggy_len # the pre-fix slice produced buggy_len + + +def test_first_chunk_start_is_not_shifted_below_zero(): + """First chunk has left_context==0 -> start = max(0, 0 - tail) = 0 (unchanged).""" + model = _ShortTrimDecoder() + codes = _make_codes(4) # single initial chunk, lc=0 + out = _decode_chunk(model, codes, lc=0) + # length == frames*up - trim, values == frame indices (GT prefix) + assert out.shape[-1] == 4 * _UP - _TRIM + torch.testing.assert_close(out, _single_shot(model, codes)) + + +def test_no_shift_when_decoder_returns_nominal_length(): + """C==0 decoder (e.g. Qwen3-TTS): tail==0 -> genuine no-op. + + With no trim, the fixed slice must equal the original + ``[left_context*up : code_seq_len*up]`` -- every chunk emits exactly + ``ctx_len*up`` samples and streaming reconstructs the full (untrimmed) audio. + """ + model = _ShortTrimDecoder(trim=0) + codes = _make_codes(116) + gt = _single_shot(model, codes) + streamed = _drive_async(model, codes) + + assert gt.shape[-1] == 116 * _UP # no trim -> nominal full length + assert streamed.shape[-1] == gt.shape[-1] + torch.testing.assert_close(streamed, gt) + + +def test_batched_streaming_only_shifts_the_boundary_row(): + """Per-row tail: in a shared window the shorter row had right context from + the longer row's frames, so it lost nothing (tail==0, no shift); only the + row at the decoded boundary (code_seq_len == window length) is shifted.""" + model = _ShortTrimDecoder() + # One window of W=30 frames, two rows. codes[:, 0, :] are absolute indices. + idx = torch.arange(30, dtype=torch.long) + window = idx.view(1, 1, 30).expand(2, _Q, 30).contiguous() + + long_frames, short_frames = 30, 20 # row0 spans the full window, row1 is shorter + lc = 5 + out = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( + model, + window, + left_context_size=[lc, lc], + seq_token_counts=[long_frames * _Q, short_frames * _Q], + ) + + batch_len = 30 * _UP - _TRIM # actual decoder output width + # Row 0 (boundary row): tail == TRIM -> start = lc*up - TRIM + assert out[0].shape[-1] == batch_len - (lc * _UP - _TRIM) + # Row 1 (shorter): code_seq_len*up < batch_len -> tail == 0 -> start = lc*up + assert out[1].shape[-1] == short_frames * _UP - lc * _UP diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py new file mode 100644 index 00000000000..e5fa38a6b4b --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for Qwen3-Omni Thinker fused-MoE LoRA setup.""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch +from peft import LoraConfig as PeftLoraConfig +from peft import get_peft_model +from torch import nn +from transformers import PretrainedConfig +from vllm.config.lora import LoRAConfig +from vllm.lora.layers.fused_moe import FusedMoE3DWithLoRA +from vllm.lora.lora_model import LoRAModel +from vllm.lora.model_manager import LoRAModelManager +from vllm.lora.peft_helper import PEFTHelper + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeThinkerForConditionalGeneration, + _ensure_thinker_architecture, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_NUM_EXPERTS = 2 +_HIDDEN_SIZE = 4 +_INTERMEDIATE_SIZE = 4 +_LORA_RANK = 1 +_EXPERTS_MODULE = "language_model.model.layers.0.mlp.experts" +_PEFT_TARGET_PARAMETERS = [ + "thinker.model.layers.0.mlp.experts.gate_up_proj", + "thinker.model.layers.0.mlp.experts.down_proj", +] + + +class _TinyExperts(nn.Module): + """Expert parameters with the same 3-D layout as Qwen3-Omni.""" + + def __init__(self): + super().__init__() + self.gate_up_proj = nn.Parameter(torch.zeros(_NUM_EXPERTS, 2 * _INTERMEDIATE_SIZE, _HIDDEN_SIZE)) + self.down_proj = nn.Parameter(torch.zeros(_NUM_EXPERTS, _HIDDEN_SIZE, _INTERMEDIATE_SIZE)) + + +class _TinyQwen3OmniMoe(nn.Module): + """Minimal module tree needed for PEFT to save expert LoRA weights.""" + + def __init__(self, model_path: Path): + super().__init__() + self.thinker = nn.ModuleDict( + { + "model": nn.ModuleDict( + {"layers": nn.ModuleList([nn.ModuleDict({"mlp": nn.ModuleDict({"experts": _TinyExperts()})})])} + ) + } + ) + self.name_or_path = str(model_path) + self.config = {"model_type": "qwen3_omni_moe"} + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return inputs + + +class _CpuFusedMoE3DWithLoRA(FusedMoE3DWithLoRA): + """CPU harness for the real 3-D MoE LoRA weight methods.""" + + def __init__(self): + nn.Module.__init__(self) + self.moe_config = SimpleNamespace( + hidden_dim=_HIDDEN_SIZE, + num_local_experts=_NUM_EXPERTS, + num_experts=_NUM_EXPERTS, + intermediate_size_per_partition=_INTERMEDIATE_SIZE, + moe_parallel_config=SimpleNamespace( + tp_size=1, + tp_rank=0, + ep_rank=0, + use_ep=False, + ), + ) + self.device = torch.device("cpu") + self.tp_size = 1 + self.tp_rank = 0 + self._w13_slices = 1 + + +def _save_peft_adapter(tmp_path: Path) -> Path: + base_model_dir = tmp_path / "base_model" + base_model_dir.mkdir() + (base_model_dir / "config.json").write_text('{"model_type": "qwen3_omni_moe"}', encoding="utf-8") + + peft_model = get_peft_model( + _TinyQwen3OmniMoe(base_model_dir), + PeftLoraConfig( + r=_LORA_RANK, + lora_alpha=_LORA_RANK, + target_modules=[], + target_parameters=_PEFT_TARGET_PARAMETERS, + ), + ) + with torch.no_grad(): + for value, parameter in enumerate( + (parameter for parameter in peft_model.parameters() if parameter.requires_grad), + start=1, + ): + parameter.fill_(value) + + adapter_dir = tmp_path / "adapter" + peft_model.save_pretrained(adapter_dir) + return adapter_dir + + +def test_peft_expert_lora_loads_into_thinker(tmp_path: Path): + adapter_dir = _save_peft_adapter(tmp_path) + peft_helper = PEFTHelper.from_local_dir(str(adapter_dir), max_position_embeddings=128) + lora_model = LoRAModel.from_local_checkpoint( + str(adapter_dir), + {"experts"}, + peft_helper=peft_helper, + lora_model_id=1, + device="cpu", + dtype=torch.float32, + weights_mapper=Qwen3OmniMoeThinkerForConditionalGeneration.hf_to_vllm_mapper, + ) + + thinker_config = PretrainedConfig() + _ensure_thinker_architecture(thinker_config) + fused_moe = _CpuFusedMoE3DWithLoRA() + fused_moe.create_lora_weights( + max_loras=1, + lora_config=LoRAConfig( + max_lora_rank=_LORA_RANK, + max_loras=1, + lora_dtype=torch.float32, + ), + model_config=thinker_config, + ) + + manager = object.__new__(LoRAModelManager) + manager.is_pooling_model = False + manager._is_3d_moe_model = Qwen3OmniMoeThinkerForConditionalGeneration.is_3d_moe_weight + manager._stack_moe_lora_weights(lora_model, fused_moe, _EXPERTS_MODULE) + + expert_lora = lora_model.get_lora(_EXPERTS_MODULE) + assert expert_lora is not None + fused_moe.set_lora(0, expert_lora.lora_a, expert_lora.lora_b) + + assert fused_moe.adapter_enabled.tolist() == [1, 0] + for weights in ( + fused_moe.w13_lora_a_stacked[0], + fused_moe.w13_lora_b_stacked[0], + fused_moe.w2_lora_a_stacked[0], + fused_moe.w2_lora_b_stacked[0], + ): + assert torch.count_nonzero(weights[0]) == weights[0].numel() + + +def test_existing_thinker_architecture_is_preserved(): + config = PretrainedConfig(architectures=["CustomThinkerArchitecture"]) + + _ensure_thinker_architecture(config) + + assert config.architectures == ["CustomThinkerArchitecture"] diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index 87281af82b2..5e31c2cbd60 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -18,8 +18,11 @@ import pytest import torch +import torch.nn.functional as F from pytest_mock import MockerFixture +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + # Direct file import to avoid vllm_omni.__init__ patch dependencies. _MODELS = os.path.join( os.path.dirname(__file__), @@ -46,6 +49,11 @@ def _load_module(name: str, filename: str): def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: """Build the dict of modules to inject into sys.modules.""" + + class NativeCustomOp(torch.nn.Module): + def forward(self, *args, **kwargs): + return self.forward_native(*args, **kwargs) + platforms_mock = mocker.MagicMock() platforms_mock.current_omni_platform.supports_torch_inductor.return_value = False platforms_mock.current_omni_platform.is_npu.return_value = False @@ -73,6 +81,8 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: vllm_parallel_mock = mocker.MagicMock() vllm_parallel_mock.VocabParallelEmbedding = torch.nn.Embedding + custom_op_mock = types.ModuleType("vllm_omni.diffusion.layers.custom_op") + custom_op_mock.CustomOp = NativeCustomOp return { "vllm_omni": mocker.MagicMock(), @@ -82,6 +92,7 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: "vllm.config.vllm": vllm_config_mod, "vllm.model_executor.model_loader.weight_utils": weight_utils_mock, "vllm.model_executor.layers.vocab_parallel_embedding": vllm_parallel_mock, + "vllm_omni.diffusion.layers.custom_op": custom_op_mock, "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"), "vllm_omni.model_executor.models": models_pkg, "vllm_omni.model_executor.models.common": common_pkg, @@ -165,6 +176,38 @@ def _make_vllm_config(mocker: MockerFixture, max_num_seqs: int = 4): return vllm_config +def test_npu_custom_ops_use_fused_norm_and_cached_rope(mocker: MockerFixture, loaded_target_classes) -> None: + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + cp_config, _ = _make_tiny_config(loaded_target_classes) + rms_calls = [] + + def npu_rms_norm(hidden_states, weight, epsilon): + rms_calls.append((hidden_states, weight, epsilon)) + return hidden_states + 1, None + + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=True) + mocker.patch.object( + common_mod, + "torch_npu", + types.SimpleNamespace(npu_rms_norm=npu_rms_norm), + create=True, + ) + + hidden_states = torch.zeros(1, 2, cp_config.hidden_size, dtype=torch.float16) + norm = common_mod._RMSNorm(cp_config.hidden_size, eps=cp_config.rms_norm_eps) + torch.testing.assert_close(norm.forward_npu(hidden_states), hidden_states + 1) + assert len(rms_calls) == 1 + assert rms_calls[0][2] == cp_config.rms_norm_eps + + rotary = common_mod._RotaryEmbedding(cp_config) + assert rotary.cos_cached.shape == (cp_config.num_code_groups + 1, cp_config.head_dim) + assert rotary.sin_cached.shape == rotary.cos_cached.shape + position_ids = torch.tensor([[0, 2, 4]]) + cos, sin = rotary.forward_npu(hidden_states, position_ids) + torch.testing.assert_close(cos, rotary.cos_cached[position_ids].to(torch.float16)) + torch.testing.assert_close(sin, rotary.sin_cached[position_ids].to(torch.float16)) + + class TestCodePredictorDtypeAlignment: """Test that code predictor buffers match model parameter dtype.""" @@ -339,6 +382,91 @@ def test_forward_generator_controls_sampling(self, mocker: MockerFixture, loaded assert not torch.equal(first[:, 1:], different[:, 1:]) +class TestCodePredictorPerRowGenerators: + """Seeded requests must stay deterministic inside a multi-row batch (#4883).""" + + def _make_predictor(self, mocker: MockerFixture, loaded_target_classes): + _, _, code_predictor_wrapper, _, _ = loaded_target_classes + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False) + cp_config, talker_config = _make_tiny_config(loaded_target_classes) + vllm_config = _make_vllm_config(mocker, max_num_seqs=4) + predictor = code_predictor_wrapper( + vllm_config=vllm_config, + config=cp_config, + talker_config=talker_config, + ) + predictor._wrapper_config.use_cuda_graphs = False + return predictor, talker_config + + def test_multinomial_per_row_matches_single_row_draws(self, mocker: MockerFixture, loaded_target_classes) -> None: + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + multinomial = common_mod.CodePredictorWrapper._multinomial + torch.manual_seed(7) + probs = torch.softmax(torch.randn(3, 64), dim=-1) + + def seeded(seed: int) -> torch.Generator: + generator = torch.Generator(device=probs.device) + generator.manual_seed(seed) + return generator + + batched = multinomial(probs, None, [seeded(11), None, seeded(22)]) + assert batched.shape == (3, 1) + # Each seeded row must reproduce a standalone draw from the same seed, + # independent of what the other rows in the batch consume. + assert torch.equal(batched[0:1], torch.multinomial(probs[0:1], num_samples=1, generator=seeded(11))) + assert torch.equal(batched[2:3], torch.multinomial(probs[2:3], num_samples=1, generator=seeded(22))) + + def test_forward_per_row_generators_are_row_independent(self, mocker: MockerFixture, loaded_target_classes) -> None: + predictor, talker_config = self._make_predictor(mocker, loaded_target_classes) + hidden = talker_config.hidden_size + torch.manual_seed(123) + row_embed = torch.randn(1, hidden) + row_hidden = torch.randn(1, hidden) + # Three identical rows: two share a seed, one differs. + layer0_code = torch.zeros(3, dtype=torch.long) + layer0_embed = row_embed.expand(3, hidden).contiguous() + last_talker_hidden = row_hidden.expand(3, hidden).contiguous() + + def seeded(seed: int) -> torch.Generator: + generator = torch.Generator(device=layer0_code.device) + generator.manual_seed(seed) + return generator + + def run(generators): + return predictor( + layer0_code=layer0_code, + layer0_embed=layer0_embed, + last_talker_hidden=last_talker_hidden, + do_sample=True, + temperature=0.9, + top_k=50, + top_p=1.0, + generators=generators, + ) + + first = run([seeded(1234), seeded(1234), seeded(4321)]) + second = run([seeded(1234), seeded(1234), seeded(4321)]) + + # Same per-row seeds -> the whole batch reproduces across calls. + assert torch.equal(first, second) + # Identical inputs + identical seeds -> identical rows within a call. + assert torch.equal(first[0], first[1]) + # A different seed on the same inputs must diverge in the residual layers. + assert not torch.equal(first[0, 1:], first[2, 1:]) + + def test_forward_rejects_mismatched_generators_length(self, mocker: MockerFixture, loaded_target_classes) -> None: + predictor, talker_config = self._make_predictor(mocker, loaded_target_classes) + hidden = talker_config.hidden_size + with pytest.raises(ValueError, match="one entry per row"): + predictor( + layer0_code=torch.zeros(2, dtype=torch.long), + layer0_embed=torch.randn(2, hidden), + last_talker_hidden=torch.randn(2, hidden), + generators=[torch.Generator()], + ) + + class TestCodePredictorModelDtype: """Test the inner model forward with different dtypes.""" @@ -457,3 +585,92 @@ def test_prefix_graph_env_requires_cuda_graphs( wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True), ) assert graph_wrapper._prefix_graphs_enabled is True + + +class TestGumbelMaxSampling: + """Sanity tests for the shared Gumbel-max sampling helper.""" + + @staticmethod + def _sample(logits: torch.Tensor, generator=None) -> torch.Tensor: + common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] + return common_mod.CodePredictorWrapper._sample_codes_gumbel(logits, generator=generator) + + def test_gumbel_max_picks_unique_high_logit(self, loaded_target_classes) -> None: + """If one logit is dominant, Gumbel-max must always pick it.""" + torch.manual_seed(0) + bsz, vocab = 4, 64 + logits = torch.full((bsz, vocab), -1e4) + logits[:, 7] = 0.0 # only finite entry + out = self._sample(logits) + assert out.shape == (bsz, 1) + assert torch.equal(out.squeeze(-1), torch.full((bsz,), 7, dtype=torch.long)) + + def test_gumbel_max_handles_single_finite_entry_rows(self, loaded_target_classes) -> None: + """Degenerate rows with one surviving finite entry should still return + a valid index instead of tripping multinomial-style asserts.""" + bsz, vocab = 2, 8 + logits = torch.full((bsz, vocab), float("-inf")) + logits[0, 3] = 0.0 + logits[1, 5] = 0.0 + out = self._sample(logits) + assert out[0].item() == 3 + assert out[1].item() == 5 + + def test_gumbel_max_generator_makes_results_reproducible(self, loaded_target_classes) -> None: + """Same ``Generator.manual_seed`` should produce identical samples.""" + torch.manual_seed(0) + logits = torch.randn(8, 32) + g1 = torch.Generator() + g1.manual_seed(123) + g2 = torch.Generator() + g2.manual_seed(123) + g3 = torch.Generator() + g3.manual_seed(456) + s1 = self._sample(logits, generator=g1) + s2 = self._sample(logits, generator=g2) + s3 = self._sample(logits, generator=g3) + assert torch.equal(s1, s2) + # Different seeds must (almost surely) yield a different sequence. + assert not torch.equal(s1, s3) + + def test_gumbel_max_per_row_generators_match_individual_rows(self, loaded_target_classes) -> None: + """A per-row generator list should preserve row independence.""" + torch.manual_seed(11) + logits = torch.randn(3, 16) + seeds = [101, 202, 303] + batched_generators = [] + expected_rows = [] + for row, seed in enumerate(seeds): + row_generator = torch.Generator() + row_generator.manual_seed(seed) + batched_generators.append(row_generator) + + solo_generator = torch.Generator() + solo_generator.manual_seed(seed) + expected_rows.append(self._sample(logits[row : row + 1], generator=solo_generator)) + + batched = self._sample(logits, generator=batched_generators) + expected = torch.cat(expected_rows, dim=0) + assert torch.equal(batched, expected) + + def test_gumbel_max_rejects_generator_length_mismatch(self, loaded_target_classes) -> None: + logits = torch.randn(2, 8) + generators = [torch.Generator()] + with pytest.raises(ValueError, match="Expected 2 per-row generators"): + self._sample(logits, generator=generators) + + def test_gumbel_max_distribution_matches_softmax(self, loaded_target_classes) -> None: + """Empirical frequency under Gumbel-max should be close to the + softmax probability mass. Loose tolerance because we draw a + finite number of samples.""" + torch.manual_seed(7) + vocab = 16 + logits = torch.randn(vocab) * 2.0 + target = F.softmax(logits, dim=-1, dtype=torch.float32) + + n_draws = 20000 + batched_logits = logits.unsqueeze(0).expand(n_draws, vocab).contiguous() + samples = self._sample(batched_logits).squeeze(-1) + empirical = torch.bincount(samples, minlength=vocab).float() / n_draws + # L1 distance under 0.03 with 20k draws is comfortable for vocab=16. + assert (empirical - target).abs().sum().item() < 0.05 diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py index 9a7460bf153..0f77455d7b8 100644 --- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py +++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py @@ -175,6 +175,77 @@ def test_padded_output_bounded(decoder, wrapper, seq_len): assert max_diff < 1.0, f"Max diff {max_diff} exceeds bound" +@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99]) +def test_padded_short_output_decode_length_matches_eager(seq_len): + """Streaming ``decode()`` of a shorter-than-nominal decoder must match eager length. + + Regression for #4466: Qwen3-Omni Code2Wav returns fewer samples than + ``seq_len * total_upsample``. The graph ``decode()`` path used to trim to the + nominal length, so for padded (non-capture-size) chunks it leaked a fixed + surplus of stale buffer-tail samples. The streaming concat turned that + per-chunk surplus into a ~46 ms "lag in the middle". The trim must instead be + relative to the captured output length, matching the eager forward exactly. + """ + short_decoder = ShortOutputDecoder().to(DEVICE).eval() + short_wrapper = CUDAGraphDecoderWrapper( + decoder=short_decoder, + capture_sizes=[25, 50, 100], + capture_batch_sizes=[1], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + short_wrapper.warmup(DEVICE) + codes = _random_codes(seq_len) + with torch.no_grad(): + eager_out = short_decoder(codes) + graph_out = short_wrapper.decode(codes) + surplus = graph_out.shape[-1] - eager_out.shape[-1] + assert surplus == 0, ( + f"graph length {graph_out.shape[-1]} != eager length {eager_out.shape[-1]} " + f"(surplus {surplus}); the padded graph output is leaking stale tail samples" + ) + # Interior content (away from the conv receptive-field boundary) must still match. + boundary = 3 * TOTAL_UPSAMPLE + if eager_out.shape[-1] > boundary: + torch.testing.assert_close(graph_out[..., :-boundary], eager_out[..., :-boundary], atol=1e-5, rtol=1e-5) + + +def test_compiled_padded_short_output_length_matches_eager(monkeypatch): + """Compiled replay path must also trim short-output decoders to eager length. + + Same regression as the graph path (#4466) but exercising the torch.compile + + CUDA Graph branch in ``_decode``. ``torch.compile`` is faked (as in the other + compiled-path tests) so the test stays fast and deterministic. + """ + + def _fake_compile(model, **_kwargs): + def _compiled(codes): + return model(codes) + + return _compiled + + monkeypatch.setattr(torch, "compile", _fake_compile) + + short_decoder = ShortOutputDecoder().to(DEVICE).eval() + compiled_wrapper = CUDAGraphDecoderWrapper( + decoder=short_decoder, + capture_sizes=[25, 50], + capture_batch_sizes=[1], + compile_shapes=[(1, 25), (1, 50)], + num_quantizers=NUM_QUANTIZERS, + enabled=True, + ) + compiled_wrapper.warmup(DEVICE) + codes = _random_codes(30) # not a capture size -> padded to bucket 50 on the compiled path + with torch.no_grad(): + eager_out = short_decoder(codes) + graph_out = compiled_wrapper.decode(codes) + surplus = graph_out.shape[-1] - eager_out.shape[-1] + assert surplus == 0, ( + f"compiled graph length {graph_out.shape[-1]} != eager length {eager_out.shape[-1]} (surplus {surplus})" + ) + + # ────────────────────────────────────────────────────────────────── # 3. Fallback to eager (size exceeds all capture sizes) → bit-identical # ────────────────────────────────────────────────────────────────── diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 6547fb29d32..4c94eb52020 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -375,6 +375,22 @@ def test_malformed_codec_length_warning_is_rate_limited(): assert out2.multimodal_outputs["model_outputs"][0].numel() == 0 +def test_forward_emits_zero_samples_for_empty_codec_payload(): + """Consumer side of #4463. + + The empty-but-finished payload that ``talker2code2wav_full_payload`` now + returns on a degenerate take carries a zero-length ``codes.audio``. Code2Wav + must turn that into a 0-sample output (finishing the request) so the Stage-1 + wait gate releases, rather than stalling the pipeline to the connector + timeout. + """ + model = _make_model() + + out = model.forward(input_ids=torch.zeros(0, dtype=torch.long)) + + assert out.multimodal_outputs["model_outputs"][0].numel() == 0 + + def test_forward_batches_equal_length_requests_in_one_decoder_call(): model = _make_model() diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py index 5fd961b4f1f..7358c789314 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py @@ -34,6 +34,7 @@ def _make_minimal_talker( """ model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration) model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16) + model._embedding_dtype = torch.bfloat16 if tts_pad_embed is None: tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16) model._tts_pad_embed = tts_pad_embed @@ -41,6 +42,7 @@ def _make_minimal_talker( def _default_raise(**_kwargs): raise AssertionError("build_prompt_embeds was not stubbed in this test") + model._embedding_dtype = torch.bfloat16 model._prompt_builder = SimpleNamespace( build_prompt_embeds=build_prompt_embeds if build_prompt_embeds is not None else _default_raise, ) @@ -81,6 +83,7 @@ def _make_minimal_builder( builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device) builder._residual_code_embeddings = lambda: [] builder._speaker_encoder = None + builder._embedding_dtype = torch.bfloat16 builder._tts_pad_embed_buffer = ( tts_pad_embed if tts_pad_embed is not None else torch.zeros((1, 4), dtype=torch.bfloat16) ) @@ -89,6 +92,7 @@ def _make_minimal_builder( ) builder._speaker_cache = None builder._text_tokenizer = None + builder._embedding_dtype = torch.bfloat16 builder._ref_audio_artifact_cache_max_entries = 256 builder._ref_audio_artifact_cache = OrderedDict() builder._resampler_cache = OrderedDict() @@ -252,6 +256,36 @@ def fake_embed_input_ids(input_ids): assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:]) +def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state(): + model = _make_minimal_talker() + + def fake_embed_input_ids(input_ids): + return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4) + + model.embed_input_ids = fake_embed_input_ids + trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4) + last_hidden = torch.full((4,), 2.0, dtype=torch.float32) + + out_ids, out_embeds, update = model.preprocess( + input_ids=torch.tensor([101, 202, 303], dtype=torch.long), + input_embeds=None, + text=["hello"], + task_type=["Base"], + hidden_states={"trailing_text": trailing_text, "last": last_hidden}, + meta={"talker_text_offset": 1, "codec_streaming": True}, + _omni_is_prefill=False, + _omni_num_computed_tokens=2400, + _omni_prompt_len=213, + ) + + assert out_ids.tolist() == [101, 202, 303] + assert torch.equal( + out_embeds.cpu(), + torch.tensor([[101.0] * 4, [202.0] * 4, [303.0] * 4], dtype=torch.bfloat16), + ) + assert update == {"meta": {"codec_streaming": True}} + + def test_decode_batch_preprocess_matches_decode_state_updates(): tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16) model = _make_minimal_talker(tts_pad_embed=tts_pad) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py index 2adfa17a410..3c145e4ca7c 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py @@ -74,3 +74,18 @@ def test_make_omni_output_omits_ref_key_when_no_request_has_one() -> None: ) assert "ref" not in out.multimodal_outputs["codes"] + + +def test_make_omni_output_keeps_hidden_states_for_replay_spans_without_audio_codes() -> None: + hidden = torch.arange(6 * 8, dtype=torch.float32).reshape(6, 8) + + out = _make_talker().make_omni_output( + hidden, + model_intermediate_buffer=[ + {"meta": {"codec_streaming": True}}, + _info(1, None), + ], + ) + + assert torch.equal(out.text_hidden_states, hidden) + assert out.multimodal_outputs["codes"]["audio"].shape == (1, _NUM_CODE_GROUPS) diff --git a/tests/model_executor/models/step_audio2/test_step_audio2_thinker.py b/tests/model_executor/models/step_audio2/test_step_audio2_thinker.py new file mode 100644 index 00000000000..8ed1ce9a4e5 --- /dev/null +++ b/tests/model_executor/models/step_audio2/test_step_audio2_thinker.py @@ -0,0 +1,151 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for Step-Audio2 thinker processor, encoder, and token handling.""" + +import numpy as np +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def create_dummy_audio(sample_rate: int = 16000, duration_sec: float = 1.0) -> tuple[np.ndarray, int]: + """Create a dummy audio signal (sine wave) for testing.""" + t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), dtype=np.float32) + audio = 0.5 * np.sin(2 * np.pi * 440 * t) + return audio, sample_rate + + +def test_processor_loads_correctly(): + """Test that Step-Audio2 processor can be imported and initialized.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + StepAudio2MultiModalProcessor, + StepAudio2ProcessingInfo, + StepAudio2Processor, + ) + + assert StepAudio2MultiModalProcessor is not None + assert StepAudio2ProcessingInfo is not None + assert StepAudio2Processor is not None + + +def test_audio_preprocessing(): + """Test that audio preprocessing produces correct tensor shapes.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + log_mel_spectrogram, + padding_mels, + ) + + audio, _sample_rate = create_dummy_audio(sample_rate=16000, duration_sec=1.0) + + mel = log_mel_spectrogram(audio) + + assert mel.ndim == 2, f"Expected 2D tensor, got {mel.ndim}D" + assert mel.shape[0] == 128, f"Expected 128 mel bins, got {mel.shape[0]}" + + mels = [mel, mel] + padded_mels, _mel_lens = padding_mels(mels) + + assert padded_mels.ndim == 3, f"Expected 3D tensor, got {padded_mels.ndim}D" + assert padded_mels.shape[0] == 2, f"Expected batch size 2, got {padded_mels.shape[0]}" + assert padded_mels.shape[1] == 128, f"Expected 128 mel bins, got {padded_mels.shape[1]}" + + +def test_feature_length_calculation(): + """Test that audio feature length calculation is correct.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + calculate_audio_feature_length, + ) + + assert calculate_audio_feature_length(1000) == 125 + + assert calculate_audio_feature_length(100) > 0 + + assert calculate_audio_feature_length(1) >= 1 + + +def test_token_config_constants(): + """Test that token configuration constants are set correctly.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( + DEFAULT_TOKEN_CONFIG, + STEP_AUDIO2_AUDIO_END, + STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID, + STEP_AUDIO2_AUDIO_START, + STEP_AUDIO2_AUDIO_VOCAB_SIZE, + STEP_AUDIO2_TEXT_MAX, + ) + + assert STEP_AUDIO2_TEXT_MAX < STEP_AUDIO2_AUDIO_START, "Text tokens should come before audio tokens" + assert STEP_AUDIO2_AUDIO_END == STEP_AUDIO2_AUDIO_START + STEP_AUDIO2_AUDIO_VOCAB_SIZE - 1 + assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID > STEP_AUDIO2_TEXT_MAX + assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID < STEP_AUDIO2_AUDIO_START + + assert DEFAULT_TOKEN_CONFIG.text_max == STEP_AUDIO2_TEXT_MAX + assert DEFAULT_TOKEN_CONFIG.audio_start == STEP_AUDIO2_AUDIO_START + + +def test_mm_field_config_structure(): + """Test that multimodal field config is properly structured for vLLM batching.""" + from vllm.multimodal.inputs import MultiModalFieldConfig + + audio_lens = torch.tensor([100, 200, 150], dtype=torch.int32) + + field_config = MultiModalFieldConfig.flat_from_sizes("audio", audio_lens) + + assert field_config is not None + + +def test_audio_encoder_output_shape(): + """Test that audio encoder produces correct output shapes.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + Adaptor, + AudioEncoder, + ) + + encoder = AudioEncoder(n_mels=128, n_ctx=1500, n_state=512, n_head=8, n_layer=6) + + batch_size = 2 + n_mels = 128 + time_steps = 400 + x = torch.randn(batch_size, n_mels, time_steps) + x_len = torch.tensor([time_steps, time_steps // 2], dtype=torch.int32) + + encoded, _encoded_lens = encoder(x, x_len) + + assert encoded.ndim == 3, f"Expected 3D tensor, got {encoded.ndim}D" + assert encoded.shape[0] == batch_size + assert encoded.shape[2] == 512 + + adapter = Adaptor(n_state=512, n_hidden=4096, kernel_size=3, stride=2) + adapted = adapter(encoded) + + assert adapted.ndim == 3 + assert adapted.shape[0] == batch_size + assert adapted.shape[2] == 4096 + + +def test_token_separation(): + """Test that token separation works correctly.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + StepAudio2ThinkerForConditionalGeneration, + ) + + token_ids = [100, 200, 151700, 300, 151800, 400] + + text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(token_ids) + + assert text_tokens == [100, 200, 300, 400] + assert audio_tokens == [4, 104] + + +def test_has_audio_output(): + """Test detection of audio tokens in output.""" + from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import ( + StepAudio2ThinkerForConditionalGeneration, + ) + + text_only = [100, 200, 300] + assert not StepAudio2ThinkerForConditionalGeneration.has_audio_output(text_only) + + with_audio = [100, 200, 151700, 300] + assert StepAudio2ThinkerForConditionalGeneration.has_audio_output(with_audio) diff --git a/tests/model_executor/models/test_aura_qwen3_vl.py b/tests/model_executor/models/test_aura_qwen3_vl.py new file mode 100644 index 00000000000..9035fcbf093 --- /dev/null +++ b/tests/model_executor/models/test_aura_qwen3_vl.py @@ -0,0 +1,69 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.aura_omni.qwen3_vl import ( + AuraQwen3VLForConditionalGeneration, + AuraQwen3VLProcessingInfo, + Qwen3VLProcessor, +) +from vllm_omni.model_executor.models.registry import OmniModelRegistry + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _FakeCtx: + def __init__(self, config): + self.config = config + self.processor_type = None + + def get_hf_config(self, typ=None): + if typ is not None: + raise TypeError("wrong config class") + return self.config + + def get_hf_processor(self, typ=None, **kwargs): + self.processor_type = typ + return SimpleNamespace(kwargs=kwargs) + + +def test_aura_qwen3_vl_processing_accepts_remote_qwen3_vl_config(): + info = object.__new__(AuraQwen3VLProcessingInfo) + info.ctx = _FakeCtx( + SimpleNamespace( + model_type="qwen3_vl", + vision_config=SimpleNamespace(spatial_merge_size=2), + text_config=SimpleNamespace(hidden_size=4096), + ) + ) + + assert info.get_hf_config().model_type == "qwen3_vl" + + +def test_aura_qwen3_vl_processing_rejects_unrelated_remote_config(): + info = object.__new__(AuraQwen3VLProcessingInfo) + info.ctx = _FakeCtx(SimpleNamespace(model_type="not_qwen3_vl")) + + with pytest.raises(TypeError, match="wrong config class"): + info.get_hf_config() + + +def test_aura_qwen3_vl_model_arch_registered(): + model_cls = OmniModelRegistry._try_load_model_cls("AuraQwen3VLForConditionalGeneration") + + assert model_cls is AuraQwen3VLForConditionalGeneration + + +def test_aura_qwen3_vl_processing_forces_upstream_processor_class(): + info = object.__new__(AuraQwen3VLProcessingInfo) + ctx = _FakeCtx(SimpleNamespace(model_type="qwen3_vl")) + info.ctx = ctx + + processor = info.get_hf_processor(extra="value") + + assert ctx.processor_type is Qwen3VLProcessor + assert processor.kwargs["use_fast"] is True + assert processor.kwargs["extra"] == "value" diff --git a/tests/model_executor/models/test_kv_cache_scale_mapper.py b/tests/model_executor/models/test_kv_cache_scale_mapper.py new file mode 100644 index 00000000000..98f016bb406 --- /dev/null +++ b/tests/model_executor/models/test_kv_cache_scale_mapper.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for vLLM's KV-cache scale mapper load path.""" + +from collections.abc import Iterable +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_STALE_API_FILES = [ + "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py", + "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py", + "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py", + "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py", +] + +_SOURCE_SCALE_NAME = "layers.0.self_attn.k_proj.output_scale" +_MAPPED_SCALE_NAME = "layers.0.self_attn.attn.k_scale" + + +class _RecordingCacheScaleMapper: + def __init__(self) -> None: + self.called = False + + def apply( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Iterable[tuple[str, torch.Tensor]]: + self.called = True + for name, weight in weights: + yield name.replace(".k_proj.output_scale", ".attn.k_scale"), weight + + +class _QuantConfigWithoutOldCacheScale: + def __init__(self) -> None: + self.mapper = _RecordingCacheScaleMapper() + + def get_cache_scale_mapper(self) -> _RecordingCacheScaleMapper: + return self.mapper + + +class _FakeModel: + def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None: + self.quant_config = quant_config + self.param = torch.nn.Parameter(torch.zeros(()), requires_grad=False) + + def named_parameters( + self, + remove_duplicate: bool = True, + ) -> Iterable[tuple[str, torch.nn.Parameter]]: + del remove_duplicate + return [(_MAPPED_SCALE_NAME, self.param)] + + +class _FakeHunyuanModel(_FakeModel): + def __init__(self, quant_config: _QuantConfigWithoutOldCacheScale) -> None: + super().__init__(quant_config) + self.config = SimpleNamespace( + tie_word_embeddings=False, + num_attention_heads=1, + num_key_value_heads=1, + use_cla=False, + ) + + def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]: + return [], {} + + def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor: + return weight + + +@pytest.mark.parametrize("rel_path", _STALE_API_FILES) +def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> None: + """Affected loaders must not call the removed vLLM get_cache_scale API.""" + source = (_REPO_ROOT / rel_path).read_text(encoding="utf-8") + assert ".get_cache_scale(" not in source + + +@pytest.mark.parametrize( + ("module_name", "class_name", "fake_model_cls"), + [ + ( + "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3", + "HunyuanModel", + _FakeHunyuanModel, + ), + ( + "vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old", + "Qwen2Model", + _FakeModel, + ), + ], +) +def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale( + module_name: str, + class_name: str, + fake_model_cls: type[_FakeModel], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Delegated loaders must not expect the removed get_cache_scale method.""" + module = __import__(module_name, fromlist=[class_name]) + monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False) + + fake_model = fake_model_cls(_QuantConfigWithoutOldCacheScale()) + load_weights = getattr(module, class_name).load_weights + + loaded = load_weights(fake_model, [(_MAPPED_SCALE_NAME, torch.ones(()))]) + + assert loaded == {_MAPPED_SCALE_NAME} + torch.testing.assert_close(fake_model.param.data, torch.ones(())) + + +@pytest.mark.parametrize( + ("module_name", "class_name"), + [ + ( + "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2", + "MammothModa2Qwen2ForCausalLM", + ), + ( + "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm", + "MiMoAudioLLMForConditionalGeneration", + ), + ], +) +def test_direct_custom_loaders_apply_cache_scale_mapper( + module_name: str, + class_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Direct custom loaders must apply the new mapper before fall-through load.""" + module = __import__(module_name, fromlist=[class_name]) + monkeypatch.setattr(module, "is_pp_missing_parameter", lambda *_: False) + + quant_config = _QuantConfigWithoutOldCacheScale() + fake_model = _FakeModel(quant_config) + load_weights = getattr(module, class_name).load_weights + + loaded = load_weights(fake_model, [(_SOURCE_SCALE_NAME, torch.ones(()))]) + + assert quant_config.mapper.called + assert loaded == {_MAPPED_SCALE_NAME} + torch.testing.assert_close(fake_model.param.data, torch.ones(())) diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py new file mode 100644 index 00000000000..8154ddd65dd --- /dev/null +++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Static regression checks for VoxCPM2 runner-assisted unified graph hooks. + +These tests intentionally avoid importing torch/vLLM so they can run in a +lightweight local checkout. Runtime audio quality still requires a CUDA test. +""" + +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[4] +RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py" +RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py" +TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py" +VOXCPM2_RUNTIME_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "runtime_config.py" +VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py" +VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py" +OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py" +DEPLOY = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm2.yaml" + + +def test_ar_runner_exposes_runner_assisted_full_metadata_hook(): + source = RUNNER.read_text() + + assert "Models without this hook keep the normal runner path" in source + assert "get_runner_assisted_full_attention_metadata_request" in source + assert "set_runner_assisted_full_attention_metadata_context" in source + assert "runner_assisted_full_attn" in source + assert "runner_assisted_full_attn_capture" in source + assert "pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn" in source + assert "for_cudagraph_capture=runner_assisted_full_attn_capture" in source + assert "cudagraph_runtime_mode=(" in source + assert "CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode" in source + refresh_source = source[source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") :] + refresh_source = refresh_source[: refresh_source.index("def _set_runner_assisted_full_attention_metadata_context")] + assert "num_computed_tokens_cpu" in refresh_source + assert "num_scheduled_tokens_np" in refresh_source + assert "optimistic_seq_lens_cpu[:num_reqs]=1" not in "".join(refresh_source.split()) + padding_source = source[ + source.index("runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request") : + ] + padding_source = padding_source[: padding_source.index("ubatch_slices, ubatch_slices_padded")] + assert "runner_assisted_full_attn_request.num_reqs_padded" in padding_source + assert "runner_assisted_full_attn_request.for_cudagraph_capture" in padding_source + assert "num_reqs_padded * max_num_scheduled_tokens" in padding_source + + +def test_ar_runner_without_model_hook_stays_on_normal_path(): + source = RUNNER.read_text() + contract_source = RUNNER_ASSISTED_METADATA.read_text() + request_source = source[source.index("def _get_runner_assisted_full_attention_metadata_request") :] + request_source = request_source[ + : request_source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") + ] + compact_request_source = "".join(request_source.split()) + + assert "Models without this hook keep the normal runner path" in request_source + assert "class RunnerAssistedAttentionMetadataProvider(Protocol)" in contract_source + assert "class RunnerAssistedFullAttentionMetadataRequest(NamedTuple)" in contract_source + assert "tuple[int, bool]" not in contract_source + assert 'hook=getattr(self.model,"get_runner_assisted_full_attention_metadata_request",None)' in ( + compact_request_source + ) + assert "ifnotcallable(hook):returnNone" in compact_request_source + assert ".tolist()" not in compact_request_source + assert "isinstance(request,RunnerAssistedFullAttentionMetadataRequest)" in compact_request_source + assert "request.num_reqs_padded" in request_source + assert "request.for_cudagraph_capture" in request_source + assert "exceptException" not in compact_request_source + + set_context_source = source[source.index("def _set_runner_assisted_full_attention_metadata_context") :] + set_context_source = set_context_source[: set_context_source.index("def _deferred_prefix_cache_mm_keys")] + assert "except Exception" not in set_context_source + + +def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise(): + source = TALKER.read_text() + runtime_config_source = VOXCPM2_RUNTIME_CONFIG.read_text() + compact_source = "".join(source.split()) + + assert "_voxcpm2_compile_without_inductor_cudagraphs" in source + assert "def _compile_without_inductor_cudagraphs" not in source + assert "_voxcpm2_unwrap_torch_compile" in source + assert "_orig_mod" in source + compile_source = source[source.index("def _voxcpm2_compile_without_inductor_cudagraphs") :] + compile_source = compile_source[: compile_source.index("def _setup_torch_compile")] + assert "torch._inductor.config" not in compile_source + assert '"triton.cudagraphs": False' in compile_source + assert '"triton.cudagraph_trees": False' in compile_source + + assert "self._enable_unified_decode_graph=self._runtime_config.unified_decode_graph_available(" in compact_source + assert "andnotself.deterministic_cfm_noise" in "".join(runtime_config_source.split()) + assert "decode_tail_graph" not in source + + unified_source = source[source.index("def _forward_unified_decode") :] + unified_source = unified_source[: unified_source.index("# -------------------- vllm hooks")] + assert "except Exception:" not in unified_source + assert "_failed_unified_graph_sizes" not in source + assert "self._enable_unified_decode_graph = False" not in unified_source + assert "_forward_unified_decode_fallback" not in source + assert "capture_failed" not in compact_source + assert "record_fallback" not in source + assert "fallbacks" not in source + assert "**info_dict: Unpack[VoxCPM2PreprocessInput]" in source + assert "**info: Unpack[VoxCPM2PostprocessInput]" in source + assert "reference_audio: Any" not in source + assert "prompt_audio: Any" not in source + assert "def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike" in source + assert "class _PrefillInputs(NamedTuple)" in source + assert "-> _PrefillInputs" in source + assert 'inputs["audio_feat"]' not in source + assert 'inputs["text_token"]' not in source + + capture_source = source[source.index("def _capture_unified_decode_graph") :] + capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")] + assert "_voxcpm2_compile_unified_capture_estimator" in capture_source + assert "_voxcpm2_compile_unified_capture_feat_encoder" in capture_source + graph_body = capture_source[capture_source.index("with torch.cuda.graph") :] + assert "g.cfm_noise.normal_()" not in graph_body + assert "g.cfm_noise.normal_()" in unified_source + + assert "_pre_capture_unified_graphs" not in source + assert "_unified_graph_pre_capture_sizes" not in source + assert "unified_decode_graph_pre_capture_sizes" not in source + + +def test_voxcpm2_batch_unified_graph_requires_runner_metadata_marker(): + source = TALKER.read_text() + compact_source = "".join(source.split()) + + assert "enable_runner_assisted_unified_decode_graph" not in source + assert "allow_unified_decode_graph_batch_attention" not in source + assert "get_runner_assisted_full_attention_metadata_request" in source + assert "set_runner_assisted_full_attention_metadata_context" in source + assert "RunnerAssistedFullAttentionMetadataRequest" in source + assert "_runner_assisted_unified_decode_graph_active" in source + assert "runner_full_metadata_missing" in source + assert "_build_unified_graph_bucket_sizes" in source + assert "_select_unified_graph_bucket_size" in source + + capture_source = source[source.index("def _capture_unified_decode_graph") :] + capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")] + assert "override_forward_context" in capture_source + assert "_nullify_volatile_metadata" in capture_source + assert "capture_context = override_forward_context(self._nullify_volatile_metadata(ctx))" in capture_source + assert "ifsize>1:continue" not in compact_source + forward_source = source[source.index("def _forward_unified_decode") :] + forward_source = forward_source[: forward_source.index("# -------------------- vllm hooks")] + assert "graph_size = self._select_unified_graph_bucket_size(num_reqs)" in forward_source + assert "self._unified_graphs[graph_size]" in forward_source + assert "g.input_embeds[num_reqs:graph_size].zero_()" not in forward_source + assert "padded rows as valid duplicates" in forward_source + assert "g.input_embeds[last : last + 1].expand" in forward_source + assert "ifnum_reqs>1andnotself._runner_assisted_unified_decode_graph_active" in compact_source + assert "andnotcfg.enable_runner_assisted_unified_decode_graph" not in compact_source + needs_source = source[source.index("def get_runner_assisted_full_attention_metadata_request") :] + needs_source = needs_source[: needs_source.index("def set_runner_assisted_full_attention_metadata_context")] + assert "_select_unified_graph_bucket_size(num_reqs)" in needs_source + assert "_should_use_decode_graph(num_reqs)" not in needs_source + assert "RunnerAssistedFullAttentionMetadataRequest(" in needs_source + + +def test_voxcpm2_unified_skip_preserves_segmented_decode_graphs(): + source = TALKER.read_text() + compact_source = "".join(source.split()) + forward_source = source[source.index("def forward(") :] + forward_source = forward_source[: forward_source.index("# -------------------- prefill / decode helpers")] + + assert "use_segmented_decode_graph" in forward_source + assert "unified_skip_reasonisnotNone" in compact_source + assert "notself._enable_unified_decode_graphorunified_skip_reasonisnotNone" in compact_source + assert "ifuse_segmented_decode_graph:" in compact_source + assert "anduse_segmented_decode_graph" in compact_source + + +def test_voxcpm2_scheduler_policy_stays_model_local(): + common_source = OMNI_AR_SCHEDULER.read_text() + voxcpm2_scheduler_source = VOXCPM2_SCHEDULER.read_text() + pipeline_source = VOXCPM2_PIPELINE.read_text() + + assert "VoxCPM2TalkerForConditionalGeneration" not in common_source + assert "voxcpm2_runtime_config" not in common_source + assert "pure_decode_graph" not in common_source + assert "_schedule_with_optional_waiting_deferral" not in common_source + assert "def _should_defer_waiting_admission(self) -> bool:" in common_source + + assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source + assert "from .runtime_config import _VoxCPM2RuntimeConfig" in voxcpm2_scheduler_source + assert "_VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config)" in voxcpm2_scheduler_source + assert "unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda())" in voxcpm2_scheduler_source + assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source + assert "def schedule(" not in voxcpm2_scheduler_source + assert "create_request_queue" not in voxcpm2_scheduler_source + assert "original_waiting.prepend_requests(deferred_waiting)" not in voxcpm2_scheduler_source + assert "return self._should_defer_waiting_for_unified_decode_graph()" in voxcpm2_scheduler_source + assert "vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler" in pipeline_source + + +def test_voxcpm2_deploy_defaults_to_full_unified_graph_only(): + source = DEPLOY.read_text() + + assert "max_num_seqs: 8" in source + assert "enable_unified_decode_graph: true" in source + assert "unified_decode_graph_max_batch_size: 8" in source + assert "unified_decode_graph_pre_capture_sizes" not in source + assert "enable_runner_assisted_unified_decode_graph" not in source + assert "allow_unified_decode_graph_batch_attention" not in source + assert "decode_graph_capture_policy" not in source diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index 28da0b1bc21..f5a94376a8a 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py @@ -52,6 +52,10 @@ def _make_bare_talker(): talker._cfm_buffers = None talker._last_audio_output_req_ids = [] talker._batched_fsq_fusion_max_batch = 32 + # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to + # default (False) to avoid AttributeError when test exercises forward + # path that checks this flag. + talker._enable_unified_decode_graph = False return talker @@ -329,23 +333,27 @@ def test_batched_fsq_fusion_matches_per_request_decode_prepare(self) -> None: for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True): torch.testing.assert_close(actual_res, expected_res) - torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"]) + torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden) assert [state.decode_step_count for state in batch_states] == [1, 1] def test_batch_decode_matches_per_request_state_updates(self) -> None: + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + _DecodeResidualMeta, + ) + seq_talker = self._make_decode_talker() batch_talker = self._make_decode_talker() seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] metas = [ - {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])}, - {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])}, + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])), + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])), ] batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]]) for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True): - seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu")) + seq_talker._finish_decode(state, meta, res_out.unsqueeze(0)) batch_talker._finish_decode_batch( [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)], @@ -360,12 +368,16 @@ def test_batch_decode_matches_per_request_state_updates(self) -> None: torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu) def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None: + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + _DecodeResidualMeta, + ) + talker = self._make_decode_talker() talker._enable_batched_cfm = False states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] metas = [ - {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])}, - {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])}, + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])), + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])), ] batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]]) diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index 3b645632e9c..cf40f9ad417 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -156,7 +156,7 @@ def compute_mm_logits( ) audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) - mm_tokens = {"audio": audio_list} + mm_tokens = {"codes": {"audio": audio_list}} return fake_eos, mm_tokens @@ -192,9 +192,9 @@ def _cfg_alpha(batch_size, value=1.2, device=DEVICE): def _unpack_audio_codes(result): - """Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes).""" + """Unpack (fake_eos, {"codes": {"audio": [list of tensors]}}) into (fake_eos, audio_codes).""" fake_eos, mm_tokens = result - audio_list = mm_tokens["audio"] + audio_list = mm_tokens["codes"]["audio"] # Each element is (1, 1, 1+n_acoustic_codebook), concat along batch dim audio_codes = torch.cat(audio_list, dim=0).squeeze(1) # (B, 1+n_acoustic) return fake_eos, audio_codes diff --git a/tests/model_executor/models/voxtral_tts/test_tts_preprocess.py b/tests/model_executor/models/voxtral_tts/test_tts_preprocess.py new file mode 100644 index 00000000000..399a7ecbdb2 --- /dev/null +++ b/tests/model_executor/models/voxtral_tts/test_tts_preprocess.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools + +import pytest +import torch +import torch.nn as nn + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@functools.lru_cache(maxsize=1) +def _voxtral_tts_model_cls(): + from tests.model_executor.helpers import bootstrap_vllm_layer_custom_op_modules + + bootstrap_vllm_layer_custom_op_modules() + import vllm.model_executor.models.utils # noqa: F401 + + from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts import ( + VoxtralTTSForConditionalGeneration, + ) + + return VoxtralTTSForConditionalGeneration + + +class FakeAudioGeneration(nn.Module): + def __init__(self): + super().__init__() + self.embed_multimodal_calls = [] + + def embed_multimodal(self, **kwargs): + self.embed_multimodal_calls.append(kwargs) + return [torch.full((1, 4), 7.0)] + + +def _make_voxtral_tts_model(): + model_cls = _voxtral_tts_model_cls() + model = model_cls.__new__(model_cls) + nn.Module.__init__(model) + model.model_stage = "audio_generation" + model.model = FakeAudioGeneration() + model._audio_token_id = 42 + return model + + +def test_tts_preprocess_consumes_nested_codes_audio_feedback(): + model = _make_voxtral_tts_model() + input_ids = torch.tensor([model._audio_token_id]) + input_embeds = torch.zeros((1, 4)) + audio_tokens = torch.tensor([[1, 2, 3, 4]]) + + _, output_embeds, _ = model.tts_preprocess( + input_ids=input_ids, + input_embeds=input_embeds, + codes={"audio": audio_tokens}, + ) + + assert len(model.model.embed_multimodal_calls) == 1 + torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens) + torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0)) + + +def test_tts_preprocess_keeps_legacy_top_level_audio_feedback(): + model = _make_voxtral_tts_model() + input_ids = torch.tensor([model._audio_token_id]) + input_embeds = torch.zeros((1, 4)) + audio_tokens = torch.tensor([[5, 6, 7, 8]]) + + _, output_embeds, _ = model.tts_preprocess( + input_ids=input_ids, + input_embeds=input_embeds, + audio=audio_tokens, + ) + + assert len(model.model.embed_multimodal_calls) == 1 + torch.testing.assert_close(model.model.embed_multimodal_calls[0]["audio_tokens"], audio_tokens) + torch.testing.assert_close(output_embeds, torch.full((1, 4), 7.0)) diff --git a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py new file mode 100644 index 00000000000..0e9f23de60c --- /dev/null +++ b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the NVFP4 W4A4 weight_scale NaN clamp installed by +``vllm_omni.patch``. + +ModelOpt 0.44's FP32->FP8 E4M3 cast of per-block weight scales can emit +literal NaN bytes (E4M3 encoding 0x7F / 0xFF) when the pre-cast scale +rounds above the FP8 max of 448 after the global-scale division. A +single NaN byte in any per-block weight_scale collapses the FlashInfer +FP4 GEMM output to NaN and the served model emits `!!!!`. The patch +overrides ``ModelOptNvFp4LinearMethod.process_weights_after_loading`` +to scan + clamp those bytes at load time. + +These tests cover the behavior of the override on a synthetic layer: + +- the NaN bytes get clamped to the FP8 E4M3 max byte (0x7E) +- a finite weight_scale is left untouched (no spurious writes) +- the install asserts itself so silent fallback is impossible +- the env-var escape hatch path is reachable + +No GPU is required — the tests construct an in-memory FP8 ``weight_scale`` +tensor with planted NaN bytes and call the patched PWAL directly. +""" + +import pytest +import torch + +# Importing vllm_omni runs vllm_omni.patch, which installs the override +# we are testing. The import must come before any direct reference to +# ``ModelOptNvFp4LinearMethod`` below. +import vllm_omni # noqa: F401 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +FP8_NAN_BYTES = (0x7F, 0xFF) +FP8_MAX_BYTE = 0x7E # encoding of finfo(float8_e4m3fn).max + + +def _fp8_tensor_with_nan_bytes(num_blocks: int, nan_indices: list[int]) -> torch.Tensor: + """Return an FP8 E4M3 tensor where every position is the FP8 max byte + except ``nan_indices``, which carry NaN bytes (alternating 0x7F/0xFF).""" + raw = bytearray([FP8_MAX_BYTE] * num_blocks) + for i, idx in enumerate(nan_indices): + raw[idx] = FP8_NAN_BYTES[i % len(FP8_NAN_BYTES)] + # Materialize on a writable storage so masked_fill_ on the uint8 view + # can write back in-place. ``frombuffer`` produces a read-only view + # that PyTorch warns about on first use. + return torch.tensor(list(raw), dtype=torch.uint8).view(torch.float8_e4m3fn) + + +@pytest.fixture +def clamp_fn(): + """Return the top-level NaN-clamp helper exposed by vllm_omni.patch. + + Skipped when vllm doesn't expose ``ModelOptNvFp4LinearMethod`` (older + pin / unsupported build) and vllm_omni.patch consequently didn't + install the helper either. + """ + try: + from vllm_omni.patch import _clamp_nvfp4_weight_scale_nans + except ImportError: + pytest.skip("vllm_omni.patch._clamp_nvfp4_weight_scale_nans not exposed") + return _clamp_nvfp4_weight_scale_nans + + +def test_clamps_nan_bytes_in_weight_scale(clamp_fn): + weight_scale = _fp8_tensor_with_nan_bytes(num_blocks=8, nan_indices=[2, 5]) + assert torch.isnan(weight_scale).sum().item() == 2 + + layer = torch.nn.Module() + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + n_clamped = clamp_fn(layer) + + assert n_clamped == 2 + assert torch.isnan(layer.weight_scale).sum().item() == 0, ( + "NaN bytes in weight_scale should have been clamped to FP8 E4M3 max" + ) + # Every byte should now be the FP8 max byte. + bytes_after = layer.weight_scale.data.view(torch.uint8).tolist() + assert bytes_after == [FP8_MAX_BYTE] * 8 + + +def test_finite_weight_scale_untouched(clamp_fn): + # All-finite weight_scale should leave bytes exactly as written. + raw = bytes([0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x7E]) + weight_scale = torch.frombuffer(raw, dtype=torch.uint8).clone().view(torch.float8_e4m3fn) + assert torch.isnan(weight_scale).sum().item() == 0 + before = weight_scale.data.view(torch.uint8).clone() + + layer = torch.nn.Module() + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + n_clamped = clamp_fn(layer) + + assert n_clamped == 0 + after = layer.weight_scale.data.view(torch.uint8) + assert torch.equal(before, after), "clamp must not modify a finite weight_scale tensor" + + +def test_clamp_is_noop_when_weight_scale_missing(clamp_fn): + # Defensive guard: subclasses without weight_scale should not crash. + layer = torch.nn.Module() + n_clamped = clamp_fn(layer) + assert n_clamped == 0 + + +def test_clamps_non_contiguous_weight_scale(clamp_fn): + # A non-contiguous weight_scale would make `.view(torch.uint8)` raise. + # The clamp must materialize a contiguous copy, write it back to the + # Parameter, and still land the in-place clamp where the kernel reads. + # NaN bytes at even indices so they survive the stride-2 slice below + # (base[::2] keeps original indices 0, 2, 4, ...). + base = _fp8_tensor_with_nan_bytes(num_blocks=16, nan_indices=[4, 12]) + # Slice every other element -> stride 2 -> non-contiguous view. + non_contig = base[::2] + assert not non_contig.is_contiguous() + assert torch.isnan(non_contig).sum().item() >= 1 + + layer = torch.nn.Module() + layer.weight_scale = torch.nn.Parameter(non_contig, requires_grad=False) + n_clamped = clamp_fn(layer) + + assert n_clamped >= 1 + assert layer.weight_scale.is_contiguous(), "weight_scale should be contiguous after clamp" + assert torch.isnan(layer.weight_scale).sum().item() == 0, ( + "NaN bytes in a non-contiguous weight_scale should have been clamped" + ) + + +def test_install_is_self_asserting(): + # If the install assert is removed, a future regression could let the + # override silently drop. There are three legitimate end states: + # (a) upstream vLLM already has its own NaN clamp and vllm_omni.patch + # intentionally skipped installing the override (self-extinguish), + # (b) upstream is unpatched and our wrapper is installed, + # (c) the env-var escape hatch is set OR the modelopt import failed, + # so nothing is installed by design — nothing to assert. + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod, + ) + except ImportError: + pytest.skip("vllm.modelopt.ModelOptNvFp4LinearMethod not available") + + from vllm_omni.patch import _already_patched_upstream, _clamp_installed + + current = ModelOptNvFp4LinearMethod.process_weights_after_loading + name = getattr(current, "__qualname__", "") or getattr(current, "__name__", "") + if _already_patched_upstream: + # Self-extinguish path: presence of our marker here would mean the + # heuristic in vllm_omni.patch failed to detect upstream's clamp. + assert "_patched_nvfp4_pwal" not in name, ( + f"vllm_omni override installed despite upstream PWAL already containing a " + f"NaN clamp ({name!r}); self-extinguish heuristic in vllm_omni.patch needs review" + ) + elif _clamp_installed: + assert "patched_nvfp4_pwal" in name or "_patched_nvfp4_pwal" in name, ( + f"NVFP4 NaN-clamp PWAL override does not appear installed (saw {name!r}); " + "another plugin may have re-patched ModelOptNvFp4LinearMethod" + ) + else: + pytest.skip("NaN-clamp not installed by design (env-var escape hatch or import failure)") + + +def test_reload_does_not_nest_wrapper(): + """Reloading vllm_omni.patch must not wrap our own wrapper a second time. + + On reload the class attribute already holds our wrapper; capturing it as + the "original" and re-wrapping would run the clamp twice per load. The + patch recovers the genuine upstream method via a sentinel, so after a + reload the installed wrapper's ``_vllm_omni_wrapped_pwal`` must point at + the real upstream method (which itself carries no sentinel), not at + another wrapper. Runs in a subprocess: importlib.reload mutates global + interpreter state we don't want leaking into other tests.""" + import subprocess + import sys + import textwrap + + code = textwrap.dedent( + """ + import importlib + import vllm_omni.patch as p + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod, + ) + except ImportError: + print("SKIP: modelopt not available") + raise SystemExit(0) + if not getattr(p, "_clamp_installed", False): + print("SKIP: clamp not installed (upstream patched or escape hatch)") + raise SystemExit(0) + importlib.reload(p) + fn = ModelOptNvFp4LinearMethod.process_weights_after_loading + wrapped = getattr(fn, "_vllm_omni_wrapped_pwal", None) + # The recovered upstream must NOT itself be one of our wrappers. + nested = hasattr(wrapped, "_vllm_omni_wrapped_pwal") + print("NESTED" if nested else "FLAT") + """ + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120) + if result.stdout.strip().startswith("SKIP:"): + pytest.skip(result.stdout.strip()) + assert "FLAT" in result.stdout, ( + f"reload nested the NVFP4 clamp wrapper — clamp would run twice per load.\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + + +@pytest.mark.parametrize( + ("env_value", "acceptable"), + [ + # Truthy values activate the escape hatch: the import short-circuits + # before the upstream-detection runs, so nothing is installed. + ("1", frozenset({"NOT_INSTALLED"})), + ("true", frozenset({"NOT_INSTALLED"})), + # Falsy values mean "do NOT use the escape hatch" — the repo-wide + # bool-env idiom in vllm_omni.patch guards against a naive + # `if os.environ.get(...)` wrongly skipping on "0"/"false". The clamp + # must therefore be in force, which is satisfied EITHER by our wrapper + # being installed OR by upstream vLLM already carrying its own clamp + # (self-extinguish, `_already_patched_upstream`), in which case skipping + # our wrapper is the correct behavior. Accept both so this test keeps + # passing once the upstream fix lands. + ("0", frozenset({"INSTALLED", "UPSTREAM_PATCHED"})), + ("false", frozenset({"INSTALLED", "UPSTREAM_PATCHED"})), + ], +) +def test_env_var_escape_hatch_disables_install(env_value, acceptable): + """Spawn a subprocess with VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP set and verify + the override install state matches the bool-env convention. The patch + module runs at ``vllm_omni`` import time, so we can't toggle it via + monkeypatch in the same process. + + The subprocess reports the authoritative state from vllm_omni.patch's + module flags (the same source of truth as ``test_install_is_self_asserting``) + rather than introspecting the method qualname — both because that is the + real contract and because ``"INSTALLED"`` is a substring of + ``"NOT_INSTALLED"`` and naive substring matching would be vacuous.""" + import os + import subprocess + import sys + import textwrap + + code = textwrap.dedent( + """ + import vllm_omni # noqa: F401 + import vllm_omni.patch as p + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod, # noqa: F401 + ) + except ImportError: + print("SKIP: modelopt not available") + raise SystemExit(0) + if getattr(p, "_already_patched_upstream", False): + # Upstream vLLM already clamps; vllm_omni.patch deliberately + # skipped installing our wrapper (self-extinguish). + print("UPSTREAM_PATCHED") + elif getattr(p, "_clamp_installed", False): + print("INSTALLED") + else: + print("NOT_INSTALLED") + """ + ) + env = os.environ.copy() + env["VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP"] = env_value + # 120s rather than 60s: cold CI containers can take 30-50s just to + # import vllm_omni → vllm → torch before the subprocess does any work. + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env) + if result.stdout.strip().startswith("SKIP:"): + pytest.skip("modelopt not available in subprocess") + # Exact token match on the final stdout line — not a substring test, since + # "INSTALLED" is a substring of "NOT_INSTALLED". + state = result.stdout.strip().splitlines()[-1] if result.stdout.strip() else "" + assert state in acceptable, ( + f"VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP={env_value!r} expected one of {set(acceptable)!r} " + f"but got {state!r}:\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) diff --git a/tests/model_executor/stage_input_processors/test_aura_omni.py b/tests/model_executor/stage_input_processors/test_aura_omni.py new file mode 100644 index 00000000000..fa7ad2e8d4d --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_aura_omni.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + PRECOMPUTED_TEXT_IDS_KEY, +) +from vllm_omni.model_executor.stage_input_processors.aura_omni import ( + SILENT_TEXT, + asr2aura, + aura2tts, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _source_output(text: str, request_id: str = "req-1", token_ids: list[int] | None = None): + output = SimpleNamespace(text=text, cumulative_token_ids=token_ids or [1, 2, 3], multimodal_output={}) + return SimpleNamespace(request_id=request_id, outputs=[output]) + + +def _source_delta_final_output(cumulative_text: str, request_id: str = "req-1"): + output = SimpleNamespace( + text="", + cumulative_text=cumulative_text, + cumulative_token_ids=[1, 2, 3], + multimodal_output={}, + ) + return SimpleNamespace(request_id=request_id, outputs=[output]) + + +def test_asr2aura_carries_video_payload_and_transcript(): + prompt = { + "multi_modal_data": {"video": ["frame-0", "frame-1"]}, + "additional_information": {"aura_system_prompt": ["system"]}, + } + + [next_input] = asr2aura([_source_output("What is happening now?")], prompt=[prompt]) + + assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]} + assert "<|video_pad|>" in next_input["prompt"] + assert "What is happening now?" in next_input["prompt"] + assert next_input["prompt"].startswith("<|im_start|>system\nsystem") + + +def test_asr2aura_drops_audio_before_qwen3_vl_stage(): + prompt = { + "multi_modal_data": { + "audio": ("wave", 16000), + "video": ["frame-0", "frame-1"], + }, + } + + [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt]) + + assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]} + assert "<|video_pad|>" in next_input["prompt"] + + +def test_asr2aura_reads_video_stashed_for_downstream_stage(): + prompt = { + "multi_modal_data": {"audio": ("wave", 16000)}, + "additional_information": { + "deferred_multi_modal_data": {"video": ["frame-0", "frame-1"]}, + }, + } + + [next_input] = asr2aura([_source_output("Check the video")], prompt=[prompt]) + + assert next_input["multi_modal_data"] == {"video": ["frame-0", "frame-1"]} + assert "<|video_pad|>" in next_input["prompt"] + + +def test_asr2aura_supports_video_only_observation(): + prompt = {"multi_modal_data": {"video": ["frame-0", "frame-1"]}} + + [next_input] = asr2aura([_source_output("")], prompt=[prompt]) + + assert "<|video_pad|>" in next_input["prompt"] + assert "<|im_start|>assistant" in next_input["prompt"] + + +def test_aura2tts_builds_qwen3_tts_prompt_information(): + prompt = { + "additional_information": { + "tts_language": ["Chinese"], + "tts_instruct": ["Calm voice."], + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert len(tts_input["prompt_token_ids"]) > 0 + assert tts_input["additional_information"]["text"] == ["Hello."] + assert PRECOMPUTED_TEXT_IDS_KEY not in tts_input["additional_information"] + assert tts_input["additional_information"]["task_type"] == ["Base"] + assert tts_input["additional_information"]["language"] == ["Chinese"] + assert tts_input["additional_information"]["ref_audio"] == ["ref.wav"] + assert tts_input["additional_information"]["ref_text"] == ["Reference transcript sample."] + assert tts_input["additional_information"]["x_vector_only_mode"] == [False] + assert tts_input["additional_information"]["instruct"] == ["Calm voice."] + + +def test_aura2tts_prefers_streaming_cumulative_text(): + prompt = { + "additional_information": { + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + } + } + + [tts_input] = aura2tts( + [_source_delta_final_output("The complete AURA reply.")], + prompt=[prompt], + ) + + assert tts_input["additional_information"]["text"] == ["The complete AURA reply."] + + +def test_aura2tts_supports_base_ref_audio_override(): + prompt = { + "additional_information": { + "tts_ref_audio": ["custom.wav"], + "tts_ref_text": ["custom transcript"], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert tts_input["additional_information"]["task_type"] == ["Base"] + assert tts_input["additional_information"]["ref_audio"] == ["custom.wav"] + assert tts_input["additional_information"]["ref_text"] == ["custom transcript"] + assert tts_input["additional_information"]["x_vector_only_mode"] == [False] + + +def test_aura2tts_supports_x_vector_only_mode_for_base(): + prompt = { + "additional_information": { + "tts_task_type": ["Base"], + "tts_x_vector_only_mode": [True], + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert tts_input["additional_information"]["x_vector_only_mode"] == [True] + + +def test_aura2tts_supports_custom_voice_mode(): + prompt = { + "additional_information": { + "tts_task_type": ["CustomVoice"], + "tts_speaker": ["vivian"], + } + } + + [tts_input] = aura2tts([_source_output("Hello.")], prompt=[prompt]) + + assert tts_input["additional_information"]["task_type"] == ["CustomVoice"] + assert tts_input["additional_information"]["speaker"] == ["Vivian"] + assert "ref_audio" not in tts_input["additional_information"] + assert len(tts_input["prompt_token_ids"]) == 14 + + +def test_aura2tts_passes_token_ids_to_qwen3_tts_when_enabled(): + prompt = { + "additional_information": { + "tts_ref_audio": ["ref.wav"], + "tts_ref_text": ["Reference transcript sample."], + "tts_pass_token_ids": [True], + } + } + + [tts_input] = aura2tts( + [ + _source_output( + "Hello.", + token_ids=[151644, 77091, 198, 108386, 1773, 151645, 198], + ) + ], + prompt=[prompt], + ) + + assert tts_input["additional_information"][PRECOMPUTED_TEXT_IDS_KEY] == [ + [151644, 77091, 198, 108386, 1773, 151645, 198, 151644, 77091, 198] + ] + assert "text" not in tts_input["additional_information"] + + +def test_aura2tts_drops_silent_response(): + assert aura2tts([_source_output(SILENT_TEXT)]) == [] diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py index c46c49d97a6..5309c7675e1 100644 --- a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py +++ b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py @@ -8,7 +8,6 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import ( talker2code2wav_async_chunk, - text2flow, text2flow_full_payload, text2flow_token_only, ) @@ -49,13 +48,13 @@ def _transfer_manager( ) -def test_text2flow_supports_batched_source_outputs(): +def test_text2flow_token_only_supports_batched_source_outputs(): source_outputs = [ _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}), _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[6, 7]])}), ] - outputs = text2flow(source_outputs=source_outputs, prompt=None) + outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None) assert len(outputs) == 2 assert outputs[0]["prompt_token_ids"] == [1, 2, 3] @@ -64,16 +63,6 @@ def test_text2flow_supports_batched_source_outputs(): assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21] -def test_text2flow_strips_reference_speech_prefix_from_cumulative_ids(): - source_outputs = [ - _source_output("req-0", [10, 11], [8, 9, 1, 2, 3], {"speech_token": torch.tensor([[8, 9]])}), - ] - - outputs = text2flow(source_outputs=source_outputs, prompt=None) - - assert outputs[0]["prompt_token_ids"] == [1, 2, 3] - - def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids(): source_outputs = [ _source_output( @@ -91,7 +80,11 @@ def test_text2flow_token_only_strips_reference_speech_prefix_from_cumulative_ids assert outputs[0]["additional_information"]["ids"]["prompt"] == [10, 11] -def test_text2flow_token_only_marks_prompt_trim_for_stop_token_completion(): +def test_text2flow_token_only_does_not_mark_prompt_trim(): + # The talker prompt is wrapped with the CosyVoice3 instruction template in + # _build_cosyvoice3_prompt, so the talker emits target-only speech and no + # prompt-trim offset is required; the flow stage trims prompt_feat itself + # (issue #4644). Confirm no talker_prefill_offset is set. source_outputs = [ _source_output( "req-stop", @@ -104,7 +97,8 @@ def test_text2flow_token_only_marks_prompt_trim_for_stop_token_completion(): outputs = text2flow_token_only(source_outputs=source_outputs, prompt=None) assert outputs[0]["prompt_token_ids"] == [1, 2, 6562] - assert outputs[0]["additional_information"]["meta"]["talker_prefill_offset"] == 2 + meta = outputs[0]["additional_information"].get("meta") or {} + assert "talker_prefill_offset" not in meta def test_text2flow_full_payload_does_not_send_codec_ids(): diff --git a/tests/model_executor/stage_input_processors/test_glm_image.py b/tests/model_executor/stage_input_processors/test_glm_image.py index 029ba99c29e..8dc333fb6f2 100644 --- a/tests/model_executor/stage_input_processors/test_glm_image.py +++ b/tests/model_executor/stage_input_processors/test_glm_image.py @@ -284,12 +284,11 @@ def test_basic_t2i(self): prompt = {"prompt": "a cat", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}} - result = ar2diffusion(source_outputs, prompt=[prompt]) - assert len(result) == 1 - assert result[0]["prompt"] == "a cat" - assert result[0]["height"] == 1024 - assert result[0]["width"] == 1024 - assert "prior_token_ids" in result[0]["extra"] + result = ar2diffusion(source_outputs, prompt=prompt) + assert result["prompt"] == "a cat" + assert result["height"] == 1024 + assert result["width"] == 1024 + assert "prior_token_ids" in result["extra"] def test_i2i_with_mm_output(self): """Test image-to-image with prior_token_image_ids from AR model.""" @@ -306,9 +305,8 @@ def test_i2i_with_mm_output(self): "multi_modal_data": {"image": img}, } - result = ar2diffusion(source_outputs, prompt=[prompt]) - assert len(result) == 1 - assert result[0]["extra"]["prior_token_image_ids"] is not None + result = ar2diffusion(source_outputs, prompt=prompt) + assert result["extra"]["prior_token_image_ids"] is not None def test_i2i_detected_via_modalities(self): """Test i2i mode detected via modalities field.""" @@ -321,11 +319,11 @@ def test_i2i_detected_via_modalities(self): "modalities": ["img2img"], } - result = ar2diffusion(source_outputs, prompt=[prompt]) - assert len(result) == 1 + result = ar2diffusion(source_outputs, prompt=prompt) + assert result["prompt"] == "edit this" - def test_empty_source_outputs_returns_empty_list(self): - assert ar2diffusion([], prompt={}) == [] + def test_empty_source_outputs_returns_none(self): + assert ar2diffusion([], prompt=None) is None def test_default_dimensions(self): """When no height/width in prompt, defaults to 1024x1024.""" @@ -333,9 +331,9 @@ def test_default_dimensions(self): source_outputs = [_source_output(token_ids)] prompt = {"prompt": "test"} - result = ar2diffusion(source_outputs, prompt=[prompt]) - assert result[0]["height"] == 1024 - assert result[0]["width"] == 1024 + result = ar2diffusion(source_outputs, prompt=prompt) + assert result["height"] == 1024 + assert result["width"] == 1024 def test_requires_multimodal_data_with_pil_image(self): """Test that pil_image is included when requires_multimodal_data=True.""" @@ -350,8 +348,8 @@ def test_requires_multimodal_data_with_pil_image(self): "multi_modal_data": {"image": img}, } - result = ar2diffusion(source_outputs, prompt=[prompt], requires_multimodal_data=True) - assert result[0]["pil_image"] is img + result = ar2diffusion(source_outputs, prompt=prompt, requires_multimodal_data=True) + assert result["pil_image"] is img def test_extra_params_passed_through(self): """Test that seed, num_inference_steps, guidance_scale, negative_prompt are passed.""" @@ -366,24 +364,22 @@ def test_extra_params_passed_through(self): "negative_prompt": "blurry", } - result = ar2diffusion(source_outputs, prompt=[prompt]) - assert result[0]["seed"] == 42 - assert result[0]["num_inference_steps"] == 50 - assert result[0]["guidance_scale"] == 7.5 - assert result[0]["negative_prompt"] == "blurry" + result = ar2diffusion(source_outputs, prompt=prompt) + assert result["prompt"] == "test" + assert result["seed"] == 42 + assert result["num_inference_steps"] == 50 + assert result["guidance_scale"] == 7.5 + assert result["negative_prompt"] == "blurry" - def test_batch_requests(self): - """Test processing multiple requests in a batch.""" + def test_multiple_source_outputs_uses_first_payload_only(self): + """Test the GLM bridge keeps a single diffusion payload for one request.""" tokens1 = list(range(256)) + list(range(1024)) + [16385] - tokens2 = list(range(256)) + list(range(1024)) + [16385] + tokens2 = [1, 2, 3] source_outputs = [_source_output(tokens1), _source_output(tokens2)] - prompts = [ - {"prompt": "first", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}}, - {"prompt": "second", "mm_processor_kwargs": {"target_h": 512, "target_w": 512}}, - ] + prompt = {"prompt": "first", "mm_processor_kwargs": {"target_h": 1024, "target_w": 1024}} - result = ar2diffusion(source_outputs, prompt=prompts) - assert len(result) == 2 - assert result[0]["prompt"] == "first" - assert result[1]["prompt"] == "second" + result = ar2diffusion(source_outputs, prompt=prompt) + assert result["prompt"] == "first" + assert result["height"] == 1024 + assert result["width"] == 1024 diff --git a/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py b/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py index 76f3e500622..19405f28fa7 100644 --- a/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py +++ b/tests/model_executor/stage_input_processors/test_hunyuan_image3_bridge.py @@ -45,6 +45,33 @@ def test_truncate_at_cot_end_strips_tail_after_recaption_marker(): assert text == "body text</recaption>" +def test_ar2diffusion_returns_one_request_payload_for_request_level_batching(): + result = ar2diffusion( + [_source_output([100], text="thought")], + prompt={"prompt": "edit"}, + ) + + assert isinstance(result, dict) + assert result["prompt"] == "edit" + + +def test_ar2diffusion_uses_parent_output_when_companions_are_present(): + result = ar2diffusion( + [ + _source_output([100], text="parent thought"), + _source_output([200], text="companion thought"), + ], + prompt={"prompt": "edit"}, + ) + + assert result is not None + assert result["extra"]["ar_generated_text"] == "parent thought" + + +def test_ar2diffusion_returns_none_without_parent_output(): + assert ar2diffusion([], prompt={"prompt": "edit"}) is None + + def test_ar2diffusion_applies_ratio_and_truncates_tail_without_tokenizer(monkeypatch: pytest.MonkeyPatch): real_import = builtins.__import__ @@ -67,10 +94,10 @@ def _block_transformers_import(name, *args, **kwargs): prompt=[{"prompt": "edit", "height": 64, "width": 64}], ) - assert len(result) == 1 - assert (result[0]["height"], result[0]["width"]) == (512, 2048) - assert result[0]["extra"]["ar_generated_text"] == "decoded without special tokens" - assert "ar_token_ids" not in result[0]["extra"] + assert result is not None + assert (result["height"], result["width"]) == (512, 2048) + assert result["extra"]["ar_generated_text"] == "decoded without special tokens" + assert "ar_token_ids" not in result["extra"] def test_ar2diffusion_forwards_custom_system_prompt_body(): @@ -88,5 +115,6 @@ def test_ar2diffusion_forwards_custom_system_prompt_body(): ], ) - assert result[0]["use_system_prompt"] == "custom" - assert result[0]["system_prompt"] == marker + assert result is not None + assert result["use_system_prompt"] == "custom" + assert result["system_prompt"] == marker diff --git a/tests/model_executor/stage_input_processors/test_indextts2.py b/tests/model_executor/stage_input_processors/test_indextts2.py new file mode 100644 index 00000000000..11468603e1b --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_indextts2.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm_omni.model_executor.stage_input_processors.indextts2 import ( + _strip_stop_token, + talker2s2mel_full_payload, + talker2s2mel_token_only, +) + +STOP_MEL_TOKEN = 8193 +LATENT_DIM = 16 + + +def _make_talker_output( + *, + mel_codes: torch.Tensor, + latent: torch.Tensor, + s_ref: torch.Tensor | None = None, + ref_mel: torch.Tensor | None = None, + style: torch.Tensor | None = None, + finished: bool = True, +): + meta = {} + if s_ref is not None: + meta["S_ref"] = s_ref + if ref_mel is not None: + meta["ref_mel"] = ref_mel + if style is not None: + meta["style"] = style + + mm = { + "codes": {"mel": mel_codes}, + "hidden_states": {"latent": latent}, + "meta": meta, + } + return SimpleNamespace( + finished=finished, + outputs=[SimpleNamespace(multimodal_output=mm)], + ) + + +# ── _strip_stop_token ── + + +def test_strip_stop_token_basic(): + codes = torch.tensor([10, 20, 30, STOP_MEL_TOKEN, 99]) + latent = torch.randn(5, LATENT_DIM) + + c, lat, lens = _strip_stop_token(codes, latent) + + assert c.shape == (1, 3) + assert lat.shape == (1, 3, LATENT_DIM) + assert lens.tolist() == [3] + assert (c[0] == torch.tensor([10, 20, 30])).all() + + +def test_strip_stop_token_no_stop(): + codes = torch.tensor([10, 20, 30]) + latent = torch.randn(3, LATENT_DIM) + + c, lat, lens = _strip_stop_token(codes, latent) + + assert c.shape == (1, 3) + assert lens.tolist() == [3] + + +def test_strip_stop_token_at_start(): + codes = torch.tensor([STOP_MEL_TOKEN, 10, 20]) + latent = torch.randn(3, LATENT_DIM) + + c, lat, lens = _strip_stop_token(codes, latent) + + assert c.shape == (1, 0) + assert lat.shape == (1, 0, LATENT_DIM) + assert lens.tolist() == [0] + + +def test_strip_stop_token_batch_with_padding(): + codes = torch.tensor( + [ + [10, 20, STOP_MEL_TOKEN, 0], + [30, 40, 50, STOP_MEL_TOKEN], + ] + ) + latent = torch.randn(2, 4, LATENT_DIM) + + c, lat, lens = _strip_stop_token(codes, latent) + + assert c.shape == (2, 3) + assert lat.shape == (2, 3, LATENT_DIM) + assert lens.tolist() == [2, 3] + assert c[0, 2].item() == 0 # padded with a valid semantic code + + +def test_strip_stop_token_all_stop(): + codes = torch.tensor([STOP_MEL_TOKEN]) + latent = torch.randn(1, LATENT_DIM) + + c, lat, lens = _strip_stop_token(codes, latent) + + assert lens.tolist() == [0] + assert c.shape[1] == 0 + assert lat.shape[1] == 0 + + +# ── full-payload connector path ── + + +def _conditioning_tensors(): + return { + "S_ref": torch.randn(1, 3, LATENT_DIM), + "ref_mel": torch.randn(1, 80, 5), + "style": torch.randn(1, 4), + } + + +def test_talker2s2mel_token_only_placeholder_finished_only(): + finished = _make_talker_output(mel_codes=torch.tensor([1]), latent=torch.randn(1, LATENT_DIM), finished=True) + unfinished = _make_talker_output(mel_codes=torch.tensor([2]), latent=torch.randn(1, LATENT_DIM), finished=False) + + outputs = talker2s2mel_token_only([unfinished, finished]) + + assert len(outputs) == 1 + assert outputs[0]["prompt_token_ids"] == [0] + assert outputs[0]["additional_information"] is None + assert outputs[0]["multi_modal_data"] is None + + +def test_talker2s2mel_full_payload_flat_keys_builds_s2mel_contract(): + cond = _conditioning_tensors() + latent = torch.randn(3, LATENT_DIM) + payload = { + "codes.mel": torch.tensor([[10], [20], [30]]), + "hidden_states.latent": latent, + "meta.S_ref": cond["S_ref"], + "meta.ref_mel": cond["ref_mel"], + "meta.style": cond["style"], + } + + result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-flat")) + + assert result is not None + assert result["mel_codes"].shape == (1, 3) + assert result["mel_codes"].tolist() == [[10, 20, 30]] + assert result["latent"].shape == (1, 3, LATENT_DIM) + assert result["code_lens"].tolist() == [3] + assert result["S_ref"].device.type == "cpu" + assert result["ref_mel"].device.type == "cpu" + assert result["style"].device.type == "cpu" + assert result["S_ref"].data_ptr() != cond["S_ref"].data_ptr() + + +def test_talker2s2mel_full_payload_nested_fallback_input(): + cond = _conditioning_tensors() + payload = { + "codes": {"mel": torch.tensor([[4], [5]])}, + "hidden_states": {"latent": torch.randn(2, LATENT_DIM)}, + "meta": cond, + } + + result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-nested")) + + assert result is not None + assert result["mel_codes"].tolist() == [[4, 5]] + assert result["latent"].shape == (1, 2, LATENT_DIM) + + +def test_talker2s2mel_full_payload_normalizes_one_by_t_mel_row(): + cond = _conditioning_tensors() + payload = { + "codes.mel": torch.tensor([[7, 8, 9]]), + "hidden_states.latent": torch.randn(3, LATENT_DIM), + "meta": cond, + } + + result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-row")) + + assert result is not None + assert result["mel_codes"].tolist() == [[7, 8, 9]] + + +def test_talker2s2mel_full_payload_trims_stop_token_and_matching_latent(): + cond = _conditioning_tensors() + latent = torch.arange(4 * LATENT_DIM, dtype=torch.float32).reshape(4, LATENT_DIM) + payload = { + "codes.mel": torch.tensor([[1], [2], [STOP_MEL_TOKEN], [99]]), + "hidden_states.latent": latent, + "meta": cond, + } + + result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-stop")) + + assert result is not None + assert result["mel_codes"].tolist() == [[1, 2]] + assert result["latent"].shape == (1, 2, LATENT_DIM) + assert result["latent"][0, 0].tolist() == latent[0].tolist() + assert result["latent"][0, 1].tolist() == latent[1].tolist() + assert result["code_lens"].tolist() == [2] + + +def test_talker2s2mel_full_payload_crops_mel_latent_to_common_length(): + cond = _conditioning_tensors() + payload = { + "codes.mel": torch.tensor([[1], [2], [3], [4]]), + "hidden_states.latent": torch.randn(2, LATENT_DIM), + "meta": cond, + } + + result = talker2s2mel_full_payload(None, payload, SimpleNamespace(request_id="r-crop")) + + assert result is not None + assert result["mel_codes"].tolist() == [[1, 2]] + assert result["latent"].shape == (1, 2, LATENT_DIM) + assert result["code_lens"].tolist() == [2] + + +def test_talker2s2mel_full_payload_missing_required_fields_returns_none(): + assert talker2s2mel_full_payload(None, {"hidden_states.latent": torch.randn(1, LATENT_DIM)}, None) is None + assert talker2s2mel_full_payload(None, {"codes.mel": torch.tensor([[1]])}, None) is None diff --git a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py index d1cc0dbf7b6..ae1c8e2f960 100644 --- a/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py +++ b/tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py @@ -11,59 +11,40 @@ from vllm_omni.model_executor.stage_input_processors import mimo_audio as sip from vllm_omni.model_executor.stage_input_processors.mimo_audio import ( MAX_CODE2WAV_TOKENS, - llm2code2wav, + llm2code2wav_full_payload, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -def _make_source_outputs(codec_codes: torch.Tensor, request_id: str = "req-0"): - output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}}) - """Build minimal source_outputs with one talker output carrying codec codes.""" - talker_output = SimpleNamespace(outputs=[output], request_id=request_id) - return [talker_output] - - -def test_llm2code2wav_truncates_when_flat_exceeds_max(caplog): - """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated, not passed through.""" - # prepend_and_flatten_colmajor produces 36 ids per (8, 4) codec frame: - # pad adds one row -> (9, 4) per frame, permuted and flattened. - # Pick enough frames to comfortably exceed the cap. +def test_llm2code2wav_full_payload_truncates_when_flat_exceeds_max(caplog): + """Flat codec sequences longer than MAX_CODE2WAV_TOKENS must be truncated.""" frames = (MAX_CODE2WAV_TOKENS // 36) + 100 codec_codes = torch.ones(frames, 1, 8, 4, dtype=torch.long) + request = SimpleNamespace(request_id="req-long") - source_outputs = _make_source_outputs(codec_codes, request_id="req-long") - - # Attach caplog's handler directly to the module logger so the warning is - # captured regardless of propagation (vllm's logger configuration can - # interact badly with caplog.at_level's default root-handler path). target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.mimo_audio") target_logger.addHandler(caplog.handler) prev_level = target_logger.level target_logger.setLevel(logging.WARNING) try: - prompts = llm2code2wav(source_outputs) + payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request) finally: target_logger.removeHandler(caplog.handler) target_logger.setLevel(prev_level) - assert len(prompts) == 1 - assert len(prompts[0]["prompt_token_ids"]) == MAX_CODE2WAV_TOKENS - assert any("truncating" in rec.getMessage() for rec in caplog.records), ( - f"Expected a 'truncating' warning; captured records: {[r.getMessage() for r in caplog.records]}" - ) + assert payload is not None + assert len(payload["codes"]["audio"]) == MAX_CODE2WAV_TOKENS -def test_llm2code2wav_short_sequence_unchanged(): - """Short codec sequences are returned without truncation.""" +def test_llm2code2wav_full_payload_short_sequence_unchanged(): codec_codes = torch.ones(4, 1, 8, 4, dtype=torch.long) - source_outputs = _make_source_outputs(codec_codes, request_id="req-short") + request = SimpleNamespace(request_id="req-short") - prompts = llm2code2wav(source_outputs) + payload = llm2code2wav_full_payload(None, {"codes.audio": codec_codes}, request) - assert len(prompts) == 1 - # 4 frames + 1 pad row, flattened col-major → well below the cap - assert 0 < len(prompts[0]["prompt_token_ids"]) <= MAX_CODE2WAV_TOKENS + assert payload is not None + assert 0 < len(payload["codes"]["audio"]) <= MAX_CODE2WAV_TOKENS def test_llm2code2wav_truncation_boundary_constant_matches_yaml(): diff --git a/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py new file mode 100644 index 00000000000..bff3e31e389 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_ming_tts_async_chunk.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.ming_tts.constants import LATENT_DIM, PATCH_SIZE +from vllm_omni.model_executor.stage_input_processors.ming_tts import ( + MING_EMIT_PATCH_COUNT_KEY, + llm2audio_vae_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_transfer_manager(*, chunk_size: int = 5, initial_chunk_size: int = 2): + return SimpleNamespace( + code_prompt_token_ids=defaultdict(list), + put_req_chunk=defaultdict(int), + request_payload={}, + connector=SimpleNamespace( + config={ + "extra": { + "latent_chunk_size": chunk_size, + "initial_latent_chunk_size": initial_chunk_size, + "latent_left_context": 0, + } + } + ), + ) + + +def _make_request(req_id: str = "req", *, finished: bool = False): + return SimpleNamespace( + external_req_id=req_id, + is_finished=lambda: finished, + ) + + +def _make_output(value: float): + return { + "ming_has_patch": torch.tensor([1], dtype=torch.bool), + "ming_latent_patch": torch.full((1, PATCH_SIZE, LATENT_DIM), value), + } + + +def _append_patch(tm, req_id: str, idx: int, *, finished: bool = False): + return llm2audio_vae_async_chunk( + tm, + _make_output(float(idx)), + _make_request(req_id, finished=finished), + is_finished=finished, + ) + + +def test_ming_async_chunk_emits_small_initial_chunk_then_steady_chunks(): + tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2) + + assert _append_patch(tm, "req", 0) is None + + first = _append_patch(tm, "req", 1) + assert first is not None + assert first.latent.shape == (2, PATCH_SIZE, LATENT_DIM) + assert first.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2 + assert first.meta.finished.item() is False + + for idx in range(2, 6): + assert _append_patch(tm, "req", idx) is None + + second = _append_patch(tm, "req", 6) + assert second is not None + assert second.latent.shape == (5, PATCH_SIZE, LATENT_DIM) + assert second.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 5 + + +def test_ming_async_chunk_flushes_short_final_chunk(): + tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2) + + payload = _append_patch(tm, "req", 0, finished=True) + + assert payload is not None + assert payload.latent.shape == (1, PATCH_SIZE, LATENT_DIM) + assert payload.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 1 + assert payload.meta.finished.item() is True + + +def test_ming_async_chunk_flushes_leftover_after_initial_chunk(): + tm = _make_transfer_manager(chunk_size=5, initial_chunk_size=2) + + assert _append_patch(tm, "req", 0) is None + assert _append_patch(tm, "req", 1) is not None + assert _append_patch(tm, "req", 2) is None + + final = _append_patch(tm, "req", 3, finished=True) + + assert final is not None + assert final.latent.shape == (2, PATCH_SIZE, LATENT_DIM) + assert final.kv_metadata[MING_EMIT_PATCH_COUNT_KEY] == 2 + assert final.meta.finished.item() is True 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 52ca2990563..8a6dfbe94b4 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 @@ -163,7 +163,7 @@ def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None: assert payload is not None assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22] - assert payload["code_predictor_codes"] == payload["codes"]["audio"] + assert "code_predictor_codes" not in payload def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None: @@ -204,7 +204,7 @@ def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> N assert payload is not None assert payload["codes"]["audio"] == [0, 0, 0] - assert payload["code_predictor_codes"] == payload["codes"]["audio"] + assert "code_predictor_codes" not in payload def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None: @@ -224,7 +224,7 @@ def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None: assert payload is not None assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9] - assert payload["code_predictor_codes"] == payload["codes"]["audio"] + assert "code_predictor_codes" not in payload def test_thinker2talker_full_payload_packs_complete_tensors() -> None: @@ -247,11 +247,33 @@ def test_thinker2talker_full_payload_packs_complete_tensors() -> None: assert payload["ids"]["all"] == [151644, 872, 3] assert payload["embed"]["prefill"].device.type == "cpu" assert payload["hidden_states"]["output"].device.type == "cpu" - assert payload["next_stage_prompt_len"] > 0 assert payload["embed"]["prefill"].shape[0] == 2 assert payload["hidden_states"]["output"].shape[0] == 2 +def test_thinker2talker_token_only_preserves_voice_metadata() -> None: + source_outputs = [ + SimpleNamespace( + request_id="req-1", + prompt_token_ids=[1, 2], + outputs=[SimpleNamespace(cumulative_token_ids=[3])], + ) + ] + prompt = { + "additional_information": { + "speaker": ["ethan"], + "language": ["English"], + } + } + + [talker_prompt] = q3.thinker2talker_token_only(source_outputs, prompt) + + assert talker_prompt["additional_information"] == { + "speaker": ["ethan"], + "language": ["English"], + } + + def test_accumulator_replaces_keys_in_replace_set() -> None: """REPLACE-key semantics: subsequent emissions of the same key replace, not append.""" from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin @@ -921,7 +943,6 @@ def test_cosyvoice3_full_payload_replace_keys_present() -> None: def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None: """Smoke: ming_flash_omni token-only carries voice metadata.""" from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import ( - thinker2talker, thinker2talker_token_only, ) @@ -939,15 +960,14 @@ def __init__(self, info): src = [_Wrap("hello world")] prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"}) - for func in (thinker2talker, thinker2talker_token_only): - out = func(src, prompt=prompt) - assert len(out) == 1 - assert out[0]["prompt_token_ids"] == [0] - info = out[0]["additional_information"] - assert info["text"] == "hello world" - assert info["voice_name"] == "ZH_FEMALE" - assert info["prompt_text"] == "ref text" - assert info["ming_task"] == "omni" + out = thinker2talker_token_only(src, prompt=prompt) + assert len(out) == 1 + assert out[0]["prompt_token_ids"] == [0] + 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/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index da25ee7bd6c..bc6239e293d 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -12,9 +12,9 @@ max_ic_for_chunk_size, ) from vllm_omni.model_executor.stage_input_processors.qwen3_tts import ( - talker2code2wav, talker2code2wav_async_chunk, talker2code2wav_full_payload, + talker2code2wav_token_only, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -368,7 +368,15 @@ def test_ref_code_context_can_be_buffered_before_first_emit(): assert rid in tm.request_payload -def test_non_async_processor_prepends_ref_code_and_sets_trim_context(): +def test_non_async_token_only_sizes_placeholder_for_ref_and_audio_frames(): + """``talker2code2wav_token_only`` only allocates placeholder prompt slots. + + After the connector refactor, actual codec flattening (ref prepend + + codebook-major layout) is performed by ``talker2code2wav_full_payload`` on + the worker data plane. The orchestrator hook still derives + ``left_context_size`` from stage-0 multimodal_output so Code2Wav can trim + reference frames once the connector payload arrives. + """ ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) audio_codes = torch.tensor( [ @@ -387,12 +395,38 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context(): engine_outputs=[SimpleNamespace(outputs=[output], finished=True)], ) - prompts = talker2code2wav(stage.engine_outputs) + prompts = talker2code2wav_token_only(stage.engine_outputs) assert len(prompts) == 1 prompt = prompts[0] assert prompt["additional_information"] == {"meta": {"left_context_size": 2}} - assert prompt["prompt_token_ids"] == [ + # 2 ref frames + 2 valid audio frames (zero row filtered), 4 quantizers. + assert prompt["prompt_token_ids"] == [0] * (_Q * (2 + 2)) + + +def test_full_payload_prepends_ref_code_and_flattens_codebook_major(): + """Worker producer is authoritative for ref prepend + codec flatten.""" + ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long) + audio_codes = torch.tensor( + [ + [0, 0, 0, 0], + [1, 2, 3, 4], + [5, 6, 7, 8], + ], + dtype=torch.long, + ) + pooling_output = { + "codes.audio": audio_codes, + "codes.ref": ref_code, + "meta.ref_code_len": 2, + } + request = SimpleNamespace(request_id="r", output_token_ids=list(range(3))) + + payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request) + + assert payload is not None + assert payload["meta"]["left_context_size"] == 2 + assert payload["codes"]["audio"].tolist() == [ 9, 8, 1, @@ -433,7 +467,7 @@ def test_non_async_processor_filters_out_of_range_codec_values(): engine_outputs=[SimpleNamespace(outputs=[output], finished=True)], ) - prompts = talker2code2wav(stage.engine_outputs) + prompts = talker2code2wav_token_only(stage.engine_outputs) assert len(prompts) == 1 prompt = prompts[0] @@ -486,3 +520,32 @@ def test_full_payload_omits_left_context_size_without_ref(): assert payload["meta"]["finished"].item() is True assert "left_context_size" not in payload["meta"] assert len(payload["codes"]["audio"]) == _Q * 2 + + +@pytest.mark.parametrize( + "pooling_output", + [ + pytest.param(SimpleNamespace(codes="not-a-dict"), id="non_dict_output"), + pytest.param({}, id="missing_codes_audio"), + pytest.param({"codes.audio": torch.zeros((3, _Q), dtype=torch.long)}, id="all_codes_filtered"), + ], +) +def test_full_payload_emits_empty_finished_payload_on_degenerate_take(pooling_output): + """Regression for #4463. + + A degenerate talker take used to return ``None`` from + ``talker2code2wav_full_payload``. The connector treats ``None`` as "drop the + request", but Stage-1 was already scheduled to receive it, so its wait gate + polls to ``connector_get_max_wait`` (~300s) and the orchestrator aborts — one + stuck request stalls the whole two-stage pipeline. Each degenerate case + (non-dict pooling_output, missing ``codes.audio``, all codec frames dropped by + the filter) must instead return an empty-but-finished payload so the gate + releases and the request finishes immediately with zero-length audio. + """ + request = SimpleNamespace(request_id="r", output_token_ids=[0, 1, 2]) + + payload = talker2code2wav_full_payload(transfer_manager=None, pooling_output=pooling_output, request=request) + + assert payload is not None + assert payload["meta"]["finished"].item() is True + assert payload["codes"]["audio"].numel() == 0 diff --git a/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py b/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py index 186f008907a..9349c1cc664 100644 --- a/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_step_audio2_async_chunk.py @@ -5,7 +5,6 @@ from types import SimpleNamespace import pytest -import torch from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( DEFAULT_STREAM_CONFIG, @@ -43,14 +42,14 @@ def test_step_audio2_async_chunk_uses_decode_only_tokens_not_prompt_history(): payload = thinker2token2wav_async_chunk( transfer_manager=transfer_manager, - pooling_output=None, + multimodal_output=None, request=request, ) assert payload is not None - assert payload["code_predictor_codes"] == [1, 2] - assert payload["left_context_size"] == 1 - assert payload["finished"].item() is True + assert payload.codes.audio.tolist() == [1, 2] + assert payload.meta.left_context_size == 1 + assert payload.meta.finished.item() is True def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens(): @@ -68,7 +67,7 @@ def test_step_audio2_async_chunk_returns_none_when_not_enough_tokens(): payload = thinker2token2wav_async_chunk( transfer_manager=transfer_manager, - pooling_output=None, + multimodal_output=None, request=request, ) @@ -91,14 +90,14 @@ def test_step_audio2_async_chunk_emits_non_last_chunk_and_advances_consumed_by_c payload = thinker2token2wav_async_chunk( transfer_manager=transfer_manager, - pooling_output=None, + multimodal_output=None, request=request, ) assert payload is not None - assert payload["left_context_size"] == 0 - assert payload["finished"].item() is False - assert payload["code_predictor_codes"] == list(range(required)) + assert payload.meta.left_context_size == 0 + assert payload.meta.finished.item() is False + assert payload.codes.audio.tolist() == list(range(required)) assert transfer_manager.code_prompt_token_ids["rid-ready"] == list(range(chunk_size)) @@ -121,12 +120,11 @@ def test_step_audio2_async_chunk_emits_eof_when_finished_with_no_remaining_audio payload = thinker2token2wav_async_chunk( transfer_manager=transfer_manager, - pooling_output=None, + multimodal_output=None, request=request, ) - assert payload == { - "code_predictor_codes": [], - "left_context_size": 1, - "finished": torch.tensor(True, dtype=torch.bool), - } + assert payload is not None + assert payload.codes.audio.numel() == 0 + assert payload.meta.left_context_size == 1 + assert payload.meta.finished.item() is True diff --git a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py index 5235d7c68f7..dc248f21179 100644 --- a/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py @@ -8,61 +8,12 @@ import torch from vllm_omni.model_executor.stage_input_processors.voxtral_tts import ( - generator2tokenizer, generator2tokenizer_async_chunk, ) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -# ---- Helpers for generator2tokenizer (non-streaming) ---- - - -def _make_stage(engine_outputs): - return SimpleNamespace(engine_outputs=engine_outputs) - - -def _make_output(audio_tensors): - """Build a single generator output with multimodal_output["audio"] = audio_tensors.""" - return SimpleNamespace( - outputs=[SimpleNamespace(multimodal_output={"audio": audio_tensors})], - ) - - -# ---- Tests for generator2tokenizer ---- - - -class TestGenerator2Tokenizer: - def test_single_output_single_chunk(self): - """Single output with one audio tensor produces correct prompt_token_ids.""" - audio = torch.tensor([10, 20, 30, 40]) - stage = _make_stage([_make_output([audio])]) - - result = generator2tokenizer(stage.engine_outputs) - - assert len(result) == 1 - assert result[0]["prompt_token_ids"] == [10, 20, 30, 40] - assert result[0]["multi_modal_data"] is None - - def test_2d_audio_tensors_flattened(self): - """2D audio tensors (e.g., multi-codebook frames) are flattened.""" - audio = torch.tensor([[1, 2, 3], [4, 5, 6]]) # (2, 3) - stage = _make_stage([_make_output([audio])]) - - result = generator2tokenizer(stage.engine_outputs) - - assert result[0]["prompt_token_ids"] == [1, 2, 3, 4, 5, 6] - - def test_no_outputs_yet_raises(self): - """Stage with engine_outputs=None raises RuntimeError.""" - stage = _make_stage(engine_outputs=None) - with pytest.raises(TypeError): - generator2tokenizer(stage.engine_outputs) - - -# ---- Helpers for generator2tokenizer_async_chunk (streaming) ---- - - def _req(external_req_id: str, *, finished: bool): return SimpleNamespace( external_req_id=external_req_id, @@ -92,7 +43,7 @@ def test_empty_chunk_when_not_finished(): payload = generator2tokenizer_async_chunk( transfer_manager=transfer_manager, - multimodal_output={"audio": torch.zeros((0,))}, + multimodal_output={"codes": {"audio": torch.zeros((0,))}}, request=request, ) @@ -153,7 +104,7 @@ def test_normal_chunk_emission(): # Feed 25 frames one by one for i in range(25): - multimodal_output = {"audio": torch.tensor([float(i)] * 4)} + multimodal_output = {"codes": {"audio": torch.tensor([float(i)] * 4)}} payload = generator2tokenizer_async_chunk( transfer_manager=transfer_manager, multimodal_output=multimodal_output, @@ -181,7 +132,7 @@ def test_small_initial_chunks(): # Feed 5 frames (should trigger emission because codec_chunk_frames_at_begin=5) for i in range(5): - multimodal_output = {"audio": torch.tensor([float(i + 1)] * 3)} + multimodal_output = {"codes": {"audio": torch.tensor([float(i + 1)] * 3)}} payload = generator2tokenizer_async_chunk( transfer_manager=transfer_manager, multimodal_output=multimodal_output, @@ -208,7 +159,7 @@ def test_no_emission_between_boundaries(): # Feed 3 frames (3 % 5 != 0, should not emit) for i in range(3): - multimodal_output = {"audio": torch.tensor([float(i)] * 4)} + multimodal_output = {"codes": {"audio": torch.tensor([float(i)] * 4)}} payload = generator2tokenizer_async_chunk( transfer_manager=transfer_manager, multimodal_output=multimodal_output, @@ -230,7 +181,7 @@ def test_context_handling_format(): # Feed 5 frames to trigger a chunk (chunk_size_at_begin=5) for i in range(5): - multimodal_output = {"audio": torch.tensor([float(i + 10)] * 2)} # codebook_dim=2 + multimodal_output = {"codes": {"audio": torch.tensor([float(i + 10)] * 2)}} # codebook_dim=2 payload = generator2tokenizer_async_chunk( transfer_manager=transfer_manager, multimodal_output=multimodal_output, diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py new file mode 100644 index 00000000000..76ef84e7eba --- /dev/null +++ b/tests/model_extras/test_model_extras.py @@ -0,0 +1,518 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import pytest +from PIL import Image + +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_extras import ( + build_image_to_image_prompt, + build_image_to_video_prompt, + build_text_to_image_prompt, + get_extra_body_params, + get_extra_output_params, + should_init_extra_args_for_non_diffusion_stages, +) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("BagelPipeline") == frozenset( + { + "cfg_text_scale", + "cfg_img_scale", + "cfg_interval", + "cfg_renorm_type", + "cfg_renorm_min", + "negative_prompt", + "think", + "max_think_tokens", + "do_sample", + "text_temperature", + "timestep_shift", + } + ) + assert get_extra_output_params("BagelPipeline") == frozenset({"text_output", "think_text"}) + assert should_init_extra_args_for_non_diffusion_stages("BagelPipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_sensenova_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("SenseNovaU1Pipeline") == frozenset( + { + "think", + "cfg_scale", + "cfg_norm", + "timestep_shift", + "t_eps", + "img_cfg_scale", + "max_tokens", + } + ) + assert get_extra_output_params("SenseNovaU1Pipeline") == frozenset({"think_text"}) + assert should_init_extra_args_for_non_diffusion_stages("SenseNovaU1Pipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) +def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_name: str) -> None: + assert get_extra_body_params(pipeline_name) == frozenset( + { + "flow_shift", + "max_sequence_length", + "use_resolution_template", + "use_duration_template", + "use_system_prompt", + "system_prompt", + "negative_prompt", + "guardrails", + "condition_frame_indexes_vision", + "condition_video_keep", + "generate_sound", + "sound_gen", + "sound_duration", + "audio_duration", + "action_mode", + "action", + "domain_name", + "domain_id", + "raw_action_dim", + "action_chunk_size", + "action_space", + "action_fps", + "image_height", + "image_width", + "history_length", + "conditioning_fps", + "resolution", + "image_size", + "use_state", + "observation", + "robot_obs", + "deterministic_seed", + "session_id", + } + ) + assert get_extra_output_params(pipeline_name) == frozenset( + { + "action", + "raw_action_dim", + "domain_id", + "action_mode", + } + ) + assert should_init_extra_args_for_non_diffusion_stages(pipeline_name) is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_magi_human_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("MagiHumanPipeline") == frozenset( + { + "seconds", + "audio_path", + "image_path", + "sr_height", + "sr_width", + "sr_num_inference_steps", + } + ) + assert get_extra_output_params("MagiHumanPipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("MingImagePipeline") == frozenset( + { + "height", + "width", + "steps", + "cfg", + "seed", + "byte5_text", + "negative_prompt", + } + ) + assert get_extra_output_params("MingImagePipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("MingImagePipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_ming_flash_omni_declared_extra_args_route_into_sampling_params() -> None: + params = OmniDiffusionSamplingParams() + declared = get_extra_body_params("MingImagePipeline") + apply_declared_extra_args( + params, + declared, + { + "steps": 6, + "cfg": 1.5, + "byte5_text": ["理解与生成统一"], + "negative_prompt": "ugly, blurry", + "unknown": "ignored", + }, + ) + assert params.extra_args == { + "steps": 6, + "cfg": 1.5, + "byte5_text": ["理解与生成统一"], + "negative_prompt": "ugly, blurry", + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_ming_flash_omni_text_to_image_prompt_builder() -> None: + assert build_text_to_image_prompt( + "MingImagePipeline", + prompt="Please draw a cute cat.", + negative_prompt="ugly, blurry", + height=512, + width=768, + ) == { + "prompt": "Please draw a cute cat.", + "modalities": ["image"], + "mm_processor_kwargs": { + "modalities": ["image"], + "target_h": 512, + "target_w": 768, + }, + "negative_prompt": "ugly, blurry", + } + # target_h/w are omitted when height/width are not supplied so the + # pipeline's 1024x1024 default still applies. + assert build_text_to_image_prompt( + "MingImagePipeline", + prompt="Draw a poster.", + negative_prompt=None, + ) == { + "prompt": "Draw a poster.", + "modalities": ["image"], + "mm_processor_kwargs": {"modalities": ["image"]}, + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_ming_flash_omni_image_to_image_prompt_builder() -> None: + dummy_image = Image.new("RGB", (64, 64)) + result = build_image_to_image_prompt( + "MingImagePipeline", + prompt="Change the background to a sandy beach at sunset.", + negative_prompt=None, + input_image=dummy_image, + height=256, + width=256, + ) + assert result["modalities"] == ["img2img"] + assert result["multi_modal_data"]["img2img"] is dummy_image + assert result["mm_processor_kwargs"] == { + "modalities": ["img2img"], + "target_h": 256, + "target_w": 256, + } + assert "negative_prompt" not in result + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize("pipeline_name", ["Cosmos3OmniDiffusersPipeline", "Cosmos3OmniPipeline"]) +def test_cosmos3_text_to_image_prompt_builder_selects_image_modality(pipeline_name: str) -> None: + assert build_text_to_image_prompt( + pipeline_name, + prompt="a red sports car at golden hour", + negative_prompt="blurry, distorted", + height=1024, + width=1024, + ) == { + "prompt": "a red sports car at golden hour", + "modalities": ["image"], + "negative_prompt": "blurry, distorted", + } + assert build_text_to_image_prompt( + pipeline_name, + prompt="a red sports car", + negative_prompt=None, + ) == {"prompt": "a red sports car", "modalities": ["image"]} + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_audiox_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("AudioXPipeline") == frozenset( + { + "audiox_task", + "seconds_start", + "seconds_total", + "sigma_min", + "sigma_max", + "cfg_rescale", + "video_path", + "audio_path", + } + ) + assert get_extra_output_params("AudioXPipeline") == frozenset({"audiox_task"}) + assert should_init_extra_args_for_non_diffusion_stages("AudioXPipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_audiox_declared_extra_args_route_into_sampling_params() -> None: + params = OmniDiffusionSamplingParams() + declared = get_extra_body_params("AudioXPipeline") + apply_declared_extra_args( + params, + declared, + { + "audiox_task": "t2a", + "seconds_total": 10.0, + "sigma_min": 0.03, + "unknown": "ignored", + }, + ) + assert params.extra_args == { + "audiox_task": "t2a", + "seconds_total": 10.0, + "sigma_min": 0.03, + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_helios_extra_registry_declares_request_and_response_params() -> None: + expected_body = frozenset( + { + "is_enable_stage2", + "pyramid_num_stages", + "pyramid_num_inference_steps_list", + "is_amplify_first_chunk", + "is_skip_first_chunk", + "use_cfg_zero_star", + "use_zero_init", + "zero_steps", + "image", + "video", + "add_noise_to_image_latents", + "image_noise_sigma_min", + "image_noise_sigma_max", + "add_noise_to_video_latents", + "video_noise_sigma_min", + "video_noise_sigma_max", + } + ) + # Both the base and pyramid class names resolve to the same declaration. + for cls in ("HeliosPipeline", "HeliosPyramidPipeline"): + assert get_extra_body_params(cls) == expected_body + assert get_extra_output_params(cls) == frozenset() + assert should_init_extra_args_for_non_diffusion_stages(cls) is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_vace_extra_registry_has_no_pipeline_params() -> None: + assert get_extra_body_params("WanVACEPipeline") == frozenset() + assert get_extra_output_params("WanVACEPipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("WanVACEPipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_unknown_pipeline_has_empty_extra_registry() -> None: + assert get_extra_body_params("UnknownPipeline") == frozenset() + assert get_extra_output_params("UnknownPipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("UnknownPipeline") is False + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_text_to_image_prompt_builder() -> None: + assert build_text_to_image_prompt( + "BagelPipeline", + prompt="a cat", + negative_prompt="blurry", + height=512, + width=768, + ) == { + "prompt": "<|im_start|>a cat<|im_end|>", + "modalities": ["image"], + "mm_processor_kwargs": { + "target_h": 512, + "target_w": 768, + "modalities": ["image"], + }, + "negative_prompt": "blurry", + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_bagel_image_to_image_prompt_builder() -> None: + dummy_image = Image.new("RGB", (64, 64)) + result = build_image_to_image_prompt( + "BagelPipeline", + prompt="paint it", + negative_prompt="ugly", + input_image=dummy_image, + height=256, + width=256, + ) + assert result["prompt"] == "<|fim_middle|><|im_start|>paint it<|im_end|>" + assert result["modalities"] == ["img2img"] + assert result["multi_modal_data"]["img2img"] is dummy_image + assert result["mm_processor_kwargs"]["target_h"] == 256 + assert result["mm_processor_kwargs"]["target_w"] == 256 + assert result["negative_prompt"] == "ugly" + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_unknown_pipeline_uses_default_text_to_image_prompt() -> None: + assert build_text_to_image_prompt( + "UnknownPipeline", + prompt="a cat", + negative_prompt=None, + height=512, + width=512, + ) == {"prompt": "a cat"} + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_unknown_pipeline_uses_default_image_to_image_prompt() -> None: + dummy_image = Image.new("RGB", (64, 64)) + result = build_image_to_image_prompt( + "UnknownPipeline", + prompt="edit", + negative_prompt=None, + input_image=dummy_image, + ) + assert result == { + "prompt": "edit", + "multi_modal_data": {"image": dummy_image}, + } + + +def _build_vace_prompt(media_inputs: dict[str, object], *, num_frames: int = 5) -> dict: + return build_image_to_video_prompt( + "WanVACEPipeline", + prompt="a bird flying", + negative_prompt=None, + media_inputs=media_inputs, + height=16, + width=320, + num_frames=num_frames, + ) + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + "media_inputs", + [ + {"image": Image.new("RGB", (320, 16), "red")}, + {"last_image": Image.new("RGB", (320, 16), "blue")}, + { + "image": Image.new("RGB", (320, 16), "red"), + "last_image": Image.new("RGB", (320, 16), "blue"), + }, + {"image": Image.new("RGB", (320, 16), "red"), "mask": Image.new("L", (320, 16), 0)}, + {"reference_images": [Image.new("RGB", (64, 64), "red")]}, + ], + ids=["i2v", "v2lf", "flf2v", "inpaint", "r2v"], +) +def test_vace_image_to_video_prompt_builder(media_inputs: dict[str, object]) -> None: + result = _build_vace_prompt(media_inputs) + mmd = result["multi_modal_data"] + if "reference_images" in mmd: + assert mmd["reference_images"] is media_inputs["reference_images"] + else: + assert len(mmd["video"]) == len(mmd["mask"]) == 5 + + +@pytest.mark.core_model +@pytest.mark.cpu +@pytest.mark.parametrize( + ("media_inputs", "message"), + [ + ({}, "requires a conditioning media input"), + ({"mask": Image.new("L", (320, 16))}, "mask input requires an image"), + ({"control_image": Image.new("RGB", (320, 16))}, "Unsupported VACE media input"), + ], +) +def test_vace_rejects_invalid_media_combinations(media_inputs: dict[str, object], message: str) -> None: + with pytest.raises(ValueError, match=message): + _build_vace_prompt(media_inputs) + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_declared_extra_args_apply_to_existing_sampling_params() -> None: + params = OmniDiffusionSamplingParams(extra_args={"existing": 1}) + + declared_extra_params: frozenset[str] = frozenset({"cfg_text_scale", "think"}) + apply_declared_extra_args( + params, + declared_extra_params, + { + "cfg_text_scale": 4.0, + "think": False, + "unknown": "ignored", + }, + ) + + assert params.extra_args == { + "existing": 1, + "cfg_text_scale": 4.0, + "think": False, + } + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None: + assert get_extra_body_params("MammothModa2DiTPipeline") == frozenset( + { + "text_guidance_scale", + "cfg_range", + "num_inference_steps", + } + ) + assert get_extra_output_params("MammothModa2DiTPipeline") == frozenset() + assert should_init_extra_args_for_non_diffusion_stages("MammothModa2DiTPipeline") is True + + +@pytest.mark.core_model +@pytest.mark.cpu +def test_mammothmoda2_text_to_image_prompt_builder() -> None: + # Image dims are converted to the AR grid (width/16 x height/16); the negative + # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path). + assert build_text_to_image_prompt( + "MammothModa2DiTPipeline", + prompt="a cat", + negative_prompt="blurry", + height=512, + width=768, + ) == { + "prompt": ( + "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n" + "<|im_start|>user\na cat<|im_end|>\n" + "<|im_start|>assistant\n" + "<|image start|>48*32<|image token|>" + ), + "additional_information": { + "omni_task": ["t2i"], + "ar_width": [48], + "ar_height": [32], + "image_height": [512], + "image_width": [768], + }, + } diff --git a/tests/model_tests/__init__.py b/tests/model_tests/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/model_tests/conftest.py b/tests/model_tests/conftest.py new file mode 100644 index 00000000000..ab71a02cdc9 --- /dev/null +++ b/tests/model_tests/conftest.py @@ -0,0 +1,41 @@ +from shutil import rmtree + +import pytest + +from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS + + +@pytest.fixture(scope="session") +def tiny_model_paths(request, run_level): + """Build or download the tiny models for the selected tests. + + At core_model level, builds tiny models via the builder function. + At advanced_model / full_model level, uses the real HF model path. + + NOTE: this is session scoped to avoid churn in tiny model creation, + but will ensure all the tiny models you need are created for the selected tests + before it starts to execute them.""" + model_paths = {} + built_paths = [] + print("Initializing tiny models...") + for item in request.session.items: + if not hasattr(item, "callspec"): + continue + model_name = item.callspec.params.get("model_name") + if model_name is None or model_name not in DIFFUSION_TEST_SETTINGS: + continue + if model_name not in model_paths: + settings = DIFFUSION_TEST_SETTINGS[model_name] + if run_level == "core_model": + print(f"Calling tiny model builder for: {model_name}") + path = settings.builder() + built_paths.append(path) + else: + print(f"Run level is {run_level}; {model_name} will use the full model") + path = settings.model + model_paths[model_name] = path + + yield model_paths + for path in built_paths: + print(f"Removing tiny model: {path}") + rmtree(path, ignore_errors=True) diff --git a/tests/model_tests/diffusion/__init__.py b/tests/model_tests/diffusion/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/model_tests/diffusion/case_filtering.py b/tests/model_tests/diffusion/case_filtering.py new file mode 100644 index 00000000000..a4fa5dc1319 --- /dev/null +++ b/tests/model_tests/diffusion/case_filtering.py @@ -0,0 +1,104 @@ +""" +Analogous to: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/vlm_utils/case_filtering.py + +This follows a similar pattern so that we can use `get_parametrized_options` to expand our test +configurations out, similar to the way in which vLLM expands supported input type tests. For now, +this mostly just means parametrizing over the test groups, which may be unique per model, but +doing so in this way lets us cleanly define test groups that may be heterogeneous + avoid having +excessive skip marks for features that may not be supported in every pipeline. +""" + +# Engine initialization is currently quite expensive. For efficiency, +# when we run the tests, we should try to collapse compatible accelerations, +# because even with tiny models, the overhead of spinning up a new instance +# stacks up fast. +import itertools + +import pytest + +from tests.helpers.mark import hardware_marks +from tests.model_tests.diffusion.config_types import DiffusionAccs, DiffusionModelTestOpts, get_required_device_count +from vllm_omni.platforms import current_omni_platform + +# These tests are intended to run on L4 GPUs since they test parallelism (i.e., need GPUs), +# but are intended to be smoke tests that we can properly validate without real model +# weights in the CI. +# +# As such, no DiffusionAccs may require more than 4 GPUs, since we would otherwise skip +# them in the CI. +MAX_CI_DEVICES = 4 + + +def get_test_group_marks(model_name: str, test_group: list[DiffusionAccs] | None, model_marks: list | None) -> list: + """Build the full set of pytest marks for a test group. + + For now, single device groups default to core model, and multi device groups + default to advanced_model marks. We also validate and throw if a test group exceeds + the max number of GPUs for the CI to prevent accidentally adding configurations that + would be skipped in the CI.""" + marks = list(model_marks) if model_marks is not None else [] + + required_devices = get_required_device_count(test_group) + if required_devices > MAX_CI_DEVICES: + raise ValueError( + f"Test group {test_group} for {model_name} requires {required_devices} devices, " + f"but the max CI device count is {MAX_CI_DEVICES}. " + ) + + marks.extend(hardware_marks(res={"cuda": "L4"}, num_cards=required_devices)) + # hardware_marks only adds skipif for num_cards > 1, so we currently handle the single-device case + # directly here. This should probably be handled in a more common way later on. + assert current_omni_platform is not None and current_omni_platform.device_count is not None + device_count = current_omni_platform.device_count() + if current_omni_platform.is_cuda() and device_count < required_devices: + marks.append(pytest.mark.skip(reason=f"Need {required_devices} devices, got {device_count}")) + if required_devices > 1: + marks.append(pytest.mark.full_model) + else: + marks.append(pytest.mark.core_model) + return marks + + +def get_model_parametrization(model_name: str, test_info: DiffusionModelTestOpts, online: bool): + """Given a model & its corresponding test options, build the list of pytest params + to be run for this model. The base case (no accelerations) is always included. Extra + test groups are always appended for offline tests, but are only added to the online + tests if online_base_only is overridden to False to avoid redundant testing. + """ + test_groups: list[list[DiffusionAccs] | None] = [None] + if not (online and test_info.online_base_only) and test_info.extra_test_groups: + test_groups.extend(test_info.extra_test_groups) + + return [ + pytest.param( + model_name, + test_group, + test_info.supported_tasks, + test_info.check_multi_output and test_group is None, + test_info.check_determinism and test_group is None, + id=f"{model_name}[{'+'.join(test_group)}]" if test_group else model_name, + marks=get_test_group_marks(model_name, test_group, test_info.marks), + ) + for test_group in test_groups + ] + + +def get_parametrized_options( + test_settings: dict[str, DiffusionModelTestOpts], + online: bool = False, +): + """Converts all the DiffusionModelTestOpts into an expanded list of parameters + based on which accelerations are available. + + When online=True and a model has online_base_only=True, only the base case + (no accelerations) is included for that model's parametrization to minimize + redundancy between offline / online tests. + """ + # Get a list per model type, where each entry contains a tuple of all of + # that model type's cases, then flatten them into the top level so that + # we can consume them in one mark.parametrize call. + parametrization = [ + get_model_parametrization(model_name, test_info, online=online) + for model_name, test_info in test_settings.items() + ] + return list(itertools.chain(*parametrization)) diff --git a/tests/model_tests/diffusion/config_types.py b/tests/model_tests/diffusion/config_types.py new file mode 100644 index 00000000000..4120900fe12 --- /dev/null +++ b/tests/model_tests/diffusion/config_types.py @@ -0,0 +1,187 @@ +""" +Common definitions for controlling what tests run where. +""" + +from collections.abc import Callable +from enum import StrEnum, auto +from typing import NamedTuple, TypeAlias + +from pytest import MarkDecorator + +from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.entrypoints.omni import Omni + +# All builder funcs take no params and return a path +TinyDiffusionBuilder: TypeAlias = Callable[[], str] + + +class DiffusionAccs(StrEnum): + """Supported acceleration types / test settings for Diffusion Models.""" + + HSDP = auto() + TEA_CACHE = auto() + CACHE_DIT = auto() + SEQUENCE_PARALLEL = auto() + CFG_PARALLEL = auto() + TENSOR_PARALLEL = auto() + CPU_OFFLOAD = auto() + LAYERWISE_OFFLOAD = auto() + VAE_PATCH_PARALLEL = auto() + + +class DiffusionTasks(StrEnum): + """Supported tasks for Diffusion Models.""" + + TEXT_TO_IMAGE = auto() + IMAGE_TO_IMAGE = auto() + # Text to video, text to audio, etc should be added here as needed + + +class DiffusionModelTestOpts(NamedTuple): + """Configuration for one Diffusion model's tests.""" + + # HF model name for real-weight tests (advanced_model / full_model level). + # For now, whether we use the real weights vs tiny weights in the common tests + # depends on the run level. + model: str + + # Creates a tiny model for the given architecture. We should always use tiny + # model weights for tests that do not require us to check the model quality. + builder: TinyDiffusionBuilder + + # Actual tasks which controls the tests actually run + supported_tasks: list[DiffusionTasks] + + # Additional acceleration groups to run beyond the base case (no acceleration). + # The base case is always run for every model in the test settings. None means + # we only run the base case. + extra_test_groups: list[list[DiffusionAccs]] | None = None + + # Pytest Marks for this model. This may be useful for selecting which models + # we want to run where, similar to the way vLLM's multimodal tests mark some + # as core models to always run in the CI. + # Example: https://github.com/vllm-project/vllm/blob/v0.23.0/tests/models/multimodal/generation/test_common.py#L131 + marks: list[MarkDecorator] | None = None + + # When True (default), online tests only run the base case (no accelerations). + # When False, online tests run all test_groups, same as offline. This should be + # True unless there is a good reason for it not to be, because the execution of + # the acceleration should be the same on both codepaths, and CLI parsing etc should + # be tested by tests adding the acceleration, and not per model. + online_base_only: bool = True + + # Additional checks to run for the base case. + check_multi_output: bool = True # Runs multiple generations in one request + check_determinism: bool = True # Runs 2 generations with the same seed and check determinism + + +### Mappings & utils for building offline Omni() instances given a list of enabled accelerations +ACC_OMNI_KWARGS = { + DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_use_tiling": True}, + DiffusionAccs.CPU_OFFLOAD: {"enable_cpu_offload": True}, + DiffusionAccs.LAYERWISE_OFFLOAD: {"enable_layerwise_offload": True}, + DiffusionAccs.CACHE_DIT: {"cache_backend": "cache_dit"}, + DiffusionAccs.TEA_CACHE: {"cache_backend": "tea_cache"}, +} + +ACC_PARALLEL_KWARGS = { + DiffusionAccs.HSDP: {"use_hsdp": True, "hsdp_shard_size": 2}, + DiffusionAccs.TENSOR_PARALLEL: {"tensor_parallel_size": 2}, + DiffusionAccs.CFG_PARALLEL: {"cfg_parallel_size": 2}, + DiffusionAccs.VAE_PATCH_PARALLEL: {"vae_patch_parallel_size": 2}, + # For SP, we don't run ring here to conserve devices, since it is easy + # to blow up the number of needed devices fast. Compatibility for ring + # and ulysses together should be tested generically outside of these tests. + DiffusionAccs.SEQUENCE_PARALLEL: {"ulysses_degree": 2}, +} + +ACC_DEVICE_COUNT_KEYS = { + DiffusionAccs.TENSOR_PARALLEL: "tensor_parallel_size", + DiffusionAccs.CFG_PARALLEL: "cfg_parallel_size", + DiffusionAccs.SEQUENCE_PARALLEL: "ulysses_degree", + DiffusionAccs.VAE_PATCH_PARALLEL: "vae_patch_parallel_size", + DiffusionAccs.HSDP: "hsdp_shard_size", +} + +### CLI args for launching an OmniServer subprocess per acceleration. +# Follows the same pattern as e2e online serving tests, which hand-code +# their server_args lists in OmniServerParams. +ACC_SERVER_ARGS: dict[DiffusionAccs, list[str]] = { + DiffusionAccs.HSDP: ["--use-hsdp", "--hsdp-shard-size", "2"], + DiffusionAccs.TEA_CACHE: ["--cache-backend", "tea_cache"], + DiffusionAccs.CACHE_DIT: ["--cache-backend", "cache_dit"], + DiffusionAccs.SEQUENCE_PARALLEL: ["--usp", "2"], + DiffusionAccs.CFG_PARALLEL: ["--cfg-parallel-size", "2"], + DiffusionAccs.TENSOR_PARALLEL: ["--tensor-parallel-size", "2"], + DiffusionAccs.CPU_OFFLOAD: ["--enable-cpu-offload"], + DiffusionAccs.LAYERWISE_OFFLOAD: ["--enable-layerwise-offload"], + DiffusionAccs.VAE_PATCH_PARALLEL: ["--vae-use-tiling", "--vae-patch-parallel-size", "2"], +} + +## TODO ^ These should be in the same object, it's getting messy. + + +def get_required_device_count(accelerations: list[DiffusionAccs] | None) -> int: + """Compute the minimum number of devices needed for a set of accelerations. + The total is the product of all parallel dimensions (defaulting to 1). + + If not enough devices are available for a test group's accelerations, + that test will be skipped.""" + count = 1 + if accelerations is None: + return count + + for acc in accelerations: + key = ACC_DEVICE_COUNT_KEYS.get(acc) + if key is not None: + count *= ACC_PARALLEL_KWARGS[acc][key] + return count + + +def build_parallel_config_from_diff_accelerations(accelerations: list[DiffusionAccs]) -> DiffusionParallelConfig | None: + """Given a list of accelerations pertaining to the current test group, + build the parallel config needed for the Omni() object (if any).""" + config_kwargs = {} + for acc in accelerations: + update_dict = ACC_PARALLEL_KWARGS.get(acc, {}) + config_kwargs.update(update_dict) + if config_kwargs: + return DiffusionParallelConfig(**config_kwargs) + return None + + +### Offline Omni() object builder +def build_omni_from_diff_accelerations(accelerations: list[DiffusionAccs] | None, **kwargs) -> Omni: + """Given one or more acceleration types, build the corresponding Omni() object.""" + # Coerce to a list and build the parallel config, since that depends on the accelerations + if accelerations is None: + accelerations = [] + parallel_config = build_parallel_config_from_diff_accelerations(accelerations) + + # Then add anything else that's a top-level kwarg + acc_kwargs = {} + if parallel_config is not None: + acc_kwargs["parallel_config"] = parallel_config + for acc in accelerations: + update_dict = ACC_OMNI_KWARGS.get(acc, {}) + acc_kwargs.update(update_dict) + + # Keys passed through should mostly be things like enforce_eager; + # if there's overlap, it's probably due to a misconfiguration + shared_keys = acc_kwargs.keys() & kwargs.keys() + if shared_keys: + raise ValueError(f"Explicit Omni kwargs and inferred Omni kwargs for accelerations overlap: {shared_keys}") + omni_kwargs = {**acc_kwargs, **kwargs} + return Omni(**omni_kwargs) + + +### Online server flag builder +def build_server_args_from_diff_accelerations(accelerations: list[DiffusionAccs] | None) -> list[str]: + """Given one or more acceleration types, build the corresponding CLI args + for launching an OmniServer subprocess.""" + if accelerations is None: + return [] + args = [] + for acc in accelerations: + args.extend(ACC_SERVER_ARGS.get(acc, [])) + return args diff --git a/tests/model_tests/diffusion/diff_model_builders.py b/tests/model_tests/diffusion/diff_model_builders.py new file mode 100644 index 00000000000..9bee71f9bd1 --- /dev/null +++ b/tests/model_tests/diffusion/diff_model_builders.py @@ -0,0 +1,63 @@ +import os +import tempfile + +import torch +from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 +from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel +from diffusers.pipelines.flux2.pipeline_flux2_klein import Flux2KleinPipeline +from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from transformers import AutoTokenizer, Qwen3Config, Qwen3ForCausalLM + +TINY_MODEL_DIR = os.path.join(tempfile.gettempdir(), "vllm-omni-tiny-models") + + +def _get_tiny_model_path(name: str) -> str: + path = os.path.join(TINY_MODEL_DIR, name) + os.makedirs(path, exist_ok=True) + return path + + +def tiny_flux2_klein_builder() -> str: + """Build a tiny Flux2Klein model.""" + model_dir = _get_tiny_model_path("Flux2KleinPipeline") + + pipe = Flux2KleinPipeline( + scheduler=FlowMatchEulerDiscreteScheduler(), + vae=AutoencoderKLFlux2( + down_block_types=("DownEncoderBlock2D",), + up_block_types=("UpDecoderBlock2D",), + block_out_channels=(32,), + layers_per_block=1, + latent_channels=16, + norm_num_groups=16, + ), + # NOTE: For now we need 28 layers because of hardcoded stuff in the model :( + text_encoder=Qwen3ForCausalLM( + Qwen3Config( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=28, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=16, + vocab_size=151936, + max_position_embeddings=512, + ) + ), + tokenizer=AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B"), + # NOTE: For now we need at least 2 layers for the transformer + # due to hardcoded hacks in CacheDiT for Flux2Klein specifically. + transformer=Flux2Transformer2DModel( + in_channels=64, + num_layers=2, + num_single_layers=2, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=96, + timestep_guidance_channels=32, + axes_dims_rope=(4, 4, 4, 4), + ), + ) + # Need dtypes to be consistent; for now we just put it on bfloat16 + pipe.to(torch.bfloat16).save_pretrained(model_dir) + return model_dir diff --git a/tests/model_tests/diffusion/model_settings.py b/tests/model_tests/diffusion/model_settings.py new file mode 100644 index 00000000000..3f44afcae5e --- /dev/null +++ b/tests/model_tests/diffusion/model_settings.py @@ -0,0 +1,42 @@ +from tests.model_tests.diffusion import diff_model_builders +from tests.model_tests.diffusion.config_types import ( + DiffusionAccs, + DiffusionModelTestOpts, + DiffusionTasks, +) + +# This object defines the (tiny) model configurations for common tests. +# +# TO ADD A NEW MODEL: +# Map the pipeline name (i.e., the name of the pipeline in _DIFFUSION_MODELS) +# to a DiffusionModelTestOpts and define the supported tasks. A base case with no +# accelerations is always run. Use extra_test_groups to define additional +# acceleration configurations to run a basic smoke test for (i.e., simple inference). +# +# The model will only be loaded once per test group, and will execute each of its +# supported tasks as a pytest subtest. The base case also runs additional checks +# that are independent of accelerations, e.g., for determinism and compatibility with +# multi-output, unless explicitly disabled in the test settings. +# +# TO RUN A SUBSET OF THE TESTS: +# These tests should be fast, but if you'd like to run a subset of the tests, you can do so +# with `-k`; the IDs of the tests are the name of the pipeline joined by +. +# +# Example: +# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline] +# ^ Runs only the base (no accelerations) case for Flux2KleinPipeline +# +# $ pytest test_common_offline.py -k test_pipeline_on_supported_tasks[Flux2KleinPipeline +# ^ Runs all test groups for Flux2KleinPipeline only +DIFFUSION_TEST_SETTINGS = { + "Flux2KleinPipeline": DiffusionModelTestOpts( + model="black-forest-labs/FLUX.2-klein-4B", + builder=diff_model_builders.tiny_flux2_klein_builder, + supported_tasks=[DiffusionTasks.TEXT_TO_IMAGE, DiffusionTasks.IMAGE_TO_IMAGE], + extra_test_groups=[ + [DiffusionAccs.HSDP, DiffusionAccs.TEA_CACHE], + [DiffusionAccs.SEQUENCE_PARALLEL, DiffusionAccs.CACHE_DIT, DiffusionAccs.LAYERWISE_OFFLOAD], + [DiffusionAccs.CFG_PARALLEL, DiffusionAccs.TENSOR_PARALLEL, DiffusionAccs.CPU_OFFLOAD], + ], + ) +} diff --git a/tests/model_tests/diffusion/task_runners.py b/tests/model_tests/diffusion/task_runners.py new file mode 100644 index 00000000000..6b9c7ff9907 --- /dev/null +++ b/tests/model_tests/diffusion/task_runners.py @@ -0,0 +1,172 @@ +""" +Definitions running and validating individual tasks, e.g., text to image, +image to image, and so on. These are called by the core test runner. +""" + +import base64 +import io +from dataclasses import replace + +import numpy as np +from PIL import Image + +from tests.helpers.runtime import DiffusionResponse, OmniServer, OpenAIClientHandler, dummy_messages_from_mix_data +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +PROMPT = "Dummy prompt" +IMAGE_DIMS = (512, 512) +HEIGHT, WIDTH = IMAGE_DIMS +INPUT_IMAGE = Image.new("RGB", IMAGE_DIMS) + +# Offline sampling params +IMAGE_GEN_SAMPLING_PARAMS = OmniDiffusionSamplingParams( + num_inference_steps=4, + height=HEIGHT, + width=WIDTH, + seed=42, +) + +# Online extra_body for diffusion requests +IMAGE_GEN_EXTRA_BODY = { + "height": HEIGHT, + "width": WIDTH, + "num_inference_steps": 4, + "seed": 42, +} + + +### Shared validation +def _validate_images(images: list[Image.Image], expected_n: int = 1): + """Given a set of images, ensure we got the expected count of images, + and that all provided images match the expected dimensions.""" + assert len(images) == expected_n + for img in images: + assert isinstance(img, Image.Image) + assert img.size == IMAGE_DIMS + return images + + +def _validate_image_gen_determinism(images_a: list[Image.Image], images_b: list[Image.Image]): + """Ensure that two image sets are valid and that the results match.""" + _validate_images(images_a) + _validate_images(images_b) + assert np.array_equal(np.array(images_a[0]), np.array(images_b[0])) + + +### Output extractor utils for offline / online paths respectively +def _get_offline_images(outputs: list[OmniRequestOutput]) -> list[Image.Image]: + """Extract the images from an Omni .generate() call.""" + assert len(outputs) == 1 + return outputs[0].images + + +def _get_online_images(responses: list[DiffusionResponse]) -> list[Image.Image]: + """Extract the images from a server response.""" + assert len(responses) == 1 + images = responses[0].images + assert images is not None + return images + + +### Offline helpers +def _run_offline_t2i(omni: Omni, params: OmniDiffusionSamplingParams = IMAGE_GEN_SAMPLING_PARAMS): + return omni.generate({"prompt": PROMPT}, params) + + +def _run_offline_i2i(omni: Omni): + return omni.generate( + {"prompt": PROMPT, "multi_modal_data": {"image": INPUT_IMAGE}}, + IMAGE_GEN_SAMPLING_PARAMS, + ) + + +### Online helpers +def _build_online_image_data_url() -> str: + """Get a valid base 64 encoded data URL corresponding to an image.""" + buf = io.BytesIO() + INPUT_IMAGE.save(buf, format="PNG") + b64 = base64.b64encode(buf.getvalue()).decode() + return f"data:image/png;base64,{b64}" + + +def _run_online_t2i( + server: OmniServer, client: OpenAIClientHandler, extra_body: dict | None = None +) -> list[DiffusionResponse]: + """Run a text to image request through the server.""" + messages = dummy_messages_from_mix_data(content_text=PROMPT) + request_config = { + "model": server.model, + "messages": messages, + "extra_body": extra_body or IMAGE_GEN_EXTRA_BODY, + } + return client.send_diffusion_request(request_config) + + +def _run_online_i2i(server: OmniServer, client: OpenAIClientHandler) -> list[DiffusionResponse]: + """Run an image to image request through the server.""" + image_data_url = _build_online_image_data_url() + messages = dummy_messages_from_mix_data( + content_text=PROMPT, + image_data_url=image_data_url, + ) + request_config = { + "model": server.model, + "messages": messages, + "extra_body": IMAGE_GEN_EXTRA_BODY, + } + return client.send_diffusion_request(request_config) + + +### Offline task runners +def run_and_validate_text_to_image_request(omni: Omni): + """Run and validate a text to image request.""" + _validate_images(_get_offline_images(_run_offline_t2i(omni))) + + +def run_and_validate_image_to_image_request(omni: Omni): + """Run and validate an image to image request.""" + _validate_images(_get_offline_images(_run_offline_i2i(omni))) + + +def run_and_validate_text_to_image_determinism(omni: Omni): + """Checks for determinism; for now we just keep this for TTI.""" + _validate_image_gen_determinism( + _get_offline_images(_run_offline_t2i(omni)), + _get_offline_images(_run_offline_t2i(omni)), + ) + + +def run_and_validate_text_to_image_multi_output(omni: Omni): + """Checks for multi-output; for now we just keep this for TTI.""" + params = replace(IMAGE_GEN_SAMPLING_PARAMS, num_outputs_per_prompt=2) + _validate_images(_get_offline_images(_run_offline_t2i(omni, params)), expected_n=2) + + +### Online task runners +def run_and_validate_online_text_to_image_request(server: OmniServer, client: OpenAIClientHandler): + """Run and validate a text to image request through the server.""" + _validate_images(_get_online_images(_run_online_t2i(server, client))) + + +def run_and_validate_online_image_to_image_request(server: OmniServer, client: OpenAIClientHandler): + """Run and validate an image to image request through the server.""" + _validate_images(_get_online_images(_run_online_i2i(server, client))) + + +def run_and_validate_online_text_to_image_determinism(server: OmniServer, client: OpenAIClientHandler): + """Checks for determinism through the server; for now we just keep this for TTI.""" + _validate_image_gen_determinism( + _get_online_images(_run_online_t2i(server, client)), + _get_online_images(_run_online_t2i(server, client)), + ) + + +def run_and_validate_online_text_to_image_multi_output(server: OmniServer, client: OpenAIClientHandler): + """Checks for multi-output through the server; for now we just keep this for TTI.""" + extra_body = {**IMAGE_GEN_EXTRA_BODY, "num_outputs_per_prompt": 2} + _validate_images( + _get_online_images(_run_online_t2i(server, client, extra_body=extra_body)), + expected_n=2, + ) diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py new file mode 100644 index 00000000000..c9564414301 --- /dev/null +++ b/tests/model_tests/diffusion/test_alignment.py @@ -0,0 +1,100 @@ +""" +Tests for alignment between _DIFFUSION_MODELS and DIFFUSION_TEST_SETTINGS; if +tests in this file are failing, you are probably adding a new model, and need +to account for it in DIFFUSION_TEST_SETTINGS +""" + +import pytest + +from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS +from vllm_omni.diffusion.registry import _DIFFUSION_MODELS + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +# These diffusion models currently do not have any added (tiny) model tests. +# In general, pipelines should only be in the list below if they have not been +# migrated yet or there is a compelling reason that they should not be (which for +# now may just be CI pressure etc due to the volume of models and the cost of +# initializing the server). +# +# The tests here validate that all keys in the diffusion registry are accounted for; +# if you are adding a new model and see the tests below fail, please follow the pattern +# for adding a new tiny model builder and corresponding entry in DIFFUSION_TEST_SETTINGS. +EXCLUDED_MODELS = [ + "QwenImagePipeline", + "QwenImageEditPipeline", + "QwenImageEditPlusPipeline", + "QwenImageLayeredPipeline", + "GlmImagePipeline", + "ZImagePipeline", + "OvisImagePipeline", + "WanPipeline", + "WanVACEPipeline", + "LTX2Pipeline", + "LTX2ImageToVideoPipeline", + "LTX2TwoStagesPipeline", + "LTX2ImageToVideoTwoStagesPipeline", + "LTX2T2VDMD2Pipeline", + "LTX2I2VDMD2Pipeline", + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + "StableAudioPipeline", + "WanImageToVideoPipeline", + "WanS2VPipeline", + "WanT2VDMD2Pipeline", + "WanI2VDMD2Pipeline", + "LongCatImagePipeline", + "BagelPipeline", + "LancePipeline", + "MingImagePipeline", + "InternVLAA1Pipeline", + "LongCatImageEditPipeline", + "StableDiffusion3Pipeline", + "FluxKontextPipeline", + "HunyuanImage3ForCausalMM", + "ErnieImagePipeline", + "NextStep11Pipeline", + "FluxPipeline", + "FluxDMD2Pipeline", + "Krea2Pipeline", + "QwenImageDMD2Pipeline", + "OmniGen2Pipeline", + "HeliosPipeline", + "HeliosPyramidPipeline", + "Flux2Pipeline", + "DreamIDOmniPipeline", + "SenseNovaU1Pipeline", + "AudioXPipeline", + "HunyuanVideo15Pipeline", + "HunyuanVideo15ImageToVideoPipeline", + "MagiHumanPipeline", + "OmniVoicePipeline", + "OmniVoice", + "Cosmos3OmniDiffusersPipeline", + "Cosmos3OmniPipeline", + "DiffusersAdapterPipeline", + "HiDreamImagePipeline", + "DreamZeroPipeline", + "StableDiffusionXLPipeline", + "Gr00tN1d7Pipeline", + "SoulXSingerPipeline", + "SoulXSingerSVCPipeline", +] + + +def test_all_non_excluded_pipelines_have_tests(): + """Ensure that every pipeline in the diffusion registry either has a test + configuration in DIFFUSION_TEST_SETTINGS, or is explicitly excluded for a + known reason.""" + non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS) + missing_tests = non_excluded_diff_pipes - DIFFUSION_TEST_SETTINGS.keys() + assert len(missing_tests) == 0, f"Pipelines missing test settings: {missing_tests}" + + +def test_no_excluded_models_have_test_settings(): + """Ensure that no models in the exclude list have a test configuration in + DIFFUSION_TEST_SETTINGS (i.e., if they have test settings, we should not + exclude them).""" + non_excluded_diff_pipes = _DIFFUSION_MODELS.keys() - set(EXCLUDED_MODELS) + missing_pipes = DIFFUSION_TEST_SETTINGS.keys() - set(non_excluded_diff_pipes) + assert len(missing_pipes) == 0 diff --git a/tests/model_tests/diffusion/test_common_offline.py b/tests/model_tests/diffusion/test_common_offline.py new file mode 100644 index 00000000000..f24ed378023 --- /dev/null +++ b/tests/model_tests/diffusion/test_common_offline.py @@ -0,0 +1,69 @@ +import pytest + +from tests.model_tests.diffusion.case_filtering import get_parametrized_options +from tests.model_tests.diffusion.config_types import ( + DiffusionAccs, + DiffusionTasks, + build_omni_from_diff_accelerations, +) +from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS +from tests.model_tests.diffusion.task_runners import ( + run_and_validate_image_to_image_request, + run_and_validate_text_to_image_determinism, + run_and_validate_text_to_image_multi_output, + run_and_validate_text_to_image_request, +) + +# NOTE : Hardware marks are added dynamically based on test requirements +pytestmark = [pytest.mark.diffusion] + + +@pytest.mark.parametrize( + "model_name,accelerations,supported_tasks,check_multioutput,check_determinism", + get_parametrized_options(DIFFUSION_TEST_SETTINGS), +) +def test_pipeline_on_supported_tasks( + model_name, + accelerations: list[DiffusionAccs] | None, + supported_tasks: list[DiffusionTasks], + check_multioutput: bool, + check_determinism: bool, + tiny_model_paths: dict[str, str], + subtests, +): + """Run a smoke test on all of the pipelines supported tasks using a set of enabled accelerations.""" + assert len(supported_tasks) > 0 + # We initialize the Omni object before running the tasks, then run each task as a pytest subtask. + # This lets us init the model once, but display separate failures in pytest, and avoid halting the + # checks on other tasks if one fails. + # + # This allows us to have some degree of test isolation without the cost of redundant initialization, + # since starting the server can take 10+ seconds, even for tiny models. + # + # NOTE: Be sure to install pytest-subtests if you're running on pytest < 9 + omni = build_omni_from_diff_accelerations( + accelerations=accelerations, + model=tiny_model_paths[model_name], + enforce_eager=True, + ) + try: + for task_type in supported_tasks: + with subtests.test(msg=task_type): + if task_type == DiffusionTasks.TEXT_TO_IMAGE: + run_and_validate_text_to_image_request(omni) + elif task_type == DiffusionTasks.IMAGE_TO_IMAGE: + run_and_validate_image_to_image_request(omni) + else: + raise ValueError(f"Task type {task_type} is not yet supported") + + # NOTE: For now, we only check determinism + multi output for the base case, + # since checking it on every extra acceleration configuration is redundant + # (see case_filtering). + if check_determinism: + with subtests.test(msg="determinism"): + run_and_validate_text_to_image_determinism(omni) + if check_multioutput: + with subtests.test(msg="multi_output"): + run_and_validate_text_to_image_multi_output(omni) + finally: + omni.close() diff --git a/tests/model_tests/diffusion/test_common_online.py b/tests/model_tests/diffusion/test_common_online.py new file mode 100644 index 00000000000..e4c6be0d4c7 --- /dev/null +++ b/tests/model_tests/diffusion/test_common_online.py @@ -0,0 +1,75 @@ +""" +Analogous to test_common_offline, but for server tests. Validates the full +online serving stack (CLI arg parsing, subprocess, API routing, response +encoding) using tiny models. +""" + +import pytest + +from tests.helpers.runtime import OmniServer, OpenAIClientHandler +from tests.model_tests.diffusion.case_filtering import get_parametrized_options +from tests.model_tests.diffusion.config_types import ( + DiffusionAccs, + DiffusionTasks, + build_server_args_from_diff_accelerations, +) +from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS +from tests.model_tests.diffusion.task_runners import ( + run_and_validate_online_image_to_image_request, + run_and_validate_online_text_to_image_determinism, + run_and_validate_online_text_to_image_multi_output, + run_and_validate_online_text_to_image_request, +) + +# NOTE : Hardware marks are added dynamically based on test requirements +pytestmark = [pytest.mark.diffusion] + + +@pytest.mark.parametrize( + "model_name,accelerations,supported_tasks,check_multioutput,check_determinism", + get_parametrized_options(DIFFUSION_TEST_SETTINGS, online=True), +) +def test_online_on_supported_tasks( + model_name: str, + accelerations: list[DiffusionAccs] | None, + supported_tasks: list[DiffusionTasks], + check_multioutput: bool, + check_determinism: bool, + tiny_model_paths: dict[str, str], + run_level: str, + subtests, +): + """Smoke test: start a tiny model server and run each supported task via the API.""" + model_path = tiny_model_paths[model_name] + server_args = build_server_args_from_diff_accelerations(accelerations) + server_args.append("--enforce-eager") + + with OmniServer(model_path, server_args) as server: + # TODO: We may want to revisit run_level validation here, + # because checks for things like image size etc should not + # depend on whether or not the weights are real or random + client = OpenAIClientHandler( + host=server.host, + port=server.port, + api_key="EMPTY", + run_level=run_level, + log_stats=server.log_stats, + ) + for task_type in supported_tasks: + with subtests.test(msg=task_type): + if task_type == DiffusionTasks.TEXT_TO_IMAGE: + run_and_validate_online_text_to_image_request(server, client) + elif task_type == DiffusionTasks.IMAGE_TO_IMAGE: + run_and_validate_online_image_to_image_request(server, client) + else: + raise ValueError(f"Task type {task_type} is not yet supported") + + # NOTE: For now, we only check determinism + multi output for the base case, + # since checking it on every extra acceleration configuration is redundant + # (see case_filtering). + if check_determinism: + with subtests.test(msg="determinism"): + run_and_validate_online_text_to_image_determinism(server, client) + if check_multioutput: + with subtests.test(msg="multi_output"): + run_and_validate_online_text_to_image_multi_output(server, client) diff --git a/tests/model_tests/diffusion/utils.py b/tests/model_tests/diffusion/utils.py new file mode 100644 index 00000000000..35577b68f28 --- /dev/null +++ b/tests/model_tests/diffusion/utils.py @@ -0,0 +1,33 @@ +""" +Utilities for resolving real models to their tiny model equivalents. +""" + +import logging + +from tests.model_tests.diffusion.model_settings import DIFFUSION_TEST_SETTINGS +from vllm_omni.diffusion.data import resolve_model_class_name + +logger = logging.getLogger(__name__) + + +def resolve_tiny_model_path(model: str) -> str: + """Given a real model name/path, resolve it to a tiny model path. + + Raises ValueError if the pipeline class cannot be determined (invalid + model). Returns the original model path if no tiny builder exists yet.""" + pipeline_class = resolve_model_class_name(model) + if pipeline_class is None: + raise ValueError( + f"Cannot resolve pipeline class for model: {model}. The model path may be invalid or its config unreadable." + ) + + test_opts = DIFFUSION_TEST_SETTINGS.get(pipeline_class) + if test_opts is None: + logger.warning( + "No tiny model builder for pipeline %s (model: %s). Using original model.", + pipeline_class, + model, + ) + return model + + return test_opts.builder() diff --git a/tests/platforms/npu/test_310p_patches.py b/tests/platforms/npu/test_310p_patches.py index 255b092c860..e138a6d63da 100644 --- a/tests/platforms/npu/test_310p_patches.py +++ b/tests/platforms/npu/test_310p_patches.py @@ -67,10 +67,37 @@ class FakeCodePredictorDecoderLayer(torch.nn.Module): class FakeCodePredictorBaseModel(torch.nn.Module): pass - class FakeMimiEuclideanCodebook(torch.nn.Module): - @property - def embed(self): - return self._embed + class FakeProjection(torch.nn.Linear): + def __init__(self): + super().__init__(4, 4, bias=False) + self.call_shapes: list[tuple[int, ...]] = [] + + def forward(self, hidden_states): + self.call_shapes.append(tuple(hidden_states.shape)) + return super().forward(hidden_states) + + class FakeCodePredictorWrapper(torch.nn.Module): + def __init__(self, *args, **kwargs): + del args, kwargs + super().__init__() + self.model = torch.nn.Module() + self.model.codec_embedding = torch.nn.ModuleList([torch.nn.Embedding(8, 4), torch.nn.Embedding(8, 4)]) + self.model.linear = torch.nn.Linear(4, 4, bias=False) + self.lm_head = torch.nn.ModuleList([torch.nn.Linear(4, 8, bias=False), torch.nn.Linear(4, 8, bias=False)]) + self.small_to_mtp_projection = FakeProjection() + with torch.no_grad(): + self.small_to_mtp_projection.weight.copy_(torch.diag(torch.tensor([1.0, 2.0, 3.0, 4.0]))) + self._wrapper_config = SimpleNamespace(use_parallel_embedding=False, sampling_mode="per_call") + self._projected_codec_embed_weight = None + + def load_weights(self, weights): + del weights + return {"loaded"} + + class FakeCode2WavBase(torch.nn.Module): + def __init__(self, *args, **kwargs): + del args, kwargs + super().__init__() class FakeEncoder: def __init__(self): @@ -136,8 +163,21 @@ class FakePromptEmbedsBuilder: CodePredictorAttention=FakeCodePredictorAttention, CodePredictorDecoderLayer=FakeCodePredictorDecoderLayer, CodePredictorBaseModel=FakeCodePredictorBaseModel, + CodePredictorWrapper=FakeCodePredictorWrapper, _rotate_half=lambda x: x, ) + fake_qwen3_tts_code_predictor_vllm = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm", + Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM=FakeCodePredictorWrapper, + Qwen3TTSTalkerCodePredictorModelVLLM=FakeCodePredictorBaseModel, + CodePredictorWrapper=FakeCodePredictorWrapper, + ) + fake_qwen3_tts_code2wav = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav", + Qwen3TTSCode2Wav=FakeCode2WavBase, + ) fake_prompt_builder = _install_fake_module( monkeypatch, "vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder", @@ -150,22 +190,40 @@ class FakePromptEmbedsBuilder: Qwen3TTSTalkerForConditionalGeneration=FakeTalkerBase, Qwen3TTSPromptEmbedsBuilder=FakePromptEmbedsBuilder, ) - fake_modeling_mimi = _install_fake_module( - monkeypatch, - "transformers.models.mimi.modeling_mimi", - MimiEuclideanCodebook=FakeMimiEuclideanCodebook, - ) - fake_mimi = _install_fake_module( - monkeypatch, - "transformers.models.mimi", - modeling_mimi=fake_modeling_mimi, - ) _install_fake_module(monkeypatch, "vllm") _install_fake_module(monkeypatch, "vllm.multimodal") _install_fake_module(monkeypatch, "vllm.multimodal.audio", AudioResampler=FakeAudioResampler) - _install_fake_module(monkeypatch, "transformers") - _install_fake_module(monkeypatch, "transformers.models", mimi=fake_mimi) + _install_fake_module(monkeypatch, "torch_npu", npu_format_cast=lambda weight, _fmt: weight) + _install_fake_module(monkeypatch, "vllm_ascend") + _install_fake_module(monkeypatch, "vllm_ascend._310p") + _install_fake_module(monkeypatch, "vllm_ascend._310p.attention") + _install_fake_module( + monkeypatch, + "vllm_ascend._310p.attention.attention_mask", + AttentionMaskBuilder310=SimpleNamespace( + gen_causal_additive_mask=lambda max_seq, device: torch.zeros( + max_seq, + max_seq, + device=device, + ) + ), + ) + _install_fake_module(monkeypatch, "vllm_ascend.sample") + _install_fake_module( + monkeypatch, + "vllm_ascend.sample.sampler", + apply_top_k_top_p=lambda logits, **_kwargs: logits, + random_sample=lambda probs, _generators: probs.argmax(dim=-1, keepdim=True), + ) + _install_fake_module( + monkeypatch, + "vllm_ascend.utils", + ACL_FORMAT_FRACTAL_NZ=29, + aligned_16=lambda tensor: tensor, + maybe_trans_nz=lambda weight: weight, + nd_to_nz_2d=lambda tensor: tensor, + ) _install_fake_module(monkeypatch, "vllm_omni") _install_fake_module(monkeypatch, "vllm_omni.model_executor") _install_fake_module(monkeypatch, "vllm_omni.model_executor.models") @@ -178,9 +236,17 @@ class FakePromptEmbedsBuilder: monkeypatch, "vllm_omni.model_executor.models.qwen3_tts", prompt_embeds_builder=fake_prompt_builder, + qwen3_tts_code2wav=fake_qwen3_tts_code2wav, + qwen3_tts_code_predictor_vllm=fake_qwen3_tts_code_predictor_vllm, qwen3_tts_talker=fake_talker, ) - return fake_qwen3_code_predictor, fake_prompt_builder, fake_talker + return ( + fake_qwen3_code_predictor, + fake_qwen3_tts_code_predictor_vllm, + fake_qwen3_tts_code2wav, + fake_prompt_builder, + fake_talker, + ) def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch): @@ -193,7 +259,7 @@ def _load_qwen3_tts_patch(monkeypatch: pytest.MonkeyPatch): def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest.MonkeyPatch) -> None: registry_path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "__init__.py" registry = _load_source_module("vllm_omni_test_310p_patch_registry", registry_path) - calls = {"worker": 0, "talker": 0} + calls = {"worker": 0, "talker": 0, "code2wav": 0} _install_fake_module( monkeypatch, @@ -204,52 +270,35 @@ def test_registry_applies_worker_once_and_model_patch_lazily(monkeypatch: pytest monkeypatch, "vllm_omni.platforms.npu._310p.patch.qwen3_tts", apply_talker_patches=lambda: calls.__setitem__("talker", calls["talker"] + 1), + apply_code2wav_patches=lambda: calls.__setitem__("code2wav", calls["code2wav"] + 1), ) registry.apply_patches() registry.apply_patches() registry.apply_model_patches(SimpleNamespace(model_arch="OtherModel")) registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSTalkerForConditionalGeneration")) + registry.apply_model_patches(SimpleNamespace(model_arch="Qwen3TTSCode2Wav")) - assert calls == {"worker": 1, "talker": 1} - - -def test_worker_patch_replaces_base_and_runs_disable_jit(monkeypatch: pytest.MonkeyPatch) -> None: - calls: list[str] = [] - - class FakeOmniNPUWorkerBase: - def _init_device(self): - calls.append("parent") - return "npu:0" - - fake_worker_base = _install_fake_module( - monkeypatch, - "vllm_omni.platforms.npu.worker.base", - OmniNPUWorkerBase=FakeOmniNPUWorkerBase, - ) - _install_fake_module(monkeypatch, "vllm_omni") - _install_fake_module(monkeypatch, "vllm_omni.platforms") - _install_fake_module(monkeypatch, "vllm_omni.platforms.npu") - _install_fake_module( - monkeypatch, - "vllm_omni.platforms.npu._310p", - disable_jit_compile=lambda: calls.append("disable_jit"), - ) - _install_fake_module(monkeypatch, "vllm_omni.platforms.npu.worker", base=fake_worker_base) - - path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "_310p" / "patch" / "worker.py" - module = _load_source_module("vllm_omni_test_310p_worker_patch", path) - module.apply_patch() + assert calls == {"worker": 1, "talker": 1, "code2wav": 1} - assert fake_worker_base.OmniNPUWorkerBase is module._OmniNPUWorkerBase310P - assert fake_worker_base.OmniNPUWorkerBase()._init_device() == "npu:0" - assert calls == ["parent", "disable_jit"] +def test_qwen3_tts_patch_registers_target_classes(monkeypatch: pytest.MonkeyPatch) -> None: + ( + module, + ( + fake_code_predictor, + fake_code_predictor_vllm, + fake_code2wav, + fake_prompt_builder, + fake_talker, + ), + ) = _load_qwen3_tts_patch(monkeypatch) -def test_qwen3_tts_patch_replaces_target_classes(monkeypatch: pytest.MonkeyPatch) -> None: - module, (fake_code_predictor, fake_prompt_builder, fake_talker) = _load_qwen3_tts_patch(monkeypatch) + original_common_wrapper = fake_code_predictor.CodePredictorWrapper + original_vllm_wrapper = fake_code_predictor_vllm.CodePredictorWrapper module.apply_talker_patches() + module.apply_code2wav_patches() assert fake_talker.Qwen3TTSTalkerForConditionalGeneration is module._Qwen3TTSTalker310P assert fake_talker.Qwen3TTSPromptEmbedsBuilder is module._Qwen3TTSPromptEmbedsBuilder310P @@ -257,6 +306,112 @@ def test_qwen3_tts_patch_replaces_target_classes(monkeypatch: pytest.MonkeyPatch assert fake_code_predictor.CodePredictorAttention is module._Qwen3CodePredictorAttention310P assert fake_code_predictor.CodePredictorDecoderLayer is module._Qwen3CodePredictorDecoderLayer310P assert fake_code_predictor.CodePredictorBaseModel is module._Qwen3CodePredictorBaseModel310P + assert ( + fake_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM + is module._Qwen3TTSTalkerCodePredictor310P + ) + assert fake_code_predictor.CodePredictorWrapper is original_common_wrapper + assert fake_code_predictor_vllm.CodePredictorWrapper is original_vllm_wrapper + assert fake_code2wav.Qwen3TTSCode2Wav is module._Qwen3TTSCode2Wav310P + + code2wav = module._Qwen3TTSCode2Wav310P( + vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))) + ) + + assert code2wav._npu_decoder_runtime_dtype(torch.device("cpu")) is torch.float16 + + +def test_qwen3_tts_code_predictor_forward_uses_projected_embedding_and_sampling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module, _ = _load_qwen3_tts_patch(monkeypatch) + predictor = module._Qwen3TTSTalkerCodePredictor310P( + vllm_config=object(), + config=object(), + talker_config=object(), + ) + codec_weight = predictor.model.codec_embedding[0].weight.detach().clone() + projection_weight = predictor.small_to_mtp_projection.weight.detach().clone() + + assert predictor.load_weights(iter(())) == {"loaded"} + torch.testing.assert_close( + predictor._projected_codec_embed_weight[0], + torch.nn.functional.linear(codec_weight, projection_weight), + ) + assert predictor.small_to_mtp_projection.call_shapes == [(8, 4), (8, 4)] + predictor.small_to_mtp_projection.call_shapes.clear() + predictor._num_groups = 3 + predictor._model_dtype = torch.float32 + predictor._prefix_graphs_enabled = False + predictor._bucket_pos_ids = {} + predictor._device_graphs = {} + predictor._lm_heads_list = [ + lambda hidden: torch.nn.functional.one_hot( + torch.full((hidden.shape[0],), 2, dtype=torch.long), + num_classes=8, + ).to(torch.float32), + lambda hidden: torch.nn.functional.one_hot( + torch.full((hidden.shape[0],), 3, dtype=torch.long), + num_classes=8, + ).to(torch.float32), + ] + predictor._setup_compile = lambda: None + predictor._padded_bsz = lambda bsz: bsz + + def ensure_buffers(device, dtype, padded_bsz): + predictor._proj_buf = torch.zeros(padded_bsz, 4, 4, device=device, dtype=dtype) + + predictor._ensure_buffers = ensure_buffers + predictor._compiled_model_fwd = lambda embeds, _positions: torch.zeros_like(embeds) + + codes = predictor.forward( + torch.tensor([1]), + torch.zeros(1, 4), + torch.zeros(1, 4), + do_sample=False, + ) + + assert codes.tolist() == [[1, 2, 3]] + assert predictor.small_to_mtp_projection.call_shapes == [(1, 2, 4)] + torch.testing.assert_close( + predictor._proj_buf[0, 2], + predictor._projected_codec_embed_weight[0, 2], + ) + + filter_calls = [] + sample_calls = [] + + def apply_top_k_top_p(logits, *, p, k, top_k): + filter_calls.append((p, k, top_k)) + return logits + + def random_sample(_probs, generators): + sample_calls.append(generators) + return torch.tensor([[4]]) + + monkeypatch.setattr(module, "apply_top_k_top_p", apply_top_k_top_p) + monkeypatch.setattr(module, "random_sample", random_sample) + predictor._num_groups = 2 + predictor._wrapper_config.sampling_mode = "stored" + predictor._top_k = 2 + predictor._top_p = 0.8 + predictor._lm_heads_list = predictor._lm_heads_list[:1] + generator = torch.Generator().manual_seed(1234) + + sampled_codes = predictor.forward( + torch.tensor([1]), + torch.zeros(1, 4), + torch.zeros(1, 4), + generator=generator, + ) + + assert sampled_codes.tolist() == [[1, 4]] + assert len(filter_calls) == 1 + top_p_tensor, top_k_tensor, top_k_hint = filter_calls[0] + torch.testing.assert_close(top_p_tensor, torch.tensor([0.8])) + assert top_k_tensor.tolist() == [2] + assert top_k_hint == 2 + assert sample_calls == [{0: generator}] def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.MonkeyPatch) -> None: @@ -265,12 +420,14 @@ def test_qwen3_tts_talker_patch_uses_fp16_runtime_dtype(monkeypatch: pytest.Monk assert talker._embedding_dtype is torch.float16 assert talker._prompt_builder._embedding_dtype is torch.float16 + assert talker.talker_mtp_graph_safe is False + assert talker.talker_mtp_accepts_per_row_generators is True assert talker.load_weights([]) == {"loaded"} - assert talker.encoder.to_calls[-1] == {"dtype": torch.float16} + assert talker.encoder.to_calls[-1] == {"device": torch.device("cpu"), "dtype": torch.float32} codes = talker._encode_ref_audio_batch([np.zeros(8, dtype=np.float32)], 24000, device=torch.device("cpu")) - assert talker.encoder.last_input_dtype is torch.float16 + assert talker.encoder.last_input_dtype is torch.float32 assert len(codes) == 1 assert codes[0].dtype is torch.long assert codes[0].shape == (4, 2) @@ -311,28 +468,154 @@ def forward(self, mels): assert speaker.dtype is torch.float16 -def test_qwen3_tts_mimi_codebook_quantize_uses_cpu_fp32_cdist(monkeypatch: pytest.MonkeyPatch) -> None: - module, _ = _load_qwen3_tts_patch(monkeypatch) - real_cdist = torch.cdist - captured = {} +def test_qwen3_tts_tokenizer_npu_patch_dispatches_fused_ops(monkeypatch: pytest.MonkeyPatch) -> None: + rotary_calls = [] + rms_calls = [] + + def rotary_mul(hidden_states, cos, sin): + rotary_calls.append((hidden_states, cos, sin)) + return hidden_states + 1 + + def rms_norm(hidden_states, weight, *, epsilon): + rms_calls.append((hidden_states, weight, epsilon)) + return hidden_states * weight, None + + _install_fake_module( + monkeypatch, + "torch_npu", + npu_rotary_mul=rotary_mul, + npu_rms_norm=rms_norm, + ) + _install_fake_module(monkeypatch, "vllm") + _install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: SimpleNamespace(debug=lambda *_: None)) + + class FakeRMSNorm(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.tensor([2.0, 3.0])) + self.variance_epsilon = 1e-5 + + def forward(self, hidden_states): + return hidden_states + + def original_rope(q, k, cos, sin): + del cos, sin + return q, k + + tokenizer = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2", + Qwen3TTSTokenizerV2DecoderRMSNorm=FakeRMSNorm, + apply_rotary_pos_emb=original_rope, + ) + tokenizer_package = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz", + modeling_qwen3_tts_tokenizer_v2=tokenizer, + ) + _install_fake_module(monkeypatch, "vllm_omni") + _install_fake_module(monkeypatch, "vllm_omni.model_executor") + _install_fake_module(monkeypatch, "vllm_omni.model_executor.models") + _install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts") + monkeypatch.setitem( + sys.modules, + "vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz", + tokenizer_package, + ) + + path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_tokenizer_v2.py" + module = _load_source_module("vllm_omni_test_qwen3_tts_tokenizer_npu_patch", path) + module.apply_qwen3_tts_tokenizer_v2_patch() + + q = torch.zeros(1, 2, 3, 2) + k = torch.ones_like(q) + cos = torch.zeros(1, 3, 2) + sin = torch.ones_like(cos) + q_out, k_out = tokenizer.apply_rotary_pos_emb(q, k, cos, sin) + norm = FakeRMSNorm() + norm_out = norm(torch.ones(1, 2)) + + assert len(rotary_calls) == 2 + assert rotary_calls[0][1].shape == (1, 1, 3, 2) + torch.testing.assert_close(q_out, q + 1) + torch.testing.assert_close(k_out, k + 1) + assert len(rms_calls) == 1 + assert rms_calls[0][2] == pytest.approx(1e-5) + torch.testing.assert_close(norm_out, torch.tensor([[2.0, 3.0]])) + + +def test_qwen3_tts_code2wav_npu_patch_prepares_loaded_decoder(monkeypatch: pytest.MonkeyPatch) -> None: + linear_weights = [] + conv_weights = [] + + def maybe_trans_nz(weight): + linear_weights.append(weight) + return weight + + def format_cast(weight, fmt): + conv_weights.append((weight, fmt)) + return weight + + class FakeDecoder(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + self.conv = torch.nn.Conv1d(4, 4, 3) + self.deconv = torch.nn.ConvTranspose1d(4, 4, 4) + self.grouped_conv = torch.nn.Conv1d(4, 4, 3, groups=2) + self.cache_precompute_calls = 0 + + def precompute_snake_caches(self): + self.cache_precompute_calls += 1 - def fake_cdist(x1, x2, p=2): - captured["x1_device"] = x1.device - captured["x2_device"] = x2.device - captured["x1_dtype"] = x1.dtype - captured["x2_dtype"] = x2.dtype - return real_cdist(x1, x2, p=p) + class FakeCode2Wav: + def __init__(self, *, vllm_config, prefix=""): + self.vllm_config = vllm_config + self.prefix = prefix + self.decoder = FakeDecoder() - monkeypatch.setattr(torch, "cdist", fake_cdist) - codebook = object.__new__(module._MimiEuclideanCodebook310P) - codebook._embed = torch.tensor([[0.0, 0.0], [2.0, 0.0]], dtype=torch.float16) + def _npu_decoder_runtime_dtype(self, _device): + return torch.float16 - indices = codebook.quantize(torch.tensor([[0.1, 0.0], [1.8, 0.0]], dtype=torch.float16)) + def load_weights(self, weights): + assert list(weights) == [] + return {"loaded"} - assert captured == { - "x1_device": torch.device("cpu"), - "x2_device": torch.device("cpu"), - "x1_dtype": torch.float32, - "x2_dtype": torch.float32, + logger = SimpleNamespace(info=lambda *_: None, debug=lambda *_: None) + current_platform = SimpleNamespace(is_npu=lambda: False) + target = _install_fake_module( + monkeypatch, + "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav", + Qwen3TTSCode2Wav=FakeCode2Wav, + ) + _install_fake_module(monkeypatch, "torch_npu", npu_format_cast=format_cast) + _install_fake_module(monkeypatch, "vllm") + _install_fake_module(monkeypatch, "vllm.config", VllmConfig=object) + _install_fake_module(monkeypatch, "vllm.logger", init_logger=lambda _name: logger) + _install_fake_module(monkeypatch, "vllm_ascend") + _install_fake_module(monkeypatch, "vllm_ascend.utils", maybe_trans_nz=maybe_trans_nz) + _install_fake_module(monkeypatch, "vllm_omni") + _install_fake_module(monkeypatch, "vllm_omni.platforms", current_omni_platform=current_platform) + _install_fake_module(monkeypatch, "vllm_omni.model_executor") + _install_fake_module(monkeypatch, "vllm_omni.model_executor.models") + _install_fake_module(monkeypatch, "vllm_omni.model_executor.models.qwen3_tts") + + path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "models" / "qwen3_tts_code2wav.py" + module = _load_source_module("vllm_omni_test_qwen3_tts_code2wav_npu_patch", path) + module.apply_qwen3_tts_code2wav_patch() + + model = target.Qwen3TTSCode2Wav( + vllm_config=SimpleNamespace(device_config=SimpleNamespace(device=torch.device("cpu"))), + prefix="stage1", + ) + assert model.load_weights(iter(())) == {"loaded"} + + assert model.prefix == "stage1" + assert model.decoder.linear.weight.dtype is torch.float16 + assert [weight.data_ptr() for weight in linear_weights] == [model.decoder.linear.weight.data_ptr()] + assert {weight.data_ptr() for weight, _ in conv_weights} == { + model.decoder.conv.weight.data_ptr(), + model.decoder.deconv.weight.data_ptr(), } - assert indices.tolist() == [0, 1] + assert all(fmt == module._ACL_FORMAT_FRACTAL_Z for _, fmt in conv_weights) + assert model.decoder.cache_precompute_calls == 1 diff --git a/tests/platforms/test_cuda_ir_op_priority.py b/tests/platforms/test_cuda_ir_op_priority.py new file mode 100644 index 00000000000..8542ea63d9c --- /dev/null +++ b/tests/platforms/test_cuda_ir_op_priority.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +from vllm.config import CompilationConfig, CompilationMode, DeviceConfig, VllmConfig +from vllm.platforms import current_platform + +# Importing CudaOmniPlatform pulls in vllm.platforms.cuda, which imports the +# CUDA-only ``vllm._C_stable_libtorch`` extension at module top level. That +# extension is absent on non-CUDA builds (e.g. XPU), so skip the whole module +# there before the import can crash collection. ``is_cuda()`` resolves the +# platform without importing cuda.py. +if not current_platform.is_cuda(): + pytest.skip("CUDA-only IR op priority tests", allow_module_level=True) + +from vllm_omni.platforms.cuda.platform import CudaOmniPlatform # noqa: E402 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _vllm_config(*, backend: str = "inductor", mode: CompilationMode) -> VllmConfig: + return VllmConfig( + device_config=DeviceConfig(device="cpu"), + compilation_config=CompilationConfig(backend=backend, mode=mode), + ) + + +@pytest.mark.parametrize( + "mode", + [CompilationMode.NONE, CompilationMode.VLLM_COMPILE, CompilationMode.STOCK_TORCH_COMPILE], +) +def test_cuda_default_ir_op_priority_prefers_vllm_c_when_inductor_backend(mode: CompilationMode) -> None: + """Regression for #4964: inductor-active configs must not switch to native-only.""" + priority = CudaOmniPlatform.get_default_ir_op_priority(_vllm_config(mode=mode)) + assert priority.rms_norm == ["vllm_c", "native"] + assert priority.fused_add_rms_norm == ["vllm_c", "native"] + + +def test_cuda_default_ir_op_priority_with_oink(monkeypatch: pytest.MonkeyPatch) -> None: + import vllm.envs as envs + + monkeypatch.setattr(envs, "VLLM_USE_OINK_OPS", True) + priority = CudaOmniPlatform.get_default_ir_op_priority( + _vllm_config(mode=CompilationMode.VLLM_COMPILE), + ) + assert priority.rms_norm == ["oink", "vllm_c", "native"] + assert priority.fused_add_rms_norm == ["oink", "vllm_c", "native"] diff --git a/tests/plugin/__init__.py b/tests/plugin/__init__.py new file mode 100644 index 00000000000..9881313609a --- /dev/null +++ b/tests/plugin/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/plugin/test_gguf_diffusion.py b/tests/plugin/test_gguf_diffusion.py new file mode 100644 index 00000000000..44de4a1fb0c --- /dev/null +++ b/tests/plugin/test_gguf_diffusion.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for GGUF quantization on diffusion models. + +Validates that GGUF-quantized diffusion models generate valid images and +use less peak GPU memory than BF16 baseline. + +Requires vllm-omni to be installed alongside the plugin. + +Usage: + pytest tests/diffusion/test_gguf_memory.py -v +""" + +from __future__ import annotations + +from typing import NamedTuple + +import numpy as np +import pytest +import torch + +from tests.helpers.env import DeviceMemoryMonitor +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniRunner +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + + +class DiffusionGGUFTestConfig(NamedTuple): + artifact_prefix: str + gguf_model: str + hf_model: str + min_cosine_similarity: float = 0.95 + + +Z_IMAGE_CONFIG = DiffusionGGUFTestConfig( + artifact_prefix="zimage", + hf_model="Tongyi-MAI/Z-Image-Turbo", + gguf_model="unsloth/Z-Image-Turbo-GGUF:Q4_0", +) + + +FLUX_CONFIG = DiffusionGGUFTestConfig( + artifact_prefix="flux2_klein", + hf_model="black-forest-labs/FLUX.2-klein-4B", + gguf_model="unsloth/FLUX.2-klein-4B-GGUF:Q8_0", +) + + +def _image_cosine_similarity(hf_image, gguf_image) -> float: + hf_tensor = torch.as_tensor(np.array(hf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten() + gguf_tensor = torch.as_tensor(np.array(gguf_image.convert("RGB"), copy=True), dtype=torch.float32).flatten() + return torch.nn.functional.cosine_similarity(hf_tensor, gguf_tensor, dim=0).item() + + +def _generate_single_stage_image( + model: str, + height: int = 256, + width: int = 256, + num_inference_steps: int = 20, + seed: int = 42, + **extra_kwargs, +) -> tuple[list, float]: + """Generate an image with a single-stage diffusion model. + + Returns (images, peak_memory_gib). + """ + omni_kwargs = dict(extra_kwargs) + + memory_monitor = DeviceMemoryMonitor(device_index=0) + memory_monitor.start() + try: + with OmniRunner(model, **omni_kwargs) as runner: + generator = torch.Generator( + device=current_omni_platform.device_type, + ).manual_seed(seed) + outputs = runner.omni.generate( + "a photo of a cat sitting on a laptop keyboard", + OmniDiffusionSamplingParams( + height=height, + width=width, + num_inference_steps=num_inference_steps, + guidance_scale=0.0, + generator=generator, + ), + ) + finally: + memory_monitor.stop() + peak_mem = memory_monitor.peak_used_mb / 1024 + + first_output = outputs[0] + assert first_output.final_output_type == "image" + if hasattr(first_output, "images") and first_output.images: + images = first_output.images + else: + assert hasattr(first_output, "request_output") and first_output.request_output + request_output = first_output.request_output + if isinstance(request_output, list): + req_out = request_output[0] + else: + req_out = request_output + assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images") + images = req_out.images + assert len(images) >= 1 + assert images[0].width == width + assert images[0].height == height + + return images, peak_mem + + +@hardware_test(res={"cuda": "L4"}) +@pytest.mark.full_model +@pytest.mark.diffusion +@pytest.mark.slow +@pytest.mark.parametrize("model", [Z_IMAGE_CONFIG, FLUX_CONFIG], ids=["Z-Image-Turbo", "FLUX.2-klein"]) +def test_single_stage_diffusion_gguf(model: DiffusionGGUFTestConfig, monkeypatch: pytest.MonkeyPatch) -> None: + """Z-Image-Turbo GGUF generates valid images and uses less memory than BF16.""" + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + + # BF16 baseline + hf_images, mem_bf16 = _generate_single_stage_image( + model=model.hf_model, + ) + + # GGUF + images, mem_gguf = _generate_single_stage_image( + model=model.hf_model, + diffusion_quantization_config={ + "method": "gguf", + "gguf_model": model.gguf_model, + }, + ) + + hf_image_path = f"test_{model.artifact_prefix}_hf.png" + gguf_image_path = f"test_{model.artifact_prefix}_gguf.png" + assert len(hf_images) >= 1 + hf_images[0].save(hf_image_path) + assert len(images) >= 1 + images[0].save(gguf_image_path) + print(f"Saved HF image: {hf_image_path}") + print(f"Saved GGUF image: {gguf_image_path}") + + image_cosine_similarity = _image_cosine_similarity(hf_images[0], images[0]) + print(f"{model.artifact_prefix} image cosine similarity: {image_cosine_similarity:.4f}") + assert image_cosine_similarity >= model.min_cosine_similarity, ( + f"GGUF image cosine similarity ({image_cosine_similarity:.4f}) should be >= {model.min_cosine_similarity:.4f}" + ) + + print(f"{model.artifact_prefix} BF16 peak VRAM delta: {mem_bf16:.2f} GiB") + print(f"{model.artifact_prefix} GGUF peak VRAM delta: {mem_gguf:.2f} GiB") + reduction = (mem_bf16 - mem_gguf) / mem_bf16 * 100 + print(f"VRAM reduction: {reduction:.1f}%") + assert mem_gguf < mem_bf16, f"GGUF ({mem_gguf:.2f} GiB) should use less VRAM than BF16 ({mem_bf16:.2f} GiB)" diff --git a/tests/quantization/test_quack_fp8_gate.py b/tests/quantization/test_quack_fp8_gate.py new file mode 100644 index 00000000000..4a8848ffaf4 --- /dev/null +++ b/tests/quantization/test_quack_fp8_gate.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for the quack FP8 auto-enable gate. + +quack's CuteDSL FP8 GEMM uses the 5th-gen ``tcgen05`` tensor-core MMA, which +exists only on datacenter Blackwell (``sm_100a`` / ``sm_101a`` / ``sm_103a``, +compute capability ``10.x``). It must NOT auto-enable on workstation/consumer +Blackwell (``sm_120`` / ``sm_121``, cc ``12.x``, e.g. RTX PRO 6000 / RTX +50-series), where ``tcgen05`` is absent and every FP8 GEMM would fall back to +FlashInfer one call at a time (catastrophically slow). See +``vllm_omni/quantization/quack_fp8.py``. + +Regression guard for the ``>= 10`` gate that matched cc ``12.x`` too. These +tests are hardware-free: the ``torch.cuda`` probes are monkeypatched, so no GPU +is required. +""" + +import pytest +import torch + +from vllm_omni.quantization import quack_fp8 + +_ENV = "VLLM_OMNI_USE_QUACK_FP8" + + +def _fake_cuda(monkeypatch: pytest.MonkeyPatch, capability: tuple[int, int]) -> None: + """Pretend a CUDA device with the given (major, minor) capability is present.""" + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability) + + +@pytest.mark.parametrize( + "capability, expected", + [ + ((10, 0), True), # sm_100 datacenter Blackwell — tcgen05 present + ((10, 3), True), # sm_103 datacenter Blackwell + ((12, 0), False), # sm_120 workstation Blackwell — the regression + ((12, 1), False), # sm_121 consumer Blackwell + ((9, 0), False), # Hopper — CUTLASS already fuses bias, quack unused + ((8, 9), False), # Ada + ], +) +def test_auto_enable_only_on_datacenter_blackwell( + monkeypatch: pytest.MonkeyPatch, + capability: tuple[int, int], + expected: bool, +) -> None: + monkeypatch.delenv(_ENV, raising=False) + _fake_cuda(monkeypatch, capability) + assert quack_fp8._is_quack_capable() is expected + assert quack_fp8.quack_enabled() is expected + + +def test_no_cuda_disables_quack(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv(_ENV, raising=False) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + assert quack_fp8._is_quack_capable() is False + assert quack_fp8.quack_enabled() is False + + +@pytest.mark.parametrize("value", ["1", "true", "yes", "on", "TRUE", "On"]) +def test_env_override_forces_on_even_on_sm120(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + # sm_120 auto-disables, but an explicit truthy override forces quack on + # (e.g. once CuteDSL ships sm_120a support). + _fake_cuda(monkeypatch, (12, 0)) + monkeypatch.setenv(_ENV, value) + assert quack_fp8.quack_enabled() is True + + +@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""]) +def test_env_override_forces_off_even_on_datacenter_blackwell(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + # sm_100 auto-enables, but an explicit non-truthy override forces quack off. + _fake_cuda(monkeypatch, (10, 0)) + monkeypatch.setenv(_ENV, value) + assert quack_fp8.quack_enabled() is False diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 77bdfb39391..e083491406c 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -5,23 +5,24 @@ """ import importlib -import inspect import warnings from dataclasses import dataclass, fields from pathlib import Path from unittest.mock import patch import pytest -from transformers import PretrainedConfig +from transformers import PretrainedConfig, Qwen3OmniMoeConfig -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability +from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline, resolve_pipeline_config from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, + _DEPLOY_DIR, DeployConfig, PipelineConfig, StageConfig, - StageConfigFactory, StageDeployConfig, StageExecutionType, StagePipelineConfig, @@ -32,7 +33,7 @@ build_stage_runtime_overrides, load_deploy_config, merge_pipeline_deploy, - register_pipeline, + pipeline_cfg_resolver, strip_parent_engine_args, ) from vllm_omni.engine.arg_utils import SHARED_FIELDS, EngineArgs, internal_blacklist_keys @@ -40,15 +41,25 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -# TODO (Alex) - remove this in the next release, this is only used in the -# patch since we didn't carry the eager registry change forward yet. -@pytest.fixture -def clean_pipeline_registry(): - """Ensure the OMNI_PIPELINES are in a clean state for a test that mutates it.""" - snapshot = dict(_PIPELINE_REGISTRY) +@pytest.fixture(autouse=True) +def _stable_test_platform(monkeypatch): + from vllm_omni import platforms + + platform = platforms.current_omni_platform + monkeypatch.setattr(platform, "device_name", "cpu", raising=False) + monkeypatch.setattr(platform, "device_type", "cpu", raising=False) + + +@pytest.fixture(autouse=True) +def clear_config_factory_caches(): + """Clear cached classmethods from the StageConfigFactory to prevent test pollution.""" yield - _PIPELINE_REGISTRY.clear() - _PIPELINE_REGISTRY.update(snapshot) + StageConfigFactory.get_hf_config.cache_clear() + StageConfigFactory.try_infer_model_type.cache_clear() + + +Q3_OMNI_ALL_STAGES_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=True) +Q3_OMNI_THINKER_HF_CONFIG = Qwen3OmniMoeConfig(enable_audio_output=False) class TestStageType: @@ -522,26 +533,61 @@ def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self): class TestPipelineDiscovery: - """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``).""" + """Tests for the central pipeline registry (``OMNI_PIPELINES``).""" def test_registry_has_known_models(self): - """Built-in pipelines are lazy-loaded from the central declaration - on first access; no eager import or discovery walk needed.""" - # ``in`` triggers the lazy-map lookup without forcing a load. - assert "qwen2_5_omni" in _PIPELINE_REGISTRY - assert "qwen3_omni_moe" in _PIPELINE_REGISTRY - assert "qwen3_tts" in _PIPELINE_REGISTRY - - def test_registry_loads_pipeline_on_getitem(self): - """Looking up a registered model_type returns the matching PipelineConfig.""" - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] + """Check that specific models are in OMNI_PIPELINES.""" + assert "qwen2_5_omni" in OMNI_PIPELINES + assert "qwen3_omni_moe" in OMNI_PIPELINES + assert "qwen3_tts" in OMNI_PIPELINES + + def test_registry_resolver_qwen3_omni_all_stages(self): + """Test that providing the HF config for qwen3 omni with audio enabled uses all stages.""" + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) assert pipeline.model_type == "qwen3_omni_moe" + assert pipeline.default_deploy_config_name == "qwen3_omni_moe.yaml" assert len(pipeline.stages) == 3 # thinker + talker + code2wav + def test_registry_resolver_qwen3_omni_thinker_only(self): + """Test that providing the HF config for qwen3 omni without audio is thinker only.""" + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_THINKER_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + assert pipeline.model_type == "qwen3_omni_moe_thinker_only" + assert pipeline.default_deploy_config_name is None + assert len(pipeline.stages) == 1 # thinker only + + @pytest.mark.parametrize( + "pipeline", + [registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)], + ids=lambda pipeline: pipeline.model_type, + ) + def test_registered_default_deploy_config_exists(self, pipeline): + if pipeline.default_deploy_config_name is None: + return + assert Path(get_deploy_config_path(pipeline.default_deploy_config_name)).is_file() + + @pytest.mark.parametrize( + "deploy_path", + sorted(_DEPLOY_DIR.glob("*.yaml")), + ids=lambda deploy_path: deploy_path.name, + ) + def test_explicit_deploy_pipeline_is_registered(self, deploy_path): + deploy = load_deploy_config(deploy_path) + if deploy.pipeline is not None: + assert deploy.pipeline in OMNI_PIPELINES + def test_registry_returns_none_for_unknown(self): - """Unknown model_types aren't found; ``get()`` returns None.""" - assert "definitely_not_a_real_model" not in _PIPELINE_REGISTRY - assert _PIPELINE_REGISTRY.get("definitely_not_a_real_model") is None + """Unknown model_types aren't found and resolve to `None`.""" + assert "definitely_not_a_real_model" not in OMNI_PIPELINES + assert OMNI_PIPELINES.get("definitely_not_a_real_model") is None + assert resolve_pipeline_config("definitely_not_a_real_model") is None def test_pipeline_config_supports_hf_architectures(self): """PipelineConfig accepts hf_architectures for HF-arch fallback @@ -566,6 +612,7 @@ def test_defaults(self): assert s.final_output is False assert s.sampling_constraints == {} assert s.engine_output_type is None + assert s.scheduler_cls is None class TestPipelineConfigNew: @@ -590,43 +637,413 @@ def test_validate_no_stages(self): assert any("no stages" in e.lower() for e in p.validate()) -class TestResolveScheduler: - def test_all_execution_types_handled(self): - for et in StageExecutionType: - _resolve_scheduler(et) +class TestPipelineRegistration: + def test_resolve_pipeline_prefers_deploy_pipeline_key(self, clean_pipeline_registry, tmp_path): + deploy_key = "deploy_selected_pipeline" + model_type_key = "hf_model_type_pipeline" + deploy_pipe = PipelineConfig(model_type=deploy_key) + model_type_pipe = PipelineConfig(model_type=model_type_key) + register_pipeline(deploy_pipe) + register_pipeline(model_type_pipe) - def test_ar_sync_when_false(self): - cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=False) - assert cls is not None - assert "Async" not in cls.__name__ + deploy_path = tmp_path / "deploy.yaml" + deploy_path.write_text(f"pipeline: {deploy_key}\n", encoding="utf-8") - def test_ar_async_when_true(self): - cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=True) - assert cls is not None - assert "Async" in cls.__name__ + class FakeConfig(PretrainedConfig): + model_type = model_type_key + + fake_config = FakeConfig() + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + pipeline_cfg = StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=str(deploy_path), + ) - def test_generation(self): - cls = _resolve_scheduler(StageExecutionType.LLM_GENERATION) - assert cls is not None - assert "Generation" in cls.__name__ + assert pipeline_cfg is deploy_pipe - def test_diffusion_returns_none(self): - assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None + def test_resolve_pipeline_rejects_invalid_deploy_pipeline_key( + self, + clean_pipeline_registry, + tmp_path, + ): + model_type_key = "registered_model_type_pipeline" + register_pipeline(PipelineConfig(model_type=model_type_key)) + deploy_path = tmp_path / "deploy.yaml" + deploy_path.write_text("pipeline: missing_pipeline\n", encoding="utf-8") -class TestPipelineRegistry: - def test_register_and_lookup(self): - p = PipelineConfig( - model_type="__test_only__", - model_arch="A", - stages=(StagePipelineConfig(stage_id=0, model_stage="a"),), + class FakeConfig(PretrainedConfig): + model_type = model_type_key + + with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): + with pytest.raises(KeyError, match="missing_pipeline"): + StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=str(deploy_path), + ) + with pytest.raises(KeyError, match="missing_pipeline"): + StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=str(deploy_path), + ) + with pytest.raises(KeyError, match="missing_pipeline"): + StageConfigFactory.create_legacy_stage_configs_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=str(deploy_path), + ) + + def test_resolve_pipeline_matches_hf_architecture_fallback(self, clean_pipeline_registry): + pipeline_key = "architecture_fallback_pipeline" + pipe_cfg = PipelineConfig( + model_type=pipeline_key, + hf_architectures=("ArchitectureFallbackForTest",), ) - register_pipeline(p) - assert _PIPELINE_REGISTRY["__test_only__"] is p - del _PIPELINE_REGISTRY["__test_only__"] + register_pipeline(pipe_cfg) + + class FakeConfig(PretrainedConfig): + model_type = "unregistered_model_type" + + fake_config = FakeConfig() + fake_config.architectures = ["ArchitectureFallbackForTest"] + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + pipeline_cfg = StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=None, + ) + assert pipeline_cfg is pipe_cfg + + def test_resolve_pipeline_architecture_fallback_respects_predicates(self, clean_pipeline_registry): + shared_arch = "SharedPredicateArchitectureForTest" + + def rejecting_predicate(_hf_config): + return False + + def raising_predicate(_hf_config): + raise ValueError("predicate failure") + + reject_cfg = PipelineConfig( + model_type="rejected_by_predicate", + hf_architectures=(shared_arch,), + hf_config_predicate=rejecting_predicate, + ) + raises_cfg = PipelineConfig( + model_type="raises_in_predicate", + hf_architectures=(shared_arch,), + hf_config_predicate=raising_predicate, + ) + accept_cfg = PipelineConfig( + model_type="accepted_by_predicate", + hf_architectures=(shared_arch,), + ) + register_pipeline(reject_cfg) + register_pipeline(raises_cfg) + register_pipeline(accept_cfg) + + class FakeConfig(PretrainedConfig): + model_type = "unregistered_model_type" + + fake_config = FakeConfig() + fake_config.architectures = [shared_arch] + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + pipeline_cfg = StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=None, + ) + + assert pipeline_cfg is accept_cfg + + def test_resolve_pipeline_architecture_fallback_rejects_non_matching_architecture( + self, + clean_pipeline_registry, + ): + predicate_calls = [] + + def predicate(_hf_config): + predicate_calls.append(_hf_config) + return True + + pipe_cfg = PipelineConfig( + model_type="predicate_without_arch_match", + hf_architectures=("DifferentArchitectureForTest",), + hf_config_predicate=predicate, + ) + register_pipeline(pipe_cfg) + + class FakeConfig(PretrainedConfig): + model_type = "unregistered_model_type" + + fake_config = FakeConfig() + fake_config.architectures = ["RequestedArchitectureForTest"] + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + pipeline_cfg = StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=None, + ) + + assert pipeline_cfg is None + assert predicate_calls == [fake_config] + + def test_resolve_pipeline_architecture_fallback_resolves_callable_pipeline(self, clean_pipeline_registry): + pipeline_key = "callable_architecture_fallback" + resolved_cfg = PipelineConfig( + model_type="callable_resolved_pipeline", + hf_architectures=("CallableArchitectureForTest",), + ) + seen_hf_configs = [] + + def resolver(hf_config): + seen_hf_configs.append(hf_config) + return resolved_cfg + + register_pipeline(resolver, model_type=pipeline_key) + + class FakeConfig(PretrainedConfig): + model_type = "unregistered_model_type" + + fake_config = FakeConfig() + fake_config.architectures = ["CallableArchitectureForTest"] + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + pipeline_cfg = StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=None, + ) + + assert pipeline_cfg is resolved_cfg + assert seen_hf_configs == [fake_config] + + def test_resolve_pipeline_returns_none_without_registry_or_architecture_match(self, clean_pipeline_registry): + class FakeConfig(PretrainedConfig): + model_type = "unregistered_model_type" + + fake_config = FakeConfig() + fake_config.architectures = ["UnmatchedArchitectureForTest"] + with patch("vllm_omni.config.config_factory.get_config", return_value=fake_config): + pipeline_cfg = StageConfigFactory.get_pipeline_config( + "fake/model", + trust_remote_code=True, + deploy_config_path=None, + ) + + assert pipeline_cfg is None + + def test_create_from_model_returns_structured_omni_config(self): + class FakeConfig(PretrainedConfig): + model_type = "qwen3_tts" + + with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): + omni_config = StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=None, + ) + + assert isinstance(omni_config, VllmOmniConfig) + assert omni_config.pipeline_config is OMNI_PIPELINES["qwen3_tts"] + assert len(omni_config.stage_configs) == 2 + + def test_create_from_model_preserves_model_on_structured_diffusion_stage(self): + class FakeConfig(PretrainedConfig): + model_type = "dreamzero" + + with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): + omni_config = StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=None, + ) + + assert isinstance(omni_config, VllmOmniConfig) + assert omni_config.stage_by_id(0).diffusion_config.model == "fake/model" + + def test_pipeline_registration(self, clean_pipeline_registry): + """Ensure that we can register and create a custom pipeline config.""" + new_model_type = "new_model_type" + pipe_cfg = PipelineConfig(model_type=new_model_type) + + # Register the new PipelineConfig + assert new_model_type not in OMNI_PIPELINES + register_pipeline(pipe_cfg) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is pipe_cfg + + class FakeConfig(PretrainedConfig): + model_type = new_model_type + + # Create the model + fake_config = FakeConfig() + with ( + patch("vllm_omni.config.config_factory.get_config", return_value=fake_config), + patch.object(VllmOmniConfig, "from_pipeline_config") as mock_create, + ): + StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=None, + ) + mock_create.assert_called_once() + assert mock_create.call_args.args == (pipe_cfg,) + assert mock_create.call_args.kwargs["cli_overrides"] == { + "trust_remote_code": False, + "model": "fake/model", + } + assert pipe_cfg.model_type == new_model_type + + def test_resolver_registration(self, clean_pipeline_registry): + """Ensure that we can register and create a custom resolver.""" + new_model_type = "new_model_type" + resolved_type = "resolved_type" + + class FakeConfig(PretrainedConfig): + model_type = new_model_type + + @pipeline_cfg_resolver(config_type=FakeConfig) + def custom_resolver( + hf_config: FakeConfig, + ) -> PipelineConfig: + return PipelineConfig(model_type=resolved_type) + + # Register the new PipelineConfig + assert new_model_type not in OMNI_PIPELINES + + register_pipeline(custom_resolver, model_type=new_model_type) + assert new_model_type in OMNI_PIPELINES + assert OMNI_PIPELINES[new_model_type] is custom_resolver + + # Create the model + fake_config = FakeConfig() + with ( + patch("vllm_omni.config.config_factory.get_config", return_value=fake_config) as mock_get_config, + patch.object(VllmOmniConfig, "from_pipeline_config") as mock_create, + ): + StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=True, + cli_overrides={}, + deploy_config_path=None, + ) + mock_create.assert_called_once() + assert mock_create.call_args.args[0].model_type == resolved_type + assert mock_create.call_args.kwargs["cli_overrides"] == { + "trust_remote_code": True, + "model": "fake/model", + } + mock_get_config.assert_called_once_with("fake/model", trust_remote_code=True) + + with patch.object(StageConfigFactory, "_create_legacy_from_registry", return_value=([], None)) as mock_legacy: + StageConfigFactory.create_legacy_stage_configs_from_model( + "fake/model", + trust_remote_code=True, + cli_overrides={}, + deploy_config_path=None, + ) + + assert mock_legacy.call_args.args[0].model_type == resolved_type + assert mock_legacy.call_args.args[1]["trust_remote_code"] is True + + def test_resolve_when_autodetect_resolves_none(self): + """Regression test for: https://github.com/vllm-project/vllm-omni/issues/4726""" + deploy_path = get_deploy_config_path("ming_tts.yaml") + resolved_config = StageConfigFactory.create_from_model( + model="inclusionAI/Ming-omni-tts-0.5B", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=deploy_path, + ) + assert resolved_config is not None + assert len(resolved_config.stage_configs) > 0 + + def test_legacy_and_structured_paths_share_deploy_pipeline_override(self, clean_pipeline_registry, tmp_path): + pipeline_key = "deploy_only_pipeline" + pipe_cfg = PipelineConfig( + model_type=pipeline_key, + model_arch="DeployOnlyArch", + stages=(StagePipelineConfig(stage_id=0, model_stage="diffusion"),), + ) + register_pipeline(pipe_cfg) + + deploy_path = tmp_path / "deploy_only.yaml" + deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8") + + class FakeConfig(PretrainedConfig): + model_type = "dense" + architectures = ("UnregisteredArch",) + + with patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()): + omni_config = StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=str(deploy_path), + ) + legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=str(deploy_path), + ) + + assert omni_config is not None + assert omni_config.pipeline_config is pipe_cfg + assert legacy_configs is not None + assert len(legacy_configs) == 1 + assert legacy_configs[0].yaml_engine_args["model_arch"] == "DeployOnlyArch" + + def test_structured_path_loads_explicit_deploy_config_once(self, clean_pipeline_registry, tmp_path): + pipeline_key = "single_load_pipeline" + register_pipeline(PipelineConfig(model_type=pipeline_key)) + deploy_path = tmp_path / "single_load.yaml" + deploy_path.write_text(f"pipeline: {pipeline_key}\n", encoding="utf-8") + + class FakeConfig(PretrainedConfig): + model_type = "unregistered" + + with ( + patch("vllm_omni.config.config_factory.get_config", return_value=FakeConfig()), + patch( + "vllm_omni.config.config_factory.load_deploy_config", + wraps=load_deploy_config, + ) as mock_load, + ): + StageConfigFactory.create_from_model( + "fake/model", + trust_remote_code=False, + cli_overrides={}, + deploy_config_path=str(deploy_path), + ) + + mock_load.assert_called_once_with(deploy_path) + + def test_legacy_default_deploy_name_ignores_cwd(self, monkeypatch, tmp_path): + default_name = "pipeline_default.yaml" + (tmp_path / default_name).write_text("stages: []\n", encoding="utf-8") + monkeypatch.chdir(tmp_path) + pipeline = PipelineConfig( + model_type="pipeline_with_default", + default_deploy_config_name=default_name, + ) + + with patch( + "vllm_omni.config.config_factory.load_deploy_config", + return_value=DeployConfig(), + ) as mock_load: + StageConfigFactory._create_legacy_from_registry(pipeline, {}) + + mock_load.assert_called_once_with(_DEPLOY_DIR / default_name) -class TestEndpointRejection: def test_deploy_override_uses_correct_endpoint_restrictions(self, clean_pipeline_registry, tmp_path): """Ensure endpoint restrictions must come from the final pipeline after deploy config overrides, not the auto-detected pipeline. @@ -653,7 +1070,7 @@ class FakeConfig(PretrainedConfig): # type + patching the config for the detected type. Ensure that the endpoint # restrictions correspond to the type in the deploy config. with patch( - "vllm.transformers_utils.config.get_config", + "vllm_omni.config.config_factory.get_config", return_value=FakeConfig(), ): restrictions = StageConfigFactory.get_pipeline_endpoint_restrictions( @@ -665,6 +1082,30 @@ class FakeConfig(PretrainedConfig): assert restrictions == () +class TestResolveScheduler: + def test_all_execution_types_handled(self): + for et in StageExecutionType: + _resolve_scheduler(et) + + def test_ar_sync_when_false(self): + cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=False) + assert cls is not None + assert "Async" not in cls.__name__ + + def test_ar_async_when_true(self): + cls = _resolve_scheduler(StageExecutionType.LLM_AR, async_scheduling=True) + assert cls is not None + assert "Async" in cls.__name__ + + def test_generation(self): + cls = _resolve_scheduler(StageExecutionType.LLM_GENERATION) + assert cls is not None + assert "Generation" in cls.__name__ + + def test_diffusion_returns_none(self): + assert _resolve_scheduler(StageExecutionType.DIFFUSION) is None + + class TestDeployConfigLoading: def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): deploy_path = tmp_path / "qwen3_tts_custom_voice.yaml" @@ -683,8 +1124,9 @@ def test_custom_voice_dir_is_pipeline_wide_engine_arg(self, tmp_path): ) deploy = load_deploy_config(deploy_path) - pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + pipeline = resolve_pipeline_config("qwen3_tts") stages = merge_pipeline_deploy(pipeline, deploy) + assert isinstance(pipeline, PipelineConfig) assert deploy.custom_voice_dir == str(custom_voice_dir) assert {s.yaml_engine_args.get("custom_voice_dir") for s in stages} == {str(custom_voice_dir)} @@ -713,12 +1155,123 @@ def test_load_ming_flash_omni_deploy_config_schema_fields(self): assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}} assert "compilation_config" not in deploy.stages[0].engine_extras + def test_load_voxcpm2_deploy_config_preserves_engine_extras(self): + deploy_path = get_deploy_config_path("voxcpm2.yaml") + raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0) + expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"] + + deploy = load_deploy_config(deploy_path) + runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"] + assert runtime_config == expected_runtime_config + + @pytest.mark.parametrize( + ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"), + [ + ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"), + ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"), + ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"), + ("mimo_audio.yaml", "mimo_audio", 2, "audio", None), + ("step_audio_2.yaml", "step_audio_2", 2, "audio", None), + ("step_audio_2_asr.yaml", "step_audio_2_asr", 1, "text", "step_audio_2_asr"), + ("step_audio_2_async_chunk.yaml", "step_audio_2", 2, "audio", None), + ("hunyuan_video_15.yaml", "hunyuan_video_15", 1, "video", None), + ("wan2_2_ti2v.yaml", "wan2_2_ti2v", 1, "video", None), + ], + ) + def test_load_new_registry_backed_deploy_configs( + self, + deploy_name: str, + pipeline_name: str, + stage_count: int, + final_output_type: str, + declared_pipeline: str | None, + ): + deploy_path = Path(get_deploy_config_path(deploy_name)) + deploy = load_deploy_config(deploy_path) + assert deploy.pipeline == declared_pipeline + + with patch("vllm_omni.platforms.current_omni_platform") as platform: + platform.device_name = "cuda" + stages = merge_pipeline_deploy(OMNI_PIPELINES[pipeline_name], deploy) + assert len(stages) == stage_count + assert stages[-1].final_output is True + assert stages[-1].final_output_type == final_output_type + if deploy.trust_remote_code is not None: + assert {s.yaml_engine_args.get("trust_remote_code") for s in stages} == {deploy.trust_remote_code} + + @pytest.mark.parametrize( + ("config_json", "model_index", "expected_pipeline"), + [ + ({"model_type": "step_audio_2"}, None, "step_audio_2"), + (None, {"_class_name": "HunyuanVideo15Pipeline"}, "hunyuan_video_15"), + (None, {"_class_name": "WanPipeline"}, "wan2_2_ti2v"), + ], + ) + def test_migrated_models_are_discovered_without_explicit_deploy( + self, + config_json, + model_index, + expected_pipeline, + ): + def get_model_file(filename, _model, revision=None): + del revision + if filename == "config.json": + return config_json + if filename == "model_index.json": + return model_index + return None + + with ( + patch.object(StageConfigFactory, "get_hf_config", return_value=None), + patch("vllm_omni.config.config_factory.get_hf_file_to_dict", side_effect=get_model_file), + ): + pipeline = StageConfigFactory.get_pipeline_config( + model="/models/unrelated-checkpoint-name", + trust_remote_code=False, + ) + + assert pipeline is not None + assert pipeline.model_type == expected_pipeline + + @pytest.mark.parametrize("deploy_name", ["step_audio_2.yaml", "step_audio_2_async_chunk.yaml"]) + def test_step_audio2_deploy_configs_fit_two_gpus(self, deploy_name): + deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name))) + + assert deploy.stages[0].devices == "0,1" + assert deploy.stages[0].tensor_parallel_size == 2 + assert deploy.stages[1].devices == "1" + + def test_step_audio2_dispatches_sync_and_async_chunk_processors(self): + pipeline = resolve_pipeline_config("step_audio_2") + assert isinstance(pipeline, PipelineConfig) + + sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False)) + assert "custom_process_next_stage_input_func" not in sync_stages[0].yaml_engine_args + assert sync_stages[1].custom_process_input_func.endswith("thinker2token2wav") + + async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True)) + assert ( + async_stages[0] + .yaml_engine_args["custom_process_next_stage_input_func"] + .endswith("thinker2token2wav_async_chunk") + ) + assert async_stages[1].custom_process_input_func is None + + def test_no_bundled_legacy_stage_config_yamls(self): + stage_config_dir = Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "stage_configs" + assert not list(stage_config_dir.glob("*.yaml")) + def test_merge_pipeline_deploy(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + deploy = load_deploy_config(deploy_path) stages = merge_pipeline_deploy(pipeline, deploy) @@ -730,7 +1283,12 @@ def test_merge_pipeline_deploy(self): assert s0.yaml_extras["default_sampling_params"]["detokenize"] is True def test_merge_pipeline_deploy_preserves_num_replicas(self, tmp_path): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) + base = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not base.exists(): pytest.skip("Deploy config not found") @@ -764,6 +1322,26 @@ def test_merge_pipeline_deploy_preserves_requires_multimodal_data(self): assert stages[0].yaml_runtime["requires_multimodal_data"] is True + def test_merge_pipeline_deploy_preserves_pipeline_scheduler_cls(self): + scheduler_cls = "tests.fake.CustomScheduler" + pipeline = PipelineConfig( + model_type="test_scheduler", + model_arch="TestModel", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar", + execution_type=StageExecutionType.LLM_AR, + scheduler_cls=scheduler_cls, + ), + ), + ) + deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(stage_id=0)]) + + stages = merge_pipeline_deploy(pipeline, deploy) + + assert stages[0].scheduler_cls == scheduler_cls + def test_mixed_schema_preserves_flat_fields(self): """Ensure flat fields are not dropped when engine_args are present.""" fake_config = { @@ -892,6 +1470,54 @@ def test_mixed_engine_extras_deep_merges_dicts(self): assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3} + def test_explicit_engine_extras_merge_with_flat_and_engine_args(self): + """Explicit engine_extras should be preserved with existing pass-through styles.""" + fake_config = { + "stages": [ + { + "stage_id": 0, + "engine_extras": { + "hf_overrides": { + "runtime_config": { + "explicit_only": True, + "shared": "explicit", + "nested": {"a": 1}, + } + } + }, + "hf_overrides": { + "runtime_config": { + "flat_only": True, + "shared": "flat", + "nested": {"b": 2}, + } + }, + "engine_args": { + "hf_overrides": { + "runtime_config": { + "engine_args_only": True, + "shared": "engine_args", + "nested": {"c": 3}, + } + } + }, + }, + ] + } + + with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): + deploy = load_deploy_config("dummy.yaml") + + assert deploy.stages[0].engine_extras["hf_overrides"] == { + "runtime_config": { + "engine_args_only": True, + "explicit_only": True, + "flat_only": True, + "nested": {"a": 1, "b": 2, "c": 3}, + "shared": "engine_args", + } + } + def test_deep_merge_does_not_mutate_inputs(self): """Merging engine_args must not mutate the base stage dict.""" fake_config = { @@ -944,14 +1570,24 @@ def test_mixed_schema_engine_args_wins_scalars(self): class TestQwen3OmniPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("qwen3_omni_moe") - assert p is not None + p = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen3OmniMoeForConditionalGeneration" assert len(p.stages) == 3 assert p.validate() == [] def test_thinker(self): - s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(0) + p = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -959,29 +1595,50 @@ def test_thinker(self): assert s.sampling_constraints["detokenize"] is True def test_talker(self): - s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(1) + p = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [2150] - assert s.custom_process_input_func is not None + # thinker2talker was removed: sync_process_input_func always wins. + assert s.custom_process_input_func is None + assert s.sync_process_input_func is not None + assert s.sync_process_input_func.endswith("thinker2talker_token_only") assert s.custom_process_next_stage_input_func is not None def test_code2wav(self): - s = _PIPELINE_REGISTRY["qwen3_omni_moe"].get_stage(2) + p = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(p, PipelineConfig) + + s = p.get_stage(2) + assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" - assert s.custom_process_input_func is not None + assert s.engine_output_type == "audio" class TestQwen2_5OmniPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("qwen2_5_omni") - assert p is not None + p = resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) assert p.model_arch == "Qwen2_5OmniForConditionalGeneration" assert len(p.stages) == 3 assert p.validate() == [] def test_thinker(self): - s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(0) + p = resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -989,13 +1646,25 @@ def test_thinker(self): assert s.requires_multimodal_data is True def test_talker(self): - s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(1) + p = resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.input_sources == (0,) assert s.sampling_constraints["stop_token_ids"] == [8294] - assert s.custom_process_input_func is not None + # thinker2talker was removed: qwen2_5_omni has no async_chunk support, + # so sync_process_input_func always wins and custom_process_input_func + # was dead code. + assert s.custom_process_input_func is None + assert s.sync_process_input_func is not None def test_code2wav(self): - s = _PIPELINE_REGISTRY["qwen2_5_omni"].get_stage(2) + p = resolve_pipeline_config("qwen2_5_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(2) + assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" @@ -1003,14 +1672,19 @@ def test_code2wav(self): class TestQwen3TTSPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("qwen3_tts") + p = resolve_pipeline_config("qwen3_tts") + assert isinstance(p, PipelineConfig) assert p is not None assert p.model_arch == "Qwen3TTSTalkerForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_talker_stage(self): - s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(0) + p = resolve_pipeline_config("qwen3_tts") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "qwen3_tts" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -1020,7 +1694,11 @@ def test_talker_stage(self): assert s.model_arch is None def test_code2wav_stage_has_per_stage_model_arch(self): - s = _PIPELINE_REGISTRY["qwen3_tts"].get_stage(1) + p = resolve_pipeline_config("qwen3_tts") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.final_output_type == "audio" assert s.engine_output_type == "audio" @@ -1036,7 +1714,7 @@ def test_per_stage_model_arch_flows_through_merge(self, tmp_path): pytest.skip("qwen3_tts deploy yaml not found") deploy = load_deploy_config(deploy_path) - pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + pipeline = OMNI_PIPELINES["qwen3_tts"] stages = merge_pipeline_deploy(pipeline, deploy) # Stage 0 inherits pipeline-level model_arch @@ -1075,14 +1753,18 @@ def test_subtalker_sampling_params_deep_merge_preserves_base_keys(self): class TestMingFlashOmniPipeline: def test_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni") - assert p is not None + p = resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_thinker_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(0) + p = resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.owns_tokenizer is True @@ -1092,7 +1774,11 @@ def test_thinker_stage(self): assert s.sampling_constraints["detokenize"] is True def test_talker_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1) + p = resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "ming_tts" assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.input_sources == (0,) @@ -1102,28 +1788,40 @@ def test_talker_stage(self): # Per-stage model_arch override (Ming talker has its own self-contained LLM) assert s.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert s.tokenizer_subdir == "talker/llm" - assert s.custom_process_input_func is not None + # thinker2talker was removed: ming_flash_omni has no async_chunk support + # and both thinker2talker / thinker2talker_token_only called _build_talker_inputs + # identically, so custom_process_input_func was dead code. + assert s.custom_process_input_func is None + assert s.sync_process_input_func is not None def test_talker_stage_processor_wiring_resolves(self): - """The custom_process_input_func string must point to a real callable. + """The sync_process_input_func string must point to a real callable. Lazy string references only fail at first inference otherwise — this catches typos in the pipeline declaration at import / registration time. """ - s = _PIPELINE_REGISTRY["ming_flash_omni"].get_stage(1) - module_path, _, attr = s.custom_process_input_func.rpartition(".") + p = resolve_pipeline_config("ming_flash_omni") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(1) + assert isinstance(s, StagePipelineConfig) + module_path, _, attr = s.sync_process_input_func.rpartition(".") module = importlib.import_module(module_path) assert callable(getattr(module, attr)) def test_tts_pipeline_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni_tts") - assert p is not None + p = resolve_pipeline_config("ming_flash_omni_tts") + assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniTalkerForConditionalGeneration" assert len(p.stages) == 1 assert p.validate() == [] def test_tts_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_tts"].get_stage(0) + p = resolve_pipeline_config("ming_flash_omni_tts") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "ming_tts" assert s.execution_type == StageExecutionType.LLM_GENERATION assert s.input_sources == () @@ -1146,7 +1844,8 @@ def test_full_yaml_loads_and_merges(self): # We won't test stage 0/1 colocation contract here, # as there could exist more variant of custom device setup - pipeline = _PIPELINE_REGISTRY["ming_flash_omni"] + pipeline = resolve_pipeline_config("ming_flash_omni") + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 2 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" @@ -1162,20 +1861,24 @@ def test_tts_yaml_loads_and_merges(self): assert len(deploy.stages) == 1 assert deploy.pipeline == "ming_flash_omni_tts" - pipeline = _PIPELINE_REGISTRY["ming_flash_omni_tts"] + pipeline = OMNI_PIPELINES["ming_flash_omni_tts"] stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniTalkerForConditionalGeneration" def test_thinker_only_pipeline_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni_thinker_only") - assert p is not None + p = resolve_pipeline_config("ming_flash_omni_thinker_only") + assert isinstance(p, PipelineConfig) assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 1 assert p.validate() == [] def test_thinker_only_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"].get_stage(0) + p = resolve_pipeline_config("ming_flash_omni_thinker_only") + assert isinstance(p, PipelineConfig) + + s = p.get_stage(0) + assert isinstance(s, StagePipelineConfig) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () @@ -1196,20 +1899,21 @@ def test_thinker_only_yaml_loads_and_merges(self): assert len(deploy.stages) == 1 assert deploy.pipeline == "ming_flash_omni_thinker_only" - pipeline = _PIPELINE_REGISTRY["ming_flash_omni_thinker_only"] + pipeline = resolve_pipeline_config("ming_flash_omni_thinker_only") + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 1 assert stages[0].yaml_engine_args["model_arch"] == "MingFlashOmniForConditionalGeneration" def test_image_pipeline_registered(self): - p = _PIPELINE_REGISTRY.get("ming_flash_omni_image") + p = OMNI_PIPELINES.get("ming_flash_omni_image") assert p is not None assert p.model_arch == "MingFlashOmniForConditionalGeneration" assert len(p.stages) == 2 assert p.validate() == [] def test_image_thinker_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0) + s = resolve_pipeline_config("ming_flash_omni_image").get_stage(0) assert s.model_stage == "thinker" assert s.execution_type == StageExecutionType.LLM_AR assert s.input_sources == () @@ -1223,7 +1927,7 @@ def test_image_thinker_stage(self): assert s.prompt_expand_func is not None def test_image_dit_stage(self): - s = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1) + s = resolve_pipeline_config("ming_flash_omni_image").get_stage(1) assert s.model_stage == "dit" assert s.execution_type == StageExecutionType.DIFFUSION assert s.input_sources == (0,) @@ -1235,8 +1939,11 @@ def test_image_dit_stage(self): def test_image_processor_wiring_resolves(self): """The prompt_expand_func and custom_process_input_func strings must point to real callables.""" - thinker = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(0) - dit = _PIPELINE_REGISTRY["ming_flash_omni_image"].get_stage(1) + pipeline = resolve_pipeline_config("ming_flash_omni_image") + assert isinstance(pipeline, PipelineConfig) + + thinker = pipeline.get_stage(0) + dit = pipeline.get_stage(1) for ref in (thinker.prompt_expand_func, dit.custom_process_input_func): module_path, _, attr = ref.rpartition(".") module = importlib.import_module(module_path) @@ -1255,7 +1962,8 @@ def test_image_yaml_loads_and_merges(self): assert deploy.connectors is not None assert "shared_memory_connector" in deploy.connectors - pipeline = _PIPELINE_REGISTRY["ming_flash_omni_image"] + pipeline = resolve_pipeline_config("ming_flash_omni_image") + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) assert len(stages) == 2 # Stage 0 thinker: AR worker that emits latents. @@ -1406,12 +2114,16 @@ class TestCLIOverrideFlow: """Test --stage-overrides JSON merge into StageConfig.""" def test_stage_overrides_merge(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) # Simulate --stage-overrides '{"0": {"gpu_memory_utilization": 0.5}}' @@ -1420,12 +2132,16 @@ def test_stage_overrides_merge(self): assert stages[0].runtime_overrides["gpu_memory_utilization"] == 0.5 def test_global_override_applies_to_all(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) overrides = {"enforce_eager": True} @@ -1434,15 +2150,242 @@ def test_global_override_applies_to_all(self): assert s.runtime_overrides["enforce_eager"] is True +class TestAuraOmniDeploy: + def test_aura_omni_deploy_forces_pipeline_override(self): + deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml" + deploy = load_deploy_config(deploy_path) + + assert deploy.pipeline == "aura_omni" + + def test_aura_omni_deploy_resolves_four_native_stages(self): + pipeline_cfg = resolve_pipeline_config("aura_omni") + + stages, _ = StageConfigFactory._create_legacy_from_registry( + pipeline_cfg, + cli_overrides={}, + deploy_config_path=str(Path(__file__).parent.parent / "vllm_omni" / "deploy" / "aura_omni.yaml"), + ) + + assert [stage.model_stage for stage in stages] == [ + "asr", + "aura", + "qwen3_tts", + "code2wav", + ] + assert [stage.final_output for stage in stages] == [False, True, False, True] + assert [stage.final_output_type for stage in stages] == [None, "text", None, "audio"] + assert stages[0].yaml_engine_args["model_arch"] == "Qwen3ASRForConditionalGeneration" + assert stages[1].yaml_engine_args["model_arch"] == "AuraQwen3VLForConditionalGeneration" + assert stages[2].yaml_engine_args["model_arch"] == "Qwen3TTSTalkerForConditionalGeneration" + assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" + + +class TestDeployCliOverrideFlow: + """Test deploy-YAML baselines overridden by CLI runtime overrides.""" + + def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "diffusion_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "4,5" + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + engine_args: + cache_backend: cache_dit + diffusion_attention_backend: FLASH_ATTN + diffusion_kv_cache_dtype: auto + step_execution: false + vae_use_tiling: false + enable_cpu_offload: false + max_generated_image_size: 1048576 + tts_max_instructions_length: 1000 +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_diffusion", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False + assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1 + assert stage.yaml_engine_args["cache_backend"] == "cache_dit" + assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False + assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1 + assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1 + assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN" + assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto" + assert stage.yaml_engine_args["step_execution"] is False + assert stage.yaml_engine_args["vae_use_tiling"] is False + assert stage.yaml_engine_args["enable_cpu_offload"] is False + assert stage.yaml_engine_args["max_generated_image_size"] == 1048576 + assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000 + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "pipeline_parallel_size": 2, + "data_parallel_size": 3, + "tensor_parallel_size": 4, + "enable_expert_parallel": True, + "sequence_parallel_size": 24, + "ulysses_degree": 2, + "ring_degree": 4, + "cfg_parallel_size": 2, + "vae_patch_parallel_size": 2, + "cache_backend": "tea_cache", + "use_hsdp": True, + "hsdp_shard_size": 8, + "hsdp_replicate_size": 2, + "diffusion_attention_backend": "SAGE_ATTN", + "diffusion_kv_cache_dtype": "fp8", + "step_execution": True, + "vae_use_tiling": True, + "enable_cpu_offload": True, + "max_generated_image_size": 2097152, + "tts_max_instructions_length": 2000, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.cache_backend == "tea_cache" + assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN" + assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8" + assert omega_config.engine_args.step_execution is True + assert omega_config.engine_args.vae_use_tiling is True + assert omega_config.engine_args.enable_cpu_offload is True + assert omega_config.engine_args.max_generated_image_size == 2097152 + assert omega_config.engine_args.tts_max_instructions_length == 2000 + assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2 + assert omega_config.engine_args.parallel_config.data_parallel_size == 3 + assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4 + assert omega_config.engine_args.parallel_config.enable_expert_parallel is True + assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24 + assert omega_config.engine_args.parallel_config.ulysses_degree == 2 + assert omega_config.engine_args.parallel_config.ring_degree == 4 + assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2 + assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2 + assert omega_config.engine_args.parallel_config.use_hsdp is True + assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8 + assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2 + + def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "llm_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "0,1" + engine_args: + tensor_parallel_size: 1 + enable_expert_parallel: false + gpu_memory_utilization: 0.5 + max_num_seqs: 16 + max_num_batched_tokens: 1024 + max_model_len: 4096 + enforce_eager: false +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_llm", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + final_output=True, + final_output_type="text", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["enable_expert_parallel"] is False + assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5 + assert stage.yaml_engine_args["max_num_seqs"] == 16 + assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024 + assert stage.yaml_engine_args["max_model_len"] == 4096 + assert stage.yaml_engine_args["enforce_eager"] is False + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "tensor_parallel_size": 2, + "enable_expert_parallel": True, + "gpu_memory_utilization": 0.9, + "max_num_seqs": 32, + "max_num_batched_tokens": 2048, + "max_model_len": 8192, + "enforce_eager": True, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.tensor_parallel_size == 2 + assert omega_config.engine_args.enable_expert_parallel is True + assert omega_config.engine_args.gpu_memory_utilization == 0.9 + assert omega_config.engine_args.max_num_seqs == 32 + assert omega_config.engine_args.max_num_batched_tokens == 2048 + assert omega_config.engine_args.max_model_len == 8192 + assert omega_config.engine_args.enforce_eager is True + + class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" def _stages(self, cli_overrides): - return StageConfigFactory._create_from_registry( - "qwen3_omni_moe", + model_type = "qwen3_omni_moe" + pipeline_cfg = resolve_pipeline_config( + model_type, + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + stages, _ = StageConfigFactory._create_legacy_from_registry( + pipeline_cfg, cli_overrides=cli_overrides, ) + return stages def test_typed_kwarg_overrides_yaml(self): stages = self._stages({"max_num_seqs": 999}) @@ -1496,16 +2439,12 @@ def test_omni_with_vars_args_anti_pattern_is_safe(self): for stage in stages: assert stage.runtime_overrides == {} - def test_create_from_registry_no_cli_explicit_keys_param(self): - sig = inspect.signature(StageConfigFactory._create_from_registry) - named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD] - assert "cli_explicit_keys" not in {p.name for p in named} - def test_async_chunk_dispatches_processors(self): """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end processors based on ``deploy.async_chunk``, without needing a separate variant pipeline registration.""" - pipeline = _PIPELINE_REGISTRY["qwen3_tts"] + pipeline = resolve_pipeline_config("qwen3_tts") + assert isinstance(pipeline, PipelineConfig) # async_chunk=True → stage 0's per-chunk processor wires up, stage 1 # has no sync input processor. @@ -1567,24 +2506,23 @@ def test_async_chunk_dispatches_qwen3_omni_processors(self): ) def test_ming_flash_omni_topology(self): - """Guard ming_flash_omni's PR3 cleanup: stage 0 has no full-payload - producer hook (the connector path was removed as fake -- arch is not - in ``_FULL_PAYLOAD_INPUT_STAGES``), and stage 1 still wires the - legacy ``thinker2talker`` (custom_process_input_func) plus the - ``thinker2talker_token_only`` placeholder (sync_process_input_func). + """Guard ming_flash_omni's SIP cleanup: stage 0 has no full-payload + producer hook (arch is not in ``_FULL_PAYLOAD_INPUT_STAGES``), and + stage 1 uses only ``thinker2talker_token_only`` (sync_process_input_func). + The dead ``thinker2talker`` (custom_process_input_func) was removed + because ming_flash_omni has no async_chunk support and both functions + called ``_build_talker_inputs`` identically. Merge under either async_chunk mode must not re-introduce a stage-0 full-payload hook.""" - from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy - - pipeline = _PIPELINE_REGISTRY["ming_flash_omni"] + pipeline = resolve_pipeline_config("ming_flash_omni") + assert isinstance(pipeline, PipelineConfig) stage0, stage1 = pipeline.stages assert stage0.custom_process_next_stage_input_func is None, ( "ming_flash_omni stage 0 must not declare a full-payload producer " "(connector path is not active for this arch)." ) - assert stage1.custom_process_input_func is not None - assert stage1.custom_process_input_func.endswith("thinker2talker") + assert stage1.custom_process_input_func is None assert stage1.sync_process_input_func is not None assert stage1.sync_process_input_func.endswith("thinker2talker_token_only") @@ -1609,15 +2547,30 @@ class TestSamplingConstraintsPrecedence: """Test that pipeline sampling_constraints override deploy defaults.""" def test_constraints_win(self): - pipeline = _PIPELINE_REGISTRY["qwen3_omni_moe"] deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): pytest.skip("Deploy config not found") deploy = load_deploy_config(deploy_path) + pipeline = resolve_pipeline_config( + "qwen3_omni_moe", + Q3_OMNI_ALL_STAGES_HF_CONFIG, + ) + assert isinstance(pipeline, PipelineConfig) stages = merge_pipeline_deploy(pipeline, deploy) # Pipeline says detokenize=True for thinker, deploy can't override assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True # Pipeline says stop_token_ids=[2150] for talker assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150] + + +class TestPipelineConfigResolvers: + @pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)]) + def test_all_resolvers_reject_bad_types(self, resolver): + """Ensure that all resolvers registered reject incorrect config types.""" + + class NotTheRightHfConfig(PretrainedConfig): + pass + + assert resolver(NotTheRightHfConfig()) is None diff --git a/tests/test_diffusion_config_propagation.py b/tests/test_diffusion_config_propagation.py index 0991c6601d3..10fe001688d 100644 --- a/tests/test_diffusion_config_propagation.py +++ b/tests/test_diffusion_config_propagation.py @@ -10,7 +10,7 @@ import pytest import torch -from vllm_omni.config.stage_config import StageConfigFactory +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.diffusion.data import ( DiffusionParallelConfig, OmniDiffusionConfig, diff --git a/tests/test_omni_request.py b/tests/test_omni_request.py new file mode 100644 index 00000000000..09f893a0399 --- /dev/null +++ b/tests/test_omni_request.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for OmniRequest substitutability with the base Request. + +vllm-omni rebinds ``vllm.v1.request.Request`` to ``OmniRequest`` at import +time (see ``vllm_omni/patch.py``). vLLM core constructs ``Request`` positionally +in some paths (notably ``vllm/v1/worker/gpu/warmup.py::warmup_kernels`` for V2 +model-runner architectures like Qwen3ForCausalLM). The omni-specific params +must therefore be keyword-only so positional construction stays Liskov- +substitutable with the base class — including base-style calls that pass +``prompt_embeds`` (itself a positional-capable base param) positionally. +""" + +import inspect + +import numpy as np +import pytest +import torch +from vllm.sampling_params import SamplingParams +from vllm.v1.request import Request + +from vllm_omni.engine import PromptEmbedsPayload +from vllm_omni.request import OmniRequest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_omni_params_are_keyword_only(): + """The three omni params must be keyword-only after ``*args``. + + Guards against re-introducing the import-time rebind bug where omni params + came first positionally and broke positional ``Request(...)`` construction. + """ + params = inspect.signature(OmniRequest.__init__).parameters + for name in ("prompt_embeds", "external_req_id", "additional_information"): + assert params[name].kind is inspect.Parameter.KEYWORD_ONLY, name + + +def test_positional_construction_matches_base_request(): + """Reproduces the warmup_kernels call: Request(id, tokens, sp, pp).""" + req = OmniRequest("req-0", [1, 2, 3], SamplingParams(), None) + + assert isinstance(req, Request) + assert req.request_id == "req-0" + assert req.prompt_token_ids == [1, 2, 3] + # Omni params default cleanly when constructed positionally. + assert req.external_req_id is None + assert req.additional_information is None + assert req.prompt_embeds_payload is None + + +def test_positional_prompt_embeds_does_not_collide(): + """Base-style positional call that includes ``prompt_embeds``. + + ``prompt_embeds`` is a positional-capable base param (after + ``arrival_time``). A base-style call must not trip + ``got multiple values for argument 'prompt_embeds'`` — the subclass must + not also inject its keyword override when the value arrived positionally. + Positional order: request_id, prompt_token_ids, sampling_params, + pooling_params, client_index, arrival_time, prompt_embeds. + """ + embeds = torch.arange(6, dtype=torch.float32).reshape(2, 3) + req = OmniRequest("req-pe", [1, 2], SamplingParams(), None, 0, None, embeds) + + assert isinstance(req, Request) + assert torch.equal(req.prompt_embeds, embeds) + # Positional tensor is not a serialized payload. + assert req.prompt_embeds_payload is None + assert req.external_req_id is None + + +def test_keyword_omni_params_round_trip(): + """Keyword omni params are preserved; serialized embeds are decoded.""" + arr = np.arange(6, dtype=np.float32).reshape(2, 3) + payload = PromptEmbedsPayload(data=arr.tobytes(), shape=[2, 3], dtype="float32") + + req = OmniRequest( + request_id="req-1", + prompt_token_ids=[7, 8], + sampling_params=SamplingParams(), + pooling_params=None, + prompt_embeds=payload, + external_req_id="ext-1", + ) + + assert req.external_req_id == "ext-1" + # The serialized payload is retained and decoded into a tensor on the base. + assert req.prompt_embeds_payload is payload + assert torch.equal(req.prompt_embeds, torch.from_numpy(arr)) diff --git a/tests/test_version.py b/tests/test_version.py index 07e622a7d15..911a79611c3 100644 --- a/tests/test_version.py +++ b/tests/test_version.py @@ -56,3 +56,26 @@ def test_warning_contains_version_strings(): msg = str(record[0].message) assert "0.19.0rc2.dev21" in msg assert "0.18.0" in msg + + +def test_omni_plugin_loader_invokes_vllm_general_plugins_first(monkeypatch): + import vllm.plugins as vllm_plugins + + import vllm_omni.plugins as omni_plugins + + calls = [] + monkeypatch.setattr(omni_plugins, "omni_plugins_loaded", False) + monkeypatch.setattr(vllm_plugins, "load_general_plugins", lambda: calls.append("vllm")) + monkeypatch.setattr( + omni_plugins, + "load_omni_plugins_by_group", + lambda group: {"omni": lambda: calls.append(group)}, + ) + + omni_plugins.load_omni_general_plugins() + + assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP] + + omni_plugins.load_omni_general_plugins() + + assert calls == ["vllm", omni_plugins.OMNI_DEFAULT_PLUGINS_GROUP] diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py old mode 100644 new mode 100755 index 33be067db3f..b92b0095b5c --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py @@ -7,6 +7,7 @@ """ import asyncio +import json import time from collections import OrderedDict, defaultdict from types import SimpleNamespace @@ -19,7 +20,7 @@ class TestHiggsAudioV3Config: def test_default_config_loads(self): - from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -34,7 +35,7 @@ def test_default_config_loads(self): 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 ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -51,7 +52,7 @@ def test_custom_audio_encoder_config(self): 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 ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -59,7 +60,7 @@ def test_invalid_num_codebooks_rejected(self): 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 ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -67,7 +68,7 @@ def test_negative_num_codebooks_rejected(self): 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 ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -77,7 +78,7 @@ def test_special_token_ids_initially_none(self): 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 ( + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) @@ -797,12 +798,12 @@ def test_mixed_prefill_decode_receives_feedback_only_on_decode_rows(self): 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, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) config = HiggsAudioV3Config() c2w = HiggsAudioV3Code2Wav(config=config) @@ -821,12 +822,12 @@ def test_bundled_missing_quantizer_key_raises(self): 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, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) config = HiggsAudioV3Config() c2w = HiggsAudioV3Code2Wav(config=config) @@ -844,12 +845,12 @@ def test_bundled_missing_fc2_raises(self): 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, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) class FakeDAC(torch.nn.Module): def __init__(self): @@ -871,12 +872,12 @@ def __init__(self): 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, ) + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( + HiggsAudioV3Config, + ) class FakeDAC(torch.nn.Module): def __init__(self): @@ -1199,9 +1200,9 @@ def test_code2wav_registered(self): assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS def test_pipeline_registered(self): - from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES + from vllm_omni.config.pipeline_registry import OMNI_PIPELINES - assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES + assert "higgs_multimodal_qwen3" in OMNI_PIPELINES def test_deploy_yaml_exists(self): import os @@ -1309,7 +1310,10 @@ def test_audio_tokenizer_dir_env_accepts_parent_or_subdir(self, tmp_path, monkey parent = tmp_path / "OmniVoice" subdir = parent / "audio_tokenizer" subdir.mkdir(parents=True) - (subdir / "config.json").write_text("{}", encoding="utf-8") + (subdir / "config.json").write_text( + json.dumps({"model_type": "higgs_audio_v2_tokenizer"}), + encoding="utf-8", + ) monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent)) assert tok._resolve_audio_tokenizer_dir() == str(subdir) @@ -1317,6 +1321,32 @@ def test_audio_tokenizer_dir_env_accepts_parent_or_subdir(self, tmp_path, monkey monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir)) assert tok._resolve_audio_tokenizer_dir() == str(subdir) + def test_audio_tokenizer_dir_rejects_omnivoice_config(self, tmp_path, monkeypatch): + import huggingface_hub + import huggingface_hub.constants as hub_constants + + 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( + json.dumps({"model_type": "omnivoice"}), + encoding="utf-8", + ) + + monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent)) + monkeypatch.setenv("HIGGS_AUDIO_V2_TOKENIZER_PATH", "") + monkeypatch.setattr( + huggingface_hub, + "try_to_load_from_cache", + lambda **_: str(subdir / "config.json"), + ) + monkeypatch.setattr(hub_constants, "HF_HUB_CACHE", str(tmp_path / "empty_cache")) + + assert tok._normalize_audio_tokenizer_dir(str(parent)) is None + assert tok._resolve_audio_tokenizer_dir() is None + def test_higgs_v3_ref_code_cache_returns_clone(self): from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech diff --git a/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py b/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py new file mode 100644 index 00000000000..77b68ca3f1f --- /dev/null +++ b/tests/unit/mammoth_moda2/test_mammoth_moda2_config.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Unit tests for MammothModa2 config classes. + +These tests verify the correct behavior of the Mammothmoda2Config class, +particularly the get_text_config() method which is called during parent +class initialization (PretrainedConfig.__init__ -> validation -> get_text_config). + +Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config' +""" + +import pytest + +from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( + Mammothmoda2Config, + Mammothmoda2Qwen2_5_VLConfig, + Mammothmoda2Qwen2_5_VLTextConfig, +) + + +@pytest.mark.cpu +class TestMammothmoda2Config: + """Tests for Mammothmoda2Config class.""" + + def test_get_text_config_returns_nested_text_config(self): + """ + Verify that get_text_config() returns the nested text_config from llm_config. + + This is a regression test for the AttributeError that occurred when + get_text_config() was called during parent class initialization before + llm_config was set. + """ + config = Mammothmoda2Config( + llm_config={ + "model_type": "mammothmoda2_qwen2_5_vl", + } + ) + + assert config.llm_config is not None + assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig) + + text_config = config.get_text_config() + assert text_config is not None + assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig) + + # The text_config should be the same object as llm_config.text_config + assert text_config is config.llm_config.text_config + + def test_get_text_config_returns_none_when_llm_config_is_none(self): + """Verify that get_text_config() returns None when llm_config is None.""" + config = Mammothmoda2Config(llm_config=None) + + assert config.llm_config is None + assert config.get_text_config() is None + + def test_init_sets_llm_config_before_parent_init(self): + """ + Verify that llm_config is set before parent __init__ is called. + + This is critical because the parent PretrainedConfig.__init__ calls + validation which in turn calls get_text_config(). If llm_config is not + set before super().__init__(), an AttributeError will be raised. + """ + # This should not raise an AttributeError + config = Mammothmoda2Config( + llm_config={ + "model_type": "mammothmoda2_qwen2_5_vl", + } + ) + assert config.llm_config is not None + + +@pytest.mark.cpu +class TestMammothmoda2Qwen2_5_VLConfig: + """Tests for Mammothmoda2Qwen2_5_VLConfig class.""" + + def test_has_text_config_attribute(self): + """Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute.""" + config = Mammothmoda2Qwen2_5_VLConfig() + assert hasattr(config, "text_config") + assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig) + + def test_text_config_has_expected_attributes(self): + """Verify that the nested text_config has expected attributes for token validation.""" + config = Mammothmoda2Qwen2_5_VLConfig() + text_config = config.text_config + + assert hasattr(text_config, "vocab_size") + assert hasattr(text_config, "image_token_id") + assert hasattr(text_config, "video_token_id") diff --git a/tests/utils/test_forced_aligner.py b/tests/utils/test_forced_aligner.py new file mode 100644 index 00000000000..4964f3900c5 --- /dev/null +++ b/tests/utils/test_forced_aligner.py @@ -0,0 +1,165 @@ +import numpy as np +import pytest + +from vllm_omni.utils import forced_aligner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_decode_timestamps_maps_boundary_bins_to_words(): + logits = np.zeros((4, 5), dtype=np.float32) + logits[0, 0] = 1.0 + logits[1, 2] = 1.0 + logits[2, 2] = 1.0 + logits[3, 4] = 1.0 + + timestamps = forced_aligner._decode_timestamps( + logits=logits, + words=["hello", "world"], + timestamp_positions=[0, 1, 2, 3], + classify_num=5, + audio_duration_ms=1000, + ) + + assert timestamps == [ + forced_aligner.WordTimestamp("hello", 0, 400), + forced_aligner.WordTimestamp("world", 400, 800), + ] + + +def test_decode_timestamps_repairs_non_monotonic_bins(): + # Word 1's end bin (1) dips below its start bin (2); fix_timestamp should + # snap it back so the frame never carries end_ms < start_ms. + logits = np.zeros((4, 5), dtype=np.float32) + logits[0, 0] = 1.0 # word0 start -> bin 0 + logits[1, 2] = 1.0 # word0 end -> bin 2 + logits[2, 1] = 1.0 # word1 start -> bin 1 (out of order) + logits[3, 3] = 1.0 # word1 end -> bin 3 + + timestamps = forced_aligner._decode_timestamps( + logits=logits, + words=["hello", "world"], + timestamp_positions=[0, 1, 2, 3], + classify_num=5, + timestamp_segment_time_ms=200, + audio_duration_ms=1000, + ) + + for ts in timestamps: + assert ts.end_ms >= ts.start_ms + # bins [0, 2, 1, 3] -> repaired [0, 2, 2, 3] -> ms x200 + assert timestamps == [ + forced_aligner.WordTimestamp("hello", 0, 400), + forced_aligner.WordTimestamp("world", 400, 600), + ] + + +def test_decode_timestamps_clamps_to_audio_duration_and_stays_ordered(): + # Bins run past the 1000 ms audio; the decoder must cap each bound at the + # audio length and keep words non-overlapping (start <= end, monotonic). + logits = np.zeros((4, 5), dtype=np.float32) + logits[0, 0] = 1.0 # word0 start -> bin 0 + logits[1, 3] = 1.0 # word0 end -> bin 3 (1200 ms, past audio) + logits[2, 3] = 1.0 # word1 start -> bin 3 (1200 ms, past audio) + logits[3, 4] = 1.0 # word1 end -> bin 4 (1600 ms, past audio) + + timestamps = forced_aligner._decode_timestamps( + logits=logits, + words=["a", "b"], + timestamp_positions=[0, 1, 2, 3], + classify_num=5, + timestamp_segment_time_ms=400, + audio_duration_ms=1000, + ) + + assert timestamps == [ + forced_aligner.WordTimestamp("a", 0, 1000), + forced_aligner.WordTimestamp("b", 1000, 1000), + ] + + +def test_decode_timestamps_rejects_marker_count_mismatch(): + logits = np.zeros((2, 5), dtype=np.float32) + + timestamps = forced_aligner._decode_timestamps( + logits=logits, + words=["hello", "world"], + timestamp_positions=[0, 1], + classify_num=5, + audio_duration_ms=1000, + ) + + assert timestamps == [] + + +def test_build_config_from_yaml(tmp_path): + cfg = tmp_path / "forced_aligner.yaml" + cfg.write_text( + """ +forced_aligner: + model: Qwen/Qwen3-ForcedAligner-0.6B + gpu_memory_utilization: 0.42 + dtype: float16 + max_model_len: 2048 + trust_remote_code: false +""", + encoding="utf-8", + ) + args = type("Args", (), {"forced_aligner": None, "forced_aligner_config": str(cfg)})() + + out = forced_aligner.build_forced_aligner_config(args) + + assert out == forced_aligner.ForcedAlignerConfig( + model="Qwen/Qwen3-ForcedAligner-0.6B", + runner="pooling", + architecture="Qwen3ASRForcedAlignerForTokenClassification", + pooling_task="token_classify", + gpu_memory_utilization=0.42, + dtype="float16", + max_model_len=2048, + trust_remote_code=False, + ) + + +def test_build_config_cli_model_overrides_yaml(tmp_path): + cfg = tmp_path / "forced_aligner.yaml" + cfg.write_text( + "forced_aligner:\n model: old\n gpu_memory_utilization: 0.2\n dtype: float16\n", + encoding="utf-8", + ) + args = type( + "Args", + (), + { + "forced_aligner": "new", + "forced_aligner_config": str(cfg), + }, + )() + + out = forced_aligner.build_forced_aligner_config(args) + + assert out is not None + # --forced-aligner overrides the YAML model; gpu_memory_utilization/dtype + # come from the user YAML (there is no longer a CLI flag for gpu mem). + assert out.model == "new" + assert out.gpu_memory_utilization == 0.2 + assert out.dtype == "float16" + assert out.runner == "pooling" + + +def test_build_config_from_cli_model_uses_default_yaml(): + args = type( + "Args", + (), + { + "forced_aligner": "local-aligner", + "forced_aligner_config": None, + }, + )() + + out = forced_aligner.build_forced_aligner_config(args) + + assert out is not None + assert out.model == "local-aligner" + assert out.runner == "pooling" + assert out.architecture == "Qwen3ASRForcedAlignerForTokenClassification" diff --git a/tests/utils/test_mm_outputs_partition.py b/tests/utils/test_mm_outputs_partition.py new file mode 100644 index 00000000000..f9affa47131 --- /dev/null +++ b/tests/utils/test_mm_outputs_partition.py @@ -0,0 +1,44 @@ +import torch + +from vllm_omni.utils.mm_outputs import partition_flat_payload, partition_payload_list + + +def test_partition_thinker_latent_payload(): + payload = { + "hidden_states.layer_0": torch.zeros(2, 4), + "hidden_states.layer_24": torch.zeros(2, 8), + "embed.tts_bos": [torch.zeros(1, 1, 4)], + } + inter, client = partition_flat_payload(payload) + assert inter == payload + assert client == {} + + +def test_partition_talker_intermediate_codes(): + payload = { + "codes.audio": torch.zeros(3, 2), + "hidden": torch.zeros(3, 16), + } + inter, client = partition_flat_payload(payload) + assert inter == payload + assert client == {} + + +def test_partition_code2wav_client_audio(): + payload = { + "model_outputs": torch.zeros(1, 2400), + "sr": torch.tensor(24000, dtype=torch.int32), + } + inter, client = partition_flat_payload(payload) + assert inter == {} + assert client == payload + + +def test_partition_payload_list_preserves_request_alignment(): + payloads = [ + {"hidden_states.layer_0": torch.zeros(1, 2)}, + {"model_outputs": torch.zeros(1, 10)}, + ] + inter_list, client_list = partition_payload_list(payloads) + assert inter_list == [payloads[0], None] + assert client_list == [None, payloads[1]] diff --git a/tests/utils/test_qwen3_force_align_processor.py b/tests/utils/test_qwen3_force_align_processor.py new file mode 100644 index 00000000000..c9af7167e33 --- /dev/null +++ b/tests/utils/test_qwen3_force_align_processor.py @@ -0,0 +1,65 @@ +import pytest + +from vllm_omni.utils import qwen3_force_align_processor as processor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_build_prompt_has_boundary_timestamp_markers(): + prompt = processor.build_prompt(["hello", "world"]) + + assert prompt.count("<timestamp>") == 4 + assert "hello<timestamp><timestamp>world" in prompt + assert prompt.endswith("<|im_start|>assistant\n") + + +@pytest.mark.parametrize( + "text,expected", + [ + # Punctuation inside a whitespace segment is stripped, not split on: + # the hand-rolled splitter used to break these into many fake words. + ("U.S.A", ["USA"]), + ("hello, world!", ["hello", "world"]), + ("don't stop", ["don't", "stop"]), + # CJK characters peel off as individual tokens, Latin runs stay whole. + ("你好world", ["你", "好", "world"]), + ("我爱 China", ["我", "爱", "China"]), + (" spaced out ", ["spaced", "out"]), + ], +) +def test_tokenize_space_lang_matches_official_segmentation(text, expected): + assert processor._tokenize_space_lang(text) == expected + + +def test_segment_words_falls_back_to_port_without_qwen_asr(monkeypatch): + # With qwen_asr unavailable, segmentation must use the built-in port and + # still produce the faithful result for the common (non-JP/KO) case. + monkeypatch.setattr(processor, "_get_official_processor", lambda: None) + + assert processor.segment_words("U.S.A is here", "auto") == ["USA", "is", "here"] + + +def test_fix_timestamp_repairs_dip_onto_monotonic_sequence(): + # 200 dips below the preceding 400; it snaps back up to 400. + assert processor.fix_timestamp([0, 400, 200, 800]) == [0, 400, 400, 800] + + +def test_fix_timestamp_passes_through_monotonic(): + assert processor.fix_timestamp([0, 100, 100, 250]) == [0, 100, 100, 250] + + +def test_resolve_timestamp_token_id_defaults_to_marker_token(): + # Regression: the default must resolve the same <timestamp> marker that + # build_prompt inserts, not None (which would degrade every request). + seen = {} + + class FakeTokenizer: + def convert_tokens_to_ids(self, token): + seen["token"] = token + return 151705 + + tid = processor.resolve_timestamp_token_id(FakeTokenizer()) + + assert tid == 151705 + assert seen["token"] == processor.TIMESTAMP_TOKEN + assert processor.TIMESTAMP_TOKEN in processor.build_prompt(["hello", "world"]) diff --git a/tests/utils/test_tracking_parser.py b/tests/utils/test_tracking_parser.py index 53745466184..00882744cb7 100644 --- a/tests/utils/test_tracking_parser.py +++ b/tests/utils/test_tracking_parser.py @@ -8,15 +8,20 @@ import yaml from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES from vllm_omni.config.stage_config import ( - _PIPELINE_REGISTRY, DeployConfig, PipelineConfig, - StageConfigFactory, StageDeployConfig, StagePipelineConfig, ) -from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace +from vllm_omni.utils.tracking_parser import ( + UNSET, + TrackingArgumentParser, + TrackingNamespace, + build_shadow_kwargs, +) ### Fake pipeline/deploy config for integration tests @@ -44,9 +49,14 @@ @pytest.fixture() def mock_stages(monkeypatch): """Register a fake pipeline and mock deploy YAML loading.""" - monkeypatch.setitem(_PIPELINE_REGISTRY._loaded, _TEST_MODEL, _TEST_PIPELINE) + from vllm_omni import platforms + + platform = platforms.current_omni_platform + monkeypatch.setattr(platform, "device_name", "cpu", raising=False) + monkeypatch.setattr(platform, "device_type", "cpu", raising=False) + monkeypatch.setitem(OMNI_PIPELINES, _TEST_MODEL, _TEST_PIPELINE) monkeypatch.setattr( - "vllm_omni.config.stage_config.load_deploy_config", + "vllm_omni.config.config_factory.load_deploy_config", lambda _path: _TEST_DEPLOY, ) return __file__ @@ -326,6 +336,132 @@ def test_omitted_nargs(): assert ns.items is None +### Tests for in-place mutating actions (append / extend / count). +# These actions mutate their default in place, which would crash on the bare +# ``UNSET`` sentinel. The shadow parser remaps them to a non-mutating store-style +# action (see build_shadow_kwargs in tracking_parser), so the real namespace +# still accumulates while explicit-arg tracking keeps working. +def test_append_action_omitted(): + """Omitted append args parse without error and aren't marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("--tag", action="append") + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.tag is None + + +def test_append_action_explicit(): + """Repeated append args are collected and marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("--tag", action="append") + ns = p.parse_args(["--tag", "a", "--tag", "b"]) + assert ns.explicit_keys == {"tag"} + assert isinstance(ns, TrackingNamespace) + assert ns.tag == ["a", "b"] + + +def test_append_const_action_omitted(): + """Omitted append_const args aren't marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("--dbg", action="append_const", const="x", dest="flags") + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.flags is None + + +def test_append_const_action_explicit(): + """Repeated append_const args accumulate the const and are explicit.""" + p = TrackingArgumentParser() + p.add_argument("--dbg", action="append_const", const="x", dest="flags") + ns = p.parse_args(["--dbg", "--dbg"]) + assert ns.explicit_keys == {"flags"} + assert isinstance(ns, TrackingNamespace) + assert ns.flags == ["x", "x"] + + +def test_extend_action_omitted(): + """Omitted extend args parse without error and aren't marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("--items", action="extend", nargs="+") + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.items is None + + +def test_extend_action_explicit(): + """Repeated extend args are flattened into one list and marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("--items", action="extend", nargs="+") + ns = p.parse_args(["--items", "a", "b", "--items", "c"]) + assert ns.explicit_keys == {"items"} + assert isinstance(ns, TrackingNamespace) + assert ns.items == ["a", "b", "c"] + + +def test_count_action_omitted(): + """Omitted count args keep their default and aren't marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("-v", "--verbose", action="count", default=0) + ns = p.parse_args([]) + assert ns.explicit_keys == set() + assert isinstance(ns, TrackingNamespace) + assert ns.verbose == 0 + + +def test_count_action_explicit(): + """Repeated count flags increment and are marked explicit.""" + p = TrackingArgumentParser() + p.add_argument("-v", "--verbose", action="count", default=0) + ns = p.parse_args(["-vv"]) + assert ns.explicit_keys == {"verbose"} + assert isinstance(ns, TrackingNamespace) + assert ns.verbose == 2 + + +def test_group_append_action_explicit(): + """append args added via a group go through the same shadow-default path.""" + p = TrackingArgumentParser() + g = p.add_argument_group("TestGroup") + g.add_argument("--tag", action="append") + ns = p.parse_args(["--tag", "a", "--tag", "b"]) + assert ns.explicit_keys == {"tag"} + assert isinstance(ns, TrackingNamespace) + assert ns.tag == ["a", "b"] + + +@pytest.mark.parametrize( + ("action", "expected"), + [ + ("append", "store"), + ("extend", "store"), + ("append_const", "store_const"), + ("count", "store_const"), + ], +) +def test_build_shadow_kwargs_remaps_mutating_actions(action, expected): + """Mutating actions are remapped to a store-style action with an UNSET default.""" + shadow = build_shadow_kwargs({"action": action, "const": "c"}) + assert shadow["action"] == expected + assert shadow["default"] is UNSET + + +def test_build_shadow_kwargs_count_gets_marker_const(): + """count carries no const, so the remapped store_const gets a non-UNSET marker.""" + shadow = build_shadow_kwargs({"action": "count"}) + assert shadow["action"] == "store_const" + assert shadow["const"] is not UNSET + + +def test_build_shadow_kwargs_leaves_other_actions_untouched(): + """Non-mutating actions keep their action and only get the UNSET default.""" + shadow = build_shadow_kwargs({"action": "store_true"}) + assert shadow["action"] == "store_true" + assert shadow["default"] is UNSET + + def test_parse_known_args_tracking(): """Ensure parse_known_args is also trackable""" p = TrackingArgumentParser() @@ -453,6 +589,12 @@ def test_cli_overrides_config(tmp_path): ### Integration tests for arg resolution through StageConfigFactory +# +# These tests still target the transitional legacy StageConfig path because the +# current runtime reads explicit CLI values from ``stage.runtime_overrides``. +# After RFC #4021 migrates engine startup to consume ``VllmOmniConfig`` directly, +# these assertions should move to the structured config/runtime boundary and the +# ``_create_legacy_from_registry`` calls should be removed. def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages): """Explicitly passed CLI values reach runtime_overrides on all stages.""" p = TrackingArgumentParser() @@ -460,8 +602,8 @@ def test_explicit_cli_arg_reaches_runtime_overrides(mock_stages): ns = p.parse_args(["--max-num-seqs", "999"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( - _TEST_MODEL, + stages, _ = StageConfigFactory._create_legacy_from_registry( + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -476,8 +618,8 @@ def test_omitted_default_not_in_runtime_overrides(mock_stages): ns = p.parse_args([]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( - _TEST_MODEL, + stages, _ = StageConfigFactory._create_legacy_from_registry( + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -497,8 +639,8 @@ def test_config_file_args_reach_runtime_overrides(mock_stages): ns = p.parse_args(["--config", "fake.yaml"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( - _TEST_MODEL, + stages, _ = StageConfigFactory._create_legacy_from_registry( + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -514,8 +656,8 @@ def test_per_stage_override_routes_correctly(mock_stages): ns = p.parse_args(["--stage-0-gpu-memory-utilization", "0.42"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( - _TEST_MODEL, + stages, _ = StageConfigFactory._create_legacy_from_registry( + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) @@ -535,8 +677,8 @@ def test_explicit_args_omitted_from_yaml(mock_stages): ns = p.parse_args(["--enforce-eager"]) explicit_kwargs = ns.get_explicit_kwargs_dict() - stages = StageConfigFactory._create_from_registry( - _TEST_MODEL, + stages, _ = StageConfigFactory._create_legacy_from_registry( + _TEST_PIPELINE, explicit_kwargs, deploy_config_path=mock_stages, ) diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 3dbbc4ac45a..90cf3afe197 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -1,9 +1,17 @@ +from collections.abc import Sequence from types import SimpleNamespace +import numpy as np import pytest import torch -from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, GPUARModelRunner +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker.gpu_ar_model_runner import ( + ExecuteModelState, + GPUARModelRunner, + OmniAsyncGPUModelRunnerOutput, +) +from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -52,6 +60,528 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker(): assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"] +def test_runner_assisted_full_attention_metadata_request_is_opt_in(): + runner = object.__new__(GPUARModelRunner) + runner.model = object() + runner.scheduler_config = SimpleNamespace(max_num_seqs=16) + + request = runner._get_runner_assisted_full_attention_metadata_request( + req_ids=["r1", "r2"], + num_reqs=2, + num_reqs_padded=4, + num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32), + num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32), + max_num_scheduled_tokens=1, + ) + + assert request is None + + +def test_runner_assisted_full_attention_metadata_request_and_context_hooks(): + calls = [] + + class Model: + def get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_scheduled_tokens: Sequence[int], + num_computed_tokens: Sequence[int], + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest: + calls.append( + ( + "request", + { + "req_ids": list(req_ids), + "num_reqs": num_reqs, + "num_scheduled_tokens": [int(n) for n in num_scheduled_tokens], + "num_computed_tokens": [int(n) for n in num_computed_tokens], + "max_num_scheduled_tokens": max_num_scheduled_tokens, + }, + ) + ) + return RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=12, + for_cudagraph_capture=True, + ) + + def set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> None: + calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs})) + + runner = object.__new__(GPUARModelRunner) + runner.model = Model() + runner.scheduler_config = SimpleNamespace(max_num_seqs=8) + + request = runner._get_runner_assisted_full_attention_metadata_request( + req_ids=["r1", "r2"], + num_reqs=2, + num_reqs_padded=4, + num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32), + num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32), + max_num_scheduled_tokens=1, + ) + context_enabled = runner._set_runner_assisted_full_attention_metadata_context( + enabled=True, + num_reqs=2, + ) + context_disabled = runner._set_runner_assisted_full_attention_metadata_context(enabled=False) + + assert request == RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=8, + for_cudagraph_capture=True, + ) + assert context_enabled + assert context_disabled + assert calls == [ + ( + "request", + { + "req_ids": ["r1", "r2"], + "num_reqs": 2, + "num_scheduled_tokens": [1, 1], + "num_computed_tokens": [5, 6], + "max_num_scheduled_tokens": 1, + }, + ), + ("context", {"enabled": True, "num_reqs": 2}), + ("context", {"enabled": False, "num_reqs": 0}), + ] + + +def test_omni_async_gpu_model_runner_output_builds_lazily_once(): + async_output = object.__new__(OmniAsyncGPUModelRunnerOutput) + calls = [] + sync_calls = [] + + def builder(): + calls.append("build") + return OmniModelRunnerOutput(req_ids=["r1"], req_id_to_index={"r1": 0}) + + async_output._model_runner_output = None + async_output._model_runner_output_builder = builder + async_output._invalid_req_indices = [] + async_output.sampled_token_ids_cpu = torch.tensor([[7]], dtype=torch.long) + async_output.async_copy_ready_event = SimpleNamespace(synchronize=lambda: sync_calls.append("sync")) + async_output._sampled_token_ids = torch.tensor([[7]], dtype=torch.long) + async_output._logprobs_tensors = None + async_output._logprobs_tensors_cpu = None + async_output._routed_experts = None + async_output._routed_experts_cpu = None + async_output.vocab_size = 10 + + output = async_output.get_output() + + assert calls == ["build"] + assert sync_calls == ["sync"] + assert async_output._model_runner_output_builder is None + assert output.req_ids == ["r1"] + assert output.sampled_token_ids == [[7]] + + +def test_omni_async_gpu_model_runner_output_reraises_background_exception(): + async_output = object.__new__(OmniAsyncGPUModelRunnerOutput) + joined = [] + + class FakeThread: + def join(self): + joined.append("join") + + async_output._background_thread = FakeThread() + async_output._background_exception = RuntimeError("background failed") + + with pytest.raises(RuntimeError, match="background failed"): + async_output.get_output() + + assert joined == ["join"] + assert async_output._background_thread is None + + +def _make_async_output_runner(engine_output_type: str = "audio"): + runner = object.__new__(GPUARModelRunner) + model_config = SimpleNamespace( + engine_output_type=engine_output_type, + async_chunk=True, + enable_return_routed_experts=False, + ) + runner.vllm_config = SimpleNamespace(model_config=model_config) + runner.model_config = model_config + runner._async_chunk = True + runner.omni_prefix_cache = None + runner.requests = {"r1": object(), "r2": object()} + runner.supports_mm_inputs = False + runner.routed_experts_initialized = False + runner.model = SimpleNamespace(has_postprocess=False) + runner.model_intermediate_buffer = {} + runner.input_batch = SimpleNamespace( + req_ids=["mutated"], + req_id_to_index={"mutated": 0}, + ) + return runner + + +def test_build_omni_output_uses_snapshots_and_connector_after_accumulation(monkeypatch): + runner = _make_async_output_runner() + events = [] + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("audio", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: True) + monkeypatch.setattr( + GPUARModelRunner, + "accumulate_full_payload_output", + lambda self, rid, payload, request: events.append(f"accumulate:{rid}"), + ) + monkeypatch.setattr( + GPUARModelRunner, + "get_omni_connector_output", + lambda self: events.append("connector") or "connector-output", + ) + + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=3, + num_scheduled_tokens={"r1": 1, "r2": 2}, + ), + hidden_states=torch.tensor([[1.0], [2.0], [3.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={"foo": torch.tensor([10.0, 20.0, 30.0])}, + req_ids_output_copy=["r1", "r2"], + req_id_to_index_output_copy={"r1": 0, "r2": 1}, + valid_sampled_token_ids=[[101], [102]], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=["r2"], + num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long), + ) + + assert output.req_ids == ["r1", "r2"] + assert output.inter_stage_outputs is not None + assert torch.equal(output.inter_stage_outputs[0]["hidden"], torch.tensor([[1.0]])) + assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]])) + assert output.multimodal_outputs is None + assert output.kv_extracted_req_ids == ["r2"] + assert output.omni_connector_output == "connector-output" + assert events == ["accumulate:r1", "accumulate:r2", "connector"] + + +def test_build_omni_output_copies_hidden_for_partial_downstream_batch(monkeypatch): + runner = _make_async_output_runner(engine_output_type="latent") + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", ["r2"]), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=6, + num_scheduled_tokens={"r1": 1, "r2": 2, "r3": 3}, + ), + hidden_states=torch.tensor([[1.0], [2.0], [3.0], [4.0], [5.0], [6.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={}, + req_ids_output_copy=["r1", "r2", "r3"], + req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2}, + valid_sampled_token_ids=[[], [], []], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([1, 2, 3], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1, 3], dtype=torch.long), + ) + + assert output.inter_stage_outputs is not None + assert output.inter_stage_outputs[0] is None + assert torch.equal(output.inter_stage_outputs[1]["hidden"], torch.tensor([[2.0], [3.0]])) + assert output.inter_stage_outputs[2] is None + assert output.multimodal_outputs is None + + +def test_process_additional_information_uses_snapshot_request_order(monkeypatch): + runner = _make_async_output_runner() + seen = [] + + class PostprocessModel: + has_postprocess = True + + def postprocess(self, hidden_states, **kwargs): + seen.append(hidden_states.clone()) + return {} + + runner.model = PostprocessModel() + runner.model_intermediate_buffer = {"r1": {}, "r2": {}} + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("audio", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_update_intermediate_buffer", lambda *args, **kwargs: None) + + GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=3, + num_scheduled_tokens={"r1": 1, "r2": 2}, + ), + hidden_states=torch.tensor([[1.0], [2.0], [3.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={}, + req_ids_output_copy=["r1", "r2"], + req_id_to_index_output_copy={"r1": 0, "r2": 1}, + valid_sampled_token_ids=[[], []], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=torch.tensor([1, 2], dtype=torch.int32).numpy(), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long), + ) + + assert len(seen) == 2 + assert torch.equal(seen[0], torch.tensor([[1.0]])) + assert torch.equal(seen[1], torch.tensor([[2.0], [3.0]])) + + +def test_async_omni_output_guard_requires_safe_conditions(): + runner = _make_async_output_runner() + runner.use_async_scheduling = True + runner.speculative_config = None + runner.model.use_async_omni_output = True + + assert GPUARModelRunner._should_use_async_omni_output(runner) + + runner.omni_prefix_cache = object() + assert not GPUARModelRunner._should_use_async_omni_output(runner) + + runner.omni_prefix_cache = None + runner.model.has_postprocess = True + assert not GPUARModelRunner._should_use_async_omni_output(runner) + + runner.model.eager_omni_postprocess_before_async_output = True + assert GPUARModelRunner._should_use_async_omni_output(runner) + + +def test_build_omni_output_skips_hidden_when_model_opts_out(monkeypatch): + runner = _make_async_output_runner(engine_output_type="latent") + runner.model.omni_pooler_payload_include_hidden = False + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=2, + num_scheduled_tokens={"r1": 2}, + ), + hidden_states=torch.tensor([[1.0], [2.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={"codes": {"audio": torch.tensor([[7, 8], [9, 10]], dtype=torch.long)}}, + req_ids_output_copy=["r1"], + req_id_to_index_output_copy={"r1": 0}, + valid_sampled_token_ids=[[101]], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([2], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0], dtype=torch.long), + ) + + assert output.inter_stage_outputs is not None + assert len(output.inter_stage_outputs) == 1 + assert "hidden" not in output.inter_stage_outputs[0] + assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], torch.tensor([[7, 8], [9, 10]], dtype=torch.long)) + assert output.multimodal_outputs is None + + +def test_build_omni_output_splits_mm_by_scheduled_tokens_when_hidden_is_tail_only(monkeypatch): + runner = _make_async_output_runner(engine_output_type="latent") + runner.model.omni_pooler_payload_include_hidden = False + runner._async_chunk = False + runner.requests = {"r1": object(), "r2": object(), "r3": object()} + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + + codes = torch.arange(48, dtype=torch.long).reshape(3, 16) + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=3, + num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1}, + ), + hidden_states=torch.tensor([[1.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={"codes": {"audio": codes}}, + req_ids_output_copy=["r1", "r2", "r3"], + req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2}, + valid_sampled_token_ids=[[101], [102], [103]], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long), + ) + + assert output.inter_stage_outputs is not None + assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1]) + assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2]) + assert torch.equal(output.inter_stage_outputs[2]["codes.audio"], codes[2:3]) + + +def test_build_omni_output_splits_mm_by_hidden_len_when_scheduled_is_padded(monkeypatch): + """Thinker mm rows align to hidden_states.shape[0], not padded scheduled count.""" + runner = _make_async_output_runner(engine_output_type="latent") + runner.model.omni_pooler_payload_include_hidden = False + runner._async_chunk = False + runner.requests = {"r1": object(), "r2": object(), "r3": object()} + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + + layers = torch.arange(96, dtype=torch.long).reshape(3, 32) + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=5, + num_scheduled_tokens={"r1": 1, "r2": 1, "r3": 1}, + ), + hidden_states=torch.randn(3, 4), + staged_hidden_states_cpu=None, + multimodal_outputs={"hidden_states": {"layers": {0: layers}}}, + req_ids_output_copy=["r1", "r2", "r3"], + req_id_to_index_output_copy={"r1": 0, "r2": 1, "r3": 2}, + valid_sampled_token_ids=[[101], [102], [103]], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([1, 1, 1], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.long), + ) + + assert output.inter_stage_outputs is not None + assert torch.equal(output.inter_stage_outputs[0]["hidden_states.layer_0"], layers[0:1]) + assert torch.equal(output.inter_stage_outputs[1]["hidden_states.layer_0"], layers[1:2]) + assert torch.equal(output.inter_stage_outputs[2]["hidden_states.layer_0"], layers[2:3]) + + +def test_async_snapshot_payload_omits_hidden_when_model_opts_out(): + runner = _make_async_output_runner() + runner.model.omni_pooler_payload_include_hidden = False + + payload = GPUARModelRunner._build_omni_async_snapshot_payload( + runner, + hidden_states=torch.tensor([[1.0], [2.0]]), + staged_hidden_states_cpu=torch.tensor([[3.0]]), + multimodal_outputs={"codes": {"audio": torch.tensor([[1]], dtype=torch.long)}}, + ) + + assert set(payload.keys()) == {"multimodal_outputs"} + assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]] + + +def test_runner_assisted_full_attention_metadata_refresh_pads_buffers(): + class QueryStartLoc: + def __init__(self): + self.np = np.full(5, -1, dtype=np.int32) + self.copied = False + + def copy_to_gpu(self): + self.copied = True + + class BlockTable: + def __init__(self): + self.commits = [] + + def commit_block_table(self, num_reqs_padded): + self.commits.append(num_reqs_padded) + + runner = object.__new__(GPUARModelRunner) + block_table = BlockTable() + runner.input_batch = SimpleNamespace( + num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32), + block_table=block_table, + ) + runner.optimistic_seq_lens_cpu = torch.zeros(4, dtype=torch.int32) + runner.seq_lens = torch.empty(4, dtype=torch.int32) + runner.query_pos = SimpleNamespace(np=np.empty(4, dtype=np.int32)) + runner.query_start_loc = QueryStartLoc() + runner._get_cumsum_and_arange = lambda scheduled, _query_pos: np.cumsum(scheduled, dtype=np.int32) + + runner._refresh_runner_assisted_full_attention_metadata_buffers( + num_reqs=2, + num_reqs_padded=4, + num_scheduled_tokens_np=np.array([1, 2], dtype=np.int32), + ) + + assert runner.optimistic_seq_lens_cpu.tolist() == [11, 22, 0, 0] + assert runner.seq_lens.tolist() == [11, 22, 0, 0] + assert runner.query_start_loc.np.tolist() == [0, 1, 3, 3, 3] + assert runner.query_start_loc.copied + assert block_table.commits == [4] + + @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"]) def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr): runner = object.__new__(GPUARModelRunner) @@ -95,6 +625,7 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk runner.vllm_config = SimpleNamespace( model_config=SimpleNamespace(engine_output_type="audio"), ) + runner._async_chunk = False monkeypatch.setattr( GPUARModelRunner, "_sample", lambda self, logits, spec_decode_metadata: SimpleNamespace(sampled_token_ids=[]) @@ -128,8 +659,65 @@ def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monk output = GPUARModelRunner.sample_tokens(runner, grammar_output=None) + # Non-async-chunk now ships the full payload to the next stage, so + # inter_stage_outputs mirrors multimodal_outputs (PR #4792). + assert output.inter_stage_outputs is not None + assert output.multimodal_outputs is not None + assert torch.equal(output.inter_stage_outputs[0]["hidden"], output.multimodal_outputs[0]["hidden"]) + assert torch.equal(output.inter_stage_outputs[1]["hidden"], output.multimodal_outputs[1]["hidden"]) assert torch.equal(output.multimodal_outputs[0]["hidden"], torch.tensor([[1.0, 10.0]])) assert torch.equal( output.multimodal_outputs[1]["hidden"], torch.tensor([[2.0, 20.0], [3.0, 30.0]]), ) + + +def test_build_omni_output_falls_back_to_mm_cpu_without_prefix_merge(monkeypatch): + """Tail-only prefix-cache models still need per-step mm passthrough (e.g. codes.audio).""" + runner = _make_async_output_runner(engine_output_type="latent") + runner.omni_prefix_cache = object() + + monkeypatch.setattr( + GPUARModelRunner, + "_resolve_pooler_payload_req_ids", + lambda self, req_ids: ("latent", req_ids), + ) + monkeypatch.setattr(GPUARModelRunner, "_model_needs_full_prefix_hidden_states", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "_deferred_prefix_cache_mm_keys", lambda self: {"codes.audio"}) + monkeypatch.setattr(GPUARModelRunner, "_stage_deferred_prefix_cache_mm_outputs", lambda *args, **kwargs: None) + monkeypatch.setattr( + GPUARModelRunner, + "_prepare_prefix_cache_pooler_payload_sources", + lambda *args, **kwargs: (None, None, None), + ) + monkeypatch.setattr(GPUARModelRunner, "_process_additional_information_updates", lambda *args, **kwargs: None) + monkeypatch.setattr(GPUARModelRunner, "_should_accumulate_full_payload_output", lambda self: False) + monkeypatch.setattr(GPUARModelRunner, "get_omni_connector_output", lambda self: None) + + codes = torch.tensor([[11.0, 12.0], [21.0, 22.0]], dtype=torch.float32) + output = GPUARModelRunner._build_omni_model_runner_output_from_snapshot( + runner, + scheduler_output=SimpleNamespace( + total_num_scheduled_tokens=2, + num_scheduled_tokens={"r1": 1, "r2": 1}, + ), + hidden_states=torch.tensor([[1.0], [2.0]]), + staged_hidden_states_cpu=None, + multimodal_outputs={"codes.audio": codes}, + req_ids_output_copy=["r1", "r2"], + req_id_to_index_output_copy={"r1": 0, "r2": 1}, + valid_sampled_token_ids=[[], []], + logprobs_lists=None, + prompt_logprobs_dict={}, + num_nans_in_logits=None, + kv_connector_output=None, + ec_connector_output=None, + cudagraph_stats=None, + kv_extracted_req_ids=None, + num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32), + query_start_loc_cpu=torch.tensor([0, 1], dtype=torch.long), + ) + + assert torch.equal(output.inter_stage_outputs[0]["codes.audio"], codes[0:1]) + assert torch.equal(output.inter_stage_outputs[1]["codes.audio"], codes[1:2]) + assert output.multimodal_outputs is None diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index be7fadcb643..0cd9aeb2683 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -40,6 +40,7 @@ def _make_runner(multimodal_outputs): runner.supports_mm_inputs = False runner.speculative_config = None runner.routed_experts_initialized = False + runner._async_chunk = False return runner diff --git a/tests/worker/test_min_tokens_stop_ids_guard.py b/tests/worker/test_min_tokens_stop_ids_guard.py new file mode 100644 index 00000000000..5d7591c993f --- /dev/null +++ b/tests/worker/test_min_tokens_stop_ids_guard.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Guard against out-of-vocabulary stop ids in min-tokens masking (#4962). + +vLLM's input processor folds the stage tokenizer's EOS id into +``SamplingParams.all_stop_token_ids``. For AR stages whose lm_head is +narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS: +3072 logits vs text EOS 151645), any ``min_tokens >= 1`` makes +``MinTokensLogitsProcessor`` write ``-inf`` at an out-of-range index and +the engine dies with a CUDA device-side assert. +""" + +import pytest +import torch +from vllm import SamplingParams +from vllm.v1.sample.logits_processor import ( + BatchUpdate, + LogitsProcessors, + MinTokensLogitsProcessor, +) + +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +TALKER_VOCAB = 3072 +CODEC_EOS = 2150 +TEXT_EOS = 151645 + + +def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor: + """Build a MinTokensLogitsProcessor with one active request, mirroring + the production path (engine folds the tokenizer EOS via + ``update_from_generation_config``).""" + params = SamplingParams(min_tokens=2, stop_token_ids=stop_token_ids) + if eos_token_id is not None: + params.update_from_generation_config({}, eos_token_id) + proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False) + proc.update_state( + BatchUpdate( + batch_size=1, + removed=[], + added=[(0, params, None, [])], + moved=[], + ) + ) + return proc + + +def test_oob_stop_id_crashes_without_guard(): + """Document the failure mode: the folded text EOS is out of range for a + narrow codec head and index_put_ raises (device-side assert on CUDA).""" + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS) + logits = torch.zeros(1, TALKER_VOCAB) + with pytest.raises((IndexError, RuntimeError)): + proc.apply(logits) + + +def test_guard_filters_oob_and_keeps_in_range_mask(): + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS) + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + + logits = torch.zeros(1, TALKER_VOCAB) + out = proc.apply(logits) + + assert out[0, CODEC_EOS] == -float("inf") + masked = (out == -float("inf")).nonzero() + assert masked.tolist() == [[0, CODEC_EOS]] + + +def test_guard_noop_when_all_stop_ids_in_range(): + """Full-width heads (e.g. CosyVoice3 talker) must be untouched: no + tensor rebuild, mask unchanged.""" + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=None) + slice_before = proc.logits_slice + + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + + assert proc.logits_slice is slice_before + out = proc.apply(torch.zeros(1, TALKER_VOCAB)) + assert out[0, CODEC_EOS] == -float("inf") + + +def test_guard_sanitizes_params_only_once(): + """The stop-id set is shared with the request's SamplingParams, so the + second call finds nothing to drop (no repeated rebuilds/warnings).""" + proc = _make_min_tokens_proc([CODEC_EOS], eos_token_id=TEXT_EOS) + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + slice_after_first = proc.logits_slice + + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + + assert proc.logits_slice is slice_after_first + + +def test_guard_handles_multiple_requests(): + params_oob = SamplingParams(min_tokens=4, stop_token_ids=[CODEC_EOS]) + params_oob.update_from_generation_config({}, TEXT_EOS) + params_ok = SamplingParams(min_tokens=4, stop_token_ids=[7]) + proc = MinTokensLogitsProcessor(None, device=torch.device("cpu"), is_pin_memory=False) + proc.update_state( + BatchUpdate( + batch_size=2, + removed=[], + added=[(0, params_oob, None, []), (1, params_ok, None, [])], + moved=[], + ) + ) + + sanitize_min_tokens_stop_ids(LogitsProcessors([proc]), TALKER_VOCAB) + out = proc.apply(torch.zeros(2, TALKER_VOCAB)) + + assert out[0, CODEC_EOS] == -float("inf") + assert out[1, 7] == -float("inf") + assert int((out == -float("inf")).sum()) == 2 diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index 6f4994b2cde..48c0bc6deab 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -4,7 +4,7 @@ import pytest import torch -from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner, _filter_mrope_kwargs_for_model from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -81,6 +81,7 @@ def forward( top_k=None, top_p=None, generator=None, + generators=None, ): self.calls.append( { @@ -90,18 +91,51 @@ def forward( "top_k": top_k, "top_p": top_p, "generator": generator, + "generators": generators, } ) codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64) return req_embeds, codes +class StrictMRoPEModel: + def get_mrope_input_positions(self, input_tokens, mm_features): + raise NotImplementedError + + +class FlexibleMRoPEModel: + def get_mrope_input_positions(self, input_tokens, mm_features=None, **kwargs): + raise NotImplementedError + + @contextmanager def _noop_forward_context(*args, **kwargs): """A no-op context manager to replace vLLM forward context in CPU tests.""" yield +def test_filter_mrope_kwargs_for_strict_model_signature(): + kwargs = { + "mm_features": ["audio"], + "hf_config": object(), + "image_grid_thw": [], + } + + assert _filter_mrope_kwargs_for_model(StrictMRoPEModel(), kwargs) == { + "mm_features": ["audio"], + } + + +def test_filter_mrope_kwargs_preserves_flexible_model_kwargs(): + kwargs = { + "mm_features": ["video"], + "hf_config": object(), + "video_grid_thw": [[1, 2, 3]], + } + + assert _filter_mrope_kwargs_for_model(FlexibleMRoPEModel(), kwargs) is kwargs + + def _make_runner(req_ids=("r1", "r2"), hidden_size=4): # Create an instance without calling OmniGPUModelRunner.__init__ runner = object.__new__(OmniGPUModelRunner) @@ -243,7 +277,7 @@ def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker runner = _make_runner(req_ids=("r1",), hidden_size=4) runner.requests["r1"].sampling_params = SimpleNamespace( seed=42, - extra_args={"qwen3_tts_request_seed": 42}, + extra_args={"tts_local_seed": 42}, ) runner.talker_mtp = CaptureTalkerMTP() runner.vllm_config = SimpleNamespace( @@ -274,6 +308,7 @@ def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_ "top_k": 9, "top_p": 0.55, "generator": runner.talker_mtp.calls[0]["generator"], + "generators": None, } ] assert runner.talker_mtp.calls[0]["generator"] is not None @@ -287,11 +322,11 @@ def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch): runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) runner.requests["r1"].sampling_params = SimpleNamespace( seed=11, - extra_args={"qwen3_tts_request_seed": 11}, + extra_args={"tts_local_seed": 11}, ) runner.requests["r2"].sampling_params = SimpleNamespace( seed=22, - extra_args={"qwen3_tts_request_seed": 22}, + extra_args={"tts_local_seed": 22}, ) runner.talker_mtp = CaptureTalkerMTP() runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={})) @@ -318,6 +353,57 @@ def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_ assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds) +def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch): + """Models with talker_mtp_accepts_per_row_generators get one batched call (#4883).""" + import vllm_omni.worker.gpu_model_runner as mod + + monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context) + + runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4) + runner.requests["r1"].sampling_params = SimpleNamespace( + seed=11, + extra_args={"tts_local_seed": 11}, + ) + runner.requests["r2"].sampling_params = SimpleNamespace( + seed=22, + extra_args={"tts_local_seed": 22}, + ) + runner.talker_mtp = CaptureTalkerMTP() + runner.model = SimpleNamespace( + talker_mtp_output_key=("codes", "audio"), + talker_mtp_accepts_per_row_generators=True, + ) + runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={})) + + def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn): + batch_desc = SimpleNamespace(num_tokens=int(num_tokens)) + return (False, batch_desc, None, None, None) + + monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner))) + + inputs_embeds = torch.zeros((6, 4), dtype=torch.float32) + OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds) + + # One batched call with distinct per-row generators, not two scalar calls. + assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2] + row_generators = runner.talker_mtp.calls[0]["generators"] + assert runner.talker_mtp.calls[0]["generator"] is None + assert len(row_generators) == 2 + assert all(generator is not None for generator in row_generators) + assert row_generators[0] is not row_generators[1] + + # The per-request generator stream persists across steps... + OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds) + assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0] + assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1] + + # ...and is evicted once its request finishes. + del runner.requests["r2"] + OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds) + assert set(runner._talker_mtp_generators) == {"r1"} + assert runner.talker_mtp.calls[2]["generator"] is row_generators[0] + + def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch): """Validate that _update_intermediate_buffer writes to model_intermediate_buffer (forward path) and mirrors to additional_information_cpu setattr (backward compat).""" diff --git a/tools/nightly/run_nightly_jobs.sh b/tools/nightly/run_nightly_jobs.sh index b115a9ee24a..a95eb8876db 100644 --- a/tools/nightly/run_nightly_jobs.sh +++ b/tools/nightly/run_nightly_jobs.sh @@ -2,14 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # -# Extract steps from .buildkite/test-nightly.yml that contain pytest, synthesize +# Extract steps from .buildkite/cuda/test-nightly.yml that contain pytest, synthesize # small bash wrappers (exports + pytest), run them, and tee output to logs named # after each step's Buildkite "key" when present (otherwise a slug of the label). # YAML steps whose label contains "Perf Test" run first, then # python tools/nightly/generate_nightly_perf_excel.py runs (even if some perf jobs failed; # the script only aggregates JSON already present under tests/dfx/perf/results). # Excel and generate_nightly_perf_excel.log are written under $REPO_ROOT/logs/; -# per-job tee logs use LOG_DIR (see --log-dir). Then remaining jobs run. +# per-job tee logs use LOG_DIR (see --log-dir; default is timestamped under $REPO_ROOT/logs/). # # Test kind (--test-type / TEST_TYPE), multiple allowed (OR semantics): # Repeat the flag and/or use comma-separated values, e.g. @@ -21,6 +21,8 @@ # function — label has neither "Perf Test" nor "Accuracy Test" (incl. Doc, Multi-Replica, etc.) # stability— fixed dfx stability scripts under tests/dfx/stability/scripts/ (see below) # local — pytest -sv -m "<MODEL_TYPE markers> and local_model" from repo root (no YAML step extract) +# When LABEL_SUBSTR is set, also runs matching tests/**/test_*.py and perf JSON configs under +# tests/dfx/perf/tests/*.json via run_benchmark.py / run_diffusion_benchmark.py. # all — no test-kind filter for YAML steps (any leaf step with pytest in test-nightly.yml) # # Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics): @@ -36,10 +38,12 @@ # local (when included in TEST_TYPE): # From repo root: pytest -sv -m "<markers> and local_model" (markers from MODEL_TYPE: omni, tts, # diffusion; all → "(omni or tts or diffusion) and local_model"). Not filtered by nightly YAML. +# LABEL_SUBSTR: if set, restrict to tests/**/test_*.py basenames and tests/dfx/perf/tests/*.json +# whose filename contains the substring (benchmark runner chosen by JSON family). # # stability (when included in TEST_TYPE): -# From repo root: pytest -s -v tests/dfx/stability/scripts/test_stability_*.py -# model_type: omni → qwen3_omni; tts → qwen3_tts; diffusion → qwen_image + wan22; all → all four +# From repo root: pytest -s -v --run-level full_model tests/dfx/stability/scripts/test_stability_*.py +# model_type: omni → qwen3_omni; tts → qwen3_tts + voxcpm2; diffusion → qwen_image + wan22 + hunyuan_image; all → all six # LABEL_SUBSTR: if set, script path / job key / filename must contain it # # Requirements: bash, python3, PyYAML (pip install pyyaml) @@ -47,22 +51,31 @@ # Usage: # bash path/to/run_nightly_jobs.sh # REPO_ROOT=/path/to/vllm-omni bash path/to/run_nightly_jobs.sh --test-type perf,acc --dry-run -# YML=/path/to/vllm-omni/.buildkite/test-nightly.yml bash path/to/run_nightly_jobs.sh +# YML=/path/to/vllm-omni/.buildkite/cuda/test-nightly.yml bash path/to/run_nightly_jobs.sh # # Repository / YAML (no dependency on where this script lives): -# • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/test-nightly.yml +# • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/cuda/test-nightly.yml # • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory # • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD, -# then from the script's directory, until .buildkite/test-nightly.yml exists +# then from the script's directory, until .buildkite/cuda/test-nightly.yml exists # # Optional environment: # REPO_ROOT - vllm-omni root (working directory for pytest); see above -# YML - path to test-nightly.yml (default: $REPO_ROOT/.buildkite/test-nightly.yml) -# LOG_DIR - logs + generated job scripts (default: $REPO_ROOT/logs/nightly_jobs) +# YML - path to test-nightly.yml (default: $REPO_ROOT/.buildkite/cuda/test-nightly.yml) +# LOG_DIR - logs + generated job scripts; when unset, a timestamped directory under +# $REPO_ROOT/logs/ is created: +# nightly_jobs_YYYYMMDD-HHMMSS (default / YAML nightly steps) +# nightly_local_jobs_YYYYMMDD-HHMMSS (--test-type local only) +# nightly_stability_jobs_YYYYMMDD-HHMMSS (--test-type stability only) +# per-job *.log plus timing_summary.log after the run; +# perf JSON from tests/dfx/perf/results/ (produced this run) -> $LOG_DIR/perf_results/ # TEST_TYPE - comma-separated and/or repeated flags (default: all); see above # MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above -# LABEL_SUBSTR - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename +# LABEL_SUBSTR - YAML mode: substring of Buildkite step label; stability: substring of path/key/filename; +# local: substring of test_*.py basename or tests/dfx/perf/tests/*.json filename # DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest +# RUN_JOB_TIMEOUT_KILL_AFTER - seconds after SIGTERM before SIGKILL on inline pytest +# timeout (default: 60; baked into generated job scripts) # set -euo pipefail @@ -212,7 +225,29 @@ done TEST_TYPE="$(_finalize_test_type_csv)" MODEL_TYPE="$(_finalize_model_type_csv)" -BUILDKITE_REL=".buildkite/test-nightly.yml" +# Default log directory basename: date + local timestamp; prefix depends on TEST_TYPE. +_nightly_log_dir_basename() { + local ts has_local=0 has_stability=0 has_yaml=0 _t + ts="$(date +%Y%m%d-%H%M%S)" + IFS=',' read -ra _TTOK <<< "${TEST_TYPE}" + for _t in "${_TTOK[@]}"; do + _t="$(printf '%s' "${_t}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" + case "${_t}" in + local) has_local=1 ;; + stability) has_stability=1 ;; + perf | acc | function | all) has_yaml=1 ;; + esac + done + if [[ "${has_local}" -eq 1 && "${has_yaml}" -eq 0 && "${has_stability}" -eq 0 ]]; then + printf 'nightly_local_jobs_%s' "${ts}" + elif [[ "${has_stability}" -eq 1 && "${has_yaml}" -eq 0 && "${has_local}" -eq 0 ]]; then + printf 'nightly_stability_jobs_%s' "${ts}" + else + printf 'nightly_jobs_%s' "${ts}" + fi +} + +BUILDKITE_REL=".buildkite/cuda/test-nightly.yml" _find_repo_containing_nightly() { local dir="${1:-}" @@ -232,8 +267,14 @@ _derive_repo_root_from_yml() { local yml="$1" local d d="$(cd "$(dirname "${yml}")" && pwd)" - [[ "$(basename "${d}")" == ".buildkite" ]] || return 1 - printf '%s\n' "$(dirname "${d}")" + while [[ "${d}" != "/" ]]; do + if [[ "$(basename "${d}")" == ".buildkite" ]]; then + printf '%s\n' "$(dirname "${d}")" + return 0 + fi + d="$(dirname "${d}")" + done + return 1 } # Resolve REPO_ROOT, YML (no relative path between script and repo) @@ -243,7 +284,7 @@ if [[ -n "${YML:-}" && -n "${REPO_ROOT:-}" ]]; then elif [[ -n "${YML:-}" ]]; then YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" if ! REPO_ROOT="$(_derive_repo_root_from_yml "${YML}")"; then - echo "Could not derive REPO_ROOT from YML=${YML} (expected file at <repo>/.buildkite/test-nightly.yml)." >&2 + echo "Could not derive REPO_ROOT from YML=${YML} (expected file at <repo>/.buildkite/cuda/test-nightly.yml)." >&2 echo "Set REPO_ROOT explicitly (or pass --repo-root) for pytest working directory." >&2 exit 2 fi @@ -269,7 +310,11 @@ else YML="${REPO_ROOT}/${BUILDKITE_REL}" fi -LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/nightly_jobs}" +if [[ -z "${LOG_DIR:-}" ]]; then + LOG_DIR="${REPO_ROOT}/logs/$(_nightly_log_dir_basename)" +fi + +echo "Log directory: ${LOG_DIR}" >&2 # Require nightly YAML only when at least one non-stability test kind needs it (validated in Python too). _needs_nightly_yml=0 @@ -310,6 +355,7 @@ if [[ "${DRY_RUN}" != "1" ]]; then rm -f "${_stale_logs[@]}" fi rm -f "${LOG_DIR}/jobs/.perf_job_keys" + rm -f "${LOG_DIR}/jobs/.job_timeouts" rm -f "${REPO_ROOT}/logs/nightly_perf_manual.xlsx" rm -f "${LOG_DIR}/nightly_perf_manual.xlsx" fi @@ -320,6 +366,7 @@ from __future__ import annotations import os import re +import shlex import stat import sys from pathlib import Path @@ -491,12 +538,60 @@ def job_key_from_step(step: dict, label: str) -> str: return slug +def step_timeout_minutes(step: dict) -> int | None: + raw = step.get("timeout_in_minutes") + if raw is None: + return None + try: + minutes = int(raw) + except (TypeError, ValueError): + return None + return minutes if minutes > 0 else None + + +_INLINE_TIMEOUT_PREFIX = re.compile(r"^timeout\s+(?:-\S+\s+)*\S+\s+", re.IGNORECASE) + + +def timeout_kill_after_seconds() -> int: + raw = os.environ.get("RUN_JOB_TIMEOUT_KILL_AFTER", "60") + try: + seconds = int(raw) + except (TypeError, ValueError): + return 60 + return seconds if seconds > 0 else 60 + + +def prepend_timeout_to_pytest(pytest_line: str, timeout_min: int | None) -> str: + """Prepend GNU timeout directly before pytest (Buildkite inline style).""" + if timeout_min is None: + return pytest_line.strip() + line = _INLINE_TIMEOUT_PREFIX.sub("", pytest_line.strip()) + kill_after = timeout_kill_after_seconds() + return ( + f"timeout --foreground --verbose --kill-after={kill_after} " + f"{timeout_min}m {line}" + ) + + +def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -> None: + manifest_path = jobs_dir / ".job_timeouts" + if job_timeouts: + manifest_path.write_text( + "\n".join(f"{k}={v}" for k, v in sorted(job_timeouts.items())) + "\n", + encoding="utf-8", + ) + elif manifest_path.is_file(): + manifest_path.unlink() + + # Fixed stability scripts (when stability is selected); MODEL_TYPES narrows which run. STABILITY_CASES: list[tuple[str, str, tuple[str, ...]]] = [ ("stability_qwen3_omni", "tests/dfx/stability/scripts/test_stability_qwen3_omni.py", ("omni",)), ("stability_qwen3_tts", "tests/dfx/stability/scripts/test_stability_qwen3_tts.py", ("tts",)), + ("stability_voxcpm2", "tests/dfx/stability/scripts/test_stability_voxcpm2.py", ("tts",)), ("stability_qwen_image", "tests/dfx/stability/scripts/test_stability_qwen_image.py", ("diffusion",)), ("stability_wan22", "tests/dfx/stability/scripts/test_stability_wan22.py", ("diffusion",)), + ("stability_hunyuan_image", "tests/dfx/stability/scripts/test_stability_hunyuan_image.py", ("diffusion",)), ] @@ -533,14 +628,148 @@ def local_pytest_marker_expr(model_types: list[str]) -> str: return f"({' or '.join(families)}) and local_model" -def run_local_mode(jobs_dir: Path, model_types: list[str]) -> int: - """Single job: pytest -sv -m '<model markers> and local_model' (testpaths from pytest config).""" +def local_test_files_by_filename(substr: str) -> list[str]: + """Return repo-relative posix paths under tests/ whose basename contains substr.""" + tests_root = REPO_ROOT / "tests" + if not tests_root.is_dir(): + return [] + matches: list[str] = [] + for path in sorted(tests_root.rglob("test_*.py")): + if substr in path.name: + matches.append(path.relative_to(REPO_ROOT).as_posix()) + return matches + + +PERF_TESTS_REL = Path("tests/dfx/perf/tests") +RUN_BENCHMARK_REL = Path("tests/dfx/perf/scripts/run_benchmark.py") +RUN_DIFFUSION_BENCHMARK_REL = Path("tests/dfx/perf/scripts/run_diffusion_benchmark.py") +TTS_PERF_JSON_HINTS = ("test_tts", "voxcpm", "higgs_audio") + + +def perf_json_model_family(json_basename: str) -> str: + """Classify a perf JSON config as omni, tts, or diffusion (mirrors nightly YAML runners).""" + name = json_basename.lower() + if name.startswith("test_qwen3_omni"): + return "omni" + if any(hint in name for hint in TTS_PERF_JSON_HINTS): + return "tts" + return "diffusion" + + +def perf_json_runner(json_basename: str) -> Path: + if perf_json_model_family(json_basename) in ("omni", "tts"): + return RUN_BENCHMARK_REL + return RUN_DIFFUSION_BENCHMARK_REL + + +def perf_json_matches_model_type(json_basename: str, model_types: list[str]) -> bool: + if "all" in model_types: + return True + return perf_json_model_family(json_basename) in model_types + + +def local_perf_json_configs(substr: str, model_types: list[str]) -> list[tuple[str, str]]: + """Return (json_rel_posix, runner_rel_posix) for perf configs whose filename contains substr.""" + perf_dir = REPO_ROOT / PERF_TESTS_REL + if not perf_dir.is_dir(): + return [] + matches: list[tuple[str, str]] = [] + for path in sorted(perf_dir.glob("*.json")): + if substr not in path.name: + continue + if not perf_json_matches_model_type(path.name, model_types): + continue + runner = perf_json_runner(path.name) + if not (REPO_ROOT / runner).is_file(): + print(f"# skip (missing runner): {runner}", file=sys.stderr) + continue + matches.append( + ( + path.relative_to(REPO_ROOT).as_posix(), + runner.as_posix(), + ) + ) + return matches + + +def _local_job_slug(label: str) -> str: + slug = re.sub(r"[^\w\-.]+", "_", label, flags=re.UNICODE) + slug = re.sub(r"_+", "_", slug).strip("_") + return slug or "filter" + + +def run_local_mode( + jobs_dir: Path, + model_types: list[str], + perf_job_keys: list[str], +) -> int: + """pytest -sv -m '<model markers> and local_model'; optional LABEL_SUBSTR filters.""" marker_expr = local_pytest_marker_expr(model_types) - key = "local_pytest" + matched = 0 + + if LABEL_SUBSTR: + rel_paths = local_test_files_by_filename(LABEL_SUBSTR) + perf_pairs = local_perf_json_configs(LABEL_SUBSTR, model_types) + + if rel_paths: + paths_arg = " ".join(shlex.quote(p) for p in rel_paths) + pytest_line = f'pytest -sv -m "{marker_expr}" {paths_arg}' + slug = _local_job_slug(LABEL_SUBSTR) + key = f"local_pytest_{slug}" + header = ( + f"# Local marker tests — MODEL_TYPE={','.join(model_types)}, " + f"LABEL_SUBSTR={LABEL_SUBSTR!r} (test_*.py), files: {', '.join(rel_paths)}" + ) + script_lines = [ + "#!/usr/bin/env bash", + header, + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + pytest_line, + ] + _write_job_script(key, script_lines, jobs_dir) + matched += 1 + + for json_rel, runner_rel in perf_pairs: + json_name = Path(json_rel).name + pytest_line = ( + f'pytest -s -v -m "{marker_expr}" ' + f"{shlex.quote(runner_rel)} " + f"--test-config-file {shlex.quote(json_rel)}" + ) + slug = _local_job_slug(Path(json_rel).stem) + key = f"local_perf_{slug}" + header = ( + f"# Local perf benchmark — MODEL_TYPE={','.join(model_types)}, " + f"config={json_name!r}, runner={Path(runner_rel).name}" + ) + script_lines = [ + "#!/usr/bin/env bash", + header, + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + pytest_line, + ] + _write_job_script(key, script_lines, jobs_dir) + if key not in perf_job_keys: + perf_job_keys.append(key) + matched += 1 + + if matched == 0: + print( + f"# skip local: no tests/**/test_*.py or " + f"tests/dfx/perf/tests/*.json filename matches " + f"LABEL_SUBSTR={LABEL_SUBSTR!r}", + file=sys.stderr, + ) + return matched + pytest_line = f'pytest -sv -m "{marker_expr}"' + key = "local_pytest" + header = f"# Local marker tests — MODEL_TYPE={','.join(model_types)}" script_lines = [ "#!/usr/bin/env bash", - f"# Local marker tests — MODEL_TYPE={','.join(model_types)}", + header, "set -euo pipefail", f'cd "{REPO_ROOT}"', pytest_line, @@ -566,7 +795,7 @@ def run_stability_mode(jobs_dir: Path, model_types: list[str]) -> int: if not rel_path.is_file(): print(f"# skip (missing file): {rel_path}", file=sys.stderr) continue - pytest_line = f"pytest -s -v {rel_posix}" + pytest_line = f"pytest -s -v --run-level full_model {rel_posix}" matched += 1 script_lines = [ "#!/usr/bin/env bash", @@ -594,13 +823,15 @@ def main() -> None: matched_stability = 0 matched_yaml = 0 + matched_local = 0 perf_job_keys: list[str] = [] + job_timeouts: dict[str, int] = {} if want_stability: matched_stability = run_stability_mode(jobs_dir, model_types) if want_local: - run_local_mode(jobs_dir, model_types) + matched_local = run_local_mode(jobs_dir, model_types, perf_job_keys) if needs_yaml: if not YML.is_file(): @@ -634,8 +865,12 @@ def main() -> None: "set -euo pipefail", f'cd "{REPO_ROOT}"', ] + timeout_min = step_timeout_minutes(step) + if timeout_min is not None: + script_lines.append(f"# Buildkite timeout_in_minutes: {timeout_min}") + job_timeouts[key] = timeout_min script_lines.extend(exports) - script_lines.extend(pys) + script_lines.extend(prepend_timeout_to_pytest(p, timeout_min) for p in pys) _write_job_script(key, script_lines, jobs_dir) if "Perf Test" in label and key not in perf_job_keys: perf_job_keys.append(key) @@ -646,6 +881,7 @@ def main() -> None: manifest_path.write_text("\n".join(perf_job_keys) + "\n", encoding="utf-8") elif manifest_path.is_file(): manifest_path.unlink() + _write_job_timeouts_manifest(jobs_dir, job_timeouts) errs: list[str] = [] if want_stability and matched_stability == 0: @@ -653,6 +889,17 @@ def main() -> None: f"No stability jobs matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} " f"LABEL_SUBSTR={LABEL_SUBSTR!r} under {REPO_ROOT}" ) + if want_local and matched_local == 0: + if LABEL_SUBSTR: + errs.append( + f"No local jobs matched LABEL_SUBSTR={LABEL_SUBSTR!r} " + f"(tests/**/test_*.py or tests/dfx/perf/tests/*.json) " + f"under {REPO_ROOT}" + ) + else: + errs.append( + f"Local mode produced no jobs TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r}" + ) if needs_yaml and matched_yaml == 0: errs.append( f"No YAML steps matched TEST_TYPE={test_types!r} MODEL_TYPE={model_types!r} " @@ -671,29 +918,70 @@ if [[ "${DRY_RUN}" == "1" ]]; then exit 0 fi +# shellcheck source=tools/run_jobs_common.sh +source "${SCRIPT_DIR}/../run_jobs_common.sh" + +PERF_RESULTS_SRC="${REPO_ROOT}/tests/dfx/perf/results" + +# Copy perf benchmark JSON written under tests/dfx/perf/results/ during this run into LOG_DIR. +_collect_perf_result_jsons() { + local since_epoch="${1:?}" + local dest="${LOG_DIR}/perf_results" + + if [[ ! -d "${PERF_RESULTS_SRC}" ]]; then + return 0 + fi + + COLLECT_SINCE_EPOCH="${since_epoch}" \ + python3 - <<'PY' +import os +import shutil +import sys +from pathlib import Path + +repo = Path(os.environ["REPO_ROOT"]).resolve() +src = repo / "tests" / "dfx" / "perf" / "results" +dest = Path(os.environ["LOG_DIR"]).resolve() / "perf_results" +since = float(os.environ["COLLECT_SINCE_EPOCH"]) + +if not src.is_dir(): + sys.exit(0) + +copied = 0 +for path in sorted(src.rglob("*.json")): + if not path.is_file(): + continue + try: + if path.stat().st_mtime < since: + continue + except OSError: + continue + rel = path.relative_to(src) + out = dest / rel + out.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, out) + copied += 1 + +if copied: + print(f"Collected {copied} perf result JSON file(s) -> {dest}", file=sys.stderr) +PY +} + # Run jobs: YAML perf steps first (if manifest exists), then generate Excel (always when # manifest non-empty, even if some perf jobs failed; Excel uses whatever JSON exists on disk), # then remaining jobs. Paths align with nightly Buildkite (tests/dfx/perf/results). run_generated_jobs_with_tee() { set -o pipefail local any_fail=0 + local _run_start local perf_list="${LOG_DIR}/jobs/.perf_job_keys" local -a perf_keys=() + _run_jobs_reset_timing + _run_start="$(_run_jobs_epoch_seconds)" if [[ -f "${perf_list}" ]] && [[ -s "${perf_list}" ]]; then mapfile -t perf_keys < "${perf_list}" fi - _run_one_job() { - local _job="$1" - local base out job_status - base="$(basename "${_job}" .sh)" - out="${LOG_DIR}/${base}.log" - echo "==> ${_job} (tee ${out})" >&2 - (cd "${REPO_ROOT}" && bash "${_job}") 2>&1 | tee "${out}" - job_status="${PIPESTATUS[0]}" - return "${job_status}" - } - local k for k in "${perf_keys[@]}"; do [[ -n "${k}" ]] || continue @@ -702,11 +990,11 @@ run_generated_jobs_with_tee() { any_fail=1 continue fi - _run_one_job "${LOG_DIR}/jobs/${k}.sh" || any_fail=1 + _run_one_job_with_timing "${LOG_DIR}/jobs/${k}.sh" || any_fail=1 done if ((${#perf_keys[@]})); then - local excel_out excel_log excel_status excel_repo_logs + local excel_out excel_log excel_status excel_repo_logs excel_start excel_end excel_elapsed excel_repo_logs="${REPO_ROOT}/logs" mkdir -p "${excel_repo_logs}" excel_out="${excel_repo_logs}/nightly_perf_$(date -u +%Y%m%d-%H%M%S).xlsx" @@ -716,6 +1004,7 @@ run_generated_jobs_with_tee() { "(report reflects JSON already under tests/dfx/perf/results)." >&2 fi echo "==> python3 tools/nightly/generate_nightly_perf_excel.py -> ${excel_out} (tee ${excel_log})" >&2 + excel_start="$(_run_jobs_epoch_seconds)" ( cd "${REPO_ROOT}" && python3 tools/nightly/generate_nightly_perf_excel.py \ --input-dir "${REPO_ROOT}/tests/dfx/perf/results" \ @@ -723,7 +1012,13 @@ run_generated_jobs_with_tee() { --output-file "${excel_out}" ) 2>&1 | tee "${excel_log}" excel_status="${PIPESTATUS[0]}" - if [[ "${excel_status}" -ne 0 ]]; then + excel_end="$(_run_jobs_epoch_seconds)" + excel_elapsed=$((excel_end - excel_start)) + _run_jobs_record_timing "generate_nightly_perf_excel" "${excel_elapsed}" "${excel_status}" + if [[ "${excel_status}" -eq 0 ]]; then + echo " finished in $(_run_jobs_format_duration "${excel_elapsed}")" >&2 + else + echo " failed after $(_run_jobs_format_duration "${excel_elapsed}") (exit ${excel_status})" >&2 any_fail=1 echo "generate_nightly_perf_excel.py failed (exit ${excel_status}). See ${excel_log}" >&2 fi @@ -739,15 +1034,16 @@ run_generated_jobs_with_tee() { for _job in "${LOG_DIR}/jobs"/*.sh; do base="$(basename "${_job}" .sh)" [[ ${_perf_done["${base}"]+isset} ]] && continue - _run_one_job "${_job}" || any_fail=1 + _run_one_job_with_timing "${_job}" || any_fail=1 done shopt -u nullglob + _collect_perf_result_jsons "${_run_start}" + + _run_jobs_print_timing_summary "${_run_start}" "${any_fail}" if [[ "${any_fail}" -ne 0 ]]; then - echo "One or more selected jobs failed. See logs under ${LOG_DIR}." >&2 return 1 fi - echo "All selected jobs finished OK. Logs: ${LOG_DIR}/*.log" >&2 return 0 } diff --git a/tools/pre_commit/check_test_marks.py b/tools/pre_commit/check_test_marks.py new file mode 100644 index 00000000000..9a727ab721d --- /dev/null +++ b/tools/pre_commit/check_test_marks.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Run a pre-commit hook that fails if test files are modified or added +that (probably) never run in the CI. For now, this means that every tests file +needs to have a CI level marker (e.g., core_model, advanced_model, full_model, +etc) and hardware mark / helper so that we ensure mutated tests will actually +be selected as long as there are pytest commands pointing at the right paths. +""" + +import os +import re +import sys + +# CI level markers +LEVEL_MARKERS = ("core_model", "advanced_model", "full_model", "slow") + +# Hardware markers. These are platforms and resources that CI filters on. +# NOTE: If new hardware marks are added to pyproject.toml etc, we need +# to add them here as well. +HARDWARE_MARKERS = ( + "cpu", + "gpu", + "cuda", + "rocm", + "xpu", + "npu", + "musa", + "H100", + "L4", + "MI325", + "B60", + "S5000", + "A2", + "A3", +) + +# Helpers from tests/helpers/mark.py that auto-apply hardware marks. +HARDWARE_HELPERS = ("hardware_test", "hardware_marks") + +# Match mark.X since we could also do `from pytest import mark`. +# \b prevents matching prefixes (e.g., mark.slow vs mark.slow_test). +LEVEL_RE = re.compile(r"mark\.(?:" + "|".join(LEVEL_MARKERS) + r")\b") +HARDWARE_RE = re.compile(r"mark\.(?:" + "|".join(HARDWARE_MARKERS) + r")\b") +HELPER_RE = re.compile(r"(?:" + "|".join(HARDWARE_HELPERS) + r")\s*\(") + +MISSING_LEVEL_MARKER = "Level" +MISSING_HARDWARE_MARKER = "Hardware" + +# Check if a file is located under tests/ and matches test_<something>.py +# or <something>_test.py, since pytest technically collects on both. +# Note that we use the former everywhere in this repo by convention. +TEST_FILE_RE = re.compile(r"^tests/(?:.*/)?(?:test_[^/]*\.py$|[^/]*_test\.py$)") + + +def is_test_file(path: str) -> bool: + """Determine whether or not a path is pointing at a test file or not.""" + return bool(TEST_FILE_RE.match(path)) + + +def read_test_file(path: str) -> str | None: + """Read a test file's contents, or return None if it doesn't exist.""" + if not os.path.isfile(path): + return None + with open(path, encoding="utf-8") as f: + return f.read() + + +def has_level_marker(contents: str) -> bool: + """Check if file contents contain at least one CI level marker.""" + return bool(LEVEL_RE.search(contents)) + + +def has_hardware_marker(contents: str) -> bool: + """Check if file contents contain any hardware marker or helpers.""" + return bool(HARDWARE_RE.search(contents) or HELPER_RE.search(contents)) + + +def get_files_missing_markers( + staged_files: list[str], +) -> dict[str, list[str]]: + """Return a dict mapping file path to list of missing marker types.""" + results: dict[str, list[str]] = {} + for path in staged_files: + if is_test_file(path) and (contents := read_test_file(path)) is not None: + missing = [] + if not has_level_marker(contents): + missing.append(MISSING_LEVEL_MARKER) + if not has_hardware_marker(contents): + missing.append(MISSING_HARDWARE_MARKER) + if missing: + results[path] = missing + return results + + +if __name__ == "__main__": + missing = get_files_missing_markers(sys.argv[1:]) + + if missing: + file_lines = "\n".join(f" - {path} [{' and '.join(problems)}]" for path, problems in missing.items()) + print( + "\033[91merror:\033[0m test files are missing pytest marks " + "required for Buildkite CI collection.\n\n" + f"Level marks, e.g.: {', '.join(LEVEL_MARKERS[:4])}\n" + f"Hardware marks, e.g.: {', '.join(HARDWARE_MARKERS[:4])}, ...\n" + f" or helpers: {', '.join(HARDWARE_HELPERS)}\n\n" + "The following files are missing marks:\n" + f"{file_lines}\n\n" + "To skip: SKIP=check-test-ci-coverage git commit ..." + ) + sys.exit(1) diff --git a/tools/run_jobs_common.sh b/tools/run_jobs_common.sh new file mode 100644 index 00000000000..7aa0861e456 --- /dev/null +++ b/tools/run_jobs_common.sh @@ -0,0 +1,659 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Shared helpers for tools/run_*_jobs.sh (source this file; do not execute). +# +# Contents: +# 1) Timing / tee runners used by ready, merge, and nightly. +# 2) Ready/merge CLI + YAML extract entrypoint: run_yaml_ci_jobs_main +# (set BUILDKITE_REL and DEFAULT_LOG_SUBDIR in the entry script, then call it). +# +# Job timeouts: generated job scripts embed GNU timeout on the pytest line itself +# (see prepend_timeout_to_pytest in the Python heredoc below). This runner only +# executes bash job.sh and records exit 124 as TIMED OUT via jobs/.job_timeouts +# metadata. + +if [[ -n "${_RUN_JOBS_COMMON_LOADED:-}" ]]; then + return 0 2>/dev/null || exit 0 +fi +_RUN_JOBS_COMMON_LOADED=1 + +if [[ -z "${SCRIPT_DIR:-}" ]]; then + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +fi + +# --------------------------------------------------------------------------- +# Timing / tee helpers (ready, merge, nightly) +# --------------------------------------------------------------------------- + +RUN_JOB_NAMES=() +RUN_JOB_SECONDS=() +RUN_JOB_STATUSES=() + +_run_jobs_reset_timing() { + RUN_JOB_NAMES=() + RUN_JOB_SECONDS=() + RUN_JOB_STATUSES=() +} + +_run_jobs_epoch_seconds() { + date +%s +} + +_run_jobs_format_duration() { + local total="${1}" + local h m s + if [[ "${total}" -lt 60 ]]; then + printf '%ss' "${total}" + return 0 + fi + m=$((total / 60)) + s=$((total % 60)) + if [[ "${m}" -lt 60 ]]; then + if [[ "${s}" -eq 0 ]]; then + printf '%dm' "${m}" + else + printf '%dm %ss' "${m}" "${s}" + fi + return 0 + fi + h=$((m / 60)) + m=$((m % 60)) + if [[ "${m}" -eq 0 && "${s}" -eq 0 ]]; then + printf '%dh' "${h}" + elif [[ "${s}" -eq 0 ]]; then + printf '%dh %dm' "${h}" "${m}" + else + printf '%dh %dm %ss' "${h}" "${m}" "${s}" + fi +} + +_run_jobs_record_timing() { + local name="${1}" + local seconds="${2}" + local status="${3}" + RUN_JOB_NAMES+=("${name}") + RUN_JOB_SECONDS+=("${seconds}") + RUN_JOB_STATUSES+=("${status}") +} + +_run_jobs_lookup_timeout_minutes() { + local base="${1}" + local manifest="${LOG_DIR}/jobs/.job_timeouts" + local line key mins + [[ -f "${manifest}" ]] || return 1 + while IFS='=' read -r key mins; do + [[ -n "${key}" ]] || continue + if [[ "${key}" == "${base}" ]]; then + printf '%s' "${mins}" + return 0 + fi + done < "${manifest}" + return 1 +} + +_run_one_job_with_timing() { + local _job="$1" + local base out job_status start end elapsed timeout_min="" + base="$(basename "${_job}" .sh)" + out="${LOG_DIR}/${base}.log" + if timeout_min="$(_run_jobs_lookup_timeout_minutes "${base}")"; then + echo "==> ${_job} (tee ${out}, pytest inline timeout ${timeout_min}m)" >&2 + else + echo "==> ${_job} (tee ${out})" >&2 + fi + start="$(_run_jobs_epoch_seconds)" + (cd "${REPO_ROOT}" && bash "${_job}") 2>&1 | tee "${out}" + job_status="${PIPESTATUS[0]}" + end="$(_run_jobs_epoch_seconds)" + elapsed=$((end - start)) + _run_jobs_record_timing "${base}" "${elapsed}" "${job_status}" + if [[ "${job_status}" -eq 0 ]]; then + echo " finished in $(_run_jobs_format_duration "${elapsed}")" >&2 + elif [[ "${job_status}" -eq 124 && -n "${timeout_min}" ]]; then + echo " timed out (inline timeout ${timeout_min}m on pytest, exit 124)" >&2 + elif [[ "${job_status}" -eq 124 ]]; then + echo " timed out (exit 124)" >&2 + else + echo " failed after $(_run_jobs_format_duration "${elapsed}") (exit ${job_status})" >&2 + fi + return "${job_status}" +} + +_run_jobs_print_timing_summary() { + local wall_start="${1:-}" + local any_fail="${2:-0}" + local summary_path="${LOG_DIR}/timing_summary.log" + local -a lines=() + local i name secs status status_str failed_count=0 + local wall_elapsed job_count total_elapsed=0 + + lines+=("=== Job timing summary ===") + for i in "${!RUN_JOB_NAMES[@]}"; do + name="${RUN_JOB_NAMES[$i]}" + secs="${RUN_JOB_SECONDS[$i]}" + status="${RUN_JOB_STATUSES[$i]}" + if [[ "${status}" -eq 0 ]]; then + status_str="OK" + elif [[ "${status}" -eq 124 ]]; then + status_str="TIMED OUT" + failed_count=$((failed_count + 1)) + else + status_str="FAILED (exit ${status})" + failed_count=$((failed_count + 1)) + fi + lines+=(" ${name} $(_run_jobs_format_duration "${secs}") ${status_str}") + done + + job_count="${#RUN_JOB_NAMES[@]}" + if [[ -n "${wall_start}" ]]; then + wall_elapsed=$(( $(_run_jobs_epoch_seconds) - wall_start )) + lines+=("Total wall time: $(_run_jobs_format_duration "${wall_elapsed}") (${job_count} jobs)") + elif [[ "${job_count}" -gt 0 ]]; then + for secs in "${RUN_JOB_SECONDS[@]}"; do + total_elapsed=$((total_elapsed + secs)) + done + lines+=("Total job time: $(_run_jobs_format_duration "${total_elapsed}") (${job_count} jobs)") + else + lines+=("Total wall time: 0s (0 jobs)") + fi + + if [[ "${failed_count}" -gt 0 ]]; then + lines+=("Failed jobs: ${failed_count}/${job_count}") + fi + + if [[ "${any_fail}" -ne 0 ]]; then + lines+=("Result: one or more jobs failed. See logs under ${LOG_DIR}.") + else + lines+=("Result: all jobs finished OK. Logs: ${LOG_DIR}/*.log") + fi + + { + for line in "${lines[@]}"; do + printf '%s\n' "${line}" + done + } | tee "${summary_path}" >&2 +} + +# --------------------------------------------------------------------------- +# Ready / merge: CLI + YAML extract + run (tools/run_ready_jobs.sh, +# tools/run_merge_jobs.sh). Nightly does not use this entrypoint. +# +# Entry scripts must set before calling run_yaml_ci_jobs_main: +# BUILDKITE_REL - e.g. .buildkite/cuda/test-ready.yml +# DEFAULT_LOG_SUBDIR - e.g. ready_jobs +# --------------------------------------------------------------------------- + +_run_yaml_ci_jobs_usage() { + # Print the entry script header (caller is run_ready_jobs.sh / run_merge_jobs.sh). + sed -n '2,38p' "$0" | sed 's/^# \{0,1\}//' +} + +_run_jobs_split_append_csv_array() { + local -n _arr="${1}" + local _raw="${2}" + local _ifs="${IFS}" + local _part + IFS=',' + for _part in ${_raw}; do + IFS="${_ifs}" + _part="$(printf '%s' "${_part}" | tr '[:upper:]' '[:lower:]' | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')" + [[ -n "${_part}" ]] || continue + _arr+=("${_part}") + done +} + +_run_jobs_finalize_model_type_csv() { + if [[ "${MODEL_TYPE_FROM_CLI}" -eq 1 ]]; then + if ((${#MODEL_TYPE_CLI_PARTS[@]} == 0)); then + echo "--model-type requires a non-empty value" >&2 + exit 2 + fi + local -A _seen=() + local -a _out=() + local _x + for _x in "${MODEL_TYPE_CLI_PARTS[@]}"; do + if [[ "${_x}" == all ]]; then + printf '%s' "all" + return 0 + fi + [[ ${_seen["${_x}"]+isset} ]] && continue + _seen["${_x}"]=1 + _out+=("${_x}") + done + (IFS=','; printf '%s' "${_out[*]}") + else + printf '%s' "${MODEL_TYPE_ENV}" + fi +} + +_run_jobs_find_repo_containing_yml() { + local dir="${1:-}" + [[ -n "$dir" ]] || return 1 + dir="$(cd "$dir" && pwd)" || return 1 + while true; do + if [[ -f "${dir}/${BUILDKITE_REL}" ]]; then + printf '%s\n' "$dir" + return 0 + fi + [[ "${dir}" == "/" ]] && return 1 + dir="$(dirname "${dir}")" + done +} + +_run_jobs_derive_repo_root_from_yml() { + local yml="$1" + local d + d="$(cd "$(dirname "${yml}")" && pwd)" + while [[ "${d}" != "/" ]]; do + if [[ "$(basename "${d}")" == ".buildkite" ]]; then + printf '%s\n' "$(dirname "${d}")" + return 0 + fi + d="$(dirname "${d}")" + done + return 1 +} + +run_yaml_ci_jobs_main() { + if [[ -z "${BUILDKITE_REL:-}" || -z "${DEFAULT_LOG_SUBDIR:-}" ]]; then + echo "run_yaml_ci_jobs_main: set BUILDKITE_REL and DEFAULT_LOG_SUBDIR before calling." >&2 + exit 1 + fi + + LABEL_SUBSTR="${LABEL_SUBSTR:-}" + MODEL_TYPE_ENV="${MODEL_TYPE:-all}" + MODEL_TYPE_CLI_PARTS=() + MODEL_TYPE_FROM_CLI=0 + DRY_RUN="${DRY_RUN:-0}" + SKIP_SIMPLE=0 + + while [[ $# -gt 0 ]]; do + case "$1" in + -h | --help) + _run_yaml_ci_jobs_usage + exit 0 + ;; + --dry-run) + DRY_RUN=1 + shift + ;; + --log-dir) + LOG_DIR="$2" + shift 2 + ;; + --yaml) + YML="$2" + shift 2 + ;; + --repo-root) + REPO_ROOT="$2" + shift 2 + ;; + --label-substr) + LABEL_SUBSTR="$2" + shift 2 + ;; + --model-type) + MODEL_TYPE_FROM_CLI=1 + _run_jobs_split_append_csv_array MODEL_TYPE_CLI_PARTS "$2" + shift 2 + ;; + --skip-simple) + SKIP_SIMPLE=1 + shift + ;; + *) + echo "Unknown option: $1" >&2 + _run_yaml_ci_jobs_usage >&2 + exit 2 + ;; + esac + done + + MODEL_TYPE="$(_run_jobs_finalize_model_type_csv)" + + if [[ -n "${YML:-}" && -n "${REPO_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)" + YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" + elif [[ -n "${YML:-}" ]]; then + YML="$(cd "$(dirname "${YML}")" && pwd)/$(basename "${YML}")" + if ! REPO_ROOT="$(_run_jobs_derive_repo_root_from_yml "${YML}")"; then + echo "Could not derive REPO_ROOT from YML=${YML} (expected file at <repo>/${BUILDKITE_REL})." >&2 + echo "Set REPO_ROOT explicitly (or pass --repo-root) for pytest working directory." >&2 + exit 2 + fi + elif [[ -n "${REPO_ROOT:-}" ]]; then + REPO_ROOT="$(cd "${REPO_ROOT}" && pwd)" + YML="${REPO_ROOT}/${BUILDKITE_REL}" + else + REPO_ROOT="" + if command -v git >/dev/null 2>&1; then + REPO_ROOT="$(git -C "${PWD}" rev-parse --show-toplevel 2>/dev/null || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + REPO_ROOT="$(_run_jobs_find_repo_containing_yml "${PWD}" || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + REPO_ROOT="$(_run_jobs_find_repo_containing_yml "${SCRIPT_DIR}" || true)" + fi + if [[ -z "${REPO_ROOT}" ]]; then + echo "Could not locate ${BUILDKITE_REL}. Set REPO_ROOT or YML, run from inside the vllm-omni clone," >&2 + echo "or place this script (or run from a cwd) under the repository tree." >&2 + exit 2 + fi + YML="${REPO_ROOT}/${BUILDKITE_REL}" + fi + + LOG_DIR="${LOG_DIR:-${REPO_ROOT}/logs/${DEFAULT_LOG_SUBDIR}}" + + if [[ ! -f "${YML}" ]]; then + echo "YAML not found: ${YML}" >&2 + exit 2 + fi + + if ! command -v python3 >/dev/null 2>&1; then + echo "python3 is required." >&2 + exit 1 + fi + + mkdir -p "${LOG_DIR}/jobs" + LOG_DIR="$(cd "${LOG_DIR}" && pwd)" + export REPO_ROOT LOG_DIR YML LABEL_SUBSTR MODEL_TYPE DRY_RUN SKIP_SIMPLE + + if [[ "${DRY_RUN}" != "1" ]]; then + shopt -s nullglob + _stale=( "${LOG_DIR}/jobs"/*.sh ) + _stale_logs=( "${LOG_DIR}"/*.log ) + shopt -u nullglob + if ((${#_stale[@]})); then + rm -f "${_stale[@]}" + fi + if ((${#_stale_logs[@]})); then + rm -f "${_stale_logs[@]}" + fi + rm -f "${LOG_DIR}/jobs/.job_timeouts" + fi + + # shellcheck disable=SC2016,SC1078,SC1079 + python3 - <<'PY' +from __future__ import annotations + +import os +import re +import stat +import sys +from pathlib import Path + +try: + import yaml +except ImportError: + print("Missing PyYAML. Install with: pip install pyyaml", file=sys.stderr) + sys.exit(1) + +REPO_ROOT = Path(os.environ["REPO_ROOT"]).resolve() +LOG_DIR = Path(os.environ["LOG_DIR"]).resolve() +YML = Path(os.environ["YML"]).resolve() +LABEL_SUBSTR = (os.environ.get("LABEL_SUBSTR") or "").strip() +DRY_RUN = os.environ.get("DRY_RUN", "0") == "1" +SKIP_SIMPLE = os.environ.get("SKIP_SIMPLE", "0") == "1" + +ALLOWED_MODEL_TYPES = frozenset({"all", "omni", "tts", "diffusion"}) + + +def parse_model_types(raw: str) -> list[str]: + parts = [p.strip().lower() for p in (raw or "").split(",") if p.strip()] + if not parts: + parts = ["all"] + bad = [p for p in parts if p not in ALLOWED_MODEL_TYPES] + if bad: + print( + f"Invalid MODEL_TYPE / --model-type value(s): {bad!r} " + f"(allowed: {', '.join(sorted(ALLOWED_MODEL_TYPES))})", + file=sys.stderr, + ) + sys.exit(2) + if "all" in parts: + return ["all"] + out: list[str] = [] + seen: set[str] = set() + for p in parts: + if p not in seen: + seen.add(p) + out.append(p) + return out + + +PYTEST_CMD_RE = re.compile( + r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*" +) + + +def iter_leaf_steps(steps, group: str | None = None): + for raw in steps or []: + if not isinstance(raw, dict): + continue + nested = raw.get("steps") + if isinstance(nested, list) and nested: + g = raw.get("group") + next_group: str | None + if isinstance(g, str): + next_group = g + elif g is not None: + next_group = str(g) + else: + next_group = group + yield from iter_leaf_steps(nested, next_group) + continue + if raw.get("commands"): + yield raw, group + + +def raw_command_text(step: dict) -> str: + raw = step.get("commands") or [] + if isinstance(raw, str): + raw = [raw] + text = "\n".join((c.strip() if isinstance(c, str) else "") for c in raw if c) + return text.replace("$$", "$") + + +def export_lines(text: str) -> list[str]: + out: list[str] = [] + for line in text.splitlines(): + s = line.strip() + if not s or s.startswith("#"): + continue + if s.startswith("export "): + out.append(s) + return out + + +def pytest_lines(text: str) -> list[str]: + out: list[str] = [] + for m in PYTEST_CMD_RE.finditer(text): + line = m.group(0).strip() + line_start = text.rfind("\n", 0, m.start()) + 1 + before = text[line_start : m.start()] + if before.lstrip().startswith("#"): + continue + if line: + out.append(line) + return out + + +def is_simple_test_step(label: str, group: str | None) -> bool: + if group and "Simple Test" in group: + return True + return label.startswith("Simple ·") + + +def label_matches_model_type(label: str, model_types: list[str]) -> bool: + if "all" in model_types: + return True + if "omni" in model_types and "Omni ·" in label: + return True + if "tts" in model_types and "TTS ·" in label: + return True + if "diffusion" in model_types and "Diffusion" in label: + return True + return False + + +def job_key_from_step(step: dict, label: str) -> str: + k = step.get("key") + if isinstance(k, str) and k.strip(): + return k.strip() + slug = re.sub(r"[^\w\-.]+", "_", label or "job", flags=re.UNICODE) + slug = re.sub(r"_+", "_", slug).strip("_") or "job" + return slug + + +def step_timeout_minutes(step: dict) -> int | None: + raw = step.get("timeout_in_minutes") + if raw is None: + return None + try: + minutes = int(raw) + except (TypeError, ValueError): + return None + return minutes if minutes > 0 else None + + +_INLINE_TIMEOUT_PREFIX = re.compile(r"^timeout\s+(?:-\S+\s+)*\S+\s+", re.IGNORECASE) + + +def timeout_kill_after_seconds() -> int: + raw = os.environ.get("RUN_JOB_TIMEOUT_KILL_AFTER", "60") + try: + seconds = int(raw) + except (TypeError, ValueError): + return 60 + return seconds if seconds > 0 else 60 + + +def prepend_timeout_to_pytest(pytest_line: str, timeout_min: int | None) -> str: + """Prepend GNU timeout directly before pytest (Buildkite inline style).""" + if timeout_min is None: + return pytest_line.strip() + line = _INLINE_TIMEOUT_PREFIX.sub("", pytest_line.strip()) + kill_after = timeout_kill_after_seconds() + return ( + f"timeout --foreground --verbose --kill-after={kill_after} " + f"{timeout_min}m {line}" + ) + + +def _write_job_timeouts_manifest(jobs_dir: Path, job_timeouts: dict[str, int]) -> None: + manifest_path = jobs_dir / ".job_timeouts" + if job_timeouts: + manifest_path.write_text( + "\n".join(f"{k}={v}" for k, v in sorted(job_timeouts.items())) + "\n", + encoding="utf-8", + ) + elif manifest_path.is_file(): + manifest_path.unlink() + + +def _write_job_script(key: str, script_lines: list[str], jobs_dir: Path) -> None: + body = "\n".join(script_lines) + "\n" + if DRY_RUN: + print(f"=== {key} ===") + print(body) + return + job_path = jobs_dir / f"{key}.sh" + job_path.write_text(body, encoding="utf-8") + try: + mode = job_path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + job_path.chmod(mode) + except OSError: + pass + print(f"generated {job_path}", file=sys.stderr) + + +def main() -> None: + model_types = parse_model_types(os.environ.get("MODEL_TYPE", "all")) + + jobs_dir = LOG_DIR / "jobs" + if not DRY_RUN: + jobs_dir.mkdir(parents=True, exist_ok=True) + + if not YML.is_file(): + print(f"YAML not found: {YML}", file=sys.stderr) + sys.exit(1) + + data = yaml.safe_load(YML.read_text(encoding="utf-8")) + top_steps = (data or {}).get("steps") or [] + + matched_yaml = 0 + job_timeouts: dict[str, int] = {} + for step, grp in iter_leaf_steps(top_steps): + label = step.get("label") or "" + if SKIP_SIMPLE and is_simple_test_step(label, grp): + continue + if LABEL_SUBSTR and LABEL_SUBSTR not in label: + continue + if not label_matches_model_type(label, model_types): + continue + text = raw_command_text(step) + exports = export_lines(text) + pys = pytest_lines(text) + if not pys: + print(f"# skip (no pytest line): {label!r}", file=sys.stderr) + continue + + key = job_key_from_step(step, label) + matched_yaml += 1 + + script_lines = [ + "#!/usr/bin/env bash", + f'# From Buildkite label: {label.replace(chr(10), " ")}', + "set -euo pipefail", + f'cd "{REPO_ROOT}"', + ] + timeout_min = step_timeout_minutes(step) + if timeout_min is not None: + script_lines.append(f"# Buildkite timeout_in_minutes: {timeout_min}") + job_timeouts[key] = timeout_min + script_lines.extend(exports) + script_lines.extend(prepend_timeout_to_pytest(p, timeout_min) for p in pys) + _write_job_script(key, script_lines, jobs_dir) + + if not DRY_RUN: + _write_job_timeouts_manifest(jobs_dir, job_timeouts) + + if matched_yaml == 0: + print( + f"No YAML steps matched MODEL_TYPE={model_types!r} " + f"LABEL_SUBSTR={LABEL_SUBSTR!r} SKIP_SIMPLE={SKIP_SIMPLE!r} in {YML}", + file=sys.stderr, + ) + sys.exit(2) + + +main() +PY + + if [[ "${DRY_RUN}" == "1" ]]; then + exit 0 + fi + + set -o pipefail + local any_fail=0 + local _run_start + local _job + _run_jobs_reset_timing + _run_start="$(_run_jobs_epoch_seconds)" + + shopt -s nullglob + for _job in "${LOG_DIR}/jobs"/*.sh; do + _run_one_job_with_timing "${_job}" || any_fail=1 + done + shopt -u nullglob + + _run_jobs_print_timing_summary "${_run_start}" "${any_fail}" + if [[ "${any_fail}" -ne 0 ]]; then + exit 1 + fi + exit 0 +} diff --git a/tools/run_merge_jobs.sh b/tools/run_merge_jobs.sh new file mode 100644 index 00000000000..9f9ab13b02c --- /dev/null +++ b/tools/run_merge_jobs.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Extract steps from .buildkite/cuda/test-merge.yml that contain pytest, synthesize +# small bash wrappers (exports + pytest), run them, and tee output to logs named +# after each step's Buildkite "key" when present (otherwise a slug of the label). +# +# Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics): +# e.g. --model-type omni,tts +# omni — label contains "Omni ·" +# tts — label contains "TTS ·" +# diffusion— label contains "Diffusion" +# all — no model filter (default) +# +# --skip-simple — skip steps in the "Simple Test" Buildkite group and labels +# starting with "Simple ·" (L1-style unit tests in the same YAML) +# +# Requirements: bash, python3, PyYAML (pip install pyyaml) +# +# Usage: +# bash tools/run_merge_jobs.sh +# REPO_ROOT=/path/to/vllm-omni bash tools/run_merge_jobs.sh --model-type diffusion --dry-run +# YML=/path/to/vllm-omni/.buildkite/cuda/test-merge.yml bash tools/run_merge_jobs.sh +# +# Repository / YAML (no dependency on where this script lives): +# • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/cuda/test-merge.yml +# • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory +# • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD, +# then from the script's directory, until .buildkite/cuda/test-merge.yml exists +# +# Optional environment: +# REPO_ROOT - vllm-omni root (working directory for pytest); see above +# YML - path to test-merge.yml (default: $REPO_ROOT/.buildkite/cuda/test-merge.yml) +# LOG_DIR - logs + generated job scripts (default: $REPO_ROOT/logs/merge_jobs); +# per-job *.log plus timing_summary.log after the run +# MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above +# LABEL_SUBSTR - substring of Buildkite step label +# DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest +# +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BUILDKITE_REL=".buildkite/cuda/test-merge.yml" +DEFAULT_LOG_SUBDIR="merge_jobs" + +# shellcheck source=tools/run_jobs_common.sh +source "${SCRIPT_DIR}/run_jobs_common.sh" +run_yaml_ci_jobs_main "$@" diff --git a/tools/run_ready_jobs.sh b/tools/run_ready_jobs.sh new file mode 100644 index 00000000000..b854faa07e5 --- /dev/null +++ b/tools/run_ready_jobs.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Extract steps from .buildkite/cuda/test-ready.yml that contain pytest, synthesize +# small bash wrappers (exports + pytest), run them, and tee output to logs named +# after each step's Buildkite "key" when present (otherwise a slug of the label). +# +# Model area (--model-type / MODEL_TYPE), multiple allowed (OR semantics): +# e.g. --model-type omni,tts +# omni — label contains "Omni ·" +# tts — label contains "TTS ·" +# diffusion— label contains "Diffusion" +# all — no model filter (default) +# +# --skip-simple — skip steps in the "Simple Test" Buildkite group and labels +# starting with "Simple ·" (L1-style unit tests in the same YAML) +# +# Requirements: bash, python3, PyYAML (pip install pyyaml) +# +# Usage: +# bash tools/run_ready_jobs.sh +# REPO_ROOT=/path/to/vllm-omni bash tools/run_ready_jobs.sh --model-type omni --dry-run +# YML=/path/to/vllm-omni/.buildkite/cuda/test-ready.yml bash tools/run_ready_jobs.sh +# +# Repository / YAML (no dependency on where this script lives): +# • Set REPO_ROOT (or pass --repo-root) — default YAML is $REPO_ROOT/.buildkite/cuda/test-ready.yml +# • Or set YML (or --yaml) — repo root is inferred as parent of the .buildkite directory +# • Or run from inside the clone: git rev-parse --show-toplevel, else walk up from $PWD, +# then from the script's directory, until .buildkite/cuda/test-ready.yml exists +# +# Optional environment: +# REPO_ROOT - vllm-omni root (working directory for pytest); see above +# YML - path to test-ready.yml (default: $REPO_ROOT/.buildkite/cuda/test-ready.yml) +# LOG_DIR - logs + generated job scripts (default: $REPO_ROOT/logs/ready_jobs); +# per-job *.log plus timing_summary.log after the run +# MODEL_TYPE - comma-separated and/or repeated flags (default: all); see above +# LABEL_SUBSTR - substring of Buildkite step label +# DRY_RUN=1 - print extracted commands only; do not write scripts or run pytest +# +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BUILDKITE_REL=".buildkite/cuda/test-ready.yml" +DEFAULT_LOG_SUBDIR="ready_jobs" + +# shellcheck source=tools/run_jobs_common.sh +source "${SCRIPT_DIR}/run_jobs_common.sh" +run_yaml_ci_jobs_main "$@" diff --git a/vllm_omni/attention/fish_kvcache_attn.py b/vllm_omni/attention/fish_kvcache_attn.py index 9e552e76210..e7235168cc3 100644 --- a/vllm_omni/attention/fish_kvcache_attn.py +++ b/vllm_omni/attention/fish_kvcache_attn.py @@ -111,6 +111,8 @@ def can_use_fish_kvcache_attn( def _is_cuda_graph_capturing() -> bool: + if not torch.cuda.is_available(): + return False return bool(torch.cuda.is_current_stream_capturing()) diff --git a/vllm_omni/attention/fish_kvcache_backend.py b/vllm_omni/attention/fish_kvcache_backend.py index 3f3e9bdc6c7..0178f722ee1 100644 --- a/vllm_omni/attention/fish_kvcache_backend.py +++ b/vllm_omni/attention/fish_kvcache_backend.py @@ -233,7 +233,12 @@ def _forward_with_fish_kvcache( if attn_metadata is not None and not attn_metadata.use_cascade: num_actual_tokens = attn_metadata.num_actual_tokens - key_cache, value_cache = kv_cache.unbind(0) + # vLLM >=0.23.0 lays out the KV cache as + # (num_blocks, 2, block_size, num_kv_heads, head_size), so key/value + # live on dim 1 (matches FlashAttentionImpl's kv_cache.unbind(1)). + # The pre-0.23 layout had them on dim 0; unbind(0) here would split + # along num_blocks and raise "too many values to unpack". + key_cache, value_cache = kv_cache.unbind(1) q = query[:num_actual_tokens] out = output[:num_actual_tokens] active_batch = int(q.shape[0]) diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py b/vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py deleted file mode 100644 index 69fbe026bd8..00000000000 --- a/vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py +++ /dev/null @@ -1,255 +0,0 @@ -"""Daily-Omni: optional consistency check between text stream and generated speech. - -The benchmark MCQ accuracy uses ``generated_text`` only. When the omni server also -streams ``modality=audio`` (TTS), this module can transcribe the concatenated WAV -with Whisper and compare the inferred option letter to the one parsed from text. - -Requires ``openai-whisper`` (``pip install openai-whisper``). Enable via env -``DAILY_OMNI_TEXT_AUDIO_CONSISTENCY=1`` or CLI ``--daily-omni-text-audio-consistency``. - -Whisper model name defaults to ``tiny`` (override with ``DAILY_OMNI_WHISPER_MODEL``). -""" - -from __future__ import annotations - -import logging -import os -import re -import threading -from typing import Any - -from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniSampleRequest -from vllm_omni.benchmarks.data_modules.daily_omni_eval import extract_predicted_choice - -logger = logging.getLogger(__name__) - -_whisper_model = None -_whisper_model_name: str | None = None -_whisper_lock = threading.Lock() - - -def env_text_audio_check_enabled() -> bool: - return os.environ.get("DAILY_OMNI_TEXT_AUDIO_CONSISTENCY", "").lower() in ( - "1", - "true", - "yes", - ) - - -def extract_choice_from_asr_transcript(transcript: str) -> str | None: - """Parse A–D from ASR text; extends :func:`extract_predicted_choice` with spoken Chinese phrases.""" - c = extract_predicted_choice(transcript) - if c: - return c - t = transcript or "" - for pat in ( - r"(?i)选项\s*([ABCD])\b", - r"(?i)选\s*([ABCD])\b", - r"(?i)答案\s*是\s*([ABCD])\b", - r"(?i)答案\s*([ABCD])\b", - ): - m = re.search(pat, t) - if m: - return m.group(1).upper() - return None - - -def _get_whisper_model(model_name: str): - global _whisper_model, _whisper_model_name - with _whisper_lock: - if _whisper_model is None or _whisper_model_name != model_name: - import whisper - - logger.warning( - "Loading Whisper model %r for Daily-Omni text/audio consistency (one-time)...", - model_name, - ) - _whisper_model = whisper.load_model(model_name) - _whisper_model_name = model_name - return _whisper_model - - -def transcribe_wav_bytes( - wav_bytes: bytes, - *, - language: str | None = None, - model_name: str | None = None, -) -> tuple[str | None, str | None]: - """Transcribe WAV bytes. Returns ``(transcript, error)`` — one of them is set. - - Args: - wav_bytes: RIFF WAV file bytes. - language: Optional Whisper language code (e.g. ``en``, ``zh``); improves accuracy/latency. - model_name: Override model id; else ``DAILY_OMNI_WHISPER_MODEL`` or ``tiny``. - """ - if not wav_bytes: - return None, "empty_wav" - if model_name is None or not str(model_name).strip(): - model_name = os.environ.get("DAILY_OMNI_WHISPER_MODEL") or "tiny" - model_name = str(model_name).strip() or "tiny" - path: str | None = None - try: - import tempfile - - model = _get_whisper_model(model_name) - fd, path = tempfile.mkstemp(suffix=".wav") - with os.fdopen(fd, "wb") as fp: - fp.write(wav_bytes) - kwargs: dict = {} - if language: - kwargs["language"] = language - result = model.transcribe(path, **kwargs) - text = (result.get("text") or "").strip() - return (text if text else None), None - except ImportError: - return None, "openai-whisper is not installed (pip install openai-whisper)" - except Exception as e: - return None, str(e)[:500] - finally: - if path: - try: - os.unlink(path) - except OSError: - pass - - -def compute_daily_omni_text_audio_consistency_metrics( - input_requests: list[Any], - outputs: list[Any], - *, - include_per_item: bool = False, -) -> dict[str, Any] | None: - """Compare option letter from ``generated_text`` vs Whisper transcript of output audio. - - Only considers requests where ``outputs[i]`` has ``generated_audio_wav_bytes`` set - (populated by the omni benchmark when TA check is enabled). - """ - if not input_requests or len(input_requests) != len(outputs): - return None - if not all(isinstance(r, DailyOmniSampleRequest) for r in input_requests): - return None - - ta_no_wav = 0 - ta_asr_failed = 0 - ta_text_unparsed = 0 - ta_audio_unparsed = 0 - ta_consistent = 0 - ta_mismatch = 0 - ta_both_parsed = 0 - items: list[dict[str, Any]] = [] - - for req, out in zip(input_requests, outputs, strict=True): - assert isinstance(req, DailyOmniSampleRequest) - rid = req.request_id - if not getattr(out, "success", False): - if include_per_item: - items.append( - { - "request_id": rid, - "skipped": True, - "reason": "request_not_success", - } - ) - continue - - wav = getattr(out, "generated_audio_wav_bytes", None) - if not wav: - ta_no_wav += 1 - if include_per_item: - items.append( - { - "request_id": rid, - "skipped": False, - "reason": "no_output_audio", - "text_choice": extract_predicted_choice(getattr(out, "generated_text", "") or ""), - } - ) - continue - - transcript, asr_err = transcribe_wav_bytes(wav) - if asr_err: - ta_asr_failed += 1 - if include_per_item: - items.append( - { - "request_id": rid, - "asr_error": asr_err, - "text_choice": extract_predicted_choice(getattr(out, "generated_text", "") or ""), - } - ) - continue - - text_choice = extract_predicted_choice(getattr(out, "generated_text", "") or "") - audio_choice = extract_choice_from_asr_transcript(transcript or "") - - if text_choice is None: - ta_text_unparsed += 1 - if audio_choice is None: - ta_audio_unparsed += 1 - - if text_choice is not None and audio_choice is not None: - ta_both_parsed += 1 - if text_choice == audio_choice: - ta_consistent += 1 - else: - ta_mismatch += 1 - - if include_per_item: - consistent: bool | None - if text_choice is None or audio_choice is None: - consistent = None - else: - consistent = text_choice == audio_choice - items.append( - { - "request_id": rid, - "text_choice": text_choice, - "audio_choice": audio_choice, - "asr_transcript": (transcript or "")[:500], - "text_audio_consistent": consistent, - } - ) - - comparable = ta_consistent + ta_mismatch - rate = (ta_consistent / comparable) if comparable else None - - out: dict[str, Any] = { - "daily_omni_ta_enabled": True, - "daily_omni_ta_no_output_audio": ta_no_wav, - "daily_omni_ta_asr_failed": ta_asr_failed, - "daily_omni_ta_text_unparsed": ta_text_unparsed, - "daily_omni_ta_audio_unparsed": ta_audio_unparsed, - "daily_omni_ta_both_parsed": ta_both_parsed, - "daily_omni_ta_consistent": ta_consistent, - "daily_omni_ta_mismatch": ta_mismatch, - "daily_omni_ta_consistency_rate": rate, - } - if include_per_item: - out["daily_omni_ta_items"] = items - return out - - -def print_daily_omni_text_audio_summary(metrics: dict[str, Any]) -> None: - if not metrics.get("daily_omni_ta_enabled"): - return - print("{s:{c}^{n}}".format(s=" Daily-Omni text vs audio (ASR) ", n=50, c="=")) - print("{:<40} {:<10}".format("No output audio captured:", metrics.get("daily_omni_ta_no_output_audio", 0))) - print("{:<40} {:<10}".format("ASR failed:", metrics.get("daily_omni_ta_asr_failed", 0))) - print("{:<40} {:<10}".format("Both text+audio letter parsed:", metrics.get("daily_omni_ta_both_parsed", 0))) - print("{:<40} {:<10}".format("Consistent (same letter):", metrics.get("daily_omni_ta_consistent", 0))) - print("{:<40} {:<10}".format("Mismatch:", metrics.get("daily_omni_ta_mismatch", 0))) - r = metrics.get("daily_omni_ta_consistency_rate") - if r is not None: - print("{:<40} {:<10.4f}".format("Consistency rate (of both parsed):", r)) - print( - "{:<40} {:<10}".format( - "Text unparsed (among w/ audio):", - metrics.get("daily_omni_ta_text_unparsed", 0), - ) - ) - print( - "{:<40} {:<10}".format( - "Audio unparsed (among w/ audio):", - metrics.get("daily_omni_ta_audio_unparsed", 0), - ) - ) diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index 3ac8e076765..13c1f2c2177 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -33,10 +33,14 @@ Enable with ``SEED_TTS_WER_EVAL=1`` or ``--seed-tts-wer-eval``. Install optional deps:: - pip install 'vllm-omni[seed-tts-eval]' + pip install 'vllm-omni[dev]' Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_MODEL`` -defaults to ``openai/whisper-large-v3`` (override for debugging only). +defaults to ``openai/whisper-large-v3`` (override for debugging only). Set +``SEED_TTS_WER_SAVE_AUDIO_DIR`` to save the captured 24 kHz mono WAV used by +WER evaluation for each synthesized utterance. +Streaming PCM is decoded using ``VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE`` / +``VLLM_OMNI_BENCH_AUDIO_CHANNELS`` (default: 24 kHz mono). """ from __future__ import annotations @@ -50,12 +54,14 @@ import tempfile import threading import wave +from pathlib import Path from typing import Any import numpy as np from vllm.benchmarks.datasets import SampleRequest from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSSampleRequest +from vllm_omni.metrics.definitions import stream_pcm_format_from_env logger = logging.getLogger(__name__) @@ -87,6 +93,30 @@ def pcm_s16le_mono_to_wav_bytes(pcm: bytes, *, sample_rate: int = 24000) -> byte return buf.getvalue() +def _safe_filename_part(value: Any, *, default: str = "item", max_len: int = 96) -> str: + text = str(value or "").strip() or default + safe = "".join(ch if ch.isalnum() or ch in ("-", "_", ".") else "_" for ch in text) + safe = safe.strip("._") or default + return safe[:max_len] + + +def _save_seed_tts_eval_audio( + pcm: bytes, + *, + output_dir: Path | None, + index: int, + utterance_id: Any, + locale: str, +) -> str | None: + if output_dir is None: + return None + output_dir.mkdir(parents=True, exist_ok=True) + stem = f"{index:05d}_{_safe_filename_part(utterance_id)}_{_safe_filename_part(locale)}" + path = output_dir / f"{stem}.wav" + path.write_bytes(pcm_s16le_mono_to_wav_bytes(pcm, sample_rate=24000)) + return str(path) + + def _get_eval_device() -> str: explicit = os.environ.get("SEED_TTS_EVAL_DEVICE", "").strip() if explicit: @@ -146,12 +176,27 @@ def process_one_official(hypo: str, truth: str, lang: str) -> tuple[float, str, return wer, raw_truth, raw_hypo -def _pcm_s16le_to_f32_16k(pcm: bytes, pcm_sample_rate: int = 24000) -> np.ndarray: +def _pcm_s16le_to_f32_16k( + pcm: bytes, + pcm_sample_rate: int | None = None, + channels: int | None = None, +) -> np.ndarray: import scipy.signal if not pcm: return np.zeros(0, dtype=np.float32) + if pcm_sample_rate is None or channels is None: + env_sample_rate, env_channels = stream_pcm_format_from_env() + if pcm_sample_rate is None: + pcm_sample_rate = env_sample_rate + if channels is None: + channels = env_channels raw = np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0 + if channels > 1: + usable = (len(raw) // channels) * channels + if usable <= 0: + return np.zeros(0, dtype=np.float32) + raw = raw[:usable].reshape(-1, channels).mean(axis=1) target_len = int(len(raw) * 16000 / pcm_sample_rate) if target_len <= 0: return np.zeros(0, dtype=np.float32) @@ -446,7 +491,7 @@ def _missing_deps_message(lang: str) -> str | None: import jiwer # noqa: F401 from zhon.hanzi import punctuation # noqa: F401 except ImportError as e: - return f"Seed-TTS WER eval needs jiwer and zhon ({e!s}). Install: pip install 'vllm-omni[seed-tts-eval]'" + return f"Seed-TTS WER eval needs jiwer and zhon ({e!s}). Install: pip install 'vllm-omni[dev]'" try: import scipy.signal # noqa: F401 import soundfile # noqa: F401 @@ -457,13 +502,13 @@ def _missing_deps_message(lang: str) -> str | None: import torch # noqa: F401 from transformers import WhisperForConditionalGeneration # noqa: F401 except ImportError as e: - return f"English WER needs torch and transformers ({e!s}). Install: pip install 'vllm-omni[seed-tts-eval]'" + return f"English WER needs torch and transformers ({e!s}). Install: pip install 'vllm-omni[dev]'" else: try: import zhconv # noqa: F401 from funasr import AutoModel # noqa: F401 except ImportError as e: - return f"Chinese WER needs funasr and zhconv ({e!s}). Install: pip install 'vllm-omni[seed-tts-eval]'" + return f"Chinese WER needs funasr and zhconv ({e!s}). Install: pip install 'vllm-omni[dev]'" return None @@ -512,13 +557,18 @@ def compute_seed_tts_wer_metrics( sim_skipped_no_ref = 0 utmos_failed = 0 utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False) + save_audio_raw = os.environ.get("SEED_TTS_WER_SAVE_AUDIO_DIR", "").strip() + save_audio_dir = Path(save_audio_raw).expanduser() if save_audio_raw else None + saved_audio = 0 + save_audio_failed = 0 - for req, out in zip(input_requests, outputs, strict=True): + for index, (req, out) in enumerate(zip(input_requests, outputs, strict=True)): assert isinstance(req, SeedTTSSampleRequest) ref = req.prompt locale = req.seed_tts_locale or "en" row_lang = "zh" if locale.lower().startswith("zh") else "en" utmos_v: float | None = None + audio_path: str | None = None if not out.success: request_failed += 1 @@ -545,8 +595,30 @@ def compute_seed_tts_wer_metrics( } ) continue + try: + audio_path = _save_seed_tts_eval_audio( + pcm, + output_dir=save_audio_dir, + index=index, + utterance_id=req.seed_tts_utterance_id, + locale=locale, + ) + if audio_path: + saved_audio += 1 + except OSError as e: + save_audio_failed += 1 + logger.warning( + "Seed-TTS WER audio save failed for utterance=%s: %s", + req.seed_tts_utterance_id, + e, + ) - wav_16k = _pcm_s16le_to_f32_16k(pcm) + # Request functions normalize ``tts_output_pcm_bytes`` to Seed-TTS WER + # format before it reaches this evaluator: 24 kHz mono int16 PCM. + # Do not apply the streamed-response env format here again; for MOSS + # local that would reinterpret already-normalized 24 kHz mono audio as + # 48 kHz stereo and corrupt ASR/WER. + wav_16k = _pcm_s16le_to_f32_16k(pcm, pcm_sample_rate=24000, channels=1) if len(wav_16k) == 0: asr_failed += 1 if include_per_item: @@ -555,6 +627,7 @@ def compute_seed_tts_wer_metrics( "utterance_id": req.seed_tts_utterance_id, "locale": locale, "error": "empty_audio", + "audio_path": audio_path, } ) continue @@ -609,6 +682,7 @@ def compute_seed_tts_wer_metrics( "locale": locale, "error": "asr_exception", "detail": str(e)[:500], + "audio_path": audio_path, } ) continue @@ -621,6 +695,7 @@ def compute_seed_tts_wer_metrics( "utterance_id": req.seed_tts_utterance_id, "locale": locale, "error": "empty_asr", + "audio_path": audio_path, } ) continue @@ -637,6 +712,7 @@ def compute_seed_tts_wer_metrics( "locale": locale, "error": "wer_compute_failed", "detail": str(e)[:500], + "audio_path": audio_path, } ) continue @@ -673,6 +749,8 @@ def compute_seed_tts_wer_metrics( "reference_raw": raw_truth, "asr_raw": raw_hypo, } + if audio_path: + row["audio_path"] = audio_path if sim_v is not None: row["sim"] = sim_v if utmos_v is not None: @@ -697,7 +775,11 @@ def compute_seed_tts_wer_metrics( "seed_tts_utmos_mean": statistics.fmean(utmos_values) if utmos_values else None, "seed_tts_utmos_median": statistics.median(utmos_values) if utmos_values else None, "seed_tts_utmos_failed": utmos_failed, + "seed_tts_saved_audio": saved_audio, + "seed_tts_save_audio_failed": save_audio_failed, } + if save_audio_dir is not None: + result["seed_tts_save_audio_dir"] = str(save_audio_dir) if include_per_item: result["seed_tts_wer_eval_items"] = items return result @@ -729,6 +811,13 @@ def print_seed_tts_wer_summary(metrics: dict[str, Any]) -> None: print("{:<40} {:<10}".format("Request failed:", metrics.get("seed_tts_request_failed", 0))) print("{:<40} {:<10}".format("No PCM captured:", metrics.get("seed_tts_no_pcm", 0))) print("{:<40} {:<10}".format("ASR / WER failed:", metrics.get("seed_tts_asr_failed", 0))) + save_dir = metrics.get("seed_tts_save_audio_dir") + if save_dir: + print("{:<40} {:<10}".format("Saved eval WAVs:", metrics.get("seed_tts_saved_audio", 0))) + print("{:<40} {}".format("Saved eval WAV dir:", save_dir)) + failed = int(metrics.get("seed_tts_save_audio_failed", 0) or 0) + if failed: + print("{:<40} {:<10}".format("Save eval WAV errors:", failed)) if sim_ev or metrics.get("seed_tts_sim_skipped_no_ref") or metrics.get("seed_tts_sim_failed"): print("{:<40} {:<10}".format("SIM evaluated (higher ~ closer):", sim_ev)) sm = metrics.get("seed_tts_sim_mean") diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 100ab19effe..824b810f3ef 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/vllm_omni/benchmarks/metrics/metrics.py @@ -63,11 +63,9 @@ (defs.VLLM_TPOTS, _FLOAT_LIST_TYPE, field(default_factory=list)), (defs.VLLM_ITLS, _FLOAT_LIST_TYPE, field(default_factory=list)), (defs.AUDIO_TTFPS, _FLOAT_LIST_TYPE, field(default_factory=list)), - (defs.AUDIO_RTFS, _FLOAT_LIST_TYPE, field(default_factory=list)), (defs.AUDIO_DURATIONS, _FLOAT_LIST_TYPE, field(default_factory=list)), (defs.AUDIO_FRAMES, _INT_LIST_TYPE, field(default_factory=list)), (defs.MISSING_AUDIO_DURATION_COUNT, int, field(default=0)), - (defs.MISSING_AUDIO_RTF_COUNT, int, field(default=0)), (defs.STAGE_GEN_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)), (defs.POSTPROCESS_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)), ("output_unit_type", str, field(default="")), @@ -443,12 +441,6 @@ def _print_audio_stage_metrics( getattr(sm, defs.AUDIO_TTFPS), selected_percentiles, ) - if defs.AUDIO_RTF in selected_percentile_metrics: - audio_rtfs = getattr(sm, defs.AUDIO_RTFS) - if audio_rtfs: - _print_percentile_metric("Real Time Factor", "AUDIO_RTF", audio_rtfs, selected_percentiles, to_ms=False) - elif getattr(sm, defs.MISSING_AUDIO_RTF_COUNT) > 0: - print("{:<40} {:<10}".format("AUDIO_RTF skipped:", "missing stage-local audio duration/sample rate")) if defs.AUDIO_DURATION in selected_percentile_metrics: audio_durations = getattr(sm, defs.AUDIO_DURATIONS) if audio_durations: @@ -643,11 +635,9 @@ def _build_stage_metrics_from_outputs( total_output = sum(int((info or {}).get(defs.NUM_TOKENS_OUT) or 0) for _, info in rows) audio_ttfps: list[float] = [] - audio_rtfs: list[float] = [] audio_durations: list[float] = [] audio_frames: list[int] = [] missing_audio_duration_count = 0 - missing_audio_rtf_count = 0 if is_audio_stage: for _, info in rows: audio_ttfp_ms = float((info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) or 0.0) @@ -665,12 +655,6 @@ def _build_stage_metrics_from_outputs( else: missing_audio_duration_count += 1 - rtf = float((info or {}).get(defs.AUDIO_RTF) or 0.0) - if rtf > 0: - audio_rtfs.append(rtf) - else: - missing_audio_rtf_count += 1 - result.append( StageBenchmarkMetrics( **{ @@ -685,11 +669,9 @@ def _build_stage_metrics_from_outputs( defs.VLLM_TPOTS: vllm_tpots, defs.VLLM_ITLS: vllm_itls, defs.AUDIO_TTFPS: audio_ttfps, - defs.AUDIO_RTFS: audio_rtfs, defs.AUDIO_DURATIONS: audio_durations, defs.AUDIO_FRAMES: audio_frames, defs.MISSING_AUDIO_DURATION_COUNT: missing_audio_duration_count, - defs.MISSING_AUDIO_RTF_COUNT: missing_audio_rtf_count, defs.STAGE_GEN_TIMES_MS: stage_gen_times_ms, defs.POSTPROCESS_TIMES_MS: postprocess_times_ms, "output_unit_type": output_unit_type, @@ -768,11 +750,18 @@ def calculate_metrics( total_input += outputs[i].prompt_len tpot = 0 if output_len > 1: - try: - latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft - except Exception: - latency_minus_ttft = outputs[i].latency - outputs[i].ttft - tpot = latency_minus_ttft / (output_len - 1) + if outputs[i].itl: + # Use mean(ITL) directly so per-request TPOT == mean(ITL). + # The ITL list records one entry per SSE chunk; server may + # bundle multiple tokens per chunk, so len(itl)+1 != output_len. + # Using mean(itl) keeps TPOT and ITL on the same footing. + tpot = sum(outputs[i].itl) / len(outputs[i].itl) + else: + try: + latency_minus_ttft = outputs[i].text_latency - outputs[i].ttft + except Exception: + latency_minus_ttft = outputs[i].latency - outputs[i].ttft + tpot = latency_minus_ttft / (output_len - 1) tpots.append(tpot) # Note: if output_len <= 1, we regard tpot as 0 for goodput all_tpots.append(tpot) diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 73720f4bd65..e4f11b1e739 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -35,8 +35,6 @@ from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike -logger = init_logger(__name__) - from vllm_omni.benchmarks.audio_continuity import compute_continuity_stats from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset, DailyOmniSampleRequest from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset @@ -51,6 +49,8 @@ from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset from vllm_omni.metrics import definitions as defs +logger = init_logger(__name__) + _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S" RETURN_STAGE_METRICS_FIELD = "return_stage_metrics" _IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"}) @@ -109,6 +109,30 @@ def _audio_continuity_threshold_s() -> float: return max(value, 0.0) +def _pcm_s16le_to_seed_tts_wer_bytes( + pcm_bytes: bytes, + *, + sample_rate: int, + channels: int, +) -> bytes: + """Normalize streamed raw PCM to the 24 kHz mono PCM used by Seed-TTS WER.""" + if not pcm_bytes: + return b"" + channels = max(1, int(channels)) + pcm = np.frombuffer(pcm_bytes, dtype=np.int16) + if channels > 1: + usable = (pcm.size // channels) * channels + pcm = pcm[:usable].reshape(-1, channels).mean(axis=1) + pcm_f32 = pcm.astype(np.float32) / 32767.0 + if int(sample_rate) != 24000 and pcm_f32.size: + from vllm.multimodal.audio import AudioResampler + + resampler = AudioResampler(target_sr=24000) + pcm_f32 = resampler.resample(pcm_f32, orig_sr=int(sample_rate)) + pcm_f32 = np.clip(pcm_f32, -1.0, 1.0) + return (pcm_f32 * 32767).astype(np.int16).tobytes() + + get_samples_old = datasets.get_samples _DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni" @@ -530,6 +554,53 @@ def _extract_output_tokens_from_metrics(metrics: dict[str, Any]) -> int | None: return max(fallback_tokens, default=None) +def _apply_usage_to_output(output: MixRequestFuncOutput, usage: dict[str, Any]) -> int | None: + """Apply OpenAI ``usage`` fields to the benchmark output.""" + if (pt := _coerce_positive_int(usage.get("prompt_tokens"))) is not None: + output.prompt_len = pt + completion_tokens = _coerce_positive_int(usage.get("completion_tokens")) + if completion_tokens is not None: + output.output_tokens = max(int(output.output_tokens or 0), completion_tokens) + return completion_tokens + + +def _resolve_token_delta_from_usage( + completion_tokens: int | None, + completion_tokens_seen: int, +) -> tuple[int, int]: + if completion_tokens is None or completion_tokens <= completion_tokens_seen: + return 0, completion_tokens_seen + delta = completion_tokens - completion_tokens_seen + return delta, completion_tokens + + +def _record_text_token_stream_intervals( + output: MixRequestFuncOutput, + *, + timestamp: float, + start_time: float, + token_delta: int, + most_recent_timestamp: float, +) -> float: + """Record TTFT/ITL for ``token_delta`` newly generated text tokens.""" + if token_delta <= 0: + return most_recent_timestamp + + if output.ttft == 0.0: + output.ttft = timestamp - start_time + output.text_latency = timestamp - start_time + most_recent_timestamp = timestamp + if token_delta > 1: + output.itl.extend([0.0] * (token_delta - 1)) + return most_recent_timestamp + + interval = max(timestamp - most_recent_timestamp, 0.0) + per_token = interval / token_delta + output.itl.extend([per_token] * token_delta) + output.text_latency = timestamp - start_time + return timestamp + + def _update_output_stage_metrics_from_payload( output: MixRequestFuncOutput, data: dict[str, Any], @@ -541,7 +612,7 @@ def _update_output_stage_metrics_from_payload( return if update_output_tokens: if (num_tokens_out := _extract_output_tokens_from_metrics(metrics)) is not None: - output.output_tokens = num_tokens_out + output.output_tokens = max(int(output.output_tokens or 0), num_tokens_out) if isinstance(sid := metrics.get("stage_id"), int): output.stage_id = sid if isinstance(final_output_type := metrics.get("final_output_type"), str): @@ -624,6 +695,14 @@ async def async_request_openai_chat_omni_completions( "stream": True, "stream_options": { "include_usage": True, + # Per-chunk completion_tokens lets _resolve_token_delta_from_usage + # compute the exact token count for each SSE flush. Without this, + # one SSE chunk can carry multiple tokens (asyncio coalescing), so + # len(itl)+1 < actual_tokens and ITL measures per-chunk latency + # (~1.74× tokens_per_chunk) rather than true per-token latency. + # NOTE: the vLLM StreamOptions field is "continuous_usage_stats", + # NOT "include_continuous_usage". + "continuous_usage_stats": True, }, } _update_payload_common(payload, request_func_input) @@ -661,7 +740,6 @@ async def async_request_openai_chat_omni_completions( first_inconsistent_wav_params: tuple[int, int, int] | None = None # For non-wav responses, accumulate encoded bytes then decode once. audio_bytes_buffer = bytearray() - ttft = 0.0 st = time.perf_counter() output.start_time = st most_recent_timestamp = st @@ -685,6 +763,7 @@ async def async_request_openai_chat_omni_completions( output.image_generation_time_ms = 0.0 output.image_pixels = 0 output.denoise_step_latency_ms = 0.0 + completion_tokens_seen = 0 try: async with session.post(url=api_url, json=payload, headers=headers) as response: if response.status == 200: @@ -712,22 +791,45 @@ async def async_request_openai_chat_omni_completions( if chunk != "[DONE]": timestamp = time.perf_counter() data = json.loads(chunk) + _update_output_stage_metrics_from_payload(output, data) + usage = data.get("usage") + completion_tokens = None + if isinstance(usage, dict): + completion_tokens = _apply_usage_to_output(output, usage) + if choices := data.get("choices"): modality = data.get("modality") - delta = choices[0].get("delta") or {} + choice = choices[0] + delta = choice.get("delta") or {} content = delta.get("content") if not content and isinstance(delta.get("audio"), dict): content = delta["audio"].get("data") if modality == "text": - # First token - if ttft == 0.0: - ttft = timestamp - st - output.ttft = ttft - else: - output.itl.append(timestamp - most_recent_timestamp) - generated_text += content or "" - most_recent_timestamp = timestamp - output.text_latency = timestamp - st + token_delta, completion_tokens_seen = _resolve_token_delta_from_usage( + completion_tokens, + completion_tokens_seen, + ) + token_ids = choice.get("token_ids") + if token_delta == 0 and token_ids: + token_delta = len(token_ids) + if completion_tokens is not None: + completion_tokens_seen = max( + completion_tokens_seen, + completion_tokens, + ) + has_text_content = bool(content) + if token_delta == 0 and has_text_content and completion_tokens is None: + token_delta = 1 + if token_delta > 0: + most_recent_timestamp = _record_text_token_stream_intervals( + output, + timestamp=timestamp, + start_time=st, + token_delta=token_delta, + most_recent_timestamp=most_recent_timestamp, + ) + if has_text_content: + generated_text += content elif modality == "audio": if output.audio_ttfp == 0.0: output.audio_ttfp = timestamp - st @@ -762,7 +864,6 @@ async def async_request_openai_chat_omni_completions( if content_image_ms > 0: output.image_generation_time_ms += content_image_ms - _update_output_stage_metrics_from_payload(output, data) ( metrics_image_count, metrics_image_ms, @@ -778,10 +879,6 @@ async def async_request_openai_chat_omni_completions( if metrics_denoise_step_ms > output.denoise_step_latency_ms: output.denoise_step_latency_ms = metrics_denoise_step_ms - if usage := data.get("usage"): - if (pt := usage.get("prompt_tokens")) is not None: - output.prompt_len = pt - if wav_inconsistent_chunk_count > 0: logger.warning( "Dropped %d wav chunks with inconsistent params during benchmark " @@ -794,6 +891,12 @@ async def async_request_openai_chat_omni_completions( output.latency = timestamp - st output.generated_text = generated_text + if output.itl: + # Align text_latency with ITL so TPOT formula and + # mean(ITL) are consistent. Do NOT infer output_tokens + # from len(itl)+1: one SSE chunk may carry multiple + # tokens, so the ITL count understates the real count. + output.text_latency = output.ttft + sum(output.itl) audio_duration_sec = 0.0 audio_frames = 0 if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None: @@ -1026,9 +1129,9 @@ async def async_request_openai_audio_speech( ) -> MixRequestFuncOutput: """Streaming request to /v1/audio/speech endpoint. - Sends ``stream=true`` with ``response_format=pcm`` so the server returns - raw PCM chunks as they are decoded. This allows measuring TTFP (time to - first audio packet) separately from E2EL. + Sends ``stream=true`` with ``stream_format=audio`` and ``response_format=pcm`` + so the server returns raw PCM chunks as they are decoded. This allows measuring + TTFP (time to first audio packet) separately from E2EL. """ api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Audio Speech API", "audio/speech") @@ -1037,12 +1140,14 @@ async def async_request_openai_audio_speech( "model": request_func_input.model_name if request_func_input.model_name else request_func_input.model, "input": request_func_input.prompt, "stream": True, + "stream_format": "audio", "response_format": "pcm", } _update_payload_common(payload, request_func_input) # Seed-TTS + WER: ``--extra-body`` may set stream=false / other formats; speech must stream PCM. if getattr(request_func_input, "seed_tts_row", False) and _seed_tts_capture_pcm_for_wer(): payload["stream"] = True + payload["stream_format"] = "audio" payload["response_format"] = "pcm" headers = { @@ -1054,10 +1159,9 @@ async def async_request_openai_audio_speech( output = MixRequestFuncOutput() output.prompt_len = request_func_input.prompt_len - # PCM format: 16-bit signed, 24 kHz, mono - sample_rate = 24000 + # PCM format: 16-bit signed; sample_rate/channels are model-dependent. + sample_rate, channels = defs.stream_pcm_format_from_env() sample_width = 2 # 16-bit = 2 bytes - channels = 1 st = time.perf_counter() output.start_time = st @@ -1107,12 +1211,20 @@ async def async_request_openai_audio_speech( output.audio_continuity_ok = continuity.is_continuous output.audio_underrun_event_count = continuity.underrun_event_count if pcm_capture is not None and pcm_capture: - output.tts_output_pcm_bytes = bytes(pcm_capture) + try: + output.tts_output_pcm_bytes = _pcm_s16le_to_seed_tts_wer_bytes( + bytes(pcm_capture), + sample_rate=sample_rate, + channels=channels, + ) + except Exception as ex: + logger.warning("Seed-TTS WER PCM normalization failed: %s", ex) + output.tts_output_pcm_bytes = bytes(pcm_capture) elif capture_wer_pcm: ct = response.headers.get("Content-Type", "") logger.warning( "Seed-TTS WER: HTTP 200 but no PCM bytes (Content-Type=%r, url=%s). " - "Check stream=true and response_format=pcm on the server.", + "Check stream=true, stream_format=audio, and response_format=pcm on the server.", ct, api_url, ) @@ -1164,6 +1276,19 @@ async def async_request_openai_audio_speech( benchmark_old = serve.benchmark +def _merge_overrides(base: dict | None, overrides: dict | None) -> dict | None: + """Merge benchmark extra_body with per-request overrides. + + vLLM 0.24 removed the private helper from ``vllm.benchmarks.serve``. + Keep the same shallow-merge behavior here, with request overrides winning. + """ + if not base and not overrides: + return None + merged = dict(base or {}) + merged.update(overrides or {}) + return merged + + async def benchmark( task_type: TaskType, endpoint_type: str, @@ -1225,6 +1350,8 @@ async def benchmark( input_requests[0].expected_output_len, input_requests[0].multi_modal_data, ) + test_extra_body = _merge_overrides(extra_body, input_requests[0].request_overrides) + test_chat_messages = input_requests[0].chat_messages assert ( test_mm_content is None @@ -1242,7 +1369,8 @@ async def benchmark( multi_modal_content=test_mm_content, ignore_eos=ignore_eos, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=test_extra_body, + chat_messages=test_chat_messages, ) _attach_daily_omni_to_request_func_input(input_requests[0], test_input) _attach_seed_tts_to_request_func_input(input_requests[0], test_input) @@ -1306,7 +1434,8 @@ async def warmup_limited_request_func(): multi_modal_content=test_mm_content, ignore_eos=ignore_eos, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=test_extra_body, + chat_messages=test_chat_messages, ) _attach_daily_omni_to_request_func_input(input_requests[0], profile_input) _attach_seed_tts_to_request_func_input(input_requests[0], profile_input) @@ -1372,6 +1501,7 @@ async def limited_request_func(request_func_input, session, pbar): request.multi_modal_data, request.request_id, ) + per_request_extra_body = _merge_overrides(extra_body, request.request_overrides) req_model_id, req_model_name = model_id, model_name if lora_modules: req_lora_module = next(lora_modules) @@ -1388,8 +1518,9 @@ async def limited_request_func(request_func_input, session, pbar): multi_modal_content=mm_content, ignore_eos=ignore_eos, extra_headers=extra_headers, - extra_body=extra_body, + extra_body=per_request_extra_body, request_id=request_id, + chat_messages=request.chat_messages, ) _attach_daily_omni_to_request_func_input(request, request_func_input) _attach_seed_tts_to_request_func_input(request, request_func_input) @@ -1561,6 +1692,8 @@ def process_one_metric( prompt_len=test_prompt_len, output_len=test_output_len, logprobs=logprobs, + extra_body=test_extra_body, + chat_messages=test_chat_messages, ) profile_output = await request_func(request_func_input=profile_input, session=session) if profile_output.success: diff --git a/vllm_omni/config/__init__.py b/vllm_omni/config/__init__.py index e5762456f0b..3069025becd 100644 --- a/vllm_omni/config/__init__.py +++ b/vllm_omni/config/__init__.py @@ -2,20 +2,38 @@ Configuration module for vLLM-Omni. """ +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.lora import LoRAConfig from vllm_omni.config.model import OmniModelConfig +from vllm_omni.config.omni_config import ( + BaseVllmOmniStageConfig, + OmniStageCacheConfig, + OmniStageConnectorConfig, + OmniStageDiffusionParallelConfig, + OmniStageLoadConfig, + OmniStageModelConfig, + OmniStageParallelConfig, + OmniStageRuntimeConfig, + OmniStageSchedulerConfig, + StageConfigType, + VllmOmniARStageConfig, + VllmOmniConfig, + VllmOmniDiffusionStageConfig, + VllmOmniGenerationStageConfig, + VllmOmniOrchestratorConfig, +) +from vllm_omni.config.pipeline_registry import register_pipeline from vllm_omni.config.stage_config import ( + PIPELINE_WIDE_ENGINE_FIELDS, DeployConfig, PipelineConfig, StageConfig, - StageConfigFactory, StageDeployConfig, StageExecutionType, StagePipelineConfig, StageType, load_deploy_config, merge_pipeline_deploy, - register_pipeline, ) from vllm_omni.config.yaml_util import ( create_config, @@ -25,19 +43,40 @@ ) __all__ = [ - "OmniModelConfig", + # Legacy model-level configs. "LoRAConfig", + "OmniModelConfig", + # Structured Omni config entry points. + "VllmOmniConfig", + "BaseVllmOmniStageConfig", + "VllmOmniARStageConfig", + "VllmOmniGenerationStageConfig", + "VllmOmniDiffusionStageConfig", + "StageConfigType", + # Structured Omni sub-configs. + "OmniStageCacheConfig", + "OmniStageConnectorConfig", + "OmniStageDiffusionParallelConfig", + "OmniStageLoadConfig", + "OmniStageModelConfig", + "VllmOmniOrchestratorConfig", + "OmniStageParallelConfig", + "OmniStageRuntimeConfig", + "OmniStageSchedulerConfig", + # Legacy pipeline/stage deploy config surface. + "PIPELINE_WIDE_ENGINE_FIELDS", + "DeployConfig", + "PipelineConfig", "StageConfig", "StageConfigFactory", + "StageDeployConfig", "StageType", "StageExecutionType", "StagePipelineConfig", - "PipelineConfig", - "StageDeployConfig", - "DeployConfig", "load_deploy_config", "merge_pipeline_deploy", "register_pipeline", + # YAML utility helpers. "create_config", "load_yaml_config", "merge_configs", diff --git a/vllm_omni/config/composable_parallel/__init__.py b/vllm_omni/config/composable_parallel/__init__.py new file mode 100644 index 00000000000..499c1454853 --- /dev/null +++ b/vllm_omni/config/composable_parallel/__init__.py @@ -0,0 +1,104 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Composable parallel strategies for vLLM-Omni. + +A small, runtime-agnostic vocabulary for declaring a stage's parallel layout +(:class:`StrategySpec` + mesh/routing/aggregation patterns), plus the logic to +translate a per-role stack of specs into concrete engine sizing and overlay it +onto merged stage configs. + +The spec/pattern *types* are deliberately pure data with no engine imports; the +translation and apply logic lives alongside them and consumes vLLM-Omni's own +config objects. +""" + +from vllm_omni.config.composable_parallel.aggregation import ( + AGGREGATION_PATTERN_VARIANTS, + AggregationConflictError, + AggregationPattern, + AllGather, + Combine, + FanInByStage, + GatherDim, + StitchPipeline, + StitchSpatial, + TakeRank, + Union, +) +from vllm_omni.config.composable_parallel.apply import ( + StrategyApplyError, + StrategyApplyResult, + apply_strategy_specs, + check_device_layout, +) +from vllm_omni.config.composable_parallel.routing import ( + ROUTING_PATTERN_VARIANTS, + Broadcast, + DuplicateWithCondTag, + PartitionByHash, + PipelineMicrobatch, + RouteByExpert, + RouteByStage, + RoutingKeyError, + RoutingPattern, + ShardSequence, + ShardSpatial, +) +from vllm_omni.config.composable_parallel.spec import ( + MESH_AXIS_KINDS, + KernelSpec, + LayerHookSpec, + MeshAxisKind, + MeshAxisSpec, + SpecMergeConflictError, + StrategySpec, +) +from vllm_omni.config.composable_parallel.translator import ( + AxisTranslationError, + OmniParallelConfig, + translate_strategy_stack, +) + +__all__ = [ + # spec types + "MeshAxisKind", + "MESH_AXIS_KINDS", + "MeshAxisSpec", + "LayerHookSpec", + "KernelSpec", + "StrategySpec", + "SpecMergeConflictError", + # routing patterns + "RoutingPattern", + "RoutingKeyError", + "Broadcast", + "PartitionByHash", + "ShardSequence", + "ShardSpatial", + "RouteByExpert", + "PipelineMicrobatch", + "DuplicateWithCondTag", + "RouteByStage", + "ROUTING_PATTERN_VARIANTS", + # aggregation patterns + "AggregationPattern", + "AggregationConflictError", + "TakeRank", + "Union", + "GatherDim", + "AllGather", + "StitchSpatial", + "StitchPipeline", + "Combine", + "FanInByStage", + "AGGREGATION_PATTERN_VARIANTS", + # translation + "OmniParallelConfig", + "translate_strategy_stack", + "AxisTranslationError", + # apply + "apply_strategy_specs", + "StrategyApplyResult", + "check_device_layout", + "StrategyApplyError", +] diff --git a/vllm_omni/config/composable_parallel/aggregation.py b/vllm_omni/config/composable_parallel/aggregation.py new file mode 100644 index 00000000000..75ed8fadd63 --- /dev/null +++ b/vllm_omni/config/composable_parallel/aggregation.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Closed aggregation-pattern hierarchy for declarative parallel strategies. + +An aggregation pattern describes how the per-worker results of one mesh axis +are recombined into a single logical result (e.g. take one tensor-parallel +rank's output, union disjoint data-parallel shards). The hierarchy is *closed*: +subclasses may only be defined in this module. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + + +@dataclass(frozen=True) +class AggregationPattern: + """Base class for closed aggregation patterns. Subclass only within this module.""" + + def __init_subclass__(cls, **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + if cls.__module__ != __name__: + raise TypeError("AggregationPattern is closed to external subclassing") + + def pattern_kind(self) -> str: + raise NotImplementedError + + +class AggregationConflictError(ValueError): + """Raised when two aggregation inputs conflict for the same request id.""" + + +@dataclass(frozen=True) +class TakeRank(AggregationPattern): + """Select one rank's result (TP — all ranks agree).""" + + rank: int = 0 + + def __post_init__(self) -> None: + if not isinstance(self.rank, int) or isinstance(self.rank, bool): + raise TypeError(f"TakeRank.rank must be int, got {type(self.rank).__name__}") + + def pattern_kind(self) -> str: + return "take_rank" + + +@dataclass(frozen=True) +class Union(AggregationPattern): + """Union of disjoint replica results (DP, HSDP).""" + + def pattern_kind(self) -> str: + return "union" + + +@dataclass(frozen=True) +class GatherDim(AggregationPattern): + """Gather along a dimension and materialize on one designated rank. + + Default designated rank is axis-rank 0 unless adapter overrides. + """ + + dim: int = 1 + + def pattern_kind(self) -> str: + return "gather_dim" + + +@dataclass(frozen=True) +class AllGather(AggregationPattern): + """All ranks materialize the same aggregate result across the axis.""" + + dim: int | None = None + + def pattern_kind(self) -> str: + return "all_gather" + + +@dataclass(frozen=True) +class StitchSpatial(AggregationPattern): + """Stitch spatial patches (VAE-PP).""" + + grid: tuple[int, int] = (1, 1) + + def pattern_kind(self) -> str: + return "stitch_spatial" + + +@dataclass(frozen=True) +class StitchPipeline(AggregationPattern): + """Stitch pipeline stage outputs (PP).""" + + def pattern_kind(self) -> str: + return "stitch_pipeline" + + +@dataclass(frozen=True) +class Combine(AggregationPattern): + """Custom reducer combination (CFG). Adapter resolves reducer_id.""" + + reducer_id: str = "default" + + def pattern_kind(self) -> str: + return "combine" + + +@dataclass(frozen=True) +class FanInByStage(AggregationPattern): + """Fan-in results routed by stage (stage replicas).""" + + def pattern_kind(self) -> str: + return "fan_in_by_stage" + + +AGGREGATION_PATTERN_VARIANTS: tuple[type[AggregationPattern], ...] = ( + TakeRank, + Union, + GatherDim, + AllGather, + StitchSpatial, + StitchPipeline, + Combine, + FanInByStage, +) + +ReducerFn = Callable[..., object] diff --git a/vllm_omni/config/composable_parallel/apply.py b/vllm_omni/config/composable_parallel/apply.py new file mode 100644 index 00000000000..879c7aec3d8 --- /dev/null +++ b/vllm_omni/config/composable_parallel/apply.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Apply declarative parallel strategies onto merged stage configs. + +This is the *override-after-merge* seam. ``merge_pipeline_deploy`` first fuses +the pipeline topology + deploy YAML into a list of ``StageConfig`` objects; this +module then overlays the engine sizing derived from a per-role +``StrategySpec`` stack onto those stages. + +Design rules (one writer per axis): + +* A strategy only writes the axes it actually declares. If a stack has no ``tp`` + axis, this module never touches ``tensor_parallel_size`` — it does not force a + default of 1. +* **Conflict-on-explicit.** If a knob was already set explicitly by the deploy + YAML and the strategy derives a different value, we raise rather than silently + override. Equal values are a no-op; unset values are filled. +* **Pre-spawn device check.** When a stage declares ``devices``, the device + count must be consistent with the engine world size (``tp * dp * pp``) and the + replica count, so misconfigurations fail at config time instead of at spawn. + +``omni_lb_policy`` is *not* a per-stage config knob: omni reads it once at +orchestrator construction (``AsyncOmniEngine``), not from stage configs. So a +stage_replica axis's derived policy is surfaced on :class:`StrategyApplyResult` +for the caller to pass to the orchestrator; it is never silently written into a +stage where it would be ignored. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any, NoReturn + +from vllm.logger import init_logger + +from vllm_omni.config.composable_parallel.spec import StrategySpec +from vllm_omni.config.composable_parallel.translator import ( + OmniParallelConfig, + translate_strategy_stack, +) + +logger = init_logger(__name__) + +# Role key -> the StageConfig fields each declared axis writes. +_ENGINE_FIELD_BY_KIND: dict[str, str] = { + "tp": "tensor_parallel_size", + "dp": "data_parallel_size", + "pp": "pipeline_parallel_size", +} + +RoleKey = str + + +class StrategyApplyError(ValueError): + """Error for applying a strategy onto stage configs. + + A single error type (rather than per-cause subclasses) keeps the public + surface small and consistent with the rest of the codebase; the specific + cause — an unmatched role, a conflict with an explicit deploy value, or a + device-count mismatch — is in the message and is logged before the raise so + it stays visible even when the type is unavailable to a caller (e.g. across a + server boundary). + """ + + +def _fail(msg: str) -> NoReturn: + """Log and raise :class:`StrategyApplyError`.""" + logger.error("[composable_parallel] %s", msg) + raise StrategyApplyError(msg) + + +@dataclass +class StrategyApplyResult: + """Outcome of applying strategies to a stage list. + + ``stages`` is the same list passed in (mutated in place and returned for + convenience). ``omni_lb_policy`` is the pipeline-wide load-balancer policy + derived from any stage_replica axes (``None`` if none declared one); the + caller wires it into the orchestrator, since omni does not read it from + stage configs. + + ``per_role_config`` is keyed by the role key the caller supplied — a + ``model_stage`` name (str). ``per_stage_config`` is the same information + keyed by the resolved integer ``stage_id`` — handy for re-validating a + stage after later layers (e.g. CLI overrides) have run. + """ + + stages: list[Any] + omni_lb_policy: str | None = None + per_role_config: dict[RoleKey, OmniParallelConfig] = field(default_factory=dict) + per_stage_config: dict[int, OmniParallelConfig] = field(default_factory=dict) + + +def _resolve_stage(stages: Sequence[Any], key: RoleKey) -> Any: + """Find the single stage matching ``key`` by ``model_stage`` name. + + Role keys are ``model_stage`` names (str) — the same human-readable names + that strategy.yaml authors use to refer to a stage (e.g. ``"thinker"``). + """ + matches = [s for s in stages if getattr(s, "model_stage", None) == key] + descriptor = f"model_stage={key!r}" + + if not matches: + available = ", ".join(f"{getattr(s, 'model_stage', '?')!r}(id={getattr(s, 'stage_id', '?')})" for s in stages) + _fail(f"strategy role {descriptor} did not match any stage; available stages: {available}") + if len(matches) > 1: + _fail(f"strategy role {descriptor} is ambiguous; it matched {len(matches)} stages") + return matches[0] + + +def _set_explicit(engine_args: dict[str, Any], field_name: str, value: Any, *, role: RoleKey) -> None: + """Write ``field_name`` with conflict-on-explicit semantics. + + ``engine_args`` is the raw mapping loaded straight from the deploy YAML + (before the engine-args dataclass is constructed), so a key is *present* only + when the YAML set it explicitly. We therefore key off membership rather than + ``.get``: an explicit ``None`` in the YAML is a real, present value and must + not be conflated with a key miss (which ``.get`` would report as ``None``). + """ + if field_name in engine_args and engine_args[field_name] != value: + existing = engine_args[field_name] + _fail( + f"role {role!r}: strategy derives {field_name}={value} but the deploy config already " + f"set {field_name}={existing}. Remove one of them (the strategy is the single writer " + "for the axes it declares)." + ) + engine_args[field_name] = value + + +def _set_num_replicas(runtime: dict[str, Any], value: int, *, role: RoleKey) -> None: + """Write ``num_replicas`` with conflict-on-explicit (treating 1 as unset).""" + existing = runtime.get("num_replicas", 1) + if existing not in (1, value): + _fail( + f"role {role!r}: strategy derives num_replicas={value} but the deploy config already " + f"set num_replicas={existing}. Remove one of them." + ) + runtime["num_replicas"] = value + + +def _parse_device_count(devices: Any) -> int | None: + """Return the number of device ids in a ``devices`` value, or None.""" + if devices is None: + return None + if isinstance(devices, (list, tuple)): + return len([d for d in devices]) + text = str(devices).strip() + if not text: + return None + return len([d for d in text.split(",") if d.strip()]) + + +def check_device_layout( + devices: Any, + *, + tensor_parallel_size: int, + data_parallel_size: int, + pipeline_parallel_size: int, + num_replicas: int, + role: RoleKey, +) -> None: + """Validate a stage's device count against its world size (× replicas in pool mode). + + Raises :class:`StrategyApplyError` when the declared ``devices`` count is + neither the per-replica world size nor the full ``num_replicas`` pool. + Exposed publicly so the resolved (post-CLI) layout can be re-checked after + later override layers, not just the strategy-derived snapshot. + """ + count = _parse_device_count(devices) + if count is None: + return + world = int(tensor_parallel_size) * int(data_parallel_size) * int(pipeline_parallel_size) + replicas = int(num_replicas or 1) + # Two valid shapes: a single per-replica template (== world) or the full + # pool across replicas (== replicas * world). This mirrors omni's own + # per-replica device splitter. + valid = {world, replicas * world} + if count not in valid: + _fail( + f"role {role!r}: declared {count} device id(s) but the strategy world size is {world} " + f"(tp={tensor_parallel_size} * dp={data_parallel_size} * pp={pipeline_parallel_size}). " + f"Provide either {world} (per-replica template) or {replicas * world} " + f"(num_replicas={replicas} pool)." + ) + + +def _check_devices(runtime: dict[str, Any], cfg: OmniParallelConfig, *, role: RoleKey) -> None: + """Pre-spawn check on the strategy-derived snapshot (see :func:`check_device_layout`).""" + check_device_layout( + runtime.get("devices"), + tensor_parallel_size=cfg.tensor_parallel_size, + data_parallel_size=cfg.data_parallel_size, + pipeline_parallel_size=cfg.pipeline_parallel_size, + num_replicas=int(runtime.get("num_replicas", 1) or 1), + role=role, + ) + + +def _apply_to_stage(stage: Any, cfg: OmniParallelConfig, *, role: RoleKey) -> None: + engine_args = stage.yaml_engine_args + runtime = stage.yaml_runtime + + declared = set(cfg.l1_owners.keys()) + for kind, field_name in _ENGINE_FIELD_BY_KIND.items(): + if kind in declared: + _set_explicit(engine_args, field_name, getattr(cfg, field_name), role=role) + if "ep" in declared and cfg.enable_expert_parallel: + _set_explicit(engine_args, "enable_expert_parallel", True, role=role) + if "stage_replica" in declared: + _set_num_replicas(runtime, cfg.stage_replica_size, role=role) + + _check_devices(runtime, cfg, role=role) + + +def apply_strategy_specs( + stages: list[Any], + strategy_specs: Mapping[RoleKey, Sequence[StrategySpec]], +) -> StrategyApplyResult: + """Overlay per-role strategy specs onto a merged stage list. + + Args: + stages: the ``list[StageConfig]`` returned by ``merge_pipeline_deploy``. + strategy_specs: maps a role (a ``model_stage`` name, e.g. ``"thinker"``) + to that stage's stack of ``StrategySpec`` (one per declared mesh + axis). Role keys are ``model_stage`` names (str). + + Returns: + A :class:`StrategyApplyResult` with the mutated stages and the derived + pipeline-wide ``omni_lb_policy`` (if any stage_replica axis set one). + + Raises: + StrategyApplyError: a role matched zero or multiple stages, a derived + value conflicts with an explicit deploy value (including two roles + deriving different ``omni_lb_policy`` values), or a stage's device + count is inconsistent with its world size. + AxisTranslationError: the spec stack is invalid/unsupported. + NotImplementedError: the spec stack requests routing not built yet. + """ + result = StrategyApplyResult(stages=stages) + lb_policy: str | None = None + lb_owner: RoleKey | None = None + + for key, specs in strategy_specs.items(): + stage = _resolve_stage(stages, key) + cfg = translate_strategy_stack(specs) + _apply_to_stage(stage, cfg, role=key) + result.per_role_config[key] = cfg + stage_id = getattr(stage, "stage_id", None) + if stage_id is not None: + result.per_stage_config[int(stage_id)] = cfg + + if cfg.stage_replica_size > 1 and cfg.omni_lb_policy is not None: + if lb_policy is not None and lb_policy != cfg.omni_lb_policy: + _fail( + f"role {key!r} derives omni_lb_policy={cfg.omni_lb_policy!r} but role " + f"{lb_owner!r} already derived {lb_policy!r}. The load-balancer policy is " + "pipeline-wide; only one value is allowed." + ) + lb_policy = cfg.omni_lb_policy + lb_owner = key + + result.omni_lb_policy = lb_policy + return result diff --git a/vllm_omni/config/composable_parallel/routing.py b/vllm_omni/config/composable_parallel/routing.py new file mode 100644 index 00000000000..38e7d4a4930 --- /dev/null +++ b/vllm_omni/config/composable_parallel/routing.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Closed routing-pattern hierarchy for declarative parallel strategies. + +A routing pattern describes how a request batch is distributed across the +workers of one mesh axis (e.g. broadcast to every tensor-parallel rank, hash +partition across data-parallel ranks). The hierarchy is *closed*: subclasses +may only be defined in this module, so the set of patterns is a fixed, +exhaustively matchable contract. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Literal + + +@dataclass(frozen=True) +class RoutingPattern: + """Base class for closed routing patterns. Subclass only within this module.""" + + def __init_subclass__(cls, **kwargs: object) -> None: + super().__init_subclass__(**kwargs) + if cls.__module__ != __name__: + raise TypeError("RoutingPattern is closed to external subclassing") + + def pattern_kind(self) -> str: + raise NotImplementedError + + +class RoutingKeyError(KeyError): + """Raised when a routing key is missing from a request.""" + + def __init__(self, request_id: object, missing_key: str) -> None: + self.request_id = request_id + self.missing_key = missing_key + super().__init__(f"Missing routing key {missing_key!r} for request_id={request_id!r}") + + +@dataclass(frozen=True) +class Broadcast(RoutingPattern): + """Every worker in the group sees the whole batch (TP, dense EP).""" + + def pattern_kind(self) -> str: + return "broadcast" + + +@dataclass(frozen=True) +class PartitionByHash(RoutingPattern): + """Stable hash partition by request key (DP, HSDP).""" + + key: str = "request_id" + + def pattern_kind(self) -> str: + return "partition_by_hash" + + +@dataclass(frozen=True) +class ShardSequence(RoutingPattern): + """Shard along a sequence dimension (SP-Ulysses, SP-Ring, CP).""" + + dim: int = 1 + + def pattern_kind(self) -> str: + return "shard_sequence" + + +@dataclass(frozen=True) +class ShardSpatial(RoutingPattern): + """Shard along a spatial grid (VAE-PP).""" + + grid: tuple[int, int] = (1, 1) + + def pattern_kind(self) -> str: + return "shard_spatial" + + +@dataclass(frozen=True) +class RouteByExpert(RoutingPattern): + """Sparse MoE expert routing.""" + + router_id: str = "default" + + def pattern_kind(self) -> str: + return "route_by_expert" + + +@dataclass(frozen=True) +class PipelineMicrobatch(RoutingPattern): + """Pipeline parallelism microbatch routing (PP).""" + + microbatch_size: int = 1 + + def pattern_kind(self) -> str: + return "pipeline_microbatch" + + +@dataclass(frozen=True) +class DuplicateWithCondTag(RoutingPattern): + """Duplicate batch with a conditional tag (CFG).""" + + flag_name: str = "guidance_branch" + + def pattern_kind(self) -> str: + return "duplicate_with_cond_tag" + + +@dataclass(frozen=True) +class RouteByStage(RoutingPattern): + """Route requests by pipeline stage policy (stage replicas). + + The ``random``/``round_robin``/``least_queue`` policies describe stateless + load balancing that a runtime (e.g. an engine's built-in DP load balancer, + or the omni StagePool balancer) can realize. ``hash`` describes + deterministic, key-stable routing that a load balancer cannot honor on its + own (see ``PartitionByHash``). + """ + + routing_policy: Literal["random", "round_robin", "least_queue", "hash"] = "round_robin" + + def pattern_kind(self) -> str: + return "route_by_stage" + + +ROUTING_PATTERN_VARIANTS: tuple[type[RoutingPattern], ...] = ( + Broadcast, + PartitionByHash, + ShardSequence, + ShardSpatial, + RouteByExpert, + PipelineMicrobatch, + DuplicateWithCondTag, + RouteByStage, +) + +# Optional callable slot for Combine-style custom routing reducers in adapters. +RouterFn = Callable[..., object] diff --git a/vllm_omni/config/composable_parallel/spec.py b/vllm_omni/config/composable_parallel/spec.py new file mode 100644 index 00000000000..299fc3ad383 --- /dev/null +++ b/vllm_omni/config/composable_parallel/spec.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Core declarative types for parallel-strategy specification. + +A :class:`StrategySpec` declares one parallelism scheme: the mesh axis it sizes +(kind + degree), how a batch is routed across that axis, and how per-worker +results are aggregated back. A *stack* of specs (one per axis) fully describes a +stage's parallel layout in a runtime-agnostic, data-only form that can be +translated into concrete engine sizing. + +The mesh-axis kinds below enumerate every parallelism dimension the contract can +describe. Only a subset is wired end-to-end today (see ``translator.py``): + +* **Wired** (translatable now): ``tp``, ``dp``, ``pp``, ``ep``, + ``stage_replica``. +* **Reserved** (declared in the type system but not yet translatable — the + translator raises ``AxisTranslationError`` for them): ``sp_ulysses``, + ``sp_ring``, ``cfg``, ``vae_pp``, ``hsdp``, ``stage_pp``, ``cp``. + +Reserved kinds fail fast at translation time rather than silently doing nothing, +so declaring one is an explicit "not yet" rather than a no-op. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Literal + +from vllm_omni.config.composable_parallel.aggregation import AggregationPattern +from vllm_omni.config.composable_parallel.routing import RoutingPattern + + +class MeshAxisKind(str, Enum): + """Enumerates every parallelism dimension the strategy contract can describe. + + Subclassing ``str`` keeps existing string comparisons and dict keys working + (``MeshAxisKind.TP == "tp"`` and both hash the same), while making the type + refactor-safe: a typo'd kind is an ``AttributeError`` on the enum rather than + a silently-wrong bare string. + """ + + TP = "tp" + DP = "dp" + PP = "pp" + EP = "ep" + SP_ULYSSES = "sp_ulysses" + SP_RING = "sp_ring" + CFG = "cfg" + VAE_PP = "vae_pp" + HSDP = "hsdp" + STAGE_PP = "stage_pp" + STAGE_REPLICA = "stage_replica" + CP = "cp" + + +# Bare string values, kept as a plain tuple so membership checks and error +# messages read as the kind names (``"tp"``) rather than enum reprs. ``str``-enum +# equality means both ``"tp"`` and ``MeshAxisKind.TP`` satisfy ``in`` checks. +MESH_AXIS_KINDS: tuple[str, ...] = tuple(k.value for k in MeshAxisKind) + +HookCategory = Literal[ + "linear", + "attention_pre", + "attention_post", + "ffn_pre", + "ffn_post", + "moe_dispatch", + "other", +] + +HOOK_CATEGORY_ORDER: dict[HookCategory, int] = { + "linear": 0, + "attention_pre": 1, + "attention_post": 2, + "ffn_pre": 3, + "ffn_post": 4, + "moe_dispatch": 5, + "other": 6, +} + + +class SpecMergeConflictError(ValueError): + """Raised when merged hook/kernel specs have conflicting declarations.""" + + +@dataclass(frozen=True) +class MeshAxisSpec: + """Declares one axis of a process mesh (kind + size).""" + + kind: MeshAxisKind + size: int + + def __post_init__(self) -> None: + if self.kind not in MESH_AXIS_KINDS: + raise ValueError(f"MeshAxisSpec.kind must be one of {MESH_AXIS_KINDS}, got {self.kind!r}") + if self.size <= 0: + raise ValueError(f"MeshAxisSpec.size must be > 0, got {self.size}") + + +@dataclass(frozen=True) +class LayerHookSpec: + """L2 hook slot referenced by StrategySpec; resolved by the model walker.""" + + hook_id: str + target: str | None = None + category: HookCategory = "other" + priority: int = 0 + axis_index: int = 0 + + +@dataclass(frozen=True) +class KernelSpec: + """L3 kernel slot referenced by StrategySpec; resolved by the kernel registry.""" + + kernel_id: str + target: str | None = None + category: HookCategory = "other" + priority: int = 0 + axis_index: int = 0 + group_axis_kind: MeshAxisKind | None = None + requires_collective: bool = False + + +@dataclass(frozen=True) +class StrategySpec: + """Declarative contract for one parallelism scheme (data + hook/kernel slots).""" + + name: str + mesh_axis: MeshAxisSpec + routing: RoutingPattern + aggregation: AggregationPattern + layer_hook_specs: tuple[LayerHookSpec, ...] = () + kernel_specs: tuple[KernelSpec, ...] = () + shard_extension: Mapping[str, Any] = field(default_factory=dict) diff --git a/vllm_omni/config/composable_parallel/strategy_loader.py b/vllm_omni/config/composable_parallel/strategy_loader.py new file mode 100644 index 00000000000..665fbe9c2f0 --- /dev/null +++ b/vllm_omni/config/composable_parallel/strategy_loader.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Load per-role parallel strategies from a strategy file. + +The strategy file is a small, optional companion to the deploy config. It maps a +logical role (``model_stage``, e.g. ``thinker``) to a list of axis declarations, +each of which becomes one :class:`StrategySpec`. YAML is the only format wired up +today, but the module/symbol names stay format-agnostic so a different source can +be added without renaming the public surface:: + + strategies: + thinker: + - axis: tp + size: 2 + talker: + - axis: stage_replica + size: 2 + routing: round_robin # random | round_robin | least_queue + code2wav: + - axis: stage_replica + size: 2 + routing: round_robin + +Per-axis fields: + +* ``axis`` (required): mesh-axis kind (``tp``, ``dp``, ``pp``, ``ep``, + ``stage_replica`` are translatable today). +* ``size`` (required): axis degree (> 0). +* ``routing`` (optional): stateless policy for ``dp`` / ``stage_replica`` axes + (``random`` / ``round_robin`` / ``least_queue``). Ignored for axes whose + routing is fixed (``tp`` broadcast, ``pp`` microbatch). +* ``l1_owner`` (optional): ``engine`` or ``delegated``; overrides the default + owner for the axis (recorded in ``shard_extension``). +* ``name`` (optional): a label for the spec; defaults to the axis kind. + +Routing/aggregation patterns are filled in from sensible per-kind defaults so +the file stays compact; the translator validates the combination. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +from vllm_omni.config.composable_parallel.aggregation import ( + AggregationPattern, + FanInByStage, + StitchPipeline, + TakeRank, + Union, +) +from vllm_omni.config.composable_parallel.routing import ( + Broadcast, + PipelineMicrobatch, + RouteByStage, + RoutingPattern, +) +from vllm_omni.config.composable_parallel.spec import MeshAxisSpec, StrategySpec +from vllm_omni.config.yaml_util import load_yaml_config, to_dict + +# Per-kind defaults for routing + aggregation, so the file only needs axis+size. +_ROUTING_POLICY_KINDS = frozenset({"dp", "stage_replica"}) + + +class StrategyLoadError(ValueError): + """Raised when a strategy file is malformed.""" + + +def _default_routing(kind: str, routing_policy: str | None) -> RoutingPattern: + if kind in _ROUTING_POLICY_KINDS: + return RouteByStage(routing_policy=routing_policy or "round_robin") # type: ignore[arg-type] + if kind == "pp": + return PipelineMicrobatch() + # tp, ep, and any other kind default to broadcast (the translator rejects + # kinds it cannot handle before it ever inspects routing). + return Broadcast() + + +def _default_aggregation(kind: str) -> AggregationPattern: + if kind == "stage_replica": + return FanInByStage() + if kind == "pp": + return StitchPipeline() + if kind in ("dp", "ep"): + return Union() + return TakeRank() + + +def _build_spec(role: str, entry: Mapping[str, Any]) -> StrategySpec: + if "axis" not in entry: + raise StrategyLoadError(f"role {role!r}: each strategy entry needs an 'axis' field") + if "size" not in entry: + raise StrategyLoadError(f"role {role!r}: strategy entry for axis {entry['axis']!r} needs a 'size'") + + kind = str(entry["axis"]) + try: + size = int(entry["size"]) + except (TypeError, ValueError) as exc: + raise StrategyLoadError(f"role {role!r}: axis {kind!r} size must be an integer, got {entry['size']!r}") from exc + + routing_policy = entry.get("routing") + if routing_policy is not None and kind not in _ROUTING_POLICY_KINDS: + raise StrategyLoadError( + f"role {role!r}: axis {kind!r} does not accept a 'routing' policy (only {sorted(_ROUTING_POLICY_KINDS)} do)" + ) + + shard_extension: dict[str, Any] = {} + l1_owner = entry.get("l1_owner") + if l1_owner is not None: + shard_extension["l1_owner"] = str(l1_owner) + + return StrategySpec( + name=str(entry.get("name", kind)), + mesh_axis=MeshAxisSpec(kind=kind, size=size), # type: ignore[arg-type] + routing=_default_routing(kind, routing_policy), + aggregation=_default_aggregation(kind), + shard_extension=shard_extension, + ) + + +def parse_strategy_specs(data: Mapping[str, Any]) -> dict[str, list[StrategySpec]]: + """Parse an already-loaded strategy mapping into per-role spec stacks.""" + strategies = data.get("strategies", data) + if not isinstance(strategies, Mapping): + raise StrategyLoadError("strategy file must contain a 'strategies' mapping of role -> list of axes") + + result: dict[str, list[StrategySpec]] = {} + for role, entries in strategies.items(): + if not isinstance(entries, (list, tuple)): + raise StrategyLoadError( + f"role {role!r}: expected a list of axis declarations, got {type(entries).__name__}" + ) + specs: list[StrategySpec] = [] + for entry in entries: + if not isinstance(entry, Mapping): + raise StrategyLoadError( + f"role {role!r}: each axis declaration must be a mapping, got {type(entry).__name__}" + ) + specs.append(_build_spec(str(role), dict(entry))) + result[str(role)] = specs + return result + + +def load_strategy_specs(path: str) -> dict[str, list[StrategySpec]]: + """Load and parse a strategy file into per-role spec stacks.""" + data = to_dict(load_yaml_config(path)) + if not isinstance(data, Mapping): + raise StrategyLoadError(f"strategy file {path!r} did not parse to a mapping") + return parse_strategy_specs(data) diff --git a/vllm_omni/config/composable_parallel/translator.py b/vllm_omni/config/composable_parallel/translator.py new file mode 100644 index 00000000000..5c394dc9b5f --- /dev/null +++ b/vllm_omni/config/composable_parallel/translator.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Translate a ``StrategySpec`` stack into vLLM-Omni parallel sizing. + +Read a stack of strategy specs (one per mesh axis) and work out the +tensor/data/pipeline parallel sizes, the dense-EP flag, and the number of stage +replicas, plus who owns each axis's request routing. The result is a plain, +CPU-computable sizing struct (:class:`OmniParallelConfig`) keyed by the real +``OmniEngineArgs`` / ``EngineArgs`` field names, so the deploy layer can splat it +onto a stage's engine args. + +The distinction this module enforces — engine data parallelism (a true vLLM +intra-engine world dimension) vs. omni stage replicas (independent engines fanned +out by omni's coordinator) — is documented on the axis validators that depend on +it (see :func:`_validate_dp` and :func:`_stage_replica_lb_policy`). Routing we do +not support yet (key-stable / affinity routing) raises ``NotImplementedError``; +any other invalid spec raises :class:`AxisTranslationError`. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Literal, NoReturn, cast, get_args + +from vllm.logger import init_logger + +from vllm_omni.config.composable_parallel.routing import ( + Broadcast, + PartitionByHash, + PipelineMicrobatch, + RouteByStage, + RoutingPattern, +) +from vllm_omni.config.composable_parallel.spec import MeshAxisKind, StrategySpec + +logger = init_logger(__name__) + +# Who owns an axis's request routing. A closed string type so an unexpected raw +# value is caught statically and at runtime (``_VALID_L1_OWNERS`` is derived from +# it, keeping one source of truth) rather than silently flowing through. +L1Owner = Literal["delegated", "engine"] + +# DP, TP, PP, (dense) EP and stage_replica translate today. DP/TP/PP are true +# world dimensions vLLM realizes intra-engine; dense EP is a flag over the +# existing TP*DP ranks; stage_replica is an omni-coordinator-level fan-out of +# independent engine replicas (NOT a vLLM world dimension). The remaining kinds +# (sequence/context parallel, CFG, VAE pipelines, ...) need per-layer hooks or +# custom collectives and arrive in later stages. +_SUPPORTED_KINDS: tuple[MeshAxisKind, ...] = ("dp", "tp", "pp", "ep", "stage_replica") + +# Which EngineArgs field each axis kind *sizes*. EP and stage_replica are +# deliberately absent: EP is not an independent world dimension but an +# ``enable_expert_parallel`` flag whose degree equals tp*dp; stage_replica is a +# per-stage deploy ``num_replicas`` count, not an intra-engine world dimension. +_AXIS_TO_ENGINE_FIELD: dict[MeshAxisKind, str] = { + "tp": "tensor_parallel_size", + "dp": "data_parallel_size", + "pp": "pipeline_parallel_size", +} + +# Default L1 owner per axis kind. ``stage_replica`` is the only delegated axis: +# omni's coordinator/StagePool owns routing across replicas. ``dp`` is engine +# data parallelism, realized by vLLM's own intra-engine DP load balancer. +_DEFAULT_L1_OWNER: dict[MeshAxisKind, L1Owner] = { + "dp": "engine", + "tp": "engine", + "pp": "engine", + "ep": "engine", + "stage_replica": "delegated", +} + +# How a RouteByStage policy maps onto omni's load balancer policy string. These +# are the stateless options omni can actually do; note there's no "hash" here, +# because omni has no key-stable balancer. +_STAGE_POLICY_TO_OMNI_LB: dict[str, str] = { + "random": "random", + "round_robin": "round-robin", + "least_queue": "least-queue-length", +} + +_VALID_L1_OWNERS: frozenset[str] = frozenset(get_args(L1Owner)) + + +class AxisTranslationError(ValueError): + """Error for an invalid or unsupported strategy spec. + + A single error type (rather than a tree of subclasses) keeps the public + surface small and consistent with the rest of the codebase; the specific + cause is in the message and is logged before the raise so it is visible even + when the type is unavailable to a caller (e.g. across a server boundary). + + Strategies that are *valid but not built yet* — key-stable / affinity routing + today — raise ``NotImplementedError`` instead, to distinguish "we haven't + implemented this" from "your config is wrong". + """ + + +def _fail(msg: str) -> NoReturn: + """Log and raise :class:`AxisTranslationError` for an invalid/unsupported spec.""" + logger.error("[composable_parallel] %s", msg) + raise AxisTranslationError(msg) + + +def _not_implemented(msg: str) -> NoReturn: + """Log and raise ``NotImplementedError`` for a valid-but-unbuilt strategy.""" + logger.error("[composable_parallel] %s", msg) + raise NotImplementedError(msg) + + +@dataclass(frozen=True) +class OmniParallelConfig: + """Result of translating a spec stack into omni parallel sizing.""" + + tensor_parallel_size: int = 1 + data_parallel_size: int = 1 + pipeline_parallel_size: int = 1 + # Dense expert parallelism: a flag over the TP*DP ranks, not a world axis. + enable_expert_parallel: bool = False + # Number of independent omni stage replicas (the ``stage_replica`` axis). + # This is NOT a vLLM world dimension; it is the per-stage deploy + # ``num_replicas`` count. 1 means a single (un-replicated) engine. + stage_replica_size: int = 1 + # The omni StagePool LB policy string for the (delegated) stage_replica + # axis, if any. Only meaningful when ``stage_replica_size > 1``. + omni_lb_policy: str | None = None + # axis kind -> resolved L1 owner ("delegated" | "engine"). + l1_owners: Mapping[MeshAxisKind, L1Owner] = field(default_factory=dict) + + @property + def world_size(self) -> int: + # EP and stage_replica are intentionally excluded. EP reuses the TP*DP + # ranks rather than adding a dimension; stage_replica spins up separate + # engines (each its own world), not extra ranks in this engine's group. + # This product is what vLLM calls ``world_size_across_dp`` (tp * pp * dp). + return self.tensor_parallel_size * self.data_parallel_size * self.pipeline_parallel_size + + @property + def delegated_axes(self) -> tuple[MeshAxisKind, ...]: + return tuple(kind for kind, owner in self.l1_owners.items() if owner == "delegated") + + def as_engine_kwargs(self) -> dict[str, object]: + """Return per-stage kwargs keyed by real OmniEngineArgs/EngineArgs field names. + + Two derived values are intentionally *not* emitted here because they are + not per-stage engine args: + + * ``stage_replica_size`` — a per-stage deploy ``num_replicas`` knob + (``StageDeployConfig``); the deploy layer consumes it separately. + * ``omni_lb_policy`` — a pipeline-wide load-balancer policy the engine + reads once at construction (see ``StrategyApplyResult.omni_lb_policy``); + it is applied at the orchestrator level, not folded into per-stage args. + """ + kwargs: dict[str, object] = { + "tensor_parallel_size": self.tensor_parallel_size, + "data_parallel_size": self.data_parallel_size, + "pipeline_parallel_size": self.pipeline_parallel_size, + } + if self.enable_expert_parallel: + kwargs["enable_expert_parallel"] = True + return kwargs + + +def _is_affinity_dp_routing(routing: RoutingPattern) -> bool: + """True when DP routing demands key-stable (hash) placement.""" + if isinstance(routing, PartitionByHash): + return True + if isinstance(routing, RouteByStage) and routing.routing_policy == "hash": + return True + return False + + +def _resolve_l1_owner(spec: StrategySpec) -> L1Owner: + kind = spec.mesh_axis.kind + declared_owner = spec.shard_extension.get("l1_owner") + owner = declared_owner + if owner is None: + owner = _DEFAULT_L1_OWNER.get(kind, "engine") + if kind not in _DEFAULT_L1_OWNER: + logger.debug( + "[composable_parallel] axis %r has unknown kind %r with no default l1_owner; falling back to %r", + spec.name, + kind, + owner, + ) + else: + logger.debug( + "[composable_parallel] axis %r (kind %r) declared no l1_owner; using default %r", + spec.name, + kind, + owner, + ) + if owner not in _VALID_L1_OWNERS: + # Distinguish a bad value supplied via shard_extension from a bad default + # indexed out of _DEFAULT_L1_OWNER, so the source of the invalid owner is + # debuggable rather than swallowed. + source = "shard_extension" if declared_owner is not None else "_DEFAULT_L1_OWNER" + logger.debug( + "[composable_parallel] axis %r (kind %r) resolved invalid l1_owner %r from %s; valid owners are %s", + spec.name, + kind, + owner, + source, + sorted(_VALID_L1_OWNERS), + ) + _fail(f"axis {kind!r} has invalid l1_owner {owner!r}; expected one of {sorted(_VALID_L1_OWNERS)}") + return cast(L1Owner, owner) + + +def _validate_dp(spec: StrategySpec, owner: L1Owner) -> None: + """Validate an engine data-parallel axis. + + DP is realized intra-engine by vLLM's own DP load balancer, so it is + engine-owned and emits no ``omni_lb_policy`` (that string configures omni's + *replica* balancer, a different layer). Key-stable (hash) request affinity is + not something vLLM's DP LB guarantees, so it is rejected rather than silently + dropped. + """ + if owner != "engine": + _fail( + f"dp axis {spec.name!r} is engine data parallelism realized intra-engine by " + f"vLLM's DP load balancer; l1_owner must be 'engine', got {owner!r}. For " + "omni-coordinator-level request fan-out across independent replicas, use a " + "'stage_replica' axis." + ) + if _is_affinity_dp_routing(spec.routing): + _not_implemented( + f"dp axis {spec.name!r} requests key-stable (hash) routing, which vLLM's " + "intra-engine DP load balancer does not guarantee — not supported yet. Use " + "RouteByStage(random|round_robin|least_queue) for stateless DP balancing." + ) + if not isinstance(spec.routing, RouteByStage): + _fail( + f"dp axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) routing, " + f"got {type(spec.routing).__name__}" + ) + if spec.routing.routing_policy not in _STAGE_POLICY_TO_OMNI_LB: + # Recognized-but-unimplemented routing (key-stable/hash) is already + # handled above via _is_affinity_dp_routing -> _not_implemented. Anything + # left here is an unknown/invalid policy value (e.g. a typo from YAML), + # i.e. invalid input -> AxisTranslationError, not NotImplementedError. + _fail( + f"dp axis {spec.name!r} has invalid routing_policy " + f"{spec.routing.routing_policy!r}; expected one of " + f"{sorted(_STAGE_POLICY_TO_OMNI_LB)}." + ) + + +def _stage_replica_lb_policy(spec: StrategySpec, owner: L1Owner) -> str: + """Return the omni StagePool LB policy for a delegated stage_replica axis. + + A ``stage_replica`` axis is *not* a vLLM world dimension: it stands up N + independent engine replicas of one pipeline stage, coordinated by the omni + coordinator and balanced over a StagePool with stateless policies only + (random / round-robin / least-queue-length). It maps to the per-stage + ``num_replicas`` count plus the pipeline-level ``omni_lb_policy`` string, so + its ``l1_owner`` must be ``"delegated"`` (omni owns the routing). Key-stable + (hash) routing is rejected because omni has no key-stable balancer yet. + """ + if owner != "delegated": + _fail( + f"stage_replica axis {spec.name!r} must be 'delegated' to omni's StagePool load " + f"balancer; got owner {owner!r}. Replica routing is owned by omni's coordinator." + ) + + routing = spec.routing + if _is_affinity_dp_routing(routing): + _not_implemented( + f"stage_replica axis {spec.name!r} requests key-stable (hash) routing, which needs a " + "dedicated load balancer — not implemented yet. Use " + "RouteByStage(random|round_robin|least_queue) to delegate to omni's load balancer." + ) + if not isinstance(routing, RouteByStage): + _fail( + f"stage_replica axis {spec.name!r} expects RouteByStage(random|round_robin|least_queue) " + f"routing, got {type(routing).__name__}" + ) + policy = _STAGE_POLICY_TO_OMNI_LB.get(routing.routing_policy) + if policy is None: + _not_implemented( + f"stage_replica axis {spec.name!r} routing_policy {routing.routing_policy!r} has no omni LB policy" + ) + return policy + + +def _validate_tp(spec: StrategySpec, owner: L1Owner) -> None: + if not isinstance(spec.routing, Broadcast): + _fail(f"tp axis {spec.name!r} expects Broadcast routing, got {type(spec.routing).__name__}") + if owner != "engine": + _fail(f"tp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}") + + +def _validate_pp(spec: StrategySpec, owner: L1Owner) -> None: + if not isinstance(spec.routing, PipelineMicrobatch): + _fail(f"pp axis {spec.name!r} expects PipelineMicrobatch routing, got {type(spec.routing).__name__}") + if owner != "engine": + _fail(f"pp axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}") + + +def _validate_ep(spec: StrategySpec, owner: L1Owner) -> None: + # Dense EP: every rank still sees the whole batch, experts are sharded + # across ranks inside the engine (sparse MoE all-to-all is a later stage). + if not isinstance(spec.routing, Broadcast): + _fail( + f"ep axis {spec.name!r} expects Broadcast routing (dense expert parallel), " + f"got {type(spec.routing).__name__}" + ) + if owner != "engine": + _fail(f"ep axis {spec.name!r} is realized intra-engine; l1_owner must be 'engine', got {owner!r}") + + +def translate_strategy_stack(specs: Sequence[StrategySpec]) -> OmniParallelConfig: + """Translate a spec stack into an ``OmniParallelConfig``. + + Supported kinds: dp (engine data parallel), tp, pp, (dense) ep, and + stage_replica (omni replicas). Raises ``NotImplementedError`` for deferred + (affinity / key-stable) routing, and :class:`AxisTranslationError` for any + other invalid spec (a kind not yet translatable, a repeated kind, an owner + incompatible with the axis kind, or unsupported routing). The EP degree must + equal tensor_parallel_size * data_parallel_size. + """ + sizes: dict[str, int] = {"tensor_parallel_size": 1, "data_parallel_size": 1, "pipeline_parallel_size": 1} + owners: dict[MeshAxisKind, L1Owner] = {} + omni_lb_policy: str | None = None + stage_replica_size = 1 + enable_expert_parallel = False + ep_size: int | None = None + + for spec in specs: + kind = spec.mesh_axis.kind + if kind not in _SUPPORTED_KINDS: + _fail( + f"axis kind {kind!r} is not translatable yet (supported: {list(_SUPPORTED_KINDS)}); " + "it is designed-for and lands in a later stage" + ) + if kind in owners: + _fail(f"axis kind {kind!r} appears more than once in the spec stack") + + owner = _resolve_l1_owner(spec) + if kind == "dp": + _validate_dp(spec, owner) + elif kind == "tp": + _validate_tp(spec, owner) + elif kind == "pp": + _validate_pp(spec, owner) + elif kind == "ep": + _validate_ep(spec, owner) + enable_expert_parallel = True + ep_size = spec.mesh_axis.size + elif kind == "stage_replica": + omni_lb_policy = _stage_replica_lb_policy(spec, owner) + stage_replica_size = spec.mesh_axis.size + + if kind in _AXIS_TO_ENGINE_FIELD: + sizes[_AXIS_TO_ENGINE_FIELD[kind]] = spec.mesh_axis.size + owners[kind] = owner + + if enable_expert_parallel: + # Dense EP shards experts across exactly the TP*DP ranks, so the declared + # EP degree must match that product (it is not its own world dimension, + # and PP is excluded). An EP axis of size 1 is a degenerate no-op that + # still sets the flag; it is only valid when TP*DP == 1. + ep_ranks = sizes["tensor_parallel_size"] * sizes["data_parallel_size"] + if ep_size != ep_ranks: + _fail( + f"ep axis size {ep_size} must equal tensor_parallel_size*data_parallel_size " + f"(={ep_ranks}); dense expert parallelism shards experts across exactly those ranks" + ) + + return OmniParallelConfig( + tensor_parallel_size=sizes["tensor_parallel_size"], + data_parallel_size=sizes["data_parallel_size"], + pipeline_parallel_size=sizes["pipeline_parallel_size"], + enable_expert_parallel=enable_expert_parallel, + stage_replica_size=stage_replica_size, + omni_lb_policy=omni_lb_policy, + l1_owners=owners, + ) diff --git a/vllm_omni/config/config_factory.py b/vllm_omni/config/config_factory.py new file mode 100644 index 00000000000..d43a7a48fd2 --- /dev/null +++ b/vllm_omni/config/config_factory.py @@ -0,0 +1,591 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Config factories for vllm-omni, e.g., StageConfigFactory.""" + +from __future__ import annotations + +import dataclasses +import functools +from collections.abc import Mapping +from dataclasses import asdict +from pathlib import Path +from typing import Any + +from transformers import PretrainedConfig +from vllm.logger import init_logger +from vllm.transformers_utils.config import get_config +from vllm.transformers_utils.repo_utils import get_hf_file_to_dict + +from vllm_omni.config.endpoint_policy import EndpointRestriction +from vllm_omni.config.omni_config import VllmOmniConfig +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, resolve_pipeline_config +from vllm_omni.config.stage_config import ( + _DEPLOY_DIR, + DeployConfig, + PipelineConfig, + StageConfig, + StageType, + build_stage_runtime_overrides, + load_deploy_config, + merge_pipeline_deploy, +) +from vllm_omni.config.yaml_util import create_config +from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero + +logger = init_logger(__name__) + + +# Default degree for any parallel axis / replica count that isn't set anywhere +# (CLI, deploy YAML, or pipeline default): a single, un-parallelized rank. +# TODO(composable_parallel): this "1" default and the per-axis device-layout +# fallbacks are currently re-derived in the merge, apply, and reconcile layers. +# Centralize them in one schema so the pre-spawn device guard and the engine-args +# defaults can't drift apart. This is the light slice; the full device-layout +# centralization is tracked as a follow-up. +_DEFAULT_PARALLEL_DEGREE = 1 + + +class StageConfigFactory: + """Factory that loads pipeline YAML and merges CLI overrides. + + Handles both single-stage and multi-stage models. + + Pipelines are declared in ``vllm_omni/config/pipeline_registry.py`` and + where keys in OMNI_PIPELINES map to either a PipelineConfig, or a callable + which accepts a Transformers config as an arg & resolves to a PipelineConfig. + + NOTE: Models with generic HF ``model_type`` collisions (e.g. MiMo Audio + reports ``qwen2``) should declare ``hf_architectures=(...)`` on their + ``PipelineConfig`` so the factory can disambiguate via ``hf_config.architectures``. + """ + + @classmethod + def get_pipeline_endpoint_restrictions( + cls, + model: str, + trust_remote_code: bool, + deploy_config_path: str | None, + ) -> tuple[EndpointRestriction, ...]: + """Given a model string, determine the corresponding endpoint restrictions. + + Args: + model: Model name or path. + trust_remote_code: Whether to trust remote code for HF config loading. + deploy_config_path: Optional path to the deploy config for the pipeline. + + Returns: + A tuple of model specific endpoint restrictions. + """ + pipeline_cfg = StageConfigFactory.get_pipeline_config( + model=model, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + ) + return pipeline_cfg.endpoint_restrictions if pipeline_cfg else () + + @classmethod + @functools.cache + def get_hf_config(cls, model: str, trust_remote_code: bool) -> PretrainedConfig | None: + """Fetch the HF config (if it exists) from the model directory. + + Args: + model: Model name or path. + trust_remote_code: Whether to trust remote code for HF config loading. + + Returns: + the model's config or None. + """ + hf_config = None + try: + return get_config(model, trust_remote_code=trust_remote_code) + except Exception as e: + logger.debug(f"`get_config` failed with exception {e}; inferred HF config is None") + return hf_config + + @classmethod + @functools.cache + def try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None: + """Auto-detect model_type from model directory and apply any model + specific patches to get the correct model_type str. If we are unable + to infer it from the model directory, we fall back to the PipelineConfig. + + Args: + model: Model name or path. + trust_remote_code: Whether to trust remote code for HF config loading. + + Returns: + model_type as a string; may be None on failure. + """ + model_type = cls._try_infer_model_type( + model=model, + trust_remote_code=trust_remote_code, + ) + if model_type == "vla": + if _looks_like_dreamzero(model): + model_type = "dreamzero" + return model_type + + @classmethod + def _try_infer_model_type(cls, model: str, trust_remote_code: bool) -> str | None: + """Auto-detect model_type from model directory. + + Args: + model: Model name or path. + trust_remote_code: Whether to trust remote code for HF config loading. + + Returns: + model_type as a string; may be None on failure. + """ + hf_config = cls.get_hf_config( + model=model, + trust_remote_code=trust_remote_code, + ) + if hf_config is not None: + return hf_config.model_type + + # Fallback: read config.json directly for custom model types that + # are not registered with transformers (e.g. qwen3_tts). + try: + config_dict = get_hf_file_to_dict("config.json", model, revision=None) + if config_dict: + if "model_type" in config_dict: + return config_dict["model_type"] + # VoxCPM2-style configs use singular ``architecture`` rather + # than HF's standard ``model_type`` / ``architectures``. Accept + # it as a fallback so the pipeline registry can still match. + if "architecture" in config_dict and isinstance(config_dict["architecture"], str): + return config_dict["architecture"] + except Exception as e: + logger.debug(f"Failed to auto-detect model type for {model}: {e}") + + # Fallback for diffusers-style models: check model_index.json. + # Some models (e.g. GLM-Image) have no root config.json but ship a + # model_index.json with _class_name that maps to a pipeline key via + # PipelineConfig.diffusers_class_name. + try: + model_index = get_hf_file_to_dict("model_index.json", model, revision=None) + if model_index and "_class_name" in model_index: + class_name = model_index["_class_name"] + for obj in OMNI_PIPELINES.values(): + # If we have a resolver, call it with the optional hf_config + # to get the default pipeline config for this key + pipeline_cfg = obj(hf_config) if callable(obj) else obj + if pipeline_cfg is not None and pipeline_cfg.diffusers_class_name == class_name: + logger.info( + "Detected pipeline %r from model_index.json (_class_name=%r)", + pipeline_cfg.model_type, + class_name, + ) + return pipeline_cfg.model_type + except Exception as e: + logger.debug(f"Failed to detect model type for diffusers-style models: {e}") + + # Final fallback: some models (e.g. CosyVoice3) ship an empty + # config.json and rely on naming conventions. Match the model path + # basename against registered pipeline keys — longest match wins + # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). + model_lower = model.lower().replace("-", "").replace("_", "") + best: str | None = None + best_len = 0 + for registered_key in OMNI_PIPELINES.keys(): + candidate = registered_key.lower().replace("-", "").replace("_", "") + if candidate and candidate in model_lower and len(candidate) > best_len: + best = registered_key + best_len = len(candidate) + if best is not None: + return best + + return None + + @classmethod + def get_pipeline_config( + cls, + model: str, + trust_remote_code: bool, + deploy_config_path: str | None = None, + user_deploy_config: DeployConfig | None = None, + ) -> PipelineConfig | None: + """Resolve the PipelineConfig for a model path/name.""" + model_type = cls.try_infer_model_type(model=model, trust_remote_code=trust_remote_code) + hf_config = cls.get_hf_config(model=model, trust_remote_code=trust_remote_code) + + # Resolve the deploy config & check if the user set the pipeline; + # If the pipeline is explicitly set, it takes highest priority + if user_deploy_config is None: + user_deploy_config = cls._load_user_deploy_config(deploy_config_path) + deploy_config_pipe = cls._get_deploy_override_pipe_config(hf_config, user_deploy_config) + if deploy_config_pipe is not None: + return deploy_config_pipe + + # Pipeline isn't set in the yaml spec, so we need infer it ourselves. + if model_type and model_type in OMNI_PIPELINES: + pipeline_cfg = resolve_pipeline_config(model_type, hf_config) + if pipeline_cfg is not None: + return pipeline_cfg + + if hf_config is not None: + if model_type is not None: + logger.warning("Inferred model type %s is not registered to an Omni pipeline", model_type) + hf_archs = set(getattr(hf_config, "architectures", []) or []) + if hf_archs: + for registered in OMNI_PIPELINES.values(): + pipeline_cfg = registered if isinstance(registered, PipelineConfig) else registered(hf_config) + if pipeline_cfg is None: + continue + predicate = pipeline_cfg.hf_config_predicate + if predicate is not None: + try: + if not predicate(hf_config): + logger.debug( + "Pipeline %r matched on architectures %s but its " + "hf_config_predicate rejected the loaded config; " + "continuing fallback search.", + pipeline_cfg.model_type, + sorted(hf_archs.intersection(pipeline_cfg.hf_architectures)), + ) + continue + except Exception: + logger.exception( + "Pipeline %r hf_config_predicate raised; skipping.", + pipeline_cfg.model_type, + ) + continue + if isinstance(pipeline_cfg, PipelineConfig) and hf_archs.intersection( + pipeline_cfg.hf_architectures + ): + return pipeline_cfg + return None + + @classmethod + def _get_deploy_override_pipe_config( + cls, + hf_config: PretrainedConfig | None, + deploy_config: DeployConfig | None, + ) -> PipelineConfig | None: + """Resolve an explicit pipeline override from a loaded deploy config.""" + if deploy_config is None or deploy_config.pipeline is None: + return None + + pipeline_cfg = resolve_pipeline_config(deploy_config.pipeline, hf_config) + if pipeline_cfg is None: + raise KeyError( + f"Pipeline {deploy_config.pipeline!r} from deploy config is not registered " + f"to OMNI_PIPELINES. Available: {sorted(OMNI_PIPELINES)}" + ) + return pipeline_cfg + + @staticmethod + def _load_user_deploy_config(deploy_config_path: str | None) -> DeployConfig | None: + """Load an explicit deploy YAML once for resolution and construction.""" + if deploy_config_path is None: + return None + deploy_path = Path(deploy_config_path) + if not deploy_path.exists() and deploy_path.parent == Path("."): + candidate = _DEPLOY_DIR / deploy_path + if candidate.exists(): + deploy_path = candidate + if not deploy_path.exists(): + raise FileNotFoundError(f"Deploy config not found: {deploy_path}") + return load_deploy_config(deploy_path) + + @classmethod + def create_from_model( + cls, + model: str, + *, + trust_remote_code: bool, + cli_overrides: dict[str, Any], + deploy_config_path: str | None, + ) -> VllmOmniConfig | None: + """Build the structured Omni config for a model/deploy pair.""" + user_deploy_config = cls._load_user_deploy_config(deploy_config_path) + pipeline_cfg = cls.get_pipeline_config( + model=model, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + user_deploy_config=user_deploy_config, + ) + if pipeline_cfg is None: + return None + + registry_cli_overrides = { + **cli_overrides, + "trust_remote_code": trust_remote_code, + "model": model, + } + return VllmOmniConfig.from_pipeline_config( + pipeline_cfg, + user_deploy_config=user_deploy_config, + deploy_config_path=deploy_config_path, + cli_overrides=registry_cli_overrides, + ) + + @classmethod + def create_legacy_stage_configs_from_model( + cls, + model: str, + *, + trust_remote_code: bool, + cli_overrides: dict[str, Any], + deploy_config_path: str | None, + strategy_specs: Mapping[Any, Any] | None = None, + ) -> tuple[list[StageConfig] | None, str | None]: + """Build current runtime stage configs from the shared resolution. + + The engine still consumes the legacy StageConfig/OmegaConf shape. + RFC #4021 will replace this transitional path as runtime consumers move + to VllmOmniConfig. + """ + user_deploy_config = cls._load_user_deploy_config(deploy_config_path) + pipeline_cfg = cls.get_pipeline_config( + model=model, + trust_remote_code=trust_remote_code, + deploy_config_path=deploy_config_path, + user_deploy_config=user_deploy_config, + ) + if pipeline_cfg is None: + return None, None + + legacy_cli_overrides = { + **cli_overrides, + "trust_remote_code": trust_remote_code, + } + return cls._create_legacy_from_registry( + pipeline_cfg, + legacy_cli_overrides, + deploy_config_path, + user_deploy_config=user_deploy_config, + strategy_specs=strategy_specs, + ) + + @classmethod + def _create_legacy_from_registry( + cls, + pipeline_cfg: PipelineConfig, + cli_overrides: dict[str, Any], + deploy_config_path: str | None = None, + user_deploy_config: DeployConfig | None = None, + strategy_specs: Mapping[Any, Any] | None = None, + ) -> tuple[list[StageConfig], str | None]: + """Create current runtime StageConfigs from registry + deploy YAML. + + Precedence: caller-typed (non-None) value > deploy YAML > + StageDeployConfig dataclass default. + + Returns ``(stages, omni_lb_policy)`` — the strategy-derived pipeline-wide + load-balancer policy (``None`` when no strategy set one) travels with the + stages instead of through a mutable out-param. + """ + if user_deploy_config is not None: + deploy_cfg = user_deploy_config + elif deploy_config_path is not None: + deploy_cfg = cls._load_user_deploy_config(deploy_config_path) + assert deploy_cfg is not None + elif pipeline_cfg.default_deploy_config_name is not None: + deploy_cfg = load_deploy_config(_DEPLOY_DIR / pipeline_cfg.default_deploy_config_name) + else: + deploy_cfg = DeployConfig() + + cli_async_chunk = cli_overrides.get("async_chunk") + if cli_async_chunk is not None: + deploy_cfg.async_chunk = bool(cli_async_chunk) + + stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides) + + # Overlay declarative parallel strategies (opt-in) before CLI overrides. + applied = cls._apply_strategy_specs(stages, strategy_specs) + + explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None} + + for stage in stages: + stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides) + + # Re-validate the resolved layout now that CLI overrides are on top. + cls._reconcile_strategy_with_cli(stages, applied) + + omni_lb_policy = applied.omni_lb_policy if applied is not None else None + return stages, omni_lb_policy + + @staticmethod + def _apply_strategy_specs( + stages: list[StageConfig], + strategy_specs: Mapping[Any, Any] | None, + ) -> Any: + """Overlay derived parallel sizing onto merged stages (opt-in). + + ``omni_lb_policy`` cannot be set from stage configs (omni reads it once + at orchestrator construction), so a derived non-default policy is logged + here and carried on the returned ``StrategyApplyResult`` for the caller + to hand to the orchestrator. + + Returns the ``StrategyApplyResult`` (or ``None`` when no strategy was + supplied) so the caller can re-validate the resolved layout once CLI + overrides have been merged on top, and read its ``omni_lb_policy``. + """ + if not strategy_specs: + return None + from vllm_omni.config.composable_parallel import apply_strategy_specs + + applied = apply_strategy_specs(stages, strategy_specs) + if applied.omni_lb_policy is not None: + logger.info( + "[composable_parallel] strategy derived omni_lb_policy=%r; it will be applied " + "to the orchestrator unless an explicit --omni-lb-policy was given.", + applied.omni_lb_policy, + ) + return applied + + @staticmethod + def _reconcile_strategy_with_cli( + stages: list[StageConfig], + applied: Any, + ) -> None: + """Reconcile CLI overrides applied *after* a strategy overlay. + + CLI overrides are applied last (in ``to_omegaconf``), so for any stage a + strategy touched we must (a) warn loudly when a CLI arg overrides a + strategy-declared axis — the strategy is meant to be the single writer + for the axes it declares, so a silent CLI win is surprising — and + (b) re-run the device-count check against the *effective* (post-CLI) + ``tp``/``dp``/``pp``/``num_replicas``/``devices`` so a CLI + ``--devices``/``--tensor-parallel-size``/``--num-replicas`` cannot slip + past the pre-spawn guard that exists to prevent silent OOMs. + """ + if applied is None: + return + from vllm_omni.config.composable_parallel import check_device_layout + + # Axis kind -> (engine-arg field, strategy-derived attribute on the cfg). + axis_fields = { + "tp": ("tensor_parallel_size", "tensor_parallel_size"), + "dp": ("data_parallel_size", "data_parallel_size"), + "pp": ("pipeline_parallel_size", "pipeline_parallel_size"), + } + + for stage in stages: + cfg = applied.per_stage_config.get(stage.stage_id) + if cfg is None: + continue + overrides = stage.runtime_overrides or {} + declared = set(cfg.l1_owners.keys()) + + for kind, (field_name, attr) in axis_fields.items(): + if kind in declared and overrides.get(field_name) is not None: + cli_val = overrides[field_name] + derived = getattr(cfg, attr) + if cli_val != derived: + logger.warning( + "[composable_parallel] stage %s: CLI %s=%s overrides the " + "strategy-derived %s=%s. The CLI value wins; remove one to avoid ambiguity.", + stage.stage_id, + field_name, + cli_val, + field_name, + derived, + ) + if "stage_replica" in declared and overrides.get("num_replicas") is not None: + cli_val = overrides["num_replicas"] + if cli_val != cfg.stage_replica_size: + logger.warning( + "[composable_parallel] stage %s: CLI num_replicas=%s overrides the " + "strategy-derived num_replicas=%s. The CLI value wins; remove one to avoid ambiguity.", + stage.stage_id, + cli_val, + cfg.stage_replica_size, + ) + + def _eff(field_name: str, fallback: Any) -> Any: + val = overrides.get(field_name) + return val if val is not None else fallback + + def _eff_degree(field_name: str, source: dict[str, Any]) -> int: + # Single place the per-axis "default to 1" fallback is applied, + # for both the override and the YAML-default sides (see the + # _DEFAULT_PARALLEL_DEGREE TODO). + value = _eff(field_name, source.get(field_name, _DEFAULT_PARALLEL_DEGREE)) + return int(value or _DEFAULT_PARALLEL_DEGREE) + + check_device_layout( + _eff("devices", stage.yaml_runtime.get("devices")), + tensor_parallel_size=_eff_degree("tensor_parallel_size", stage.yaml_engine_args), + data_parallel_size=_eff_degree("data_parallel_size", stage.yaml_engine_args), + pipeline_parallel_size=_eff_degree("pipeline_parallel_size", stage.yaml_engine_args), + num_replicas=_eff_degree("num_replicas", stage.yaml_runtime), + role=stage.model_stage, + ) + + @classmethod + def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: + """Single-stage diffusion - no YAML needed. + + Creates a default diffusion stage configuration for single-stage + diffusion models. Returns a legacy OmegaConf-compatible dict for + backward compatibility with OmniStage. + + Args: + kwargs: Engine arguments from CLI/API. + + Returns: + List containing a single config dict for the diffusion stage. + """ + # Calculate devices based on parallel config + devices = "0" + if "parallel_config" in kwargs: + num_devices = kwargs["parallel_config"].world_size + for i in range(1, num_devices): + devices += f",{i}" + + engine_args: dict[str, Any] = {} + for key, value in kwargs.items(): + if key in ("parallel_config",): + continue + engine_args[key] = value + + # Serialize parallel_config as dict for OmegaConf. Test helpers + # sometimes pass SimpleNamespace rather than a dataclass instance. + if "parallel_config" in kwargs: + parallel_config = kwargs["parallel_config"] + if dataclasses.is_dataclass(parallel_config) and not isinstance(parallel_config, type): + engine_args["parallel_config"] = asdict(parallel_config) + elif hasattr(parallel_config, "__dict__"): + engine_args["parallel_config"] = dict(vars(parallel_config)) + else: + engine_args["parallel_config"] = parallel_config + + engine_args.setdefault("cache_backend", "none") + engine_args["model_stage"] = "diffusion" + + # Convert dtype to string for OmegaConf + if "dtype" in engine_args: + engine_args["dtype"] = str(engine_args["dtype"]) + + engine_args.setdefault("max_num_seqs", 1) + + config_dict: dict[str, Any] = { + "stage_id": 0, + "stage_type": StageType.DIFFUSION.value, + "runtime": { + "process": True, + "devices": devices, + }, + "engine_args": create_config(engine_args), + "final_output": True, + "final_output_type": "image", + } + + return [config_dict] + + @classmethod + def _merge_cli_overrides( + cls, + stage: StageConfig, + cli_overrides: dict[str, Any], + ) -> dict[str, Any]: + """Merge global and per-stage (``stage_N_*``) CLI overrides. + + Orchestrator-owned keys are filtered by ``build_stage_runtime_overrides`` + using ``OrchestratorArgs`` as the single source of truth; unknown + server/uvicorn keys are dropped downstream by + ``filter_dataclass_kwargs(OmniEngineArgs, ...)``. + """ + return build_stage_runtime_overrides(stage.stage_id, cli_overrides) diff --git a/vllm_omni/config/endpoint_policy.py b/vllm_omni/config/endpoint_policy.py index d22c0acc121..5a1449614d6 100644 --- a/vllm_omni/config/endpoint_policy.py +++ b/vllm_omni/config/endpoint_policy.py @@ -21,6 +21,7 @@ class RouteTarget(NamedTuple): class OmniServingCapability(Enum): """Serving capabilities that pipelines can shut down.""" + CHAT_COMPLETIONS_BATCH = RouteTarget("/v1/chat/completions/batch", frozenset({"POST"})) COMPLETIONS = RouteTarget("/v1/completions", frozenset({"POST"})) @property @@ -38,6 +39,16 @@ class EndpointRestriction: reason: str +# Routes that are not supported for any model, but are supported in vLLM. +# This is only temporary to avoid 500s for batched chat completions. +UNSUPPORTED_ROUTES: tuple[EndpointRestriction, ...] = ( + EndpointRestriction( + OmniServingCapability.CHAT_COMPLETIONS_BATCH, + "Batched chat completions are not yet supported by vLLM Omni.", + ), +) + + def build_rejection_handler(reason: str): """Build a rejection handler for a given endpoint for the provided reason.""" @@ -60,7 +71,11 @@ def shutdown_unsupported_routes( """ from vllm_omni.entrypoints.openai.api_server import _remove_route_from_app - for end_restrict in endpoint_restrictions: + # Generally these should not overlap since there is no point. If they do, + # we use the reason message in UNSUPPORTED_ROUTES, for consistent error messages. + restricted_endpoints = (*endpoint_restrictions, *UNSUPPORTED_ROUTES) + + for end_restrict in restricted_endpoints: capability = end_restrict.capability # Remove the route from the app _remove_route_from_app(app, capability.path, capability.methods) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py new file mode 100644 index 00000000000..4c2154401a1 --- /dev/null +++ b/vllm_omni/config/omni_config.py @@ -0,0 +1,1309 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structured vLLM-Omni configuration classes. + +This module is additive for Phase 2 of RFC #4021. +``VllmOmniConfig.from_pipeline_config`` builds the structured view from an +already-resolved pipeline and deploy config so parity can be proven before +later PRs cut consumers over to these classes. +""" + +from __future__ import annotations + +import copy +from collections.abc import Mapping +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any, TypeAlias, TypedDict, cast + +from pydantic import ConfigDict, Field +from vllm.config.utils import config +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +from vllm_omni.config.stage_config import ( + _DEPLOY_DIR, + _STAGE_DEPLOY_FIELDS, + PIPELINE_WIDE_ENGINE_FIELDS, + DeployConfig, + PipelineConfig, + StageDeployConfig, + StageExecutionType, + StagePipelineConfig, + StageType, + _apply_platform_overrides, + _resolve_scheduler, + _scheduler_path, + _select_processor_funcs, + build_stage_runtime_overrides, + load_deploy_config, +) + +_EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = { + StageExecutionType.LLM_AR: (StageType.LLM, "ar"), + StageExecutionType.LLM_GENERATION: (StageType.LLM, "generation"), + StageExecutionType.DIFFUSION: (StageType.DIFFUSION, None), +} + +_PIPELINE_DEPLOY_CLI_FIELDS = PIPELINE_WIDE_ENGINE_FIELDS + +_QuantizationConfigType: TypeAlias = QuantizationConfig | str | Mapping[str, Any] | None + + +class _QuantizationEngineOverrides(TypedDict, total=False): + quantization_config: _QuantizationConfigType + quantization: str + + +class _ModelEngineOverrides(TypedDict, total=False): + active_stream_window: int + enable_sleep_mode: bool + subtalker_sampling_params: dict[str, Any] + has_sampling_extra_args: bool + custom_voice_dir: str + task_type: str + codec_frame_rate_hz: float + enforce_eager: bool + enable_flashinfer_autotune: bool + compilation_config: dict[str, Any] + enable_multithread_weight_load: bool + num_weight_load_threads: int + disable_autocast: bool + + +class _LoadEngineOverrides(TypedDict, total=False): + load_format: str + tokenizer_mode: str + config_format: str + skip_mm_profiling: bool + + +class _CacheEngineOverrides(TypedDict, total=False): + gpu_memory_utilization: float + enable_prefix_caching: bool + disable_hybrid_kv_cache_manager: bool + mm_processor_cache_gb: float + + +class _SchedulerEngineOverrides(TypedDict, total=False): + max_num_seqs: int + max_num_batched_tokens: int + max_model_len: int + enable_chunked_prefill: bool + async_scheduling: bool + + +class _RuntimeEngineOverrides(TypedDict, total=False): + devices: str + num_replicas: int + env: dict[str, Any] + num_gpus: int + log_level: str + log_stats: bool + profiler_config: dict[str, Any] + + +class _ParallelConfigEngineOverrides(TypedDict, total=False): + pipeline_parallel_size: int + data_parallel_size: int + tensor_parallel_size: int + sequence_parallel_size: int + ulysses_degree: int + ring_degree: int + ulysses_mode: str + cfg_parallel_size: int + vae_patch_parallel_size: int + use_hsdp: bool + mask_sp_padding: bool + hsdp_shard_size: int + hsdp_replicate_size: int + enable_expert_parallel: bool + + +class _ParallelEngineOverrides(_ParallelConfigEngineOverrides, total=False): + parallel_config: _ParallelConfigEngineOverrides | Mapping[str, Any] + + +@dataclass(frozen=True) +class _StageEngineValues: + """Typed projections of legacy flat per-stage ``yaml_engine_args``.""" + + quantization: _QuantizationEngineOverrides + model: _ModelEngineOverrides + load: _LoadEngineOverrides + cache: _CacheEngineOverrides + scheduler: _SchedulerEngineOverrides + runtime: _RuntimeEngineOverrides + parallel: _ParallelEngineOverrides + diffusion: _DiffusionEngineOverrides + + +@dataclass(frozen=True) +class _DiffusionEngineOverrides: + """Validated diffusion projection of legacy flat per-stage engine args.""" + + _values: dict[str, Any] + + @classmethod + def from_engine(cls, engine: Mapping[str, Any]) -> _DiffusionEngineOverrides: + return cls(_select_engine_overrides(engine, _DIFFUSION_STAGE_ENGINE_FIELDS)) + + def to_kwargs(self) -> dict[str, Any]: + return {name: _copy_value(value) for name, value in self._values.items()} + + +_IMMUTABLE_CONFIG_VALUE_TYPES = (str, int, float, bool, bytes, type(None)) + + +def _can_share_config_value(value: Any) -> bool: + if isinstance(value, _IMMUTABLE_CONFIG_VALUE_TYPES): + return True + if isinstance(value, tuple): + return all(_can_share_config_value(item) for item in value) + return False + + +def _copy_value(value: Any) -> Any: + """Copy nested config values so the structured view owns its data.""" + if _can_share_config_value(value): + return value + return copy.deepcopy(value) + + +def _config_kwargs(overrides: Mapping[str, Any]) -> dict[str, Any]: + return {name: _copy_value(value) for name, value in overrides.items() if value is not None} + + +def _first_defined(*values: Any) -> Any: + for value in values: + if value is not None: + return _copy_value(value) + return None + + +def _validate_async_chunk_support(pipeline: PipelineConfig, deploy: DeployConfig) -> None: + has_inter_stage_edges = any(stage.input_sources for stage in pipeline.stages) + if ( + deploy.async_chunk + and has_inter_stage_edges + and not any(stage.async_chunk_process_next_stage_input_func for stage in pipeline.stages) + ): + raise ValueError( + f"Pipeline {pipeline.model_type!r} has async_chunk=True in deploy but no stage " + "declares a dedicated async-chunk next-stage processor " + "(``async_chunk_process_next_stage_input_func``). " + "Either set async_chunk=False or implement an async-chunk producer on the pipeline." + ) + + +def _resolve_execution_mode(execution_type: StageExecutionType) -> tuple[StageType, str | None]: + try: + return _EXECUTION_TYPE_TO_STAGE_WORKER[execution_type] + except KeyError as exc: + raise ValueError(f"Unsupported stage execution type: {execution_type!r}") from exc + + +def _resolve_scheduler_path(execution_type: StageExecutionType, async_scheduling: bool = True) -> str | None: + return _scheduler_path(_resolve_scheduler(execution_type, async_scheduling)) + + +def _stage_cli_overrides(stage_id: int, cli_overrides: Mapping[str, Any]) -> dict[str, Any]: + runtime_overrides = build_stage_runtime_overrides(stage_id, dict(cli_overrides)) + global_stage_fields = _global_stage_cli_fields() + result: dict[str, Any] = {} + for key, value in runtime_overrides.items(): + if key in global_stage_fields or f"stage_{stage_id}_{key}" in cli_overrides: + result[key] = _copy_value(value) + return result + + +def _resolve_deploy_path(deploy_config_path: str) -> Path: + deploy_path = Path(deploy_config_path) + if not deploy_path.exists() and deploy_path.parent == Path("."): + bare_name = deploy_path.name + if not bare_name.endswith(".yaml"): + bare_name = f"{bare_name}.yaml" + candidate = _DEPLOY_DIR / bare_name + if candidate.exists(): + return candidate + return deploy_path + + +def _get_deploy_config( + pipeline_cfg: PipelineConfig, + user_deploy_config: DeployConfig | None, + deploy_config_path: str | None, +) -> tuple[DeployConfig, str | None]: + """Select user-provided, pipeline-default, or empty deploy settings.""" + if user_deploy_config is not None: + loaded_path = str(_resolve_deploy_path(deploy_config_path)) if deploy_config_path is not None else None + return copy.deepcopy(user_deploy_config), loaded_path + + if deploy_config_path is not None: + resolved_path = _resolve_deploy_path(deploy_config_path) + if not resolved_path.exists(): + raise FileNotFoundError(f"Deploy config not found: {resolved_path}") + return load_deploy_config(resolved_path), str(resolved_path) + + if pipeline_cfg.default_deploy_config_name is not None: + default_path = _DEPLOY_DIR / pipeline_cfg.default_deploy_config_name + return load_deploy_config(default_path), str(default_path) + + return DeployConfig(), None + + +@config +class OmniStageModelConfig: + """Per-stage model behavior.""" + + active_stream_window: int = Field(default=0, ge=0) + enable_sleep_mode: bool = False + default_sampling_params: dict[str, Any] | None = None + subtalker_sampling_params: dict[str, Any] | None = None + has_sampling_extra_args: bool = False + custom_voice_dir: str | None = None + task_type: str | None = None + codec_frame_rate_hz: float | None = None + enforce_eager: bool = False + enable_flashinfer_autotune: bool | None = None + compilation_config: dict[str, Any] | None = None + enable_multithread_weight_load: bool = True + num_weight_load_threads: int = Field(default=4, ge=1) + disable_autocast: bool = False + + +@config +class OmniStageLoadConfig: + """Per-stage loading behavior.""" + + load_format: str = "auto" + tokenizer_mode: str = "auto" + config_format: str | None = None + skip_mm_profiling: bool | None = None + + +@config +class OmniStageCacheConfig: + """Per-stage engine cache and memory behavior. + + This is separate from ``_DiffusionConfigProjection.cache_config``, which configures + vLLM-Omni diffusion-specific cache backends such as TeaCache and Cache-DiT. + """ + + gpu_memory_utilization: float = Field(default=0.90, gt=0.0, le=1.0) + enable_prefix_caching: bool = False + disable_hybrid_kv_cache_manager: bool = False + mm_processor_cache_gb: float | None = Field(default=None, ge=0.0) + + +@config +class OmniStageSchedulerConfig: + """Per-stage request scheduling behavior.""" + + max_num_seqs: int = Field(default=128, ge=1) + max_num_batched_tokens: int | None = Field(default=None, ge=1) + max_model_len: int | None = Field(default=None, ge=-1) + enable_chunked_prefill: bool = False + async_scheduling: bool = True + + def __post_init__(self) -> None: + if self.max_num_batched_tokens is not None and self.max_num_batched_tokens < self.max_num_seqs: + raise ValueError( + f"max_num_batched_tokens ({self.max_num_batched_tokens}) must be >= max_num_seqs ({self.max_num_seqs})" + ) + + +@config +class OmniStageConnectorConfig: + """Per-stage inter-stage connector wiring.""" + + stage_connector: dict[str, Any] = field( + default_factory=lambda: { + "name": "SharedMemoryConnector", + "extra": {}, + } + ) + output_connectors: dict[str, Any] | None = None + input_connectors: dict[str, Any] | None = None + + +@config +class OmniStageRuntimeConfig: + """Per-stage process placement and runtime behavior.""" + + devices: str | None = None + num_replicas: int = Field(default=1, ge=1) + env: dict[str, Any] | None = None + num_gpus: int = Field(default=1, ge=1) + log_level: str = "info" + log_stats: bool = False + profiler_config: dict[str, Any] | None = None + + +@config +class OmniStageParallelConfig: + """Common per-stage distributed parallelism behavior.""" + + pipeline_parallel_size: int = Field(default=1, ge=1) + data_parallel_size: int = Field(default=1, ge=1) + tensor_parallel_size: int = Field(default=1, ge=1) + enable_expert_parallel: bool = False + world_size: int = Field(default=1, ge=1, init=False) + + def __post_init__(self) -> None: + self.world_size = self.pipeline_parallel_size * self.data_parallel_size * self.tensor_parallel_size + + +@config +class OmniStageDiffusionParallelConfig(OmniStageParallelConfig): + """Diffusion-stage distributed parallelism behavior.""" + + sequence_parallel_size: int = Field(default=1, ge=1, init=False) + ulysses_degree: int = Field(default=1, ge=1) + ring_degree: int = Field(default=1, ge=1) + ulysses_mode: str = "strict" + cfg_parallel_size: int = Field(default=1, ge=1, le=3) + vae_patch_parallel_size: int = Field(default=1, ge=1) + vae_parallel_mode: str = "tile" + use_hsdp: bool = False + mask_sp_padding: bool = False + hsdp_shard_size: int = -1 + hsdp_replicate_size: int = Field(default=1, ge=1) + + def __post_init__(self) -> None: + self.sequence_parallel_size = self.ulysses_degree * self.ring_degree + if self.ulysses_mode not in {"strict", "advanced_uaa"}: + raise ValueError("ulysses_mode must be 'strict' or 'advanced_uaa'") + if self.vae_parallel_mode not in {"tile", "spatial_shard_height", "spatial_shard_width"}: + raise ValueError( + "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + f"but got {self.vae_parallel_mode!r}." + ) + + other_parallel_world_size = ( + self.pipeline_parallel_size + * self.data_parallel_size + * self.tensor_parallel_size + * self.sequence_parallel_size + * self.cfg_parallel_size + ) + if self.use_hsdp: + if self.tensor_parallel_size > 1 or self.data_parallel_size > 1: + raise ValueError( + "HSDP (use_hsdp=True) cannot be used with TP or DP " + f"(tensor_parallel_size={self.tensor_parallel_size}, " + f"data_parallel_size={self.data_parallel_size})" + ) + if self.hsdp_shard_size == -1: + if other_parallel_world_size == 1: + raise ValueError("Cannot auto-calculate hsdp_shard_size when other parallelism is all 1") + if other_parallel_world_size % self.hsdp_replicate_size != 0: + raise ValueError( + f"hsdp_replicate_size ({self.hsdp_replicate_size}) must evenly divide " + f"world_size ({other_parallel_world_size}) when hsdp_shard_size is -1" + ) + self.hsdp_shard_size = other_parallel_world_size // self.hsdp_replicate_size + self.world_size = other_parallel_world_size + else: + if self.hsdp_shard_size <= 0: + raise ValueError("hsdp_shard_size must be > 0 when use_hsdp=True") + hsdp_world_size = self.hsdp_replicate_size * self.hsdp_shard_size + if other_parallel_world_size == 1: + self.world_size = hsdp_world_size + else: + if hsdp_world_size != other_parallel_world_size: + raise ValueError( + f"HSDP dimensions ({self.hsdp_replicate_size} x {self.hsdp_shard_size} = " + f"{hsdp_world_size}) must equal world_size from other parallelism " + f"({other_parallel_world_size})" + ) + self.world_size = other_parallel_world_size + else: + self.world_size = other_parallel_world_size + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class _DiffusionConfigProjection: + """Diffusion-specific per-stage settings. + + Shared AR/diffusion fields are projected into the other sub-configs. This + class keeps the diffusion-only knobs from ``OmniDiffusionConfig`` without + running its startup-time side effects such as port probing or HF metadata + loading. + """ + + stage_id: int = 0 + model: str | None = None + model_class_name: str | None = None + model_arch: str | None = None + dtype: Any = "auto" + trust_remote_code: bool = False + revision: str | None = None + distributed_executor_backend: str = "mp" + dist_timeout: int | None = None + nccl_port: int | None = None + master_port: int | None = None + host: str | None = None + port: int | None = None + model_config: dict[str, Any] = field(default_factory=dict) + tf_model_config: Any = None + diffusion_attention_config: Any = None + cache_strategy: str = "none" + cache_backend: str = "none" + cache_config: Any = field(default_factory=dict) + enable_cache_dit_summary: bool = False + enable_prompt_embed_cache: bool = False + prompt_embed_cache_size: int = Field(default=32, ge=1) + diffusion_load_format: str = "default" + diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) + diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) + diffusers_pipeline_cls: Any = None + lora_path: str | None = None + lora_scale: float = 1.0 + max_cpu_loras: int | None = None + output_type: str = "pil" + enable_cpu_offload: bool = False + enable_layerwise_offload: bool = False + pin_cpu_memory: bool = True + vae_use_slicing: bool = False + vae_use_tiling: bool = False + mask_strategy_file_path: str | None = None + skip_time_steps: int = 15 + VSA_sparsity: float = 0.0 + moba_config_path: str | None = None + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | list[int] | tuple[int, ...] | set[int] | None = None + diffusion_kv_cache_skip_layers: str | list[int] | tuple[int, ...] | set[int] | None = None + diffusion_kv_cache_skip_step_indices: set[int] | None = None + diffusion_kv_cache_skip_layer_indices: set[int] | None = None + moe_backend: str = "auto" + force_cutlass_fp8: bool = False + enable_diffusion_pipeline_profiler: bool = False + step_execution: bool = False + supports_multimodal_inputs: bool = False + max_multimodal_image_inputs: int | None = None + model_paths: dict[str, str] = field(default_factory=dict) + model_loaded: dict[str, bool] = field( + default_factory=lambda: { + "transformer": True, + "vae": True, + } + ) + override_transformer_cls_name: str | None = None + worker_extension_cls: str | None = None + custom_pipeline_args: dict[str, Any] | None = None + additional_config: dict[str, Any] = field(default_factory=dict) + enable_stage_verification: bool = True + prompt_file_path: str | None = None + quantization_config: _QuantizationConfigType = None + extras: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_kwargs(cls, **kwargs: Any) -> _DiffusionConfigProjection: + from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs + + normalized_kwargs = normalize_omni_diffusion_kwargs(kwargs) + valid_fields = {f.name for f in fields(cls)} + return cls(**{k: v for k, v in normalized_kwargs.items() if k in valid_fields}) + + def __post_init__(self) -> None: + # Keep diffusion imports lazy so importing vllm_omni.config does not + # pull in the full diffusion stack unless a diffusion stage is built. + from vllm_omni.diffusion.data import ( + AttentionConfig, + DiffusionCacheConfig, + TransformerConfig, + build_attention_config, + parse_kv_cache_skip_selector, + ) + from vllm_omni.quantization import build_quant_config + + if self.tf_model_config is None: + self.tf_model_config = TransformerConfig() + elif isinstance(self.tf_model_config, Mapping): + self.tf_model_config = TransformerConfig.from_dict(dict(self.tf_model_config)) + + if self.additional_config is None: + self.additional_config = {} + elif isinstance(self.additional_config, Mapping): + self.additional_config = dict(self.additional_config) + else: + raise TypeError(f"additional_config must be a mapping or None, got {type(self.additional_config)!r}") + + if isinstance(self.dtype, str): + # Import torch only when string dtype normalization is needed. + import torch + + dtype_map = { + "auto": torch.bfloat16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, + "float16": torch.float16, + "fp16": torch.float16, + "half": torch.float16, + "float32": torch.float32, + "fp32": torch.float32, + "float": torch.float32, + } + self.dtype = dtype_map.get(self.dtype.lower(), torch.bfloat16) + + if isinstance(self.cache_config, Mapping): + self.cache_config = DiffusionCacheConfig.from_dict(dict(self.cache_config)) + elif not isinstance(self.cache_config, DiffusionCacheConfig): + self.cache_config = DiffusionCacheConfig() + + self._propagate_quantization_from_tf_config(self.tf_model_config) + if self.quantization_config is not None: + if isinstance(self.quantization_config, QuantizationConfig): + pass + elif isinstance(self.quantization_config, str): + self.quantization_config = build_quant_config(self.quantization_config) + elif isinstance(self.quantization_config, Mapping): + self.quantization_config = dict(self.quantization_config) + else: + raise TypeError( + "quantization_config must be str, dict, QuantizationConfig, or None, " + f"got {type(self.quantization_config)!r}" + ) + + if self.diffusion_attention_config is None or isinstance( + self.diffusion_attention_config, + (AttentionConfig, Mapping), + ): + self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config) + else: + raise TypeError( + "diffusion_attention_config must be an AttentionConfig, mapping, or None, " + f"got {type(self.diffusion_attention_config)!r}" + ) + + self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps) + self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers) + + if self.max_cpu_loras is None: + self.max_cpu_loras = 1 + elif self.max_cpu_loras < 1: + raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): + raise ValueError( + "diffusers_load_kwargs and diffusers_call_kwargs are only " + "valid together with diffusion_load_format=diffusers" + ) + + def _propagate_quantization_from_tf_config(self, tf_config: Any) -> None: + quant_config = getattr(tf_config, "quant_config", None) + if quant_config is None: + return + quant_method = getattr(tf_config, "quant_method", None) + is_checkpoint_fp8 = bool(getattr(quant_config, "is_checkpoint_fp8_serialized", False)) + is_checkpoint_nvfp4 = bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False)) + should_use_checkpoint_config = ( + self.quantization_config is None + or (is_checkpoint_fp8 and self._is_generic_fp8_quant_config(self.quantization_config)) + or (is_checkpoint_nvfp4 and self._is_generic_nvfp4_quant_config(self.quantization_config)) + ) + if should_use_checkpoint_config: + self.quantization_config = quant_config + if quant_method is not None: + self.additional_config.setdefault("auto_detected_quant_method", quant_method) + + @staticmethod + def _is_generic_fp8_quant_config(quant_config: object) -> bool: + if isinstance(quant_config, str): + return quant_config.lower() == "fp8" + if isinstance(quant_config, Mapping): + method = quant_config.get("method", quant_config.get("quant_method")) + return isinstance(method, str) and method.lower() == "fp8" + if hasattr(quant_config, "get_name"): + return quant_config.get_name() == "fp8" + return False + + @staticmethod + def _is_generic_nvfp4_quant_config(quant_config: object) -> bool: + if isinstance(quant_config, str): + return quant_config.lower() in {"fp4", "nvfp4", "modelopt_fp4"} + if isinstance(quant_config, Mapping): + method = quant_config.get("method", quant_config.get("quant_method")) + return isinstance(method, str) and method.lower() in {"fp4", "nvfp4", "modelopt_fp4"} + if hasattr(quant_config, "get_name"): + return quant_config.get_name() == "modelopt_fp4" + return False + + def set_tf_model_config(self, tf_config: Any) -> None: + self.tf_model_config = tf_config + self._propagate_quantization_from_tf_config(tf_config) + + def enrich_config(self) -> None: + from vllm_omni.diffusion.data import OmniDiffusionConfig + + omni_diffusion_fields = frozenset(f.name for f in fields(OmniDiffusionConfig)) + kwargs = { + name: _copy_value(getattr(self, name)) for name in _DIFFUSION_CONFIG_FIELDS if name in omni_diffusion_fields + } + omni_diffusion_config = OmniDiffusionConfig(**kwargs) + omni_diffusion_config.enrich_config() + for name in _DIFFUSION_CONFIG_FIELDS: + if hasattr(omni_diffusion_config, name): + setattr(self, name, _copy_value(getattr(omni_diffusion_config, name))) + + +_DIFFUSION_CONFIG_FIELDS = frozenset(f.name for f in fields(_DiffusionConfigProjection)) + +# Current OmniDiffusionConfig still contains a flat mix of shared engine, +# runtime, parallel, and diffusion-specific knobs. Keep this classification +# explicit while Phase 2 is additive; later cutover PRs can move or remove +# fields without rediscovering the current boundary. +_DIFFUSION_SHARED_CONFIG_FIELDS = frozenset( + { + "stage_id", + "model", + "model_arch", + "dtype", + "trust_remote_code", + "revision", + "distributed_executor_backend", + "dist_timeout", + "model_config", + "quantization_config", + } +) +_DIFFUSION_RUNTIME_CONFIG_FIELDS = frozenset( + { + "host", + "port", + "nccl_port", + "master_port", + "worker_extension_cls", + "enable_stage_verification", + "prompt_file_path", + } +) +_DIFFUSION_ONLY_CONFIG_FIELDS = ( + _DIFFUSION_CONFIG_FIELDS - _DIFFUSION_SHARED_CONFIG_FIELDS - _DIFFUSION_RUNTIME_CONFIG_FIELDS +) +_DIFFUSION_MOVED_SHARED_FIELDS = frozenset( + { + "parallel_config", + "num_gpus", + "log_level", + "profiler_config", + "omni_kv_config", + "cfg_kv_collect_func", + "max_num_seqs", + "enable_sleep_mode", + "enforce_eager", + "enable_multithread_weight_load", + "num_weight_load_threads", + "disable_autocast", + } +) + + +_STAGE_DEPLOY_ENGINE_FIELDS: tuple[str, ...] = tuple(_STAGE_DEPLOY_FIELDS) + +_DIFFUSION_BACKCOMPAT_ENGINE_FIELDS = frozenset( + { + "diffusion_attention_backend", + "kv_cache_dtype", + "kv_cache_skip_layers", + "kv_cache_skip_steps", + "static_lora_scale", + } +) +_DIFFUSION_STAGE_ENGINE_FIELDS = (_DIFFUSION_CONFIG_FIELDS | _DIFFUSION_BACKCOMPAT_ENGINE_FIELDS) - { + "model", + "stage_id", +} + +_QUANTIZATION_ENGINE_FIELDS = frozenset(_QuantizationEngineOverrides.__annotations__) +_MODEL_ENGINE_FIELDS = frozenset(_ModelEngineOverrides.__annotations__) +_LOAD_ENGINE_FIELDS = frozenset(_LoadEngineOverrides.__annotations__) +_CACHE_ENGINE_FIELDS = frozenset(_CacheEngineOverrides.__annotations__) +_SCHEDULER_ENGINE_FIELDS = frozenset(_SchedulerEngineOverrides.__annotations__) +_RUNTIME_ENGINE_FIELDS = frozenset(_RuntimeEngineOverrides.__annotations__) +_PARALLEL_CONFIG_ENGINE_FIELDS = frozenset(_ParallelConfigEngineOverrides.__annotations__) +_PARALLEL_ENGINE_FIELDS = _PARALLEL_CONFIG_ENGINE_FIELDS | {"parallel_config"} + + +def _global_stage_cli_fields() -> frozenset[str]: + # Lazy import avoids vllm_omni.config -> omni_config -> engine.arg_utils -> + # vllm_omni.config during package-level config imports. + from vllm_omni.engine.arg_utils import OmniEngineArgs + + return ( + frozenset(f.name for f in fields(OmniEngineArgs)) + | frozenset(_STAGE_DEPLOY_ENGINE_FIELDS) + | frozenset(_PIPELINE_DEPLOY_CLI_FIELDS) + ) - {"model", "stage_id", "stage_configs_path", "async_chunk"} + + +def _mapping_or_empty(value: Any) -> Mapping[str, Any]: + return value if isinstance(value, Mapping) else {} + + +def _select_engine_overrides(engine: Mapping[str, Any], keys: set[str] | frozenset[str]) -> dict[str, Any]: + return {name: _copy_value(engine[name]) for name in keys if name in engine and engine[name] is not None} + + +def _stage_engine_overrides(stage_deploy: StageDeployConfig | None) -> dict[str, Any]: + if stage_deploy is None: + return {} + + overrides: dict[str, Any] = {} + for name in _STAGE_DEPLOY_ENGINE_FIELDS: + value = getattr(stage_deploy, name) + if value is not None: + overrides[name] = _copy_value(value) + overrides.update(_copy_value(stage_deploy.engine_extras)) + return overrides + + +def _stage_engine_values( + stage_deploy: StageDeployConfig | None, + stage_cli_overrides: Mapping[str, Any] | None = None, +) -> _StageEngineValues: + engine = _stage_engine_overrides(stage_deploy) + if stage_cli_overrides: + engine.update(_copy_value(stage_cli_overrides)) + return _StageEngineValues( + quantization=cast( + _QuantizationEngineOverrides, + _select_engine_overrides(engine, _QUANTIZATION_ENGINE_FIELDS), + ), + model=cast(_ModelEngineOverrides, _select_engine_overrides(engine, _MODEL_ENGINE_FIELDS)), + load=cast(_LoadEngineOverrides, _select_engine_overrides(engine, _LOAD_ENGINE_FIELDS)), + cache=cast(_CacheEngineOverrides, _select_engine_overrides(engine, _CACHE_ENGINE_FIELDS)), + scheduler=cast( + _SchedulerEngineOverrides, + _select_engine_overrides(engine, _SCHEDULER_ENGINE_FIELDS), + ), + runtime=cast(_RuntimeEngineOverrides, _select_engine_overrides(engine, _RUNTIME_ENGINE_FIELDS)), + parallel=cast(_ParallelEngineOverrides, _select_engine_overrides(engine, _PARALLEL_ENGINE_FIELDS)), + diffusion=_DiffusionEngineOverrides.from_engine(engine), + ) + + +def _stage_sampling_params( + stage_deploy: StageDeployConfig | None, + topology: StagePipelineConfig, +) -> dict[str, Any] | None: + sampling: dict[str, Any] = {} + if stage_deploy is not None and stage_deploy.default_sampling_params: + sampling.update(_copy_value(stage_deploy.default_sampling_params)) + sampling.update(_copy_value(topology.sampling_constraints)) + return sampling or None + + +def _orchestrator_cli_overrides(cli_overrides: Mapping[str, Any]) -> dict[str, Any]: + overrides: dict[str, Any] = {} + for config_field in fields(VllmOmniOrchestratorConfig): + name = config_field.name + if name == "deploy_config_path": + continue + if cli_overrides.get(name) is not None: + overrides[name] = _copy_value(cli_overrides[name]) + return overrides + + +@config +class VllmOmniOrchestratorConfig: + """Configuration consumed by the orchestrator process only.""" + + stage_init_timeout: int = Field(default=300, ge=1) + init_timeout: int = Field(default=600, ge=1) + worker_backend: str = "multi_process" + ray_address: str | None = None + deploy_config_path: str | None = None + omni_master_address: str | None = None + omni_master_port: int | None = None + omni_dp_size_local: int = Field(default=1, ge=1) + omni_lb_policy: str = "random" + omni_heartbeat_timeout: float = Field(default=30.0, gt=0.0) + shm_threshold_bytes: int = Field(default=65536, ge=0) + batch_timeout: int = Field(default=10, ge=0) + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class BaseVllmOmniStageConfig: + """Common structured config contract shared by all Omni stage realizations.""" + + stage_pipeline_config: StagePipelineConfig + model_config: OmniStageModelConfig = field(default_factory=OmniStageModelConfig) + load_config: OmniStageLoadConfig = field(default_factory=OmniStageLoadConfig) + cache_config: OmniStageCacheConfig = field(default_factory=OmniStageCacheConfig) + scheduler_config: OmniStageSchedulerConfig = field(default_factory=OmniStageSchedulerConfig) + connector_config: OmniStageConnectorConfig = field(default_factory=OmniStageConnectorConfig) + runtime_config: OmniStageRuntimeConfig = field(default_factory=OmniStageRuntimeConfig) + parallel_config: OmniStageParallelConfig = field(default_factory=OmniStageParallelConfig) + quantization_config: _QuantizationConfigType = None + + @property + def stage_id(self) -> int: + return self.stage_pipeline_config.stage_id + + @property + def model_stage(self) -> str: + return self.stage_pipeline_config.model_stage + + @property + def input_sources(self) -> list[int]: + return list(self.stage_pipeline_config.input_sources) + + @property + def final_output(self) -> bool: + return self.stage_pipeline_config.final_output + + @property + def final_output_type(self) -> str | None: + return self.stage_pipeline_config.final_output_type + + @property + def hf_config_name(self) -> str | None: + return self.stage_pipeline_config.hf_config_name + + @property + def stage_type(self) -> StageType: + stage_type, _ = _resolve_execution_mode(self.stage_pipeline_config.execution_type) + return stage_type + + @property + def worker_type(self) -> str | None: + _, worker_type = _resolve_execution_mode(self.stage_pipeline_config.execution_type) + return worker_type + + @property + def scheduler_cls(self) -> str | None: + return _resolve_scheduler_path( + self.stage_pipeline_config.execution_type, + self.scheduler_config.async_scheduling, + ) + + @property + def custom_process_input_func(self) -> str | None: + return getattr( + self, + "_resolved_custom_process_input_func", + self.stage_pipeline_config.custom_process_input_func, + ) + + @property + def custom_process_next_stage_input_func(self) -> str | None: + return getattr( + self, + "_resolved_custom_process_next_stage_input_func", + self.stage_pipeline_config.custom_process_next_stage_input_func, + ) + + @property + def is_comprehension(self) -> bool: + return self.stage_pipeline_config.owns_tokenizer + + @property + def engine_output_type(self) -> str | None: + return self.stage_pipeline_config.engine_output_type + + @property + def requires_multimodal_data(self) -> bool: + return self.stage_pipeline_config.requires_multimodal_data + + @property + def prompt_expand_func(self) -> str | None: + return self.stage_pipeline_config.prompt_expand_func + + @property + def cfg_kv_collect_func(self) -> str | None: + return self.stage_pipeline_config.cfg_kv_collect_func + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniARStageConfig(BaseVllmOmniStageConfig): + """Structured config for autoregressive LLM stages.""" + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniGenerationStageConfig(BaseVllmOmniStageConfig): + """Structured config for generation LLM stages.""" + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniDiffusionStageConfig(BaseVllmOmniStageConfig): + """Structured config for diffusion stages.""" + + parallel_config: OmniStageDiffusionParallelConfig = field(default_factory=OmniStageDiffusionParallelConfig) + diffusion_config: _DiffusionConfigProjection = field(default_factory=_DiffusionConfigProjection) + + +StageConfigType: TypeAlias = VllmOmniARStageConfig | VllmOmniGenerationStageConfig | VllmOmniDiffusionStageConfig + + +def _build_common_stage_config_kwargs( + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + parallel_config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig, +) -> tuple[dict[str, Any], str | None, str | None]: + input_proc, next_stage_proc = _select_processor_funcs(topology, bool(deploy.async_chunk)) + quantization_config = _build_quantization_config(deploy, engine.quantization) + parallel_config = _build_parallel_config(deploy, engine.parallel, parallel_config_cls) + + return ( + { + "stage_pipeline_config": topology, + "model_config": _build_model_config(topology, stage_deploy, engine.model), + "load_config": _build_load_config(engine.load), + "cache_config": _build_cache_config(deploy, engine.cache), + "scheduler_config": _build_scheduler_config(deploy, engine.scheduler), + "connector_config": _build_connector_config(stage_deploy), + "runtime_config": _build_runtime_config(stage_deploy, engine.runtime, parallel_config), + "parallel_config": parallel_config, + "quantization_config": _copy_value(quantization_config), + }, + input_proc, + next_stage_proc, + ) + + +def _with_resolved_processors( + stage_config: StageConfigType, + input_proc: str | None, + next_stage_proc: str | None, +) -> StageConfigType: + stage_config._resolved_custom_process_input_func = input_proc + stage_config._resolved_custom_process_next_stage_input_func = next_stage_proc + return stage_config + + +def _build_ar_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> VllmOmniARStageConfig: + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + deploy, + topology, + stage_deploy, + engine, + ) + return cast( + VllmOmniARStageConfig, + _with_resolved_processors( + VllmOmniARStageConfig(**common_kwargs), + input_proc, + next_stage_proc, + ), + ) + + +def _build_generation_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> VllmOmniGenerationStageConfig: + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + deploy, + topology, + stage_deploy, + engine, + ) + return cast( + VllmOmniGenerationStageConfig, + _with_resolved_processors( + VllmOmniGenerationStageConfig(**common_kwargs), + input_proc, + next_stage_proc, + ), + ) + + +def _build_diffusion_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> VllmOmniDiffusionStageConfig: + common_kwargs, input_proc, next_stage_proc = _build_common_stage_config_kwargs( + deploy, + topology, + stage_deploy, + engine, + OmniStageDiffusionParallelConfig, + ) + common_kwargs["diffusion_config"] = _build_diffusion_config_projection( + pipeline, + deploy, + topology, + engine.diffusion, + model=model, + quantization_config=common_kwargs["quantization_config"], + ) + return cast( + VllmOmniDiffusionStageConfig, + _with_resolved_processors( + VllmOmniDiffusionStageConfig(**common_kwargs), + input_proc, + next_stage_proc, + ), + ) + + +_STAGE_CONFIG_BUILDERS = { + StageExecutionType.LLM_AR: _build_ar_stage_config, + StageExecutionType.LLM_GENERATION: _build_generation_stage_config, + StageExecutionType.DIFFUSION: _build_diffusion_stage_config, +} + + +def _build_stage_config( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _StageEngineValues, + *, + model: str | None, +) -> StageConfigType: + try: + builder = _STAGE_CONFIG_BUILDERS[topology.execution_type] + except KeyError as exc: + raise ValueError(f"Unsupported stage execution type: {topology.execution_type!r}") from exc + return builder( + pipeline, + deploy, + topology, + stage_deploy, + engine, + model=model, + ) + + +def _build_quantization_config( + deploy: DeployConfig, + engine: _QuantizationEngineOverrides, +) -> _QuantizationConfigType: + return _first_defined( + engine.get("quantization_config"), + engine.get("quantization"), + deploy.quantization, + ) + + +def _build_model_config( + topology: StagePipelineConfig, + stage_deploy: StageDeployConfig | None, + engine: _ModelEngineOverrides, +) -> OmniStageModelConfig: + default_sampling_params = _stage_sampling_params(stage_deploy, topology) + kwargs = _config_kwargs(engine) + if "has_sampling_extra_args" not in kwargs: + kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args")) + return OmniStageModelConfig( + default_sampling_params=default_sampling_params, + **kwargs, + ) + + +def _build_load_config(engine: _LoadEngineOverrides) -> OmniStageLoadConfig: + return OmniStageLoadConfig(**_config_kwargs(engine)) + + +def _build_cache_config( + deploy: DeployConfig, + engine: _CacheEngineOverrides, +) -> OmniStageCacheConfig: + kwargs = _config_kwargs(engine) + if "enable_prefix_caching" not in kwargs and deploy.enable_prefix_caching is not None: + kwargs["enable_prefix_caching"] = _copy_value(deploy.enable_prefix_caching) + return OmniStageCacheConfig(**kwargs) + + +def _build_scheduler_config( + deploy: DeployConfig, + engine: _SchedulerEngineOverrides, +) -> OmniStageSchedulerConfig: + kwargs = _config_kwargs(engine) + if "enable_chunked_prefill" not in kwargs and deploy.enable_chunked_prefill is not None: + kwargs["enable_chunked_prefill"] = _copy_value(deploy.enable_chunked_prefill) + return OmniStageSchedulerConfig(**kwargs) + + +def _build_connector_config(stage_deploy: StageDeployConfig | None) -> OmniStageConnectorConfig: + output_connectors = stage_deploy.output_connectors if stage_deploy is not None else None + input_connectors = stage_deploy.input_connectors if stage_deploy is not None else None + return OmniStageConnectorConfig( + output_connectors=_copy_value(output_connectors) if output_connectors else None, + input_connectors=_copy_value(input_connectors) if input_connectors else None, + ) + + +def _build_runtime_config( + stage_deploy: StageDeployConfig | None, + engine: _RuntimeEngineOverrides, + parallel_config: OmniStageParallelConfig, +) -> OmniStageRuntimeConfig: + kwargs = _config_kwargs(engine) + if "devices" not in kwargs and stage_deploy is not None and stage_deploy.devices is not None: + kwargs["devices"] = _copy_value(stage_deploy.devices) + if "num_replicas" not in kwargs and stage_deploy is not None: + kwargs["num_replicas"] = stage_deploy.num_replicas + if "env" not in kwargs and stage_deploy is not None and stage_deploy.env is not None: + kwargs["env"] = _copy_value(stage_deploy.env) + kwargs["num_gpus"] = parallel_config.world_size + return OmniStageRuntimeConfig(**kwargs) + + +def _build_parallel_config( + deploy: DeployConfig, + engine: _ParallelEngineOverrides, + config_cls: type[OmniStageParallelConfig] = OmniStageParallelConfig, +) -> OmniStageParallelConfig: + parallel_config = _mapping_or_empty(engine.get("parallel_config")) + config_fields = { + config_field.name + for config_field in fields(config_cls) + if getattr(config_field.default, "init", None) is not False + } + kwargs = { + name: _copy_value(value) for name in _PARALLEL_CONFIG_ENGINE_FIELDS if (value := engine.get(name)) is not None + } + kwargs = {name: value for name, value in kwargs.items() if name in config_fields} + kwargs.update({name: value for name, value in _config_kwargs(parallel_config).items() if name in config_fields}) + if "pipeline_parallel_size" not in kwargs and deploy.pipeline_parallel_size is not None: + kwargs["pipeline_parallel_size"] = _copy_value(deploy.pipeline_parallel_size) + if "data_parallel_size" not in kwargs and deploy.data_parallel_size is not None: + kwargs["data_parallel_size"] = _copy_value(deploy.data_parallel_size) + return config_cls(**kwargs) + + +def _build_diffusion_config_projection( + pipeline: PipelineConfig, + deploy: DeployConfig, + topology: StagePipelineConfig, + engine: _DiffusionEngineOverrides, + *, + model: str | None, + quantization_config: _QuantizationConfigType, +) -> _DiffusionConfigProjection: + diffusion_kwargs = engine.to_kwargs() + diffusion_kwargs["stage_id"] = topology.stage_id + diffusion_kwargs["model_arch"] = _first_defined( + diffusion_kwargs.get("model_arch"), + topology.model_arch, + pipeline.model_arch, + ) + if "dtype" not in diffusion_kwargs and deploy.dtype is not None: + diffusion_kwargs["dtype"] = _copy_value(deploy.dtype) + if "trust_remote_code" not in diffusion_kwargs and deploy.trust_remote_code is not None: + diffusion_kwargs["trust_remote_code"] = _copy_value(deploy.trust_remote_code) + if "distributed_executor_backend" not in diffusion_kwargs and deploy.distributed_executor_backend is not None: + diffusion_kwargs["distributed_executor_backend"] = _copy_value(deploy.distributed_executor_backend) + if model is not None: + diffusion_kwargs["model"] = model + if quantization_config is not None: + diffusion_kwargs["quantization_config"] = _copy_value(quantization_config) + + return _DiffusionConfigProjection.from_kwargs(**{k: v for k, v in diffusion_kwargs.items() if v is not None}) + + +@config(config=ConfigDict(arbitrary_types_allowed=True)) +class VllmOmniConfig: + """Top-level structured Omni config built once from registry inputs.""" + + pipeline_config: PipelineConfig + stage_configs: tuple[StageConfigType, ...] + orchestrator_config: VllmOmniOrchestratorConfig = field(default_factory=VllmOmniOrchestratorConfig) + + def stage_by_id(self, stage_id: int) -> StageConfigType: + for stage in self.stage_configs: + if stage.stage_id == stage_id: + return stage + raise KeyError(f"no stage {stage_id}") + + @classmethod + def from_pipeline_config( + cls, + pipeline_cfg: PipelineConfig, + *, + user_deploy_config: DeployConfig | None = None, + deploy_config_path: str | None = None, + cli_overrides: dict[str, Any] | None = None, + ) -> VllmOmniConfig: + """Create a structured config from a resolved pipeline and deploy YAML.""" + if cli_overrides is None: + cli_overrides = {} + + deploy, loaded_deploy_config_path = _get_deploy_config( + pipeline_cfg, + user_deploy_config, + deploy_config_path, + ) + + if cli_overrides.get("async_chunk") is not None: + deploy.async_chunk = bool(cli_overrides["async_chunk"]) + for name in _PIPELINE_DEPLOY_CLI_FIELDS: + if cli_overrides.get(name) is not None: + setattr(deploy, name, _copy_value(cli_overrides[name])) + + deploy = _apply_platform_overrides(deploy) + if len(pipeline_cfg.stages) <= 1: + deploy.async_chunk = False + _validate_async_chunk_support(pipeline_cfg, deploy) + deploy_by_id = {stage.stage_id: stage for stage in deploy.stages} + model = cli_overrides.get("model") + + stage_configs = tuple( + _build_stage_config( + pipeline_cfg, + deploy, + topology, + deploy_by_id.get(topology.stage_id), + _stage_engine_values( + deploy_by_id.get(topology.stage_id), + _stage_cli_overrides(topology.stage_id, cli_overrides), + ), + model=model, + ) + for topology in pipeline_cfg.stages + ) + + orchestrator_config = VllmOmniOrchestratorConfig( + deploy_config_path=loaded_deploy_config_path, + **_orchestrator_cli_overrides(cli_overrides), + ) + return cls( + pipeline_config=pipeline_cfg, + stage_configs=stage_configs, + orchestrator_config=orchestrator_config, + ) + + +__all__ = [ + "OmniStageCacheConfig", + "OmniStageConnectorConfig", + "BaseVllmOmniStageConfig", + "OmniStageLoadConfig", + "OmniStageModelConfig", + "VllmOmniOrchestratorConfig", + "OmniStageDiffusionParallelConfig", + "OmniStageParallelConfig", + "OmniStageRuntimeConfig", + "OmniStageSchedulerConfig", + "StageConfigType", + "VllmOmniARStageConfig", + "VllmOmniConfig", + "VllmOmniDiffusionStageConfig", + "VllmOmniGenerationStageConfig", +] diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 5e35914c91f..0ee96eca4c0 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -1,170 +1,187 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Central declarative registry of all vllm-omni pipelines. +"""Pipeline registry and factory for vllm-omni. -Mirrors the pattern in ``vllm/model_executor/models/registry.py``: each entry -is ``model_type -> (module_path, variable_name)``, and the module is imported -lazily on first lookup (see ``_LazyPipelineRegistry`` in -``vllm_omni/config/stage_config.py``). Keeping every pipeline declared in one -file makes it easy to spot a missing registration, which was the original -motivation in https://github.com/vllm-project/vllm-omni/issues/2887 (item 4). +``OMNI_PIPELINES`` maps each ``model_type`` to either a ``PipelineConfig`` +instance or a resolver callable that accepts an optional HF config and returns +a ``PipelineConfig``. -Per-model ``pipeline.py`` modules still define the ``PipelineConfig`` instance; -they just no longer need to self-register via ``register_pipeline(...)``. - -Adding a new pipeline: +To add a new pipeline: 1. Define the ``PipelineConfig`` instance as a module-level variable in ``vllm_omni/.../pipeline.py``. - 2. Add one line to ``_OMNI_PIPELINES`` below. + 2. If the model needs to support several configurations, e.g., because some + stages are optional, implement a resolver that consumes the HF config + and returns a ``PipelineConfig``. + 3. Update the registry to map the key to the new config object (in the case + of new keys) or to the resolver func. -Plain single-stage diffusion models continue to use the -``_create_default_diffusion_stage_cfg`` fallback in ``async_omni_engine.py``. -The empty ``_DIFFUSION_PIPELINES`` placeholder previously here (#2915) was -removed once #2987 (which would have populated it) was deferred. +Out of tree pipeline configs or resolvers can also be registered with register_pipeline. -``register_pipeline(config)`` in ``stage_config`` is still supported for -out-of-tree plugins and tests that create pipelines at runtime; those override -the entries declared here. +NOTE: Single-stage diffusion models continue to use the +``_create_default_diffusion_stage_cfg`` fallback in +``async_omni_engine.py``; for now we do not add them to registry. """ from __future__ import annotations +from collections.abc import Callable +from typing import TypeAlias + +from transformers import PretrainedConfig +from vllm.logger import init_logger + +from vllm_omni.config.stage_config import ( + PipelineConfig, +) +from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE +from vllm_omni.model_executor.models.bagel.pipeline import ( + BAGEL_PIPELINE, + BAGEL_SINGLE_STAGE_PIPELINE, + BAGEL_THINK_PIPELINE, +) +from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE +from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE +from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE +from vllm_omni.model_executor.models.dynin_omni.pipeline import DYNIN_OMNI_PIPELINE +from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE +from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE +from vllm_omni.model_executor.models.glm_tts.pipeline import GLM_TTS_PIPELINE +from vllm_omni.model_executor.models.gr00t.pipeline import GR00T_N1D7_PIPELINE +from vllm_omni.model_executor.models.higgs_audio_v2.pipeline import HIGGS_AUDIO_V2_PIPELINE +from vllm_omni.model_executor.models.higgs_audio_v3.pipeline import HIGGS_AUDIO_V3_PIPELINE +from vllm_omni.model_executor.models.hunyuan_image3.pipeline import ( + HUNYUAN_IMAGE3_AR_PIPELINE, + HUNYUAN_IMAGE3_DIT_PIPELINE, + HUNYUAN_IMAGE3_PIPELINE, +) +from vllm_omni.model_executor.models.hunyuan_video.pipeline import HUNYUAN_VIDEO_15_PIPELINE +from vllm_omni.model_executor.models.indextts2.pipeline import INDEXTTS2_PIPELINE +from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE +from vllm_omni.model_executor.models.mammoth_moda2.pipeline import ( + MAMMOTH_MODA2_AR_PIPELINE, + MAMMOTH_MODA2_PIPELINE, +) +from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE +from vllm_omni.model_executor.models.ming_flash_omni.pipeline import ( + MING_FLASH_OMNI_IMAGE_PIPELINE, + MING_FLASH_OMNI_PIPELINE, + MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, + MING_FLASH_OMNI_TTS_PIPELINE, +) +from vllm_omni.model_executor.models.ming_tts.pipeline import ( + MING_TTS_MOE_PIPELINE, + MING_TTS_PIPELINE, +) +from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO_4_5_PIPELINE +from vllm_omni.model_executor.models.moss_tts.pipeline import ( + MOSS_TTS_LOCAL_PIPELINE, + MOSS_TTS_PIPELINE, + MOSS_TTS_REALTIME_PIPELINE, +) +from vllm_omni.model_executor.models.moss_tts_nano.pipeline import MOSS_TTS_NANO_PIPELINE +from vllm_omni.model_executor.models.omnivoice.pipeline import OMNIVOICE_PIPELINE +from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import ( + QWEN2_5_OMNI_PIPELINE, + QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, +) +from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline +from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE +from vllm_omni.model_executor.models.soulx_singer.pipeline import ( + SOULXSINGER_SVC_PIPELINE, + SOULXSINGER_SVS_PIPELINE, +) +from vllm_omni.model_executor.models.step_audio2.pipeline import ( + STEP_AUDIO2_ASR_PIPELINE, + STEP_AUDIO2_PIPELINE, +) +from vllm_omni.model_executor.models.voxcpm2.pipeline import VOXCPM2_PIPELINE +from vllm_omni.model_executor.models.voxtral_tts.pipeline import VOXTRAL_TTS_PIPELINE +from vllm_omni.model_executor.models.wan2_2.pipeline import WAN2_2_TI2V_PIPELINE + +logger = init_logger(__name__) + +PipelineResolverFunc: TypeAlias = Callable[[PretrainedConfig | None], PipelineConfig | None] + # --- Multi-stage omni pipelines (LLM-centric; audio / video I/O) --- -_OMNI_PIPELINES: dict[str, tuple[str, str]] = { - # model_type -> (module_path, variable_name) - "qwen2_5_omni": ( - "vllm_omni.model_executor.models.qwen2_5_omni.pipeline", - "QWEN2_5_OMNI_PIPELINE", - ), - "qwen2_5_omni_thinker_only": ( - "vllm_omni.model_executor.models.qwen2_5_omni.pipeline", - "QWEN2_5_OMNI_THINKER_ONLY_PIPELINE", - ), - "qwen3_omni_moe": ( - "vllm_omni.model_executor.models.qwen3_omni.pipeline", - "QWEN3_OMNI_PIPELINE", - ), - # HACK (Alex) - this is a hack to avoid bringing in the whole resolution - # pr to support qwen3 omni models that do not enable audio output. - "qwen3_omni_moe_thinker_only": ( - "vllm_omni.model_executor.models.qwen3_omni.pipeline", - "QWEN3_OMNI_THINKER_ONLY_PIPELINE", - ), - "qwen3_tts": ( - "vllm_omni.model_executor.models.qwen3_tts.pipeline", - "QWEN3_TTS_PIPELINE", - ), - "covo_audio": ( - "vllm_omni.model_executor.models.covo_audio.pipeline", - "COVO_AUDIO_PIPELINE", - ), - "bagel": ( - "vllm_omni.model_executor.models.bagel.pipeline", - "BAGEL_PIPELINE", - ), - "bagel_think": ( - "vllm_omni.model_executor.models.bagel.pipeline", - "BAGEL_THINK_PIPELINE", - ), - "bagel_single_stage": ( - "vllm_omni.model_executor.models.bagel.pipeline", - "BAGEL_SINGLE_STAGE_PIPELINE", - ), - # Lance (ByteDance) — BAGEL-lineage unified AR+diffusion, single-stage DiT. - "lance": ( - "vllm_omni.model_executor.models.lance.pipeline", - "LANCE_PIPELINE", - ), - "dreamzero": ( - "vllm_omni.model_executor.models.dreamzero.pipeline", - "DREAMZERO_PIPELINE", - ), - "glm_image": ( - "vllm_omni.model_executor.models.glm_image.pipeline", - "GLM_IMAGE_PIPELINE", - ), - "hunyuan_image_3_moe": ( - "vllm_omni.model_executor.models.hunyuan_image3.pipeline", - "HUNYUAN_IMAGE3_PIPELINE", - ), - "hunyuan_image3_ar": ( - "vllm_omni.model_executor.models.hunyuan_image3.pipeline", - "HUNYUAN_IMAGE3_AR_PIPELINE", - ), - "hunyuan_image3_dit": ( - "vllm_omni.model_executor.models.hunyuan_image3.pipeline", - "HUNYUAN_IMAGE3_DIT_PIPELINE", - ), - "voxcpm2": ( - "vllm_omni.model_executor.models.voxcpm2.pipeline", - "VOXCPM2_PIPELINE", - ), - "cosyvoice3": ( - "vllm_omni.model_executor.models.cosyvoice3.pipeline", - "COSYVOICE3_PIPELINE", - ), - "mimo_audio": ( - "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", - ), - "glm_tts": ( - "vllm_omni.model_executor.models.glm_tts.pipeline", - "GLM_TTS_PIPELINE", - ), - "fish_qwen3_omni": ( - "vllm_omni.model_executor.models.fish_speech.pipeline", - "FISH_SPEECH_PIPELINE", - ), - "ming_flash_omni": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_PIPELINE", - ), - "ming_flash_omni_tts": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_TTS_PIPELINE", - ), - "ming_flash_omni_thinker_only": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_THINKER_ONLY_PIPELINE", - ), - "ming_flash_omni_image": ( - "vllm_omni.model_executor.models.ming_flash_omni.pipeline", - "MING_FLASH_OMNI_IMAGE_PIPELINE", - ), - "moss_tts_nano": ( - "vllm_omni.model_executor.models.moss_tts_nano.pipeline", - "MOSS_TTS_NANO_PIPELINE", - ), - "moss_tts_delay": ( - "vllm_omni.model_executor.models.moss_tts.pipeline", - "MOSS_TTS_PIPELINE", - ), - "moss_tts_realtime": ( - "vllm_omni.model_executor.models.moss_tts.pipeline", - "MOSS_TTS_REALTIME_PIPELINE", - ), - "minicpmo_4_5": ( - "vllm_omni.model_executor.models.minicpmo_4_5.pipeline", - "MINICPMO_4_5_PIPELINE", - ), - "higgs_audio_v2": ( - "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", - ), +OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { + "aura_omni": AURA_OMNI_PIPELINE, + "qwen2_5_omni": QWEN2_5_OMNI_PIPELINE, + "qwen2_5_omni_thinker_only": QWEN2_5_OMNI_THINKER_ONLY_PIPELINE, + "qwen3_omni_moe": resolve_qwen3_omni_pipeline, + "qwen3_tts": QWEN3_TTS_PIPELINE, + "step_audio_2": STEP_AUDIO2_PIPELINE, + "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE, + "covo_audio": COVO_AUDIO_PIPELINE, + "bagel": BAGEL_PIPELINE, + "bagel_think": BAGEL_THINK_PIPELINE, + "bagel_single_stage": BAGEL_SINGLE_STAGE_PIPELINE, + "lance": LANCE_PIPELINE, + "dreamzero": DREAMZERO_PIPELINE, + "Gr00tN1d7": GR00T_N1D7_PIPELINE, + "glm_image": GLM_IMAGE_PIPELINE, + "hunyuan_image_3_moe": HUNYUAN_IMAGE3_PIPELINE, + "hunyuan_image3_ar": HUNYUAN_IMAGE3_AR_PIPELINE, + "hunyuan_image3_dit": HUNYUAN_IMAGE3_DIT_PIPELINE, + "hunyuan_video_15": HUNYUAN_VIDEO_15_PIPELINE, + "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE, + "voxcpm2": VOXCPM2_PIPELINE, + "cosyvoice3": COSYVOICE3_PIPELINE, + "mimo_audio": MIMO_AUDIO_PIPELINE, + "ming_tts": MING_TTS_PIPELINE, + "ming_tts_moe": MING_TTS_MOE_PIPELINE, + "voxtral_tts": VOXTRAL_TTS_PIPELINE, + "glm_tts": GLM_TTS_PIPELINE, + "fish_qwen3_omni": FISH_SPEECH_PIPELINE, + "ming_flash_omni": MING_FLASH_OMNI_PIPELINE, + "ming_flash_omni_tts": MING_FLASH_OMNI_TTS_PIPELINE, + "ming_flash_omni_thinker_only": MING_FLASH_OMNI_THINKER_ONLY_PIPELINE, + "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE, + "moss_tts_nano": MOSS_TTS_NANO_PIPELINE, + "omnivoice": OMNIVOICE_PIPELINE, + "mammoth_moda2": MAMMOTH_MODA2_PIPELINE, + "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE, + "moss_tts_delay": MOSS_TTS_PIPELINE, + "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE, + "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE, + "minicpmo_4_5": MINICPMO_4_5_PIPELINE, + "higgs_audio_v2": HIGGS_AUDIO_V2_PIPELINE, + "higgs_multimodal_qwen3": HIGGS_AUDIO_V3_PIPELINE, + "dynin_omni": DYNIN_OMNI_PIPELINE, + "indextts2": INDEXTTS2_PIPELINE, + "soulxsinger_svc": SOULXSINGER_SVC_PIPELINE, + "soulxsinger_svs": SOULXSINGER_SVS_PIPELINE, } + + +def register_pipeline(pipeline: PipelineConfig | PipelineResolverFunc, model_type: str | None = None): + """Register an out of tree pipeline or PipelineResolverFunc to a model_type key. + If a PipelineConfig is provided, model_type is optional, and pipeline.model_type + will be used by default. If a callable is provided, model_type must be provided, + since resolvers can return multiple different PipelineConfigs depending on the + consumed config. + """ + errors: list[str] = [] + if isinstance(pipeline, PipelineConfig): + errors = pipeline.validate() + model_type = model_type if model_type is not None else pipeline.model_type + else: + if model_type is None: + raise ValueError("Model type must be explicitly provided when registering a pipeline resolver") + + if model_type in OMNI_PIPELINES: + errors.append(f"Model type {model_type} is already registered; the old mapping will be clobbered") + if errors: + logger.warning("Registration for pipeline of type %s produced the following issues: %s", model_type, errors) + OMNI_PIPELINES[model_type] = pipeline + + +def resolve_pipeline_config( + model_type: str, + hf_config: PretrainedConfig | None = None, +) -> PipelineConfig | None: + """Resolve a registry key to a concrete pipeline config.""" + if model_type not in OMNI_PIPELINES: + logger.warning("Model type %s is not registered to OMNI_PIPELINES", model_type) + return None + pipeline = OMNI_PIPELINES[model_type] + return pipeline(hf_config) if callable(pipeline) else pipeline diff --git a/vllm_omni/config/server_settings.py b/vllm_omni/config/server_settings.py new file mode 100644 index 00000000000..bb63712ba10 --- /dev/null +++ b/vllm_omni/config/server_settings.py @@ -0,0 +1,58 @@ +import os +import warnings +from typing import Literal + +from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic_settings import BaseSettings, SettingsConfigDict + + +def _resolve_deprecated_env(old_name: str, new_name: str, value): + old = os.environ.get(old_name) + if old is None: + return value + warnings.warn( + f"{old_name} is deprecated and will be removed in a future release; use {new_name} instead.", + DeprecationWarning, + stacklevel=2, + ) + return value if new_name in os.environ else old + + +class FileBackend(BaseModel): + type: Literal["file"] = "file" + path: str = Field(default="/tmp/storage", validate_default=True, description="Local path to store completed files.") + file_concurrency: int = Field( + default=4, validate_default=True, description="Maximum number of file operations permitted at a time" + ) + file_ttl: int | None = Field( + default=None, description="Optional TTL (in seconds) configuration settings for locally stored files." + ) + ttl_sweep_interval: int | None = Field( + default=None, description="Optional frequency (in seconds) to enforce file TTLs." + ) + + @field_validator("path", mode="before") + @classmethod + def _migrate_path(cls, value): + return _resolve_deprecated_env("VLLM_OMNI_STORAGE_PATH", "VLLM_OMNI_SERVER_STORAGE__PATH", value) + + @field_validator("file_concurrency", mode="before") + @classmethod + def _migrate_file_concurrency(cls, value): + return _resolve_deprecated_env( + "VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY", value + ) + + @model_validator(mode="after") + def set_default_ttl_sweep_interval(self) -> "FileBackend": + if self.file_ttl is not None and self.ttl_sweep_interval is None: + self.ttl_sweep_interval = 300 + return self + + +class ServerSettings(BaseSettings): + model_config = SettingsConfigDict(env_prefix="VLLM_OMNI_SERVER_", env_nested_delimiter="__") + storage: FileBackend = Field(default_factory=FileBackend) + + +SERVER_SETTINGS_CONFIG = ServerSettings() diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index dbc5a83633a..367e1b317aa 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -5,6 +5,7 @@ from __future__ import annotations import dataclasses +import functools import re import warnings from collections.abc import Callable @@ -13,6 +14,7 @@ from pathlib import Path from typing import Any, NamedTuple +from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler @@ -23,10 +25,26 @@ logger = init_logger(__name__) +_DEPLOY_DIR = Path(__file__).resolve().parent.parent / "deploy" _STAGE_OVERRIDE_PATTERN = re.compile(r"^stage_(\d+)_(.+)$") +def pipeline_cfg_resolver(config_type: type[PretrainedConfig]): + """Wraps a resolver such that we return None if a hf_config of the wrong type is provided.""" + + def resolver_builder(func): + @functools.wraps(func) + def wrapper(hf_config: PretrainedConfig | None): + if hf_config is None or not isinstance(hf_config, config_type): + return None + return func(hf_config) + + return wrapper + + return resolver_builder + + def build_stage_runtime_overrides( stage_id: int, cli_overrides: dict[str, Any], @@ -46,7 +64,12 @@ def build_stage_runtime_overrides( if internal_keys is None: from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys - internal_keys = internal_blacklist_keys() | SHARED_FIELDS + # Some fields are modeled as orchestrator-owned for top-level CLI + # parsing, but are also legitimate deploy-time stage overrides. Keep + # the default blacklist for true orchestrator/shared fields while + # allowing any field explicitly represented by the deploy schema to + # continue flowing into per-stage overrides. + internal_keys = (internal_blacklist_keys() | SHARED_FIELDS) - deploy_runtime_override_keys() result: dict[str, Any] = {} @@ -206,6 +229,7 @@ class StagePipelineConfig: prompt_expand_func: str | None = None cfg_kv_collect_func: str | None = None omni_kv_config: dict[str, Any] | None = None + scheduler_cls: str | None = None # Model subdirectory indirections: for multi-component HF repos where the # stage's config/tokenizer lives in a subdirectory (e.g. GLM-Image's AR # config is in ``vision_language_encoder/``). Consumed at stage-init time @@ -246,6 +270,9 @@ class PipelineConfig: # multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model. diffusers_class_name: str | None = None endpoint_restrictions: tuple[EndpointRestriction, ...] = () + # Bundled deploy defaults for this concrete pipeline topology. The file is + # loaded from vllm_omni/deploy; None uses DeployConfig defaults. + default_deploy_config_name: str | None = None def get_stage(self, stage_id: int) -> StagePipelineConfig | None: """Look up a stage by its ID.""" @@ -275,153 +302,6 @@ def validate(self) -> list[str]: return errors -class _LazyPipelineRegistry: - """Dict-like registry that lazy-loads pipelines from the central declaration. - - In-tree pipelines are declared once in - ``vllm_omni/config/pipeline_registry.py`` as - ``model_type -> (module_path, variable_name)`` entries; the module is - imported only when the pipeline is first looked up. This mirrors the - pattern in ``vllm/model_executor/models/registry.py`` and addresses - https://github.com/vllm-project/vllm-omni/issues/2887 (item 4): having - every registration in one file makes a missing entry easy to spot. - - Out-of-tree / dynamic registrations via ``register_pipeline()`` are stored - directly in ``_loaded`` and take precedence over the lazy-map entry with - the same ``model_type``. - - The class exposes the subset of ``dict`` operations the rest of this - module relies on (``__contains__``, ``__getitem__``, ``__setitem__``, - ``get``, ``keys``, ``values``, ``items``, ``__iter__``), so existing call - sites don't need to change. - """ - - def __init__(self) -> None: - self._loaded: dict[str, PipelineConfig] = {} - # Populated lazily to avoid a circular import at module init time. - self._lazy_map: dict[str, tuple[str, str]] | None = None - - def _get_lazy_map(self) -> dict[str, tuple[str, str]]: - if self._lazy_map is None: - from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES - - self._lazy_map = _OMNI_PIPELINES - return self._lazy_map - - def _load_lazy(self, model_type: str) -> PipelineConfig | None: - entry = self._get_lazy_map().get(model_type) - if entry is None: - return None - module_path, var_name = entry - import importlib - - try: - module = importlib.import_module(module_path) - pipeline = getattr(module, var_name, None) - if pipeline is None: - logger.error( - "Pipeline variable %r not found in module %r (registered for %r)", - var_name, - module_path, - model_type, - ) - return None - errors = pipeline.validate() - if errors: - logger.warning("Pipeline %s has issues: %s", pipeline.model_type, errors) - self._loaded[model_type] = pipeline - return pipeline - except Exception: - logger.exception("Failed to import pipeline module %r for %r", module_path, model_type) - return None - - def __contains__(self, model_type: str) -> bool: - if model_type in self._loaded: - return True - return model_type in self._get_lazy_map() - - def __getitem__(self, model_type: str) -> PipelineConfig: - if model_type in self._loaded: - return self._loaded[model_type] - pipeline = self._load_lazy(model_type) - if pipeline is None: - raise KeyError(model_type) - return pipeline - - def get(self, model_type: str, default: PipelineConfig | None = None) -> PipelineConfig | None: - if model_type in self._loaded: - return self._loaded[model_type] - pipeline = self._load_lazy(model_type) - return pipeline if pipeline is not None else default - - def __setitem__(self, model_type: str, pipeline: PipelineConfig) -> None: - self._loaded[model_type] = pipeline - - def __delitem__(self, model_type: str) -> None: - """Remove a dynamically-registered pipeline. - - Only the dynamic-cache side of the registry can be mutated; the - central declarative registry is immutable at runtime. Calling ``del`` - on a model_type that only exists in the central registry raises - ``KeyError``. - """ - if model_type in self._loaded: - del self._loaded[model_type] - return - if model_type in self._get_lazy_map(): - raise KeyError( - f"{model_type!r} is declared in the central pipeline_registry and " - "cannot be removed at runtime. Edit " - "vllm_omni/config/pipeline_registry.py to delete a built-in entry." - ) - raise KeyError(model_type) - - def keys(self) -> set[str]: - return set(self._get_lazy_map().keys()) | set(self._loaded.keys()) - - def _safe_get(self, key: str) -> PipelineConfig | None: - try: - return self[key] - except Exception: - logger.warning("Skipping pipeline %r because it failed to load.", key) - return None - - def values(self): - # Iterating forces a lazy import for each pipeline; failures are logged and skipped. - for key in self.keys(): - if (p := self._safe_get(key)) is not None: - yield p - - def items(self): - for key in self.keys(): - if (p := self._safe_get(key)) is not None: - yield key, p - - def __iter__(self): - return iter(self.keys()) - - -_PIPELINE_REGISTRY = _LazyPipelineRegistry() - - -def register_pipeline(pipeline: PipelineConfig) -> None: - """Register a pipeline config dynamically. - - In-tree pipelines are declared in ``pipeline_registry._OMNI_PIPELINES`` - and loaded lazily; calling ``register_pipeline`` is only needed for - out-of-tree plugins or tests that build a ``PipelineConfig`` at runtime. - A dynamic registration overrides the central-registry entry with the same - ``model_type``. - """ - errors = pipeline.validate() - if errors: - logger.warning("Pipeline %s has issues: %s", pipeline.model_type, errors) - _PIPELINE_REGISTRY[pipeline.model_type] = pipeline - - -_DEPLOY_DIR = Path(__file__).resolve().parent.parent / "deploy" - - @dataclass class StageDeployConfig: """Per-stage deployment knobs. @@ -434,7 +314,7 @@ class StageDeployConfig: the top level of ``DeployConfig`` and propagated to every stage. """ - # === Omni fields === + # === Omni stage wrapper fields === # Stage identity and Omni runtime placement. stage_id: int devices: str | None = None @@ -447,42 +327,82 @@ class StageDeployConfig: default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None - # === vLLM EngineArgs fields === - # Parallelism and scheduler/memory capacity. + # === Generic stage engine fields === + # Parallelism, scheduler, and memory-capacity controls. tensor_parallel_size: int | None = None + enable_expert_parallel: bool | None = None gpu_memory_utilization: float | None = None max_num_seqs: int | None = None max_num_batched_tokens: int | None = None max_model_len: int | None = None - # Execution, scheduling, and KV/cache behavior. + # Generic execution, scheduling, and KV/cache behavior. enforce_eager: bool | None = None async_scheduling: bool | None = None disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = None - # Diffusion parallel_config deploy override fields. - enable_expert_parallel: bool | None = None + # Generic compilation, profiling, tokenizer/config parsing, and model + # loading controls. + compilation_config: dict[str, Any] | None = None + profiler_config: dict[str, Any] | None = None + skip_mm_profiling: bool | None = None + enable_flashinfer_autotune: bool | None = None + config_format: str | None = None + load_format: str | None = None + tokenizer_mode: str | None = None + + # === Diffusion stage runtime fields === + # Diffusion parallel_config deploy/runtime override fields. ulysses_degree: int | None = None ulysses_mode: str | None = None ring_degree: int | None = None sequence_parallel_size: int | None = None cfg_parallel_size: int | None = None vae_patch_parallel_size: int | None = None + vae_parallel_mode: str | None = None use_hsdp: bool | None = None hsdp_shard_size: int | None = None hsdp_replicate_size: int | None = None - # Compilation, profiling, tokenizer/config parsing, and model loading. - compilation_config: dict[str, Any] | None = None - profiler_config: dict[str, Any] | None = None - skip_mm_profiling: bool | None = None - enable_flashinfer_autotune: bool | None = None - config_format: str | None = None - load_format: str | None = None - tokenizer_mode: str | None = None - - # Pass-through vLLM EngineArgs fields that are not represented above. + # Diffusion model loading and adapter construction. + model_class_name: str | None = None + diffusion_load_format: str | None = None + diffusers_load_kwargs: dict[str, Any] | None = None + diffusers_call_kwargs: dict[str, Any] | None = None + diffusion_quantization_config: str | None = None + diffusion_attention_backend: str | None = None + diffusion_attention_config: dict[str, Any] | None = None + + # Diffusion execution, cache, and VAE behavior. + cache_backend: str | None = None + cache_config: dict[str, Any] | None = None + enable_cache_dit_summary: bool | None = None + step_execution: bool | None = None + vae_use_slicing: bool | None = None + vae_use_tiling: bool | None = None + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | None = None + diffusion_kv_cache_skip_layers: str | None = None + auxiliary_text_encoder: str | None = None + + # Runtime optimizations used by diffusion loading/execution. + enable_multithread_weight_load: bool | None = None + num_weight_load_threads: int | None = None + enable_cpu_offload: bool | None = None + enable_layerwise_offload: bool | None = None + + # Diffusion-specific debug and observability knobs. + enable_diffusion_pipeline_profiler: bool | None = None + + # Modality/service constraints consumed outside the core engine config. + max_generated_image_size: int | None = None + tts_max_instructions_length: int | None = None + + # === Pass-through stage engine fields === + # Pass-through stage engine args that are not represented above. engine_extras: dict[str, Any] = field(default_factory=dict) @@ -539,10 +459,22 @@ class DeployConfig: _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS} +def deploy_runtime_override_keys() -> frozenset[str]: + """Return deploy-schema fields that are valid CLI/runtime overrides. + + These keys form the positive contract for stage override propagation: + stage-scoped deploy knobs plus top-level pipeline-wide engine settings. + They must remain overridable even if they are also modeled on + ``OrchestratorArgs`` for top-level CLI parsing. + """ + return frozenset(_STAGE_DEPLOY_FIELDS) | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS) + + def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" # Get the non-reserved keys for this stage flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS} + explicit_engine_extras = dict(stage_data.get("engine_extras") or {}) runtime_cfg = dict(stage_data.get("runtime", {})) devices = runtime_cfg.get("devices", stage_data.get("devices")) num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1)) @@ -570,7 +502,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") - kwargs["engine_extras"] = flat_args + kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args) return StageDeployConfig(**kwargs) @@ -733,7 +665,8 @@ def _apply_platform_overrides( if platform is None: from vllm_omni.platforms import current_omni_platform - platform = current_omni_platform.device_name.lower() + device_name = current_omni_platform.device_name + platform = device_name.lower() if device_name is not None else None if platform is None or deploy.platforms is None: return deploy platform_section = deploy.platforms.get(platform) @@ -821,6 +754,7 @@ def _select_processor_funcs( "active_stream_window", "custom_voice_dir", ) +PIPELINE_WIDE_ENGINE_FIELDS = _PIPELINE_WIDE_ENGINE_FIELDS def _build_engine_args( @@ -837,6 +771,8 @@ def _build_engine_args( ``engine_extras`` can still carry a stage-specific ``dtype``). """ engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch} + if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch: + engine_args.setdefault("model_class_name", ps.model_arch) if ps.engine_output_type: engine_args["engine_output_type"] = ps.engine_output_type if next_stage_proc: @@ -907,24 +843,31 @@ def merge_pipeline_deploy( deploy = _apply_platform_overrides(deploy) deploy_by_id = {s.stage_id: s for s in deploy.stages} + # async_chunk is irrelevant for single-stage pipelines, so we always disable it if len(pipeline.stages) <= 1: deploy.async_chunk = False - # A pipeline supports async_chunk if any stage has either an explicit - # async-chunk-only processor slot OR a custom next-stage processor (some - # pipelines like qwen3_omni wire async-chunk processing directly through - # ``custom_process_next_stage_input_func``). Only raise when neither is - # present — that's the "user enabled async_chunk but pipeline has no - # inter-stage processing at all" case. - if deploy.async_chunk and not any( - ps.async_chunk_process_next_stage_input_func or ps.custom_process_next_stage_input_func - for ps in pipeline.stages + # async_chunk only applies to multi-stage pipelines: a pipeline with no + # consumer stages (every stage has empty input_sources) has no inter-stage + # edges, so async_chunk is a no-op and we skip the check entirely. + # For pipelines that DO have inter-stage edges, require a dedicated per-step + # async producer (``async_chunk_process_next_stage_input_func``). + # ``custom_process_next_stage_input_func`` is the full-payload / connector-path + # producer and does NOT imply async_chunk support — pipelines like qwen2_5_omni + # and covo_audio have it but removed their consumer-side ``custom_process_input_func`` + # because they don't support async_chunk, so accepting them here would silently + # miswire the consumer stage instead of raising a clear error. + _has_inter_stage_edges = any(ps.input_sources for ps in pipeline.stages) + if ( + deploy.async_chunk + and _has_inter_stage_edges + and not any(ps.async_chunk_process_next_stage_input_func for ps in pipeline.stages) ): raise ValueError( f"Pipeline {pipeline.model_type!r} has async_chunk=True in deploy but no stage " - "declares a next-stage input processor " - "(``async_chunk_process_next_stage_input_func`` or ``custom_process_next_stage_input_func``). " - "Either set async_chunk=False or implement an async-chunk processor on the pipeline." + "declares a dedicated async-chunk next-stage processor " + "(``async_chunk_process_next_stage_input_func``). " + "Either set async_chunk=False or implement an async-chunk producer on the pipeline." ) result: list[StageConfig] = [] @@ -959,7 +902,7 @@ def merge_pipeline_deploy( final_output=ps.final_output, final_output_type=ps.final_output_type, worker_type=worker_type, - scheduler_cls=_scheduler_path(sched_cls), + scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls), hf_config_name=ps.hf_config_name, is_comprehension=ps.owns_tokenizer, yaml_engine_args=engine_args, @@ -1057,376 +1000,3 @@ def to_omegaconf(self) -> Any: config_dict.update(self.yaml_extras) return create_config(config_dict) - - -class StageConfigFactory: - """Factory that loads pipeline YAML and merges CLI overrides. - - Handles both single-stage and multi-stage models. - - Pipelines are declared in ``vllm_omni/config/pipeline_registry.py`` and - loaded lazily via ``_PIPELINE_REGISTRY``; no hardcoded model-type → - directory mapping is maintained here. Models with generic HF - ``model_type`` collisions (e.g. MiMo Audio reports ``qwen2``) should - declare ``hf_architectures=(...)`` on their ``PipelineConfig`` so the - factory can disambiguate via ``hf_config.architectures``. - """ - - @classmethod - def get_pipeline_endpoint_restrictions( - cls, - model: str, - trust_remote_code: bool, - deploy_config_path: str | None, - ) -> tuple[EndpointRestriction, ...]: - """Given a model string, determine the corresponding endpoint restrictions. - - Args: - model: Model name or path. - trust_remote_code: Whether to trust remote code for HF config loading. - deploy_config_path: Optional path to the deploy config for the pipeline. - - Returns: - A tuple of model specific endpoint restrictions. - """ - # Alex - this is gross, it's basically a copy+paste of what we need to resolve the model type. - # This has been refactored a lot in the upstream, so will go away after this release. - model_type, _ = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code) - if model_type == "vla": - from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero - - if _looks_like_dreamzero(model): - model_type = "dreamzero" - - # Deploy config applies over the inferred config - if deploy_config_path is None: - deploy_path = _DEPLOY_DIR / f"{model_type}.yaml" - else: - deploy_path = Path(deploy_config_path) - # Bare-name lookup: if the value isn't a real path, try the - # built-in deploy directory. This lets ``deploy_config="moss_tts"`` - # resolve to ``vllm_omni/deploy/moss_tts.yaml`` when several model - # variants share an HF ``model_type`` and need different YAMLs. - if not deploy_path.exists() and deploy_path.parent == Path("."): - bare_name = deploy_path.name - if not bare_name.endswith(".yaml"): - bare_name = f"{bare_name}.yaml" - candidate = _DEPLOY_DIR / bare_name - if candidate.exists(): - deploy_path = candidate - - if not deploy_path.exists(): - logger.warning( - "Deploy config not found: %s — using pipeline defaults only", - deploy_path, - ) - deploy_cfg = DeployConfig() - else: - deploy_cfg = load_deploy_config(deploy_path) - - pipeline_key = deploy_cfg.pipeline or model_type - if pipeline_key is None or pipeline_key not in _PIPELINE_REGISTRY: - logger.warning("Failed to resolve a registered pipeline for endpoint restrictions") - return () - pipeline_cfg = _PIPELINE_REGISTRY[pipeline_key] - return pipeline_cfg.endpoint_restrictions if pipeline_cfg else () - - @classmethod - def create_from_model( - cls, - model: str, - *, - trust_remote_code: bool = False, - cli_overrides: dict[str, Any] | None = None, - deploy_config_path: str | None = None, - **deprecated_kwargs: Any, - ) -> list[StageConfig] | None: - """Load pipeline + deploy config, merge with CLI overrides. - - Checks _PIPELINE_REGISTRY first (new path), falls back to legacy YAML. - """ - if cli_overrides is None: - cli_overrides = {} - - # --- New path: check pipeline registry by model_type first --- - model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code) - if model_type == "vla": - from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero - - if _looks_like_dreamzero(model): - model_type = "dreamzero" - if model_type and model_type in _PIPELINE_REGISTRY: - return cls._create_from_registry(model_type, cli_overrides, deploy_config_path) - - # --- HF architecture fallback: some models report a generic - # model_type that collides with another model. Match by the - # hf_architectures declared on each registered PipelineConfig. - # When a pipeline declares ``hf_config_predicate`` (e.g. to - # disambiguate generation-revision collisions like MiniCPM-o 2.6 vs - # 4.5 which share ``architectures=["MiniCPMO"]``), the predicate - # must also accept the loaded ``hf_config``. - if hf_config is not None: - hf_archs = set(getattr(hf_config, "architectures", []) or []) - if hf_archs: - for registered in _PIPELINE_REGISTRY.values(): - if not hf_archs.intersection(registered.hf_architectures): - continue - predicate = registered.hf_config_predicate - if predicate is not None: - try: - if not predicate(hf_config): - logger.debug( - "Pipeline %r matched on architectures %s but its " - "hf_config_predicate rejected the loaded config; " - "continuing fallback search.", - registered.model_type, - sorted(hf_archs.intersection(registered.hf_architectures)), - ) - continue - except Exception: - logger.exception( - "Pipeline %r hf_config_predicate raised; skipping.", - registered.model_type, - ) - continue - return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path) - - # An explicit deploy config can select a pipeline when the model's HF - # metadata is too generic for registry auto-detection. - if deploy_config_path is not None: - deploy_path = Path(deploy_config_path) - if deploy_path.exists(): - deploy_config = load_deploy_config(deploy_path) - if deploy_config.pipeline in _PIPELINE_REGISTRY: - return cls._create_from_registry( - deploy_config.pipeline, - cli_overrides, - deploy_config_path, - ) - - # Not in the pipeline registry — let the caller fall back to the - # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path). - return None - - @classmethod - def _create_from_registry( - cls, - model_type: str, - cli_overrides: dict[str, Any], - deploy_config_path: str | None = None, - ) -> list[StageConfig]: - """Create StageConfigs from pipeline registry + deploy YAML. - - Precedence: caller-typed (non-None) value > deploy YAML > - StageDeployConfig dataclass default. - """ - # Resolve deploy config path - if deploy_config_path is None: - deploy_path = _DEPLOY_DIR / f"{model_type}.yaml" - else: - deploy_path = Path(deploy_config_path) - # Bare-name lookup: if the value isn't a real path, try the - # built-in deploy directory. This lets ``deploy_config="moss_tts"`` - # resolve to ``vllm_omni/deploy/moss_tts.yaml`` when several model - # variants share an HF ``model_type`` and need different YAMLs. - if not deploy_path.exists() and deploy_path.parent == Path("."): - bare_name = deploy_path.name - if not bare_name.endswith(".yaml"): - bare_name = f"{bare_name}.yaml" - candidate = _DEPLOY_DIR / bare_name - if candidate.exists(): - deploy_path = candidate - - if not deploy_path.exists(): - logger.warning( - "Deploy config not found: %s — using pipeline defaults only", - deploy_path, - ) - deploy_cfg = DeployConfig() - else: - deploy_cfg = load_deploy_config(deploy_path) - - cli_async_chunk = cli_overrides.get("async_chunk") - if cli_async_chunk is not None: - deploy_cfg.async_chunk = bool(cli_async_chunk) - - pipeline_key = deploy_cfg.pipeline or model_type - if pipeline_key not in _PIPELINE_REGISTRY: - raise KeyError( - f"Pipeline {pipeline_key!r} not in registry " - f"(resolved from {deploy_path.name!r}). Available: " - f"{sorted(_PIPELINE_REGISTRY.keys())}" - ) - pipeline_cfg = _PIPELINE_REGISTRY[pipeline_key] - - stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides) - - explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None} - - for stage in stages: - stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides) - - return stages - - @classmethod - def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]: - """Single-stage diffusion - no YAML needed. - - Creates a default diffusion stage configuration for single-stage - diffusion models. Returns a legacy OmegaConf-compatible dict for - backward compatibility with OmniStage. - - Args: - kwargs: Engine arguments from CLI/API. - - Returns: - List containing a single config dict for the diffusion stage. - """ - # Calculate devices based on parallel config - devices = "0" - if "parallel_config" in kwargs: - num_devices = kwargs["parallel_config"].world_size - for i in range(1, num_devices): - devices += f",{i}" - - engine_args: dict[str, Any] = {} - for key, value in kwargs.items(): - if key in ("parallel_config",): - continue - engine_args[key] = value - - # Serialize parallel_config as dict for OmegaConf. Test helpers - # sometimes pass SimpleNamespace rather than a dataclass instance. - if "parallel_config" in kwargs: - parallel_config = kwargs["parallel_config"] - if dataclasses.is_dataclass(parallel_config) and not isinstance(parallel_config, type): - engine_args["parallel_config"] = asdict(parallel_config) - elif hasattr(parallel_config, "__dict__"): - engine_args["parallel_config"] = dict(vars(parallel_config)) - else: - engine_args["parallel_config"] = parallel_config - - engine_args.setdefault("cache_backend", "none") - engine_args["model_stage"] = "diffusion" - - # Convert dtype to string for OmegaConf - if "dtype" in engine_args: - engine_args["dtype"] = str(engine_args["dtype"]) - - engine_args.setdefault("max_num_seqs", 1) - - config_dict: dict[str, Any] = { - "stage_id": 0, - "stage_type": StageType.DIFFUSION.value, - "runtime": { - "process": True, - "devices": devices, - }, - "engine_args": create_config(engine_args), - "final_output": True, - "final_output_type": "image", - } - - return [config_dict] - - @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. - - Args: - model: Model name or path. - trust_remote_code: Whether to trust remote code for HF config loading. - - Returns: - Tuple of (model_type, hf_config). Both may be None on failure. - """ - try: - from vllm.transformers_utils.config import get_config - - hf_config = get_config(model, trust_remote_code=trust_remote_code) - - # Alex - beeg hack for qwen3 omni single stage models. We do this - # to avoid bringing in this PR for this release: - # https://github.com/vllm-project/vllm-omni/pull/3760 - # - # We should remove this hack after the next rebase - if ( - hf_config.model_type == "qwen3_omni_moe" - and hasattr(hf_config, "enable_audio_output") - and not hf_config.enable_audio_output - ): - logger.info("Remapping qwen3_omni_moe -> qwen3_omni_moe_thinker_only (no audio output)") - hf_config.model_type = "qwen3_omni_moe_thinker_only" - return hf_config.model_type, hf_config - except Exception as e: - logger.debug(f"`get_config` failed for {e}; Falling back to raw config.json path") - - # Fallback: read config.json directly for custom model types that - # are not registered with transformers (e.g. qwen3_tts). - try: - from vllm.transformers_utils.config import get_hf_file_to_dict - - config_dict = get_hf_file_to_dict("config.json", model, revision=None) - if config_dict: - if "model_type" in config_dict: - return config_dict["model_type"], None - # VoxCPM2-style configs use singular ``architecture`` rather - # than HF's standard ``model_type`` / ``architectures``. Accept - # it as a fallback so the pipeline registry can still match. - if "architecture" in config_dict and isinstance(config_dict["architecture"], str): - return config_dict["architecture"], None - except Exception as e: - logger.debug(f"Failed to auto-detect model type for {model}: {e}") - - # Fallback for diffusers-style models: check model_index.json. - # Some models (e.g. GLM-Image) have no root config.json but ship a - # model_index.json with _class_name that maps to a pipeline key via - # PipelineConfig.diffusers_class_name. - try: - from vllm.transformers_utils.config import get_hf_file_to_dict - - model_index = get_hf_file_to_dict("model_index.json", model, revision=None) - if model_index and "_class_name" in model_index: - class_name = model_index["_class_name"] - for pipeline_cfg in _PIPELINE_REGISTRY.values(): - if pipeline_cfg.diffusers_class_name == class_name: - logger.info( - "Detected pipeline %r from model_index.json (_class_name=%r)", - pipeline_cfg.model_type, - class_name, - ) - return pipeline_cfg.model_type, None - except Exception as e: - logger.debug(f"Failed to detect model type for diffusers-style models: {e}") - - # Final fallback: some models (e.g. CosyVoice3) ship an empty - # config.json and rely on naming conventions. Match the model path - # basename against registered pipeline keys — longest match wins - # so "cosyvoice3" (length 10) beats "cosyvoice" (length 9). - model_lower = model.lower().replace("-", "").replace("_", "") - best: str | None = None - best_len = 0 - for registered_key in _PIPELINE_REGISTRY.keys(): - candidate = registered_key.lower().replace("-", "").replace("_", "") - if candidate and candidate in model_lower and len(candidate) > best_len: - best = registered_key - best_len = len(candidate) - if best is not None: - return best, None - - return None, None - - @classmethod - def _merge_cli_overrides( - cls, - stage: StageConfig, - cli_overrides: dict[str, Any], - ) -> dict[str, Any]: - """Merge global and per-stage (``stage_N_*``) CLI overrides. - - Orchestrator-owned keys are filtered by ``build_stage_runtime_overrides`` - using ``OrchestratorArgs`` as the single source of truth; unknown - server/uvicorn keys are dropped downstream by - ``filter_dataclass_kwargs(OmniEngineArgs, ...)``. - """ - return build_stage_runtime_overrides(stage.stage_id, cli_overrides) diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py index 25acfe8edcf..17a9811402f 100644 --- a/vllm_omni/core/prefix_cache.py +++ b/vllm_omni/core/prefix_cache.py @@ -2,6 +2,8 @@ Utilities for Prefix Caching in Omni models. """ +from dataclasses import dataclass, field + import torch from vllm.logger import init_logger from vllm.v1.worker.gpu_input_batch import InputBatch @@ -11,6 +13,23 @@ logger = init_logger(__name__) +@dataclass +class _PendingAsyncWrite: + """One in-flight async write into the prefix cache. + + Holds the CPU tensors that PyTorch implicitly allocated via + ``.to("cpu", non_blocking=True)`` together with a CUDA event recorded + on the copy stream. The consumer calls ``event.synchronize()`` once + (blocking) to guarantee the D2H has completed before reading. + """ + + event: "torch.cuda.Event" + num_tokens: int + slots_cpu: torch.Tensor + hidden_cpu: torch.Tensor | None + mm_cpu: dict[str, torch.Tensor] = field(default_factory=dict) + + class OmniTensorPrefixCache: """Prefix cache for hidden states (model outputs) and model specific multimodal outputs. @@ -49,6 +68,19 @@ def __init__( self._new_req_cache_hit_ids: set[str] = set() self._deferred_mm_outputs: dict[str, dict[str, list[tuple[int, torch.Tensor]]]] = {} + # Async write pipeline state: + # - one dedicated copy stream + # - one in-flight pending write at a time (the producer is exactly + # one forward step ahead of the consumer) + # - schedule_async_write() consumes the *previous* pending write + # (via event.synchronize() + scatter) before kicking off the new + # D2H. By the time we hit synchronize, the previous step's GPU + # work has been done for ~one forward step, so the sync is + # near-zero wall-clock. + self._async_initialized: bool = False + self._async_copy_stream: torch.cuda.Stream | None = None + self._pending_write: _PendingAsyncWrite | None = None + def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int): """Given multimodal outputs from executing the model, dynamically determine which multimodal outputs are tensors depending on sequence @@ -84,14 +116,186 @@ def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: in logger.info("Initializing multimodal output cache of size %s for key: %s", list(new_tensor_shape), key) def _get_cache_tensor(self, dtype: torch.dtype, hidden_size: int | None = None) -> torch.Tensor: - """Allocate a CPU cache tensor for a specific key.""" + """Allocate a CPU cache tensor for a specific key. + + When CUDA is available the tensor is pinned: this is what lets the + async-write pipeline ride the GPU->CPU copy on the dedicated copy + stream as a true async ``cudaMemcpyAsync`` and lets the post-copy CPU + scatter avoid page-resolution stalls. On CPU-only builds pinning is + unsupported (``torch.zeros(pin_memory=True)`` raises), and the async + pipeline is disabled anyway (``schedule_async_write`` early-returns + when ``torch.cuda.is_available()`` is False), so we fall back to a + plain pageable allocation with no behavioral change. + """ actual_hidden_size = hidden_size if hidden_size is not None else self.default_hidden_size return torch.zeros( (self.num_blocks, self.block_size, actual_hidden_size), dtype=dtype, device="cpu", + pin_memory=torch.cuda.is_available(), + ) + + # ===================================================================== + # Async write pipeline: + # schedule_async_write : called per forward step, enqueues a + # non-blocking D2H + records a CUDA event + # drain_ready_async_writes : non-blocking poll, applies any writes + # whose D2H has completed + # ===================================================================== + + def _ensure_async_resources(self) -> None: + """Lazily allocate the dedicated copy stream. + + The per-step CPU staging tensors are allocated by PyTorch on demand + inside ``.to("cpu", non_blocking=True)``, which gives them + implicitly-pinned memory and lets PyTorch's allocator amortize / + reuse them. + """ + if self._async_initialized: + return + if not torch.cuda.is_available(): + return + self._async_copy_stream = torch.cuda.Stream() + self._async_initialized = True + logger.info("Omni prefix cache async-write copy stream initialized.") + + def schedule_async_write( + self, + hidden_states_gpu: torch.Tensor | None, + multimodal_outputs_gpu: dict[str, torch.Tensor] | None, + slot_mapping_gpu: torch.Tensor, + num_tokens_unpadded: int, + num_tokens_padded: int, + skip_mm_cache_keys: set[str] | None = None, + ) -> None: + """Kick off a non-blocking GPU->CPU copy for this step's writes. + + 1. **Consume the previous step's write** (if any): block on its + event and scatter the CPU tensors into the cache. By the time we + reach this point, the previous step's GPU work has been done for + ~one forward step, so the synchronize is near-zero wall-clock. + 2. **Schedule THIS step's D2H** on a dedicated copy stream using + ``.to("cpu", non_blocking=True)``. PyTorch allocates the pinned + destination implicitly. + 3. **Record an event** and stash everything for the next step's + consume. + + Holding exactly one in-flight write at a time keeps the design + simple: no ring, no event-polling, no driver-mutex pressure from + repeated event.query() calls. + """ + if not torch.cuda.is_available(): + return + + self._ensure_async_resources() + if not self._async_initialized: + return + + # Decide which mm keys we'll cache. Init the dest cache tensors + # before scheduling the copies so the apply step has somewhere + # to scatter into. + skip = skip_mm_cache_keys or set() + cacheable_mm_keys: list[str] = [] + if multimodal_outputs_gpu: + self.maybe_init_missing_mm_cache_keys(multimodal_outputs_gpu, seq_len=num_tokens_padded) + for k in self.mm_cache_keys: + if k in skip: + continue + v = multimodal_outputs_gpu.get(k) + if not isinstance(v, torch.Tensor) or v.ndim < 2: + continue + if v.shape[0] < num_tokens_unpadded: + continue + cacheable_mm_keys.append(k) + + # Step 1: consume the previous step's write FIRST. The previous + # step's event was recorded ~one forward step ago, so + # synchronize() is effectively a no-op cost-wise; the scatter + # into pinned CPU caches is also fast. + self._consume_pending_write() + + if hidden_states_gpu is None and not cacheable_mm_keys: + return # nothing to schedule this step + + # Step 2: schedule the D2H on the copy stream. PyTorch implicitly + # allocates pinned host destinations for the ``.to("cpu", + # non_blocking=True)`` results. + copy_stream = self._async_copy_stream + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(copy_stream): + copy_stream.wait_stream(default_stream) + + # slot_mapping: copy as-is; the int64 cast required by + # ``index_copy_`` is deferred to ``_consume_pending_write`` + # so we avoid launching an extra GPU kernel here that would + # contend the driver mutex with the compute stream. + slots_cpu = slot_mapping_gpu[:num_tokens_unpadded].to("cpu", non_blocking=True) + hidden_cpu: torch.Tensor | None = None + if hidden_states_gpu is not None: + hidden_cpu = hidden_states_gpu[:num_tokens_unpadded].to("cpu", non_blocking=True) + mm_cpu: dict[str, torch.Tensor] = {} + for k in cacheable_mm_keys: + mm_cpu[k] = multimodal_outputs_gpu[k][:num_tokens_unpadded].to("cpu", non_blocking=True) + + event = torch.cuda.Event() + event.record() + + # Step 3: stash for next step's consume. + self._pending_write = _PendingAsyncWrite( + event=event, + num_tokens=num_tokens_unpadded, + slots_cpu=slots_cpu, + hidden_cpu=hidden_cpu, + mm_cpu=mm_cpu, ) + def drain_ready_async_writes(self) -> int: + """Consume one in-flight write if it has finished (non-blocking). + + Exposed so the model runner can drain at the top of + ``execute_model`` regardless of whether a new write will be + scheduled this step. + """ + if self._pending_write is None: + return 0 + if not self._pending_write.event.query(): + return 0 + self._consume_pending_write() + return 1 + + def _consume_pending_write(self) -> None: + """Synchronize on the pending event and scatter the CPU tensors + into the destination caches. Drops state on completion. + """ + entry = self._pending_write + if entry is None: + return + self._pending_write = None + entry.event.synchronize() + n = entry.num_tokens + if n <= 0: + return + + slots = entry.slots_cpu + if slots.dtype != torch.int64: + slots = slots.to(torch.int64) + + # NOTE: ``index_copy_(0, idx, src)`` is the row-scatter primitive in + # PyTorch. It's semantically identical to ``flat[idx] = src`` but + # dispatches to a much simpler CPU implementation (single dim, + # no broadcasting, no accumulate), which on pinned memory at large + # row sizes is observed to be much faster than the generic + # ``aten::index_put_`` path used by advanced-indexing assignment. + if entry.hidden_cpu is not None: + flat = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) + flat.index_copy_(0, slots, entry.hidden_cpu) + for k, src_cpu in entry.mm_cpu.items(): + mm_cache = self.mm_outputs_cache.get(k) + if mm_cache is None: + continue + flat = mm_cache.view(-1, mm_cache.shape[-1]) + flat.index_copy_(0, slots, src_cpu) + def add_prefix_cached_new_req_id(self, req_id: str): """Adds a new request ID to the set of prefix cache hits on the batch.""" self._new_req_cache_hit_ids.add(req_id) @@ -173,9 +377,17 @@ def update_omni_tensor_prefix_cache( num_tokens_unpadded, hidden_states_cpu, ) - # View the cache as 2D so that we can treat our slots as row indices + # View the cache as 2D so that we can treat our slots as row indices. + # ``index_copy_`` is the row-scatter primitive — it's semantically + # ``cache[idx] = src`` but dispatches to a faster CPU path than the + # generic advanced-indexing assignment (which goes through + # ``aten::index_put_`` and is dominated by per-row dispatch + # overhead on large feature widths). flat_cache = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) - flat_cache[unpadded_slot_mapping] = hidden_states + slot_idx = unpadded_slot_mapping + if slot_idx.dtype != torch.int64: + slot_idx = slot_idx.to(torch.int64) + flat_cache.index_copy_(0, slot_idx, hidden_states) logger.debug("Writing to hidden states for %s tokens", num_tokens_unpadded) # Do the same for the stage's cached multimodal outputs @@ -195,7 +407,10 @@ def update_omni_tensor_prefix_cache( mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded] mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state) flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) - flat_cache[unpadded_slot_mapping] = mm_state + slot_idx = unpadded_slot_mapping + if slot_idx.dtype != torch.int64: + slot_idx = slot_idx.to(torch.int64) + flat_cache.index_copy_(0, slot_idx, mm_state) logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded) def stage_deferred_mm_outputs( @@ -331,7 +546,9 @@ def _write_deferred_mm_tensor( mm_state_cpu = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state[: slots.numel()]) mm_cache = self.mm_outputs_cache[mm_key] flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) - flat_cache[slots] = mm_state_cpu + if slots.dtype != torch.int64: + slots = slots.to(torch.int64) + flat_cache.index_copy_(0, slots, mm_state_cpu) def _get_slot_ids_for_token_range( self, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 984844bdbae..74d09ab6007 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -11,6 +11,7 @@ from vllm.logger import init_logger from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMScheduler from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.request_queue import create_request_queue 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, FinishReason @@ -133,6 +134,9 @@ def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool: self._omits_kv_transfer_cache[rid] = result return result + def _should_defer_waiting_admission(self) -> bool: + return False + def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool: """ Check triggers and process side effects (marking transfer). @@ -205,7 +209,7 @@ def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int return False - def schedule(self) -> SchedulerOutput: # type: ignore[override] + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # Remove FINISHED_ABORTED requests before the upstream scheduler sees # them. Upstream vllm raises RuntimeError on this status; omni allows # async abort (e.g. client disconnect during TTS streaming) to leave @@ -222,9 +226,19 @@ def schedule(self) -> SchedulerOutput: # type: ignore[override] self.waiting, self.running, scheduler_requests=self.requests ) + original_waiting = None + if self._should_defer_waiting_admission(): + original_waiting = self.waiting + self.waiting = create_request_queue(self.policy) + try: - scheduler_output = super().schedule() + scheduler_output = super().schedule(throttle_prefills) finally: + if original_waiting is not None: + deferred_waiting = list(self.waiting) + if deferred_waiting: + original_waiting.prepend_requests(deferred_waiting) + self.waiting = original_waiting if self.chunk_transfer_adapter: # Add request waiting for chunk to the waiting and running queue self.chunk_transfer_adapter.restore_queues( @@ -289,6 +303,7 @@ def update_from_output( num_scheduled_tokens = scheduler_output.num_scheduled_tokens pooler_outputs = model_runner_output.pooler_output mm_outputs = getattr(model_runner_output, "multimodal_outputs", None) + inter_stage_outputs = getattr(model_runner_output, "inter_stage_outputs", None) num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats @@ -299,13 +314,6 @@ def update_from_output( outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats: KVConnectorStats | None = ( - kv_connector_output.kv_connector_stats if kv_connector_output else None - ) - if kv_connector_stats and self.connector: - kv_stats = self.connector.get_kv_connector_stats() - if kv_stats: - kv_connector_stats = kv_connector_stats.aggregate(kv_stats) failed_kv_load_req_ids = None if kv_connector_output and kv_connector_output.invalid_block_ids: @@ -391,6 +399,7 @@ def update_from_output( new_token_ids = generated_token_ids pooler_output = pooler_outputs[req_index] if pooler_outputs else None mm_output = mm_outputs[req_index] if mm_outputs else None + inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None kv_transfer_params = None status_before_stop = request.status finish_reason = None @@ -487,12 +496,17 @@ def update_from_output( new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id, None), ) ) - if self.chunk_transfer_adapter is not None: - self.chunk_transfer_adapter.save_async(mm_output, request, is_segment_finished) else: # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors + if self.chunk_transfer_adapter is not None and (inter_stage_output is not None or is_segment_finished): + self.chunk_transfer_adapter.save_async( + inter_stage_output, + request, + is_segment_finished, + ) + # Remove the stopped requests from the running and waiting queues. if stopped_running_reqs: self.running = remove_all(self.running, stopped_running_reqs) @@ -542,6 +556,20 @@ def update_from_output( if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty(): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + # collect KV cache events from KV cache manager events = self.kv_cache_manager.take_events() diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 2e99c8e1a45..973b8c55a04 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -56,7 +56,7 @@ def __init__(self, *args, **kwargs): ) self._latest_omni_connector_output: OmniConnectorOutput | None = None - def schedule(self) -> SchedulerOutput: + def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: """Diffusion fast path: - Feed all input tokens of the request at once (if 0, allocate 1 placeholder token). @@ -217,7 +217,7 @@ def schedule(self) -> SchedulerOutput: scheduler_requests=self.requests, ) else: - res = super().schedule() + res = super().schedule(throttle_prefills) if self.input_coordinator: self.input_coordinator.restore_queues(self.waiting, self.running) return self._wrap_omni_scheduler_output(res) @@ -431,14 +431,6 @@ def update_from_output( outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats: KVConnectorStats | None = ( - kv_connector_output.kv_connector_stats if kv_connector_output else None - ) - # Merge connector-side stats (align with v0.14.0) - if kv_connector_stats and self.connector: - kv_stats = self.connector.get_kv_connector_stats() - if kv_stats: - kv_connector_stats = kv_connector_stats.aggregate(kv_stats) failed_kv_load_req_ids = None if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None): @@ -646,6 +638,20 @@ def update_from_output( if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty(): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + # Collect and publish KV cache events (align with v0.14.0) events = self.kv_cache_manager.take_events() if self.connector is not None: diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 165c04e964a..be73db9c3e8 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -26,15 +26,15 @@ # Scope: this constant only covers the full-payload coordinator path # (``input_coordinator``). The async-chunk path uses # ``chunk_transfer_adapter`` and is not affected by this constant. -_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "300") +_INPUT_WAIT_TIMEOUT_RAW = os.environ.get("VLLM_OMNI_INPUT_WAIT_TIMEOUT_S", "600") try: DEFAULT_INPUT_WAIT_TIMEOUT_S: float = float(_INPUT_WAIT_TIMEOUT_RAW) except ValueError: logger.warning( - "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 300 seconds.", + "Invalid VLLM_OMNI_INPUT_WAIT_TIMEOUT_S=%r; falling back to 600 seconds.", _INPUT_WAIT_TIMEOUT_RAW, ) - DEFAULT_INPUT_WAIT_TIMEOUT_S = 300.0 + DEFAULT_INPUT_WAIT_TIMEOUT_S = 600.0 class OmniSchedulerMixin: diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 4056fd93861..af6c73d6dc0 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -51,6 +51,9 @@ ("Qwen3TTSCode2Wav", "code2wav"), # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1). ("CosyVoice3Model", "cosyvoice3_code2wav"), + # indextts2: indextts2_talker (Stage 0) -> indextts2_s2mel_decoder + # (Stage 1). Stage 1 consumes the complete mel/latent payload. + ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"), # dynin: token2text (Stage 0) -> token2image (Stage 1) -> # token2audio (Stage 2). Producer wires via # custom_process_next_stage_input_func: *_full_payload in deploy yaml. diff --git a/vllm_omni/deploy/Gr00tN1d7.yaml b/vllm_omni/deploy/Gr00tN1d7.yaml new file mode 100644 index 00000000000..0e90d8536b7 --- /dev/null +++ b/vllm_omni/deploy/Gr00tN1d7.yaml @@ -0,0 +1,41 @@ +# GR00T N1.7 deploy: single diffusion stage. +# +# Topology is declared in vllm_omni/model_executor/models/gr00t/pipeline.py. +# This default uses one GPU with TP=1 and CFG parallel disabled. + +pipeline: Gr00tN1d7 +# Required: async_chunk defaults to True, which the deploy validator rejects for a +# single-stage pipeline (no next-stage processor). Must be false here. +async_chunk: false + +stages: + - stage_id: 0 + max_num_seqs: 1 + enforce_eager: true + model_class_name: Gr00tN1d7Pipeline + model_config: + embodiment_tag: OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT + # OpenPI client handshake — model/embodiment-specific values for GR00T-N1.7. + # Sent verbatim to the client (the pipeline never reads these); they describe the + # action space / inputs the client must produce for this checkpoint + embodiment. + policy_server_config: + image_resolution: [180, 320] + needs_session_id: true + embodiment_tag: OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT + action_horizon: 40 + action_keys: + - eef_9d + - gripper_position + - joint_position + supported_embodiments: + - OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT + - XDOF + - XDOF_SUBTASK + - REAL_G1 + - REAL_R1_PRO_SHARPA + - REAL_R1_PRO_SHARPA_HUMAN + - REAL_R1_PRO_SHARPA_MAXINSIGHTS + - REAL_R1_PRO_SHARPA_MECKA + - LIBERO_PANDA + - SIMPLER_ENV_GOOGLE + - SIMPLER_ENV_WIDOWX diff --git a/vllm_omni/deploy/aura_omni.yaml b/vllm_omni/deploy/aura_omni.yaml new file mode 100644 index 00000000000..486e44b57e6 --- /dev/null +++ b/vllm_omni/deploy/aura_omni.yaml @@ -0,0 +1,110 @@ +# AURA Omni deploy profile: Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav. +# +# Pass `--deploy-config vllm_omni/deploy/aura_omni.yaml` and point the parent +# `--model` at any local placeholder repo if needed; each stage overrides its +# own model below. +pipeline: aura_omni +# AURA has a semantic ASR -> AURA -> TTS dependency. Do not prewarm downstream +# stages before AURA text is converted into Qwen3-TTS metadata by aura2tts. +async_chunk: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + codec_chunk_frames: 25 + codec_left_context_frames: 72 + initial_codec_chunk_frames: 1 + decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325] + +stages: + - stage_id: 0 + max_num_seqs: 8 + gpu_memory_utilization: 0.1 + trust_remote_code: true + enable_prefix_caching: true + async_scheduling: true + max_model_len: 8192 + max_num_batched_tokens: 4096 + devices: "0" + model: Qwen/Qwen3-ASR-1.7B + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 256 + + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.4 + trust_remote_code: true + enable_prefix_caching: true + async_scheduling: true + max_model_len: 262144 + max_num_batched_tokens: 15360 + devices: "0" + model: aurateam/AURA + default_sampling_params: + temperature: 0.5 + top_p: 1.0 + top_k: -1 + max_tokens: 256 + repetition_penalty: 1.0 + + - stage_id: 2 + max_num_seqs: 8 + gpu_memory_utilization: 0.2 + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 512 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_3: connector_of_shared_memory + model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + # model: Qwen/Qwen3-TTS-12Hz-1.7B-Base + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.05 + subtalker_sampling_params: + do_sample: true + temperature: 0.9 + top_k: 50 + top_p: 1.0 + + - stage_id: 3 + max_num_seqs: 8 + gpu_memory_utilization: 0.2 + enforce_eager: false + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_2: connector_of_shared_memory + model: Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice + # model: Qwen/Qwen3-TTS-12Hz-1.7B-Base + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 + +platforms: + npu: + stages: + - stage_id: 2 + enforce_eager: true diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml index e77afad8409..62b3a5eed9b 100644 --- a/vllm_omni/deploy/dreamzero.yaml +++ b/vllm_omni/deploy/dreamzero.yaml @@ -12,8 +12,10 @@ stages: - stage_id: 0 devices: "0" max_num_seqs: 1 - enforce_eager: true + enforce_eager: false model_class_name: DreamZeroPipeline + # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache). + engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine model_config: default_robot_embodiment: roboarena policy_server_config: @@ -23,3 +25,8 @@ stages: needs_stereo_camera: false needs_session_id: true action_space: joint_position + cache_backend: step_cache + cache_config: + step_cache_dit_enabled: true + velocity_sim_thresholds: [0.95, 0.93] + velocity_skip_countdowns: [4, 2] diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml index ed9988c92a3..a453d43f3bf 100644 --- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml +++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml @@ -9,7 +9,7 @@ stages: - stage_id: 0 devices: "0,1" max_num_seqs: 1 - enforce_eager: true + enforce_eager: false model_class_name: DreamZeroPipeline parallel_config: tensor_parallel_size: 1 @@ -23,3 +23,8 @@ stages: needs_stereo_camera: false needs_session_id: true action_space: joint_position + cache_backend: step_cache + cache_config: + step_cache_dit_enabled: true + velocity_sim_thresholds: [0.95, 0.93] + velocity_skip_countdowns: [4, 2] diff --git a/vllm_omni/deploy/fish_qwen3_omni.yaml b/vllm_omni/deploy/fish_qwen3_omni.yaml index 5b0c44988a0..539b3777507 100644 --- a/vllm_omni/deploy/fish_qwen3_omni.yaml +++ b/vllm_omni/deploy/fish_qwen3_omni.yaml @@ -63,3 +63,15 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 + +platforms: + xpu: + stages: + # Verified on Intel Arc B60. + - stage_id: 0 + gpu_memory_utilization: 0.8 + max_model_len: 2048 + max_num_batched_tokens: 2048 + - stage_id: 1 + max_model_len: 2048 + max_num_batched_tokens: 2048 diff --git a/vllm_omni/deploy/hunyuan_image_3_moe.yaml b/vllm_omni/deploy/hunyuan_image_3_moe.yaml index 67ce8b66aab..b661674c77a 100644 --- a/vllm_omni/deploy/hunyuan_image_3_moe.yaml +++ b/vllm_omni/deploy/hunyuan_image_3_moe.yaml @@ -62,7 +62,7 @@ stages: max_tokens: 8192 detokenize: true skip_special_tokens: false - include_stop_str_in_output: true, + include_stop_str_in_output: true - stage_id: 1 max_num_seqs: 1 @@ -71,6 +71,8 @@ stages: distributed_executor_backend: "mp" omni_kv_config: need_recv_cache: true + enable_kv_async_prefetch: true + kv_prefetch_min_free_mem_ratio: 0.2 parallel_config: tensor_parallel_size: 2 enable_expert_parallel: true diff --git a/vllm_omni/deploy/hunyuan_video_15.yaml b/vllm_omni/deploy/hunyuan_video_15.yaml new file mode 100644 index 00000000000..b1f4f73517e --- /dev/null +++ b/vllm_omni/deploy/hunyuan_video_15.yaml @@ -0,0 +1,18 @@ +# HunyuanVideo-1.5 deploy. +# +# Single GPU by default. Increase tensor_parallel_size and devices for multi-GPU TP. +async_chunk: false +trust_remote_code: true +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + tensor_parallel_size: 1 + enforce_eager: true + # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of + # the full video, so tiling the decode is the main memory lever. + vae_use_tiling: true + default_sampling_params: + seed: 42 diff --git a/vllm_omni/deploy/indextts2.yaml b/vllm_omni/deploy/indextts2.yaml new file mode 100644 index 00000000000..75eb60188f5 --- /dev/null +++ b/vllm_omni/deploy/indextts2.yaml @@ -0,0 +1,57 @@ +async_chunk: false +trust_remote_code: true +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: false + 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: 4 + gpu_memory_utilization: 0.4 + enable_prefix_caching: false + skip_tokenizer_init: true + tokenizer: gpt2 + async_scheduling: true + max_num_batched_tokens: 512 + max_model_len: 2560 + devices: "0" + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.8 + top_p: 0.8 + top_k: 30 + max_tokens: 1500 + repetition_penalty: 10.0 + stop_token_ids: [8193] + detokenize: false + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.4 + enable_prefix_caching: false + skip_tokenizer_init: true + tokenizer: gpt2 + enforce_eager: true + async_scheduling: true + max_num_batched_tokens: 8192 + max_model_len: 32768 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + hf_overrides: + diffusion_steps: 12 + s2mel_dit_bf16: true + s2mel_vocoder_bf16: true + s2mel_vocoder_cuda_graph: true + s2mel_dit_cuda_graph: true + s2mel_dit_cuda_graph_max_graphs: 4 + s2mel_vocoder_capture_sizes: [128, 256, 384, 512] + s2mel_vocoder_compile_shapes: [256, 384, 512] + default_sampling_params: + temperature: 0.0 + max_tokens: 32767 diff --git a/vllm_omni/deploy/indextts2_low_latency.yaml b/vllm_omni/deploy/indextts2_low_latency.yaml new file mode 100644 index 00000000000..621116ab2cf --- /dev/null +++ b/vllm_omni/deploy/indextts2_low_latency.yaml @@ -0,0 +1,76 @@ +# IndexTTS2 low-latency deploy profile. +# +# Stage 0 (GPT2 AR talker) uses vLLM FULL_DECODE_ONLY CUDA graph + FlashInfer +# attention backend + async_scheduling=false, mirroring the Higgs low-latency +# profile. This reduces decode launch overhead and per-step scheduling +# fragmentation for c1-c4 workloads. +# +# Usage: +# vllm serve IndexTeam/IndexTTS-2 --omni --trust-remote-code \ +# --deploy-config vllm_omni/deploy/indextts2_low_latency.yaml +# +# For 2-GPU testing, copy this file and set stage 1 devices: "1". + +async_chunk: false +trust_remote_code: true +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: false + 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 + enable_prefix_caching: false + skip_tokenizer_init: true + tokenizer: gpt2 + enforce_eager: false + attention_backend: FLASHINFER + async_scheduling: false + max_num_batched_tokens: 4096 + max_model_len: 2560 + devices: "0" + compilation_config: + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_num_of_warmups: 1 + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.8 + top_p: 0.8 + top_k: 30 + max_tokens: 1500 + repetition_penalty: 10.0 + stop_token_ids: [8193] + detokenize: false + - stage_id: 1 + max_num_seqs: 4 + gpu_memory_utilization: 0.4 + enable_prefix_caching: false + skip_tokenizer_init: true + tokenizer: gpt2 + enforce_eager: true + async_scheduling: true + max_num_batched_tokens: 8192 + max_model_len: 32768 + devices: "0" + input_connectors: + from_stage_0: connector_of_shared_memory + hf_overrides: + diffusion_steps: 12 + s2mel_dit_bf16: true + s2mel_vocoder_bf16: true + s2mel_vocoder_cuda_graph: true + s2mel_dit_cuda_graph: true + s2mel_dit_cuda_graph_max_graphs: 4 + s2mel_vocoder_capture_sizes: [128, 256, 384, 512] + s2mel_vocoder_compile_shapes: [256, 384, 512] + default_sampling_params: + temperature: 0.0 + max_tokens: 32767 diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml index d1a49d662c7..afae38ee0e1 100644 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ b/vllm_omni/deploy/mammoth_moda2.yaml @@ -1,42 +1,27 @@ -# MammothModa2 AR→DiT (text-to-image). Canonical copy lives at -# ``vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml``; keep in sync. -# -# Must remain legacy ``stage_args`` layout (no top-level ``stages``): Omni resolves -# explicit ``--stage-configs-path`` via ``load_stage_configs_from_yaml``. +# MammothModa2 AR→DiT (text-to-image). -stage_args: +async_chunk: false +pipeline: mammoth_moda2 + +stages: - stage_id: 0 - runtime: - devices: "0" - engine_args: - model_stage: ar - max_num_seqs: 100 - model_arch: MammothModa2ForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 8192 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - final_output: false + devices: "0" + max_num_seqs: 100 + max_model_len: 8192 + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false - stage_id: 1 - runtime: - devices: "0" - engine_args: - model_stage: dit - max_num_seqs: 1 - model_arch: MammothModa2ForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - gpu_memory_utilization: 0.3 - enforce_eager: true - trust_remote_code: true - engine_output_type: image - enable_prefix_caching: false - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit - final_output: true - final_output_type: image + devices: "0" + max_num_seqs: 1 + gpu_memory_utilization: 0.3 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + default_sampling_params: + extra_args: + text_guidance_scale: 9.0 + cfg_range: [0.0, 1.0] + num_inference_steps: 50 diff --git a/vllm_omni/deploy/mammoth_moda2_ar.yaml b/vllm_omni/deploy/mammoth_moda2_ar.yaml new file mode 100644 index 00000000000..45e2580d959 --- /dev/null +++ b/vllm_omni/deploy/mammoth_moda2_ar.yaml @@ -0,0 +1,14 @@ +# MammothModa2 AR-only (image/video understanding and summarization). + +async_chunk: false +pipeline: mammoth_moda2_ar + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 16 + max_model_len: 8192 + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false diff --git a/vllm_omni/deploy/mimo_audio.yaml b/vllm_omni/deploy/mimo_audio.yaml index 6ebe5920957..a5063296474 100644 --- a/vllm_omni/deploy/mimo_audio.yaml +++ b/vllm_omni/deploy/mimo_audio.yaml @@ -4,8 +4,8 @@ # Default mode is async-chunk streaming on a single GPU (both stages on # device 0) through SharedMemoryConnector. For the legacy 2-GPU sync # pipeline (stage 1 on a second card), pass ``--no-async-chunk -# --stage-1-devices 1 --stage-1-max-model-len 18192 -# --stage-1-max-num-batched-tokens 18192`` to the serve CLI. +# --stage-overrides '{"1": {"devices": "1", "max_model_len": 18192, "max_num_batched_tokens": 18192}}'`` +# to the serve CLI. async_chunk: true dtype: bfloat16 diff --git a/vllm_omni/deploy/ming_tts.yaml b/vllm_omni/deploy/ming_tts.yaml index 1a2dc99eca1..684a47f989a 100644 --- a/vllm_omni/deploy/ming_tts.yaml +++ b/vllm_omni/deploy/ming_tts.yaml @@ -15,6 +15,7 @@ connectors: name: SharedMemoryConnector extra: latent_chunk_size: 25 + initial_latent_chunk_size: 4 latent_left_context: 0 stages: diff --git a/vllm_omni/deploy/ming_tts_moe.yaml b/vllm_omni/deploy/ming_tts_moe.yaml new file mode 100644 index 00000000000..76425712a85 --- /dev/null +++ b/vllm_omni/deploy/ming_tts_moe.yaml @@ -0,0 +1,69 @@ +# Ming-omni-tts-16.8B-A3B (MoE) deploy: stage-0 bailing_moe LLM+flow -> stage-1 audio VAE. +# Same two-stage pipeline as the dense ming_tts.yaml; the Stage-0 backbone is the +# 16.8B-A3B bailing_moe MoE (~31GB bf16). +# +# pipeline: HF config reports model_type="bailingmm", which is not a registry key, +# already declared by MING_TTS_MOE_PIPELINE for auto-detection. +# +# NOTE: ~31GB bf16 + the tiny audio VAE fit one 80GB+ GPU, so both stages colocate +# on GPU0 (Stage-0 0.60 util ≈ 58GB + Stage-1 0.25 util ≈ 24GB ≈ 0.85, fits 97GB +# H20). For smaller cards raise Stage-0 tensor_parallel_size or split the stages +# onto separate devices. Multi-request batching is not yet validated (max_num_seqs: 1). +pipeline: ming_tts_moe +async_chunk: true +# The MoE ships a custom tokenizer (tokenization_bailing.BailingTokenizer) that +# must be executed to load. The model config still resolves to the registered +# MingMoeConfig (vLLM _CONFIG_REGISTRY precedence over the repo auto_map). +trust_remote_code: true +dtype: bfloat16 +enable_prefix_caching: false + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + latent_chunk_size: 25 + initial_latent_chunk_size: 4 + latent_left_context: 0 + +stages: + - stage_id: 0 + max_num_seqs: 1 + tensor_parallel_size: 1 + gpu_memory_utilization: 0.60 + enforce_eager: false + compilation_config: + cudagraph_mode: PIECEWISE + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 8192 + max_model_len: 8192 + # Pin the unquantized MoE kernel (wer diff between vllm 0.22 vs 0.23) + moe_backend: triton + 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/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 9806154793b..d191dddccfa 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -1,9 +1,9 @@ -# MiniCPM-o 4.5 default deploy: 2-GPU layout. +# MiniCPM-o 4.5 default deploy: single-GPU layout. # -# Stage 0 (thinker): GPU 0 — multimodal understanding, ~70% mem. -# Stage 1 (talker) : GPU 1 — MiniCPMTTS + Token2Wav, ~75% mem -# (covers both the talker LM and the in-process -# Token2Wav vocoder that emits the waveform). +# Both stages are co-located on GPU 0: +# Stage 0 (thinker): multimodal understanding, ~80% mem. +# Stage 1 (talker) : MiniCPMTTS + in-process Token2Wav vocoder, ~10% mem +# (talker LM plus the vocoder that emits the waveform). # # Other GPU layouts: minicpmo_4_5_3gpu.yaml, minicpmo_4_5_8x4090.yaml. async_chunk: false @@ -13,9 +13,9 @@ enable_prefix_caching: false stages: - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.7 - enforce_eager: true + max_num_seqs: 16 + gpu_memory_utilization: 0.8 + enforce_eager: false tensor_parallel_size: 1 max_num_batched_tokens: 16384 devices: "0" @@ -34,10 +34,10 @@ stages: # Talker only consumes runtime_additional_information[0]; any value > 1 # makes concurrent requests share request-0's audio (mm_outputs fallback). max_num_seqs: 1 - gpu_memory_utilization: 0.75 - enforce_eager: true + gpu_memory_utilization: 0.1 + enforce_eager: false max_num_batched_tokens: 8192 - devices: "1" + devices: "0" default_sampling_params: temperature: 0.0 top_p: 1.0 @@ -46,6 +46,14 @@ stages: seed: 42 detokenize: false -# MiniCPM-o 4.5 has only been verified on CUDA devices. Add per-platform -# overrides here when other backends (ROCm/XPU/NPU) are verified end-to-end. -platforms: {} +platforms: + npu: + stages: + - stage_id: 0 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 1 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml new file mode 100644 index 00000000000..7a9c400cca6 --- /dev/null +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -0,0 +1,64 @@ +# MiniCPM-o 4.5 deploy: 2-GPU layout (thinker on GPU 0, talker on GPU 1). +# +# Stage 0 (thinker): GPU 0 — multimodal understanding, ~90% mem (thinker +# gets a dedicated card, so it can take more of +# the GPU than the single-GPU co-located layout). +# Stage 1 (talker) : GPU 1 — MiniCPMTTS + Token2Wav, ~75% mem +# (covers both the talker LM and the in-process +# Token2Wav vocoder that emits the waveform). +# +# Use this when you have two GPUs and want the talker + Token2Wav vocoder +# on its own card instead of sharing GPU 0 with the thinker (the default +# single-GPU layout co-locates both and can OOM under load). +pipeline: minicpmo_4_5 +async_chunk: false + +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + max_num_seqs: 16 + gpu_memory_utilization: 0.9 + enforce_eager: false + tensor_parallel_size: 1 + max_num_batched_tokens: 16384 + devices: "0" + limit_mm_per_prompt: + video: 1 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: true + repetition_penalty: 1.1 + + - stage_id: 1 + # Talker only consumes runtime_additional_information[0]; any value > 1 + # makes concurrent requests share request-0's audio (mm_outputs fallback). + max_num_seqs: 1 + gpu_memory_utilization: 0.75 + enforce_eager: false + max_num_batched_tokens: 8192 + devices: "1" + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 + detokenize: false + +platforms: + npu: + stages: + - stage_id: 0 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 1 + max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index f5c0b03a63c..69f77eb6d38 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -12,9 +12,9 @@ enable_prefix_caching: false stages: - stage_id: 0 - max_num_seqs: 1 + max_num_seqs: 16 gpu_memory_utilization: 0.7 - enforce_eager: true + enforce_eager: false tensor_parallel_size: 2 max_num_batched_tokens: 16384 devices: "0,1" @@ -34,7 +34,7 @@ stages: # makes concurrent requests share request-0's audio (mm_outputs fallback). max_num_seqs: 1 gpu_memory_utilization: 0.75 - enforce_eager: true + enforce_eager: false max_num_batched_tokens: 8192 devices: "2" default_sampling_params: diff --git a/vllm_omni/deploy/moss_sound_effect.yaml b/vllm_omni/deploy/moss_sound_effect.yaml index 4ef5fa7afc6..0fe616290cf 100644 --- a/vllm_omni/deploy/moss_sound_effect.yaml +++ b/vllm_omni/deploy/moss_sound_effect.yaml @@ -4,6 +4,7 @@ # no reference audio required. # Rate: ~12.5 tokens/second; pass `tokens` to control duration. async_chunk: true +trust_remote_code: true # OpenMOSS-Team/MOSS-SoundEffect ships custom modeling code connectors: shm: diff --git a/vllm_omni/deploy/moss_tts.yaml b/vllm_omni/deploy/moss_tts.yaml index dbab816eb6d..b1549d0d036 100644 --- a/vllm_omni/deploy/moss_tts.yaml +++ b/vllm_omni/deploy/moss_tts.yaml @@ -2,6 +2,7 @@ # HF: OpenMOSS-Team/MOSS-TTS # General-purpose TTS with voice cloning. Production-recommended variant. async_chunk: true +trust_remote_code: true # OpenMOSS-Team/MOSS-TTS ships custom modeling code connectors: shm: diff --git a/vllm_omni/deploy/moss_tts_80gb.yaml b/vllm_omni/deploy/moss_tts_80gb.yaml new file mode 100644 index 00000000000..e61d4cf5da3 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_80gb.yaml @@ -0,0 +1,64 @@ +# Single-80GB variant of moss_tts.yaml (see #4643): lowers stage-0 +# gpu_memory_utilization and disables codec FlashInfer autotune so both +# stages initialize on one H100/H800 80GB at full max_model_len. +# MOSS-TTS (MossTTSDelayModel, 8B, n_vq=32, 24 kHz) +# HF: OpenMOSS-Team/MOSS-TTS +# General-purpose TTS with voice cloning. Production-recommended variant. +async_chunk: true +trust_remote_code: true # OpenMOSS-Team/MOSS-TTS ships custom modeling code + +connectors: + shm: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + 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 + max_num_seqs: 4 + gpu_memory_utilization: 0.60 + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 1.7 + top_p: 0.8 + top_k: 25 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.0 + stop_token_ids: [151645] # <|im_end|> — terminates the talker AR loop + + - stage_id: 1 + max_num_seqs: 1 + gpu_memory_utilization: 0.12 + enforce_eager: true + enable_flashinfer_autotune: false + enable_prefix_caching: false + async_scheduling: true + # Talker can emit up to max_tokens=4096 frames × n_vq=32 = 131072 codec + # tokens per request; cap codec stage above that so a single long output + # doesn't trigger "could not broadcast (>65536,) into (65536,)". + max_num_batched_tokens: 196608 + max_model_len: 196608 + devices: "0" + input_connectors: + from_stage_0: shm + 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/moss_tts_local.yaml b/vllm_omni/deploy/moss_tts_local.yaml new file mode 100644 index 00000000000..1455ccc0942 --- /dev/null +++ b/vllm_omni/deploy/moss_tts_local.yaml @@ -0,0 +1,58 @@ +# MOSS-TTS-Local-Transformer-v1.5 (MossTTSLocalModel, n_vq=12, 48 kHz stereo) +# HF: OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 +# Codec: OpenMOSS-Team/MOSS-Audio-Tokenizer-v2 (48 kHz, stereo, channel-interleaved RVQ) +# Uses a 1-layer GPT2-style local depth transformer (no delay warm-up), reset every +# frame — same per-frame timing (80 ms/frame, 12.5 fps) as MOSS-TTS-Realtime. +async_chunk: true + +connectors: + shm: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: true + connector_get_sleep_s: 0.005 + connector_get_max_wait_first_chunk: 1000 + connector_get_max_wait: 300 + initial_codec_chunk_frames: 1 # 1 frame * 80 ms/frame = 80 ms first audio chunk + codec_chunk_frames: 15 # 15 frames * 80 ms/frame = 1.2 s steady chunks + codec_left_context_frames: 0 + codec_stream_slots: 8 + +stages: + - stage_id: 0 + max_num_seqs: 8 + gpu_memory_utilization: 0.60 + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 4096 + max_model_len: 4096 + devices: "0" + output_connectors: + to_stage_1: shm + default_sampling_params: + temperature: 1.0 + top_p: 0.95 + top_k: 50 + max_tokens: 4096 + seed: 42 + repetition_penalty: 1.1 + + - stage_id: 1 + max_num_seqs: 8 + gpu_memory_utilization: 0.15 + enforce_eager: true + enable_prefix_caching: false + async_scheduling: true + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: "0" + input_connectors: + from_stage_0: shm + 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/moss_tts_realtime.yaml b/vllm_omni/deploy/moss_tts_realtime.yaml index 124b3b8be06..b7909661ad8 100644 --- a/vllm_omni/deploy/moss_tts_realtime.yaml +++ b/vllm_omni/deploy/moss_tts_realtime.yaml @@ -3,6 +3,9 @@ # Real-time streaming TTS with TTFB ~180 ms and LLM+TTS first-sentence ~377 ms. # Uses local depth transformer (no delay warm-up) — lowest latency variant. async_chunk: true +# OpenMOSS-Team/MOSS-TTS-Realtime ships custom modeling code; the newer vLLM +# ModelConfig hard-rejects it at validation without this (was lenient before). +trust_remote_code: true connectors: shm: diff --git a/vllm_omni/deploy/moss_ttsd.yaml b/vllm_omni/deploy/moss_ttsd.yaml index 3907941dbfe..f0b5f6eaef8 100644 --- a/vllm_omni/deploy/moss_ttsd.yaml +++ b/vllm_omni/deploy/moss_ttsd.yaml @@ -2,6 +2,7 @@ # HF: OpenMOSS-Team/MOSS-TTSD-v1.0 # Dialogue TTS — optimised for multi-turn conversational speech generation. async_chunk: true +trust_remote_code: true # OpenMOSS-Team/MOSS-TTSD-v1.0 ships custom modeling code connectors: shm: diff --git a/vllm_omni/deploy/moss_voice_generator.yaml b/vllm_omni/deploy/moss_voice_generator.yaml index 8045dcbcfb6..010f0420504 100644 --- a/vllm_omni/deploy/moss_voice_generator.yaml +++ b/vllm_omni/deploy/moss_voice_generator.yaml @@ -2,6 +2,9 @@ # HF: OpenMOSS-Team/MOSS-VoiceGenerator # Zero-shot voice design — create a new voice from a text description. async_chunk: true +# OpenMOSS-Team/MOSS-VoiceGenerator ships custom modeling code; the newer vLLM +# ModelConfig hard-rejects it at validation without this (was lenient before). +trust_remote_code: true connectors: shm: @@ -19,6 +22,11 @@ stages: - stage_id: 0 max_num_seqs: 8 gpu_memory_utilization: 0.60 + # The MOSS talker (custom audio heads + per-request GPU buffers) is not + # CUDA-graph-safe; FULL-decode graph capture crashes the worker. Run the + # LLM stage eager, matching the other TTS talker configs (ming_tts, + # moss_tts_nano). + enforce_eager: true enable_prefix_caching: false async_scheduling: true max_num_batched_tokens: 4096 diff --git a/vllm_omni/deploy/omnivoice.yaml b/vllm_omni/deploy/omnivoice.yaml index 6491be6dc14..6769dd5e995 100644 --- a/vllm_omni/deploy/omnivoice.yaml +++ b/vllm_omni/deploy/omnivoice.yaml @@ -1,19 +1,13 @@ # OmniVoice TTS deploy: diffusion engine single stage (dit). # Text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio. -# (Migrated from model_executor/stage_configs/omnivoice.yaml for deploy/ layout.) -stage_args: +async_chunk: false +pipeline: omnivoice + +stages: - stage_id: 0 - stage_type: diffusion - is_comprehension: true - runtime: - devices: 0 - engine_args: - model_stage: dit - model_class_name: "OmniVoicePipeline" - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - dtype: "float32" - final_output: true - final_output_type: audio + devices: "0" + enforce_eager: true + trust_remote_code: true + distributed_executor_backend: "mp" + dtype: "float32" diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 06869a6577f..5bcbb0fe301 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -47,7 +47,6 @@ stages: temperature: 0.9 top_k: 50 max_tokens: 4096 - seed: 42 repetition_penalty: 1.05 - stage_id: 2 @@ -67,7 +66,6 @@ stages: top_p: 1.0 top_k: -1 max_tokens: 65536 - seed: 42 repetition_penalty: 1.1 platforms: @@ -77,13 +75,19 @@ platforms: gpu_memory_utilization: 0.6 tensor_parallel_size: 2 max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE # Qwen3-Omni not FULL-cudagraph-safe on NPU devices: "0,1" - stage_id: 1 gpu_memory_utilization: 0.6 max_num_batched_tokens: 8192 + compilation_config: + cudagraph_mode: PIECEWISE # code_predictor MaskedSelect / inner NPU graph replay devices: "2" - stage_id: 2 gpu_memory_utilization: 0.3 + compilation_config: + cudagraph_mode: PIECEWISE devices: "2" # Code2Wav decode uses UpSample1d, which calls F.conv_transpose1d. diff --git a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml b/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml deleted file mode 100644 index 4f94fea6069..00000000000 --- a/vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml +++ /dev/null @@ -1,109 +0,0 @@ -# Qwen3-Omni-MoE multi-replica deploy, verified on 3x H20-96G -# (stage 0 on cuda:0, stages 1+2 scaled out across cuda:1,2). -# -# Fields omitted from a stage fall back to StageDeployConfig defaults (see -# vllm_omni/config/stage_config.py). Notable implicit defaults for this -# layout: -# * Stages 0/1 (thinker, talker) do not set max_num_batched_tokens: -# chat-sized prefill fits in the 32768 default. -# * Stage 1 (talker) does not set enforce_eager: cudagraph runs by -# default (false). -# * Stage 2 (code2wav) sets enforce_eager=true because its generation -# loop is cudagraph-incompatible. -# * Stages 1/2 use num_replicas=2 and share devices "1,2". Replica -# routing is handled by StagePool with request-level affinity. -async_chunk: true - -connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - initial_codec_chunk_frames: 4 - codec_chunk_frames: 25 - codec_left_context_frames: 25 - -stages: - - stage_id: 0 - max_num_batched_tokens: 32768 - max_num_seqs: 64 - gpu_memory_utilization: 0.9 - trust_remote_code: true - enable_prefix_caching: false - devices: "0" - 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 - max_num_batched_tokens: 32768 - max_num_seqs: 64 - gpu_memory_utilization: 0.6 - trust_remote_code: true - enable_prefix_caching: false - devices: "1,2" - num_replicas: 2 - input_connectors: - from_stage_0: connector_of_shared_memory - default_sampling_params: - temperature: 0.9 - top_k: 50 - max_tokens: 4096 - seed: 42 - repetition_penalty: 1.05 - - - stage_id: 2 - max_num_batched_tokens: 65536 - max_num_seqs: 64 - gpu_memory_utilization: 0.1 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - async_scheduling: false - enable_chunked_prefill: false - devices: "1,2" - num_replicas: 2 - input_connectors: - from_stage_1: 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.1 - -platforms: - npu: - stages: - - stage_id: 0 - gpu_memory_utilization: 0.6 - tensor_parallel_size: 2 - max_num_batched_tokens: 8192 - devices: "0,1" - - stage_id: 1 - gpu_memory_utilization: 0.6 - max_num_batched_tokens: 8192 - devices: "2" - - stage_id: 2 - gpu_memory_utilization: 0.3 - devices: "2" - - xpu: - stages: - - stage_id: 0 - tensor_parallel_size: 4 - enforce_eager: true - max_cudagraph_capture_size: 0 - devices: "0,1,2,3" - - stage_id: 1 - enforce_eager: true - max_cudagraph_capture_size: 0 - devices: "4" - - stage_id: 2 - gpu_memory_utilization: 0.3 - max_cudagraph_capture_size: 0 - devices: "4" diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 55bac5833eb..f27da1b5be4 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -10,8 +10,7 @@ # codec prefill length (Q * num_frames) exceeds the 32k default. # * Stage 0 does not set enforce_eager — talker runs cudagraph by default. # Stage 1 defaults to false so Code2Wav can use its inner cudagraph; -# set true to disable that path in eager mode. NPU platform flips stage 0 to true where -# cudagraph is not yet verified. +# set true to disable that path in eager mode. async_chunk: true connectors: @@ -36,7 +35,7 @@ connectors: stages: - stage_id: 0 - max_num_seqs: 10 + max_num_seqs: 64 gpu_memory_utilization: 0.3 trust_remote_code: true # Prefix caching on by default. The HS opt-out (talker only consumes the @@ -44,9 +43,9 @@ stages: # PR make per-step bookkeeping negligible, and any text repetition across # requests (same voice or shared ref_audio) is a net throughput win. # Disable explicitly if memory headroom is tight on a single-GPU deploy. - enable_prefix_caching: true + enable_prefix_caching: false async_scheduling: true - max_num_batched_tokens: 512 + max_num_batched_tokens: 32768 max_model_len: 4096 devices: "0" output_connectors: @@ -55,7 +54,10 @@ stages: temperature: 0.9 top_k: 50 max_tokens: 4096 - seed: 42 + # Official qwen-tts generates the talker with min_new_tokens=2; without + # it, checkpoints that sample codec EOS on the first frame return empty + # audio (HTTP 400). See #4962. + min_tokens: 2 repetition_penalty: 1.05 subtalker_sampling_params: do_sample: true @@ -65,7 +67,7 @@ stages: - stage_id: 1 # Tuned for Base voice clone; CustomVoice / VoiceDesign are TTFA-optimal at 1. - max_num_seqs: 10 + max_num_seqs: 64 gpu_memory_utilization: 0.3 enforce_eager: false trust_remote_code: true @@ -86,11 +88,15 @@ stages: top_p: 1.0 top_k: -1 max_tokens: 65536 - seed: 42 repetition_penalty: 1.0 platforms: npu: stages: - stage_id: 0 - enforce_eager: true + enforce_eager: false + compilation_config: + cudagraph_mode: PIECEWISE + - stage_id: 1 + compilation_config: + cudagraph_mode: PIECEWISE diff --git a/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml new file mode 100644 index 00000000000..41145343d96 --- /dev/null +++ b/vllm_omni/deploy/qwen3_tts_forced_aligner.yaml @@ -0,0 +1,19 @@ +# Reference deploy knobs for Qwen/Qwen3-ForcedAligner-0.6B, the shared forced +# aligner for streaming TTS word timestamps. Pass with --forced-aligner-config; +# --forced-aligner overrides the model path without editing this file. +# +# Run as LLM(runner="pooling") with pooling_task="token_classify". Prompt tokens, +# chat template and word segmentation are Qwen-specific and live in +# qwen3_force_align_processor (not configurable here). + +forced_aligner: + model: Qwen/Qwen3-ForcedAligner-0.6B + runner: pooling + architecture: Qwen3ASRForcedAlignerForTokenClassification + pooling_task: token_classify + + gpu_memory_utilization: 0.30 + + dtype: bfloat16 + max_model_len: 4096 + trust_remote_code: true diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 9c6ec19fe0d..939396840a9 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -60,7 +60,10 @@ stages: temperature: 0.9 top_k: 50 max_tokens: 4096 - seed: 42 + # Official qwen-tts generates the talker with min_new_tokens=2; without + # it, checkpoints that sample codec EOS on the first frame return empty + # audio (HTTP 400). See #4962. + min_tokens: 2 repetition_penalty: 1.05 subtalker_sampling_params: do_sample: true @@ -85,7 +88,6 @@ stages: top_p: 1.0 top_k: -1 max_tokens: 65536 - seed: 42 repetition_penalty: 1.0 platforms: diff --git a/vllm_omni/deploy/soulxsinger_svc.yaml b/vllm_omni/deploy/soulxsinger_svc.yaml new file mode 100644 index 00000000000..a9d8158068c --- /dev/null +++ b/vllm_omni/deploy/soulxsinger_svc.yaml @@ -0,0 +1,19 @@ +# SoulX-Singer SVC: single-stage DiT (preprocess inline in pre_process_func). +# +# Use: vllm serve <SoulX-Singer-svc> --omni +pipeline: soulxsinger_svc +async_chunk: false +dtype: float16 +trust_remote_code: true +distributed_executor_backend: "mp" + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.5 + enforce_eager: false + devices: "0" + default_sampling_params: + num_inference_steps: 32 + guidance_scale: 3.0 + seed: 42 diff --git a/vllm_omni/deploy/soulxsinger_svs.yaml b/vllm_omni/deploy/soulxsinger_svs.yaml new file mode 100644 index 00000000000..3f021cb53e1 --- /dev/null +++ b/vllm_omni/deploy/soulxsinger_svs.yaml @@ -0,0 +1,19 @@ +# SoulX-Singer SVS: single-stage DiT (preprocess inline in pre_process_func). +# +# Use: vllm serve <SoulX-Singer> --omni +pipeline: soulxsinger_svs +async_chunk: false +dtype: float16 +trust_remote_code: true +distributed_executor_backend: "mp" + +stages: + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.5 + enforce_eager: false + devices: "0" + default_sampling_params: + num_inference_steps: 32 + guidance_scale: 3.0 + seed: 42 diff --git a/vllm_omni/deploy/step_audio_2.yaml b/vllm_omni/deploy/step_audio_2.yaml new file mode 100644 index 00000000000..d225e974109 --- /dev/null +++ b/vllm_omni/deploy/step_audio_2.yaml @@ -0,0 +1,39 @@ +# Step-Audio2 deploy: thinker -> token2wav. +# +# Sequential mode waits for the full Thinker output before running Token2Wav. +async_chunk: false +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + devices: "0,1" + tensor_parallel_size: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.7 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 8192 + default_sampling_params: + temperature: 0.7 + top_p: 0.9 + top_k: -1 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 1 + devices: "1" + max_num_seqs: 1 + gpu_memory_utilization: 0.2 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 4096 + hf_overrides: + architectures: ["StepAudio2Token2WavForConditionalGeneration"] + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 diff --git a/vllm_omni/deploy/step_audio_2_asr.yaml b/vllm_omni/deploy/step_audio_2_asr.yaml new file mode 100644 index 00000000000..0eb0a852352 --- /dev/null +++ b/vllm_omni/deploy/step_audio_2_asr.yaml @@ -0,0 +1,22 @@ +# Step-Audio2 ASR deploy: thinker only. +pipeline: step_audio_2_asr +async_chunk: false +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + devices: "0,1" + tensor_parallel_size: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.8 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 2048 + default_sampling_params: + temperature: 0.7 + top_p: 0.9 + top_k: -1 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.05 diff --git a/vllm_omni/deploy/step_audio_2_async_chunk.yaml b/vllm_omni/deploy/step_audio_2_async_chunk.yaml new file mode 100644 index 00000000000..75b5cf901bd --- /dev/null +++ b/vllm_omni/deploy/step_audio_2_async_chunk.yaml @@ -0,0 +1,37 @@ +# Step-Audio2 deploy: thinker -> token2wav with async chunk streaming. +async_chunk: true +trust_remote_code: true +enable_prefix_caching: false + +stages: + - stage_id: 0 + devices: "0,1" + tensor_parallel_size: 2 + max_num_seqs: 1 + gpu_memory_utilization: 0.7 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 1024 + default_sampling_params: + temperature: 0.7 + top_p: 0.9 + top_k: -1 + max_tokens: 1024 + seed: 42 + repetition_penalty: 1.05 + + - stage_id: 1 + devices: "1" + max_num_seqs: 1 + gpu_memory_utilization: 0.2 + enforce_eager: true + async_scheduling: false + max_num_batched_tokens: 2048 + hf_overrides: + architectures: ["StepAudio2Token2WavForConditionalGeneration"] + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 diff --git a/vllm_omni/deploy/voxcpm2.yaml b/vllm_omni/deploy/voxcpm2.yaml index 5d366521b0a..63996f1b779 100644 --- a/vllm_omni/deploy/voxcpm2.yaml +++ b/vllm_omni/deploy/voxcpm2.yaml @@ -8,16 +8,17 @@ dtype: bfloat16 stages: - stage_id: 0 - max_num_seqs: 32 + max_num_seqs: 8 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). + # max_num_seqs(8) * max_model_len(4096) ~= 32k tokens; 6 GiB leaves headroom + # for longer live requests while keeping VoxCPM2 TTS decode batches small + # enough to avoid the high ITL observed with large admission caps. # 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; + # Assumes prefix caching off and typical live decode concurrency <= ~8; # raise this cap if either grows or enable_prefix_caching is turned on. kv_cache_memory_bytes: 6442450944 # 6 GiB enforce_eager: true @@ -39,6 +40,9 @@ stages: enable_loc_dit_fused_mlp: true enable_loc_dit_zero_dt_cache: true enable_loc_dit_skip_qkv_contig: true + enable_loc_dit_reduce_overhead_no_cg: true + enable_unified_decode_graph: true + unified_decode_graph_max_batch_size: 8 default_sampling_params: temperature: 0.0 top_p: 1.0 diff --git a/vllm_omni/deploy/wan2_2_ti2v.yaml b/vllm_omni/deploy/wan2_2_ti2v.yaml new file mode 100644 index 00000000000..bdb5cfecb69 --- /dev/null +++ b/vllm_omni/deploy/wan2_2_ti2v.yaml @@ -0,0 +1,17 @@ +# Wan2.2 TI2V-5B deploy. +# +# Single GPU by default. Increase tensor_parallel_size and devices for A14B MoE variants. +async_chunk: false +trust_remote_code: true +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + tensor_parallel_size: 1 + enforce_eager: true + # Keep VAE tiling on by default for higher resolutions / longer clips. + vae_use_tiling: true + default_sampling_params: + seed: 42 diff --git a/vllm_omni/diffusion/attention/backends/flash_attn_hub.py b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py new file mode 100644 index 00000000000..bce607888e2 --- /dev/null +++ b/vllm_omni/diffusion/attention/backends/flash_attn_hub.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import threading +from functools import partial + +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.attention.backends.abstract import ( + AttentionBackend, + AttentionImpl, + AttentionMetadata, +) +from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import ( + piecewise_attn, +) + +logger = init_logger(__name__) + + +_hub_modules: dict[str, object] = {} +_hub_lock = threading.Lock() + + +def _load_hub_module(repo_id: str): + from kernels import get_kernel + + logger.info("Loading %s kernel from HuggingFace Hub...", repo_id) + last_error = None + for version in (1, 2, None): + try: + if version is not None: + return get_kernel(repo_id, version=version) + return get_kernel(repo_id) + except Exception as exc: + logger.info("Failed to load %s version %s: %s", repo_id, version, exc) + last_error = exc + raise RuntimeError(f"Failed to load HuggingFace Hub kernel {repo_id!r}") from last_error + + +def _get_hub_module(repo_id: str): + with _hub_lock: + if repo_id not in _hub_modules: + _hub_modules[repo_id] = _load_hub_module(repo_id) + return _hub_modules[repo_id] + + +def _run_varlen_dense( + varlen_func, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + causal: bool, + softmax_scale: float, +) -> torch.Tensor: + batch_size, q_len = query.size()[:2] + k_len = key.size(1) + cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device) + cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device) + + out = varlen_func( + q=query.flatten(0, 1), + k=key.flatten(0, 1), + v=value.flatten(0, 1), + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=q_len, + max_seqlen_k=k_len, + causal=causal, + softmax_scale=softmax_scale, + ) + if isinstance(out, tuple): + out = out[0] + return out.reshape(batch_size, q_len, *out.shape[1:]) + + +class FlashAttentionHubBackend(AttentionBackend): + accept_output_buffer: bool = True + + @classmethod + def supports_attention_mask(cls) -> bool: + return True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 96, 128, 192, 256] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN_HUB" + + @staticmethod + def get_impl_cls() -> type["FlashAttentionHubImpl"]: + return FlashAttentionHubImpl + + +class FlashAttentionHubImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + qkv_layout: str | None = None, + backend_kwargs: dict | None = None, + **extra_impl_args, + ) -> None: + self.num_heads = num_heads + self.causal = causal + self.softmax_scale = softmax_scale + self.qkv_layout = qkv_layout + if backend_kwargs: + logger.warning("FlashAttentionHubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys())) + + hub_module = _get_hub_module("kernels-community/flash-attn2") + self.flash_attn_func = getattr(hub_module, "flash_attn_func", None) + self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None) + if self.flash_attn_func is None and self.flash_attn_varlen_func is None: + raise RuntimeError("Failed to load flash-attn2 kernel from HuggingFace Hub: no functions found") + + @staticmethod + def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor: + # FA3 may return (out, lse), FA2 returns out + return out[0] if isinstance(out, tuple) else out + + def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs): + if attn_func is not None: + return self._unwrap_flash_output(attn_func(q, k, v, **kwargs)) + return _run_varlen_dense( + self.flash_attn_varlen_func, + q, + k, + v, + causal=kwargs.get("causal", self.causal), + softmax_scale=kwargs.get("softmax_scale", self.softmax_scale), + ) + + def _forward_varlen_masked( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + from vllm_omni.diffusion.attention.backends.utils.fa import ( + _pad_input, + _unpad_input, + _upad_input, + ) + + assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)" + query_length = query.size(1) + q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( + query, key, value, attention_mask, query_length, _unpad_input + ) + + out_unpad = self.flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_k=cu_seq_lens_k, + max_seqlen_q=max_length_q, + max_seqlen_k=max_length_k, + **{ + "causal": self.causal, + "softmax_scale": self.softmax_scale, + }, + ) + out_unpad = self._unwrap_flash_output(out_unpad) + return _pad_input(out_unpad, indices_q, query.size(0), query_length) + + def _forward_varlen_dense( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + return _run_varlen_dense( + self.flash_attn_varlen_func, + query, + key, + value, + causal=self.causal, + softmax_scale=self.softmax_scale, + ) + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata = None, + ) -> torch.Tensor: + attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None + full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None + + # Try piecewise attention + if full_attn_spans is not None: + logger.debug("Using piecewise Flash Attention for mixed causal/full mask") + attn_func = partial( + self._flash_wrapper, + attn_func=self.flash_attn_func, + ) + + return piecewise_attn( + query, + key, + value, + full_attn_spans, + self.softmax_scale, + attn_func, + ) + + if attention_mask is not None and torch.any(~attention_mask): + return self._forward_varlen_masked( + query, + key, + value, + attention_mask, + ) + + if self.flash_attn_func is not None: + out = self.flash_attn_func( + query, + key, + value, + causal=self.causal, + softmax_scale=self.softmax_scale, + ) + return self._unwrap_flash_output(out) + + return self._forward_varlen_dense( + query, + key, + value, + ) + + +class FlashAttention3HubBackend(AttentionBackend): + accept_output_buffer: bool = True + + @classmethod + def supports_attention_mask(cls) -> bool: + return True + + @staticmethod + def get_supported_head_sizes() -> list[int]: + return [64, 96, 128, 192, 256] + + @staticmethod + def get_name() -> str: + return "FLASH_ATTN_3_HUB" + + @staticmethod + def get_impl_cls() -> type["FlashAttention3HubImpl"]: + return FlashAttention3HubImpl + + +class FlashAttention3HubImpl(AttentionImpl): + def __init__( + self, + num_heads: int, + head_size: int, + softmax_scale: float, + causal: bool = False, + num_kv_heads: int | None = None, + prefix: str = "", + qkv_layout: str | None = None, + backend_kwargs: dict | None = None, + **extra_impl_args, + ) -> None: + self.num_heads = num_heads + self.causal = causal + self.softmax_scale = softmax_scale + self.qkv_layout = qkv_layout + if backend_kwargs: + logger.warning("FlashAttention3HubImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys())) + + hub_module = _get_hub_module("kernels-community/flash-attn3") + self.flash_attn_func = getattr(hub_module, "flash_attn_func", None) + self.flash_attn_varlen_func = getattr(hub_module, "flash_attn_varlen_func", None) + if self.flash_attn_func is None and self.flash_attn_varlen_func is None: + raise RuntimeError("Failed to load flash-attn3 kernel from HuggingFace Hub: no functions found") + + @staticmethod + def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor: + # FA3 returns (out, lse) + return out[0] if isinstance(out, tuple) else out + + def _flash_wrapper(self, q, k, v, *, attn_func, **kwargs): + if attn_func is not None: + return self._unwrap_flash_output(attn_func(q, k, v, **kwargs)) + return _run_varlen_dense( + self.flash_attn_varlen_func, + q, + k, + v, + causal=kwargs.get("causal", self.causal), + softmax_scale=kwargs.get("softmax_scale", self.softmax_scale), + ) + + def _forward_varlen_masked( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + ) -> torch.Tensor: + from vllm_omni.diffusion.attention.backends.utils.fa import ( + _pad_input, + _unpad_input, + _upad_input, + ) + + assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)" + query_length = query.size(1) + q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( + query, key, value, attention_mask, query_length, _unpad_input + ) + + out_unpad = self.flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu_seq_lens_q, + cu_seqlens_k=cu_seq_lens_k, + max_seqlen_q=max_length_q, + max_seqlen_k=max_length_k, + **{ + "causal": self.causal, + "softmax_scale": self.softmax_scale, + }, + ) + out_unpad = self._unwrap_flash_output(out_unpad) + return _pad_input(out_unpad, indices_q, query.size(0), query_length) + + def _forward_varlen_dense( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + return _run_varlen_dense( + self.flash_attn_varlen_func, + query, + key, + value, + causal=self.causal, + softmax_scale=self.softmax_scale, + ) + + def forward_cuda( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata = None, + ) -> torch.Tensor: + attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None + full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None + + # Try piecewise attention + if full_attn_spans is not None: + logger.debug("Using piecewise Flash Attention for mixed causal/full mask") + attn_func = partial( + self._flash_wrapper, + attn_func=self.flash_attn_func, + ) + + return piecewise_attn( + query, + key, + value, + full_attn_spans, + self.softmax_scale, + attn_func, + ) + + if attention_mask is not None and torch.any(~attention_mask): + return self._forward_varlen_masked( + query, + key, + value, + attention_mask, + ) + + if self.flash_attn_func is not None: + out = self.flash_attn_func( + query, + key, + value, + causal=self.causal, + softmax_scale=self.softmax_scale, + ) + return self._unwrap_flash_output(out) + + return self._forward_varlen_dense( + query, + key, + value, + ) diff --git a/vllm_omni/diffusion/attention/backends/registry.py b/vllm_omni/diffusion/attention/backends/registry.py index 8c22cc2cd7d..a822d78fb17 100644 --- a/vllm_omni/diffusion/attention/backends/registry.py +++ b/vllm_omni/diffusion/attention/backends/registry.py @@ -62,6 +62,8 @@ class MyCustomBackend: SAGE_ATTN_3 = "vllm_omni.diffusion.attention.backends.sage_attn3.SageAttention3Backend" CUDNN_ATTN = "vllm_omni.diffusion.attention.backends.cudnn_attn.CuDNNAttentionBackend" FLASHINFER_ATTN = "vllm_omni.diffusion.attention.backends.flashinfer_attn.FlashInferAttentionBackend" + FLASH_ATTN_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttentionHubBackend" + FLASH_ATTN_3_HUB = "vllm_omni.diffusion.attention.backends.flash_attn_hub.FlashAttention3HubBackend" def get_path(self, include_classname: bool = True) -> str: """Get the class path for this backend (respects overrides). diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index e0fb8d2645a..0ea1222733a 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/vllm_omni/diffusion/attention/layer.py @@ -230,6 +230,34 @@ def forward( key: torch.Tensor, value: torch.Tensor, attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + if torch.compiler.is_compiling() and is_forward_context_available(): + od_config = get_forward_context().omni_diffusion_config + parallel_config = getattr(od_config, "parallel_config", None) + if getattr(parallel_config, "use_hsdp", False): + # Keep HSDP/FSDP2 parameter all-gather outside Inductor's + # attention graph; otherwise scheduler dependency analysis can + # fail on the fused attention region. + return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata) + + return self._forward_impl(query, key, value, attn_metadata) + + @torch.compiler.disable + def _forward_hsdp_compile_boundary( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, + ) -> torch.Tensor: + return self._forward_impl(query, key, value, attn_metadata) + + def _forward_impl( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_metadata: AttentionMetadata | None = None, ) -> torch.Tensor: # Get the appropriate parallel strategy based on SP active state strategy = self._get_active_parallel_strategy() diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py index d7807622497..aafbf0ba7a3 100644 --- a/vllm_omni/diffusion/cache/base.py +++ b/vllm_omni/diffusion/cache/base.py @@ -13,6 +13,8 @@ the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit". 2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from CacheBackend. Used via cache_backend="tea_cache". +3. StepCacheBackend: Velocity cosine step-skipping for DreamZero. Inherits from + CacheBackend. Used via cache_backend="step_cache". All backends implement the same interface: - enable(pipeline): Enable cache on the pipeline diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index ff1eea6aeaa..5d63ce3c015 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -10,11 +10,19 @@ import functools from collections.abc import Callable from contextlib import ExitStack -from typing import Any, Optional +from dataclasses import dataclass +from typing import Any, Optional, TypeAlias import cache_dit import torch -from cache_dit import BlockAdapter, DBCacheConfig, ForwardPattern, ParamsModifier, TaylorSeerCalibratorConfig +from cache_dit import ( + BlockAdapter, + CalibratorConfig, + DBCacheConfig, + ForwardPattern, + ParamsModifier, + TaylorSeerCalibratorConfig, +) from cache_dit.caching.block_adapters import FakeDiffusionPipeline from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter from cache_dit.caching.cache_blocks.pattern_0_1_2 import CachedBlocks_Pattern_0_1_2 @@ -25,26 +33,94 @@ from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig, OmniDiffusionConfig -from vllm_omni.diffusion.utils.tf_utils import get_transformer_from_pipeline logger = init_logger(__name__) +RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None] + + +@dataclass +class CacheDiTAdapterConfig: + """Config for creating a Cache DiT's block adapter; to enable CacheDiT, + most models just need to define an instance of this class as a class + var in the DiT. + """ + + block_forward_patterns: dict[str, ForwardPattern] + has_separate_cfg: bool = False + cached_adapter_cls: type[CachedAdapter] | None = None + check_forward_pattern: bool = True + + +# Registry of custom cache-dit enablers for specific models +# Maps pipeline names to their cache-dit enablement functions +CUSTOM_DIT_ENABLERS: dict[str, Callable] = {} + # Small helper to centralize cache-dit summaries. def cache_summary(pipeline: Any, details: bool = True) -> None: - cache_dit.summary(pipeline.transformer, details=details) + if hasattr(pipeline, "transformer"): + cache_dit.summary(pipeline.transformer, details=details) if hasattr(pipeline, "transformer_2"): cache_dit.summary(pipeline.transformer_2, details=details) + if not hasattr(pipeline, "transformer") and not hasattr(pipeline, "transformer_2"): + logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute") -# Registry of custom cache-dit enablers for specific models -# Maps pipeline names to their cache-dit enablement functions -# Models in this registry require custom handling (e.g., dual-transformer architectures) -# Will be populated after function definitions -CUSTOM_DIT_ENABLERS: dict[str, Callable] = {} +def default_get_pipeline_transformer(pipeline: Any) -> Any: + return pipeline.transformer + + +def build_cache_context_refresh( + cache_config: DiffusionCacheConfig, + get_pipeline_transformer: Callable[[Any], Any] = default_get_pipeline_transformer, +) -> RefreshCacheContextFunc: + """Build the cache context refresh func for a single Transformer.""" + + def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh cache context for the transformer with new num_inference_steps. + + Args: + pipeline: The pipeline instance. + num_inference_steps: New number of inference steps. + """ + transformer = get_pipeline_transformer(pipeline) + + # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run) + scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6) + + if cache_config.scm_steps_mask_policy is None or not scm_supported_steps: + cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) + else: + cache_dit.refresh_context( + 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 _resolve_calibrator_config(cache_config: DiffusionCacheConfig) -> CalibratorConfig | None: + """Resolves the Calibrator subconfig For DiT cache.""" + calibrator = None + if cache_config.enable_taylorseer: + taylorseer_order = cache_config.taylorseer_order + calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) + logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + # In the future, more calibrators will likely be added to DiT Cache, + # e.g., focal; handle them generically here. + return calibrator -def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: +def _build_db_cache_config(cache_config: DiffusionCacheConfig) -> DBCacheConfig: """Build DBCacheConfig with optional SCM (Step Computation Masking) support. Args: @@ -53,9 +129,8 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: Returns: DBCacheConfig instance with SCM support if configured. """ - return DBCacheConfig( - # we will refresh the context when gets num_inference_steps in the first inference request + # we will refresh the context when we get num_inference_steps in the first inference request num_inference_steps=None, Fn_compute_blocks=cache_config.Fn_compute_blocks, Bn_compute_blocks=cache_config.Bn_compute_blocks, @@ -68,7 +143,100 @@ def _build_db_cache_config(cache_config: Any) -> DBCacheConfig: ) -def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], None]: +def enable_cache_for_dit( + pipeline: Any, + cache_config: Any, + block_adapter: BlockAdapter | None = None, + adapter_cls: type[CachedAdapter] | None = None, +) -> RefreshCacheContextFunc: + """Enable cache-dit for regular single-transformer DiT models. + + Args: + pipeline: The diffusion pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + block_adapter: Custom block adapters for specific model architectures. + adapter_cls: Custom cached adapter class for specific model architectures. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + # Build DBCacheConfig with optional SCM support + db_cache_config = _build_db_cache_config(cache_config) + + # Build calibrator config if TaylorSeer is enabled + calibrator_config = _resolve_calibrator_config(cache_config) + + logger.info( + f"Enabling cache-dit on transformer: " + f"Fn={db_cache_config.Fn_compute_blocks}, " + f"Bn={db_cache_config.Bn_compute_blocks}, " + f"W={db_cache_config.max_warmup_steps}, " + ) + + # Enable cache-dit on the transformer + transformer = default_get_pipeline_transformer(pipeline) + + # If we have a custom cached adapter subclass, call apply directly + if adapter_cls is not None: + adapter_cls.apply( + transformer if block_adapter is None else block_adapter, + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + else: + # Otherwise, call enable cache, which will call CachedAdapter.apply for us + cache_dit.enable_cache( + transformer if block_adapter is None else block_adapter, + cache_config=db_cache_config, + calibrator_config=calibrator_config, + ) + + return build_cache_context_refresh(cache_config) + + +### Complex / custom enablers for DiT cache +# NOTE (Alex): This case is rare; you should only really need to do this if you have a dual transformer +# architecture, since it hasn't been handled generically yet, or if the model class has unique attributes +# that it sets during Cache DiT enablement. +# +# For the vast majority of models, you should only need to add a _cache_dit_adapter_config attribute +# to the Transformer class, which controls the forward pattern, whether or not we have separate CFG, +# and so on. + + +# from https://github.com/vipshop/cache-dit/pull/542 +def _split_wan22_inference_steps(pipeline, num_inference_steps: int) -> tuple[int, int]: + """Split inference steps into high-noise and low-noise steps for Wan2.2. + + This is an internal helper function specific to Wan2.2's dual-transformer + architecture that uses boundary_ratio to determine the split point. + + Args: + num_inference_steps: Total number of inference steps. + + Returns: + A tuple of (num_high_noise_steps, num_low_noise_steps). + """ + if pipeline.boundary_ratio is not None: + boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps + else: + boundary_timestep = None + + # Set timesteps to calculate the split + device = next(pipeline.transformer.parameters()).device + pipeline.scheduler.set_timesteps(num_inference_steps, device=device) + + timesteps = pipeline.scheduler.timesteps + num_high_noise_steps = 0 # high-noise steps for transformer + for t in timesteps: + if boundary_timestep is None or t >= boundary_timestep: + num_high_noise_steps += 1 + # low-noise steps for transformer_2 + num_low_noise_steps = num_inference_steps - num_high_noise_steps + return num_high_noise_steps, num_low_noise_steps + + +def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: """Enable cache-dit for Wan2.2 single or dual-transformer architecture. Wan2.2 can use single or dual transformers (transformer and transformer_2) that need @@ -81,43 +249,31 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int], Returns: A refresh function that can be called to update cache context with new num_inference_steps. """ + # Build DBCacheConfig with optional SCM support + db_cache_config = _build_db_cache_config(cache_config) + calibrator_config = _resolve_calibrator_config(cache_config) if getattr(pipeline, "transformer_2", None) is None: logger.info("transformer_2 not found, enabling cache-dit for single transformer mode") - db_cache_config = _build_db_cache_config(cache_config) cache_dit.enable_cache( BlockAdapter( transformer=pipeline.transformer, + # For VACE, cache only the main denoising blocks. The + # conditioning branch (vace_blocks) has a different forward + # contract and produces per-step hints from the current latent + # plus vace_context; keeping it outside CacheDiT preserves the + # control signal while still accelerating the repeated backbone. blocks=[pipeline.transformer.blocks], forward_pattern=[ForwardPattern.Pattern_2], params_modifiers=[ - ParamsModifier(cache_config=db_cache_config), + ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config), ], has_separate_cfg=True, ), cache_config=db_cache_config, + calibrator_config=calibrator_config, ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for single 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 + return build_cache_context_refresh(cache_config) cache_dit.enable_cache( BlockAdapter( @@ -126,6 +282,9 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool pipeline.transformer_2, ], blocks=[ + # See the single-transformer branch above: VACE conditioning + # blocks are intentionally recomputed each step and are not + # wrapped by CacheDiT's main-block Pattern_2 adapter. pipeline.transformer.blocks, pipeline.transformer_2.blocks, ], @@ -140,57 +299,24 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool max_warmup_steps=cache_config.max_warmup_steps, max_cached_steps=cache_config.max_cached_steps, ), + calibrator_config=calibrator_config, ), ParamsModifier( cache_config=DBCacheConfig().reset( max_warmup_steps=2, max_cached_steps=20, ), + calibrator_config=calibrator_config, ), ], has_separate_cfg=True, ), - cache_config=DBCacheConfig( - Fn_compute_blocks=cache_config.Fn_compute_blocks, - Bn_compute_blocks=cache_config.Bn_compute_blocks, - max_warmup_steps=cache_config.max_warmup_steps, - max_cached_steps=cache_config.max_cached_steps, - max_continuous_cached_steps=cache_config.max_continuous_cached_steps, - residual_diff_threshold=cache_config.residual_diff_threshold, - num_inference_steps=None, - ), + cache_config=db_cache_config, + calibrator_config=calibrator_config, ) - # from https://github.com/vipshop/cache-dit/pull/542 - def _split_inference_steps(num_inference_steps: int) -> tuple[int, int]: - """Split inference steps into high-noise and low-noise steps for Wan2.2. - - This is an internal helper function specific to Wan2.2's dual-transformer - architecture that uses boundary_ratio to determine the split point. - - Args: - num_inference_steps: Total number of inference steps. - - Returns: - A tuple of (num_high_noise_steps, num_low_noise_steps). - """ - if pipeline.boundary_ratio is not None: - boundary_timestep = pipeline.boundary_ratio * pipeline.scheduler.config.num_train_timesteps - else: - boundary_timestep = None - - # Set timesteps to calculate the split - device = next(pipeline.transformer.parameters()).device - pipeline.scheduler.set_timesteps(num_inference_steps, device=device) - - timesteps = pipeline.scheduler.timesteps - num_high_noise_steps = 0 # high-noise steps for transformer - for t in timesteps: - if boundary_timestep is None or t >= boundary_timestep: - num_high_noise_steps += 1 - # low-noise steps for transformer_2 - num_low_noise_steps = num_inference_steps - num_high_noise_steps - return num_high_noise_steps, num_low_noise_steps + refresh_trans_one = build_cache_context_refresh(cache_config) + refresh_trans_two = build_cache_context_refresh(cache_config, lambda pipeline: pipeline.transformer_2) def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: """Refresh cache context for both transformers with new num_inference_steps. @@ -199,50 +325,14 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool pipeline: The Wan2.2 pipeline instance. num_inference_steps: New number of inference steps. """ - - num_high_noise_steps, num_low_noise_steps = _split_inference_steps(num_inference_steps) - # Refresh context for high-noise transformer - if cache_config.scm_steps_mask_policy is None: - # cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_high_noise_steps, verbose=verbose) - cache_dit.refresh_context( - pipeline.transformer, - num_inference_steps=num_high_noise_steps, - verbose=verbose, - ) - cache_dit.refresh_context( - pipeline.transformer_2, - num_inference_steps=num_low_noise_steps, - verbose=verbose, - ) - else: - cache_dit.refresh_context( - pipeline.transformer, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_high_noise_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_high_noise_steps - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) - - cache_dit.refresh_context( - pipeline.transformer_2, - cache_config=DBCacheConfig().reset( - num_inference_steps=num_low_noise_steps, - steps_computation_mask=cache_dit.steps_mask( - mask_policy=cache_config.scm_steps_mask_policy, total_steps=num_low_noise_steps - ), - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, - ) + num_high_noise_steps, num_low_noise_steps = _split_wan22_inference_steps(pipeline, num_inference_steps) + refresh_trans_one(pipeline, num_high_noise_steps, verbose) + refresh_trans_two(pipeline, num_low_noise_steps, verbose) return refresh_cache_context -def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]: +def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: """Enable cache-dit for Wan2.2 S2V. S2V uses a single transformer, but unlike the other Wan2.2 variants its @@ -316,1635 +406,550 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool return refresh_cache_context -def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for LongCatImage pipeline. - - Args: - pipeline: The LongCatImage pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. +class BagelCachedContextManager(CachedContextManager): + """ + Custom CachedContextManager for Bagel that safely handles NaiveCache objects + (mapped to encoder_hidden_states) by skipping tensor operations on them. """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on LongCatImage transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1], - params_modifiers=[modifier], - has_separate_cfg=True, - ) - ), - cache_config=db_cache_config, - ) + @torch.compiler.disable + def apply_cache( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor = None, + prefix: str = "Bn", + encoder_prefix: str = "Bn_encoder", + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # Allow Bn and Fn prefix to be used for residual cache. + if "Bn" in prefix: + hidden_states_prev = self.get_Bn_buffer(prefix) + else: + hidden_states_prev = self.get_Fn_buffer(prefix) - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. + assert hidden_states_prev is not None, f"{prefix}_buffer must be set before" - Args: - pipeline: The LongCatImage pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + if self.is_cache_residual(): + hidden_states = hidden_states_prev + hidden_states 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, - ) + # If cache is not residual, we use the hidden states directly + hidden_states = hidden_states_prev - return refresh_cache_context + hidden_states = hidden_states.contiguous() + if encoder_hidden_states is not None: + if "Bn" in encoder_prefix: + encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix) + else: + encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix) -def enable_cache_for_flux(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Flux.1-dev pipeline. + if encoder_hidden_states_prev is not None: + if self.is_encoder_cache_residual(): + # FIX: Check if encoder_hidden_states is a tensor before adding + if isinstance(encoder_hidden_states, torch.Tensor) and isinstance( + encoder_hidden_states_prev, torch.Tensor + ): + encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states + else: + # If encoder cache is not residual, we use the encoder hidden states directly + encoder_hidden_states = encoder_hidden_states_prev - Args: - pipeline: The Flux pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) + # FIX: Check if encoder_hidden_states is a tensor before calling contiguous + if isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states.contiguous() - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + return hidden_states, encoder_hidden_states - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - logger.info( - f"Enabling cache-dit on Flux transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) +class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2): + """ + Custom CachedBlocks for Bagel that safely handles NaiveCache objects + by adding isinstance checks in call_Mn_blocks and compute_or_prune. + """ - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1], - params_modifiers=[modifier], - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The Flux pipeline instance. - num_inference_steps: New number of inference steps. - """ - 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, + def call_Mn_blocks( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + *args, + **kwargs, + ): + original_hidden_states = hidden_states + original_encoder_hidden_states = encoder_hidden_states + for block in self._Mn_blocks(): + hidden_states = block( + hidden_states, + encoder_hidden_states, + *args, + **kwargs, ) + hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states) - return refresh_cache_context + # compute hidden_states residual + hidden_states = hidden_states.contiguous() + hidden_states_residual = hidden_states - original_hidden_states -def enable_cache_for_flux2_klein(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for FLUX.2-klein-4B pipeline. + if ( + encoder_hidden_states is not None + and original_encoder_hidden_states is not None + and isinstance(encoder_hidden_states, torch.Tensor) # FIX: Added Check + ): + encoder_hidden_states = encoder_hidden_states.contiguous() + encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states + else: + encoder_hidden_states_residual = None - Args: - pipeline: The FLUX.2-klein-4B pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - # The Fn_compute_blocks = 2 override is the most important decision here, - # and the rationale (quality degradation at Fn=1) only lives in flux2_klein. - db_cache_config.Fn_compute_blocks = 2 + return ( + hidden_states, + encoder_hidden_states, + hidden_states_residual, + encoder_hidden_states_residual, + ) - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + def compute_or_prune( + self, + block_id: int, # Block index in the transformer blocks + # Below are the inputs to the block + block, # The transformer block to be executed + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + *args, + **kwargs, + ): + # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case. + # Copy-pasted from original but adding checks. - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) + original_hidden_states = hidden_states + original_encoder_hidden_states = encoder_hidden_states - logger.info( - f"Enabling cache-dit on Flux2-Klein transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) + can_use_prune = self._maybe_prune( + block_id, + hidden_states, + prefix=f"{self.cache_prefix}_{block_id}_Fn_original", + ) - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2], - params_modifiers=[modifier], + torch._dynamo.graph_break() + if can_use_prune: + self.context_manager.add_pruned_step() + hidden_states, encoder_hidden_states = self.context_manager.apply_prune( + hidden_states, + encoder_hidden_states, + prefix=( + f"{self.cache_prefix}_{block_id}_Bn_residual" + if self.context_manager.is_cache_residual() + else f"{self.cache_prefix}_Bn_hidden_states" + ), + encoder_prefix=( + f"{self.cache_prefix}_{block_id}_Bn_encoder_residual" + if self.context_manager.is_encoder_cache_residual() + else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states" + ), ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The FLUX.2-klein-4B pipeline instance. - num_inference_steps: New number of inference steps. - """ - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) + torch._dynamo.graph_break() 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, - ), - Fn_compute_blocks=db_cache_config.Fn_compute_blocks, - steps_computation_policy=cache_config.scm_steps_policy, - ), - verbose=verbose, + # Normal steps: Compute the block and cache the residuals. + hidden_states = block( + hidden_states, + encoder_hidden_states, + *args, + **kwargs, ) + hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states) + if not self._skip_prune(block_id): + hidden_states = hidden_states.contiguous() + hidden_states_residual = hidden_states - original_hidden_states - return refresh_cache_context - - -def enable_cache_for_stable_audio_open(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Stable Audio Open pipeline. - - Args: - pipeline: The StableAudioPipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. + if ( + encoder_hidden_states is not None + and original_encoder_hidden_states is not None + and isinstance(encoder_hidden_states, torch.Tensor) # FIX: Added Check + ): + encoder_hidden_states = encoder_hidden_states.contiguous() + encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states + else: + encoder_hidden_states_residual = None - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - db_cache_config = _build_db_cache_config(cache_config) + self.context_manager.set_Fn_buffer( + original_hidden_states, + prefix=f"{self.cache_prefix}_{block_id}_Fn_original", + ) + if self.context_manager.is_cache_residual(): + self.context_manager.set_Bn_buffer( + hidden_states_residual, + prefix=f"{self.cache_prefix}_{block_id}_Bn_residual", + ) + else: + self.context_manager.set_Bn_buffer( + hidden_states, + prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states", + ) + if encoder_hidden_states_residual is not None: + if self.context_manager.is_encoder_cache_residual(): + self.context_manager.set_Bn_encoder_buffer( + encoder_hidden_states_residual, + prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual", + ) + else: + self.context_manager.set_Bn_encoder_buffer( + encoder_hidden_states_residual, + prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states", + ) + torch._dynamo.graph_break() - 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}") + return hidden_states, encoder_hidden_states - # StableAudio is officially registered in CacheDiT as Pattern_3: - # https://github.com/vipshop/cache-dit/blob/69e82bd1/src/cache_dit/caching/block_adapters/__init__.py#L562 - # - # Pattern_3 is required because StableAudioDiT uses cross-attention - # with static encoder_hidden_states that do not change inside the - # transformer block loop. - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_3, - params_modifiers=[ - ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - ], - ), - cache_config=db_cache_config, - ) - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. +class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5): + """CacheDiT blocks wrapper that preserves S2V per-layer audio injection.""" - Args: - pipeline: The StableAudioPipeline instance. - num_inference_steps: New number of inference steps. - verbose: Whether to log refresh operations. - """ - # Bypass SCM for step counts that don't support predefined masks (e.g., vLLM's 1-step dummy run) - scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6) + 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 - if cache_config.scm_steps_mask_policy is None or not scm_supported_steps: - cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - updated_scm_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, + 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 - cache_dit.refresh_context( - pipeline.transformer, - cache_config=updated_scm_config, - verbose=verbose, - ) - - return refresh_cache_context - - -def enable_cache_for_sd3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for StableDiffusion3Pipeline. - - Args: - pipeline: The StableDiffusion3 pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on StableDiffusion3 transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_1, - params_modifiers=[modifier], - ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The LongCatImage pipeline instance. - num_inference_steps: New number of inference steps. - """ - 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_ltx2(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for LTX2 pipelines (audio-video transformer blocks).""" - transformer = get_transformer_from_pipeline(pipeline) - - 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}") - - blocks = transformer.transformer_blocks - - logger.info( - f"Enabling cache-dit on LTX2 transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=transformer, - blocks=blocks, - # LTX2 blocks return (hidden_states, audio_hidden_states) - forward_pattern=ForwardPattern.Pattern_0, - # Treat audio_hidden_states as encoder_hidden_states in Pattern_0 - check_forward_pattern=False, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = get_transformer_from_pipeline(pipeline) - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - 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_dit(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for regular single-transformer DiT models. - - Args: - pipeline: The diffusion pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - # Build DBCacheConfig with optional SCM support - db_cache_config = _build_db_cache_config(cache_config) - - # Build calibrator config if TaylorSeer is enabled - 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}") - - logger.info( - f"Enabling cache-dit on transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit on the transformer - transformer = get_transformer_from_pipeline(pipeline) - cache_dit.enable_cache( - transformer, - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The diffusion pipeline instance. - num_inference_steps: New number of inference steps. - """ - transformer = get_transformer_from_pipeline(pipeline) - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - 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_hunyuan_image3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for HunyuanImage3 pipeline. - - HunyuanImage3 stores its main transformer stack at ``pipeline.model`` with - decoder blocks in ``pipeline.model.layers``. - """ - 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}") - - transformer = getattr(pipeline, "model", None) - if transformer is None or not hasattr(transformer, "layers"): - raise ValueError( - f"HunyuanImage3 cache-dit enabler expects pipeline.model.layers, got pipeline={pipeline.__class__.__name__}" - ) - - logger.info( - f"Enabling cache-dit on HunyuanImage3 model: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - # HunyuanImage3 decoder layers are single-stream `hidden_states` blocks - # (`HunyuanImage3DecoderLayer.forward(hidden_states, ...)`) while caller - # still reads outputs as tuple (`layer_outputs[0]`). Pattern_4 matches this: - # hidden-only input with hidden-first tuple output. - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=transformer, - blocks=transformer.layers, - forward_pattern=ForwardPattern.Pattern_4, - params_modifiers=[modifier], - check_forward_pattern=False, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = getattr(pipeline, "model", None) - if transformer is None: - raise ValueError("HunyuanImage3 cache-dit refresh expects pipeline.model to exist.") - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - 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 - - -class BagelCachedContextManager(CachedContextManager): - """ - Custom CachedContextManager for Bagel that safely handles NaiveCache objects - (mapped to encoder_hidden_states) by skipping tensor operations on them. - """ - - @torch.compiler.disable - def apply_cache( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor = None, - prefix: str = "Bn", - encoder_prefix: str = "Bn_encoder", - ) -> tuple[torch.Tensor, torch.Tensor | None]: - # Allow Bn and Fn prefix to be used for residual cache. - if "Bn" in prefix: - hidden_states_prev = self.get_Bn_buffer(prefix) - else: - hidden_states_prev = self.get_Fn_buffer(prefix) - - assert hidden_states_prev is not None, f"{prefix}_buffer must be set before" - - if self.is_cache_residual(): - hidden_states = hidden_states_prev + hidden_states - else: - # If cache is not residual, we use the hidden states directly - hidden_states = hidden_states_prev - - hidden_states = hidden_states.contiguous() - - if encoder_hidden_states is not None: - if "Bn" in encoder_prefix: - encoder_hidden_states_prev = self.get_Bn_encoder_buffer(encoder_prefix) - else: - encoder_hidden_states_prev = self.get_Fn_encoder_buffer(encoder_prefix) - - if encoder_hidden_states_prev is not None: - if self.is_encoder_cache_residual(): - # FIX: Check if encoder_hidden_states is a tensor before adding - if isinstance(encoder_hidden_states, torch.Tensor) and isinstance( - encoder_hidden_states_prev, torch.Tensor - ): - encoder_hidden_states = encoder_hidden_states_prev + encoder_hidden_states - else: - # If encoder cache is not residual, we use the encoder hidden states directly - encoder_hidden_states = encoder_hidden_states_prev - - # FIX: Check if encoder_hidden_states is a tensor before calling contiguous - if isinstance(encoder_hidden_states, torch.Tensor): - encoder_hidden_states = encoder_hidden_states.contiguous() - - return hidden_states, encoder_hidden_states - - -class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2): - """ - Custom CachedBlocks for Bagel that safely handles NaiveCache objects - by adding isinstance checks in call_Mn_blocks and compute_or_prune. - """ - - def call_Mn_blocks( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - *args, - **kwargs, - ): - original_hidden_states = hidden_states - original_encoder_hidden_states = encoder_hidden_states - for block in self._Mn_blocks(): - hidden_states = block( - hidden_states, - encoder_hidden_states, - *args, - **kwargs, - ) - hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states) - - # compute hidden_states residual - hidden_states = hidden_states.contiguous() - - hidden_states_residual = hidden_states - original_hidden_states - - if ( - encoder_hidden_states is not None - and original_encoder_hidden_states is not None - and isinstance(encoder_hidden_states, torch.Tensor) # FIX: Added Check - ): - encoder_hidden_states = encoder_hidden_states.contiguous() - encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states - else: - encoder_hidden_states_residual = None - - return ( - hidden_states, - encoder_hidden_states, - hidden_states_residual, - encoder_hidden_states_residual, - ) - - def compute_or_prune( - self, - block_id: int, # Block index in the transformer blocks - # Below are the inputs to the block - block, # The transformer block to be executed - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - *args, - **kwargs, - ): - # NOTE: Although Bagel likely won't use pruning, implementing safe version just in case. - # Copy-pasted from original but adding checks. - - original_hidden_states = hidden_states - original_encoder_hidden_states = encoder_hidden_states - - can_use_prune = self._maybe_prune( - block_id, - hidden_states, - prefix=f"{self.cache_prefix}_{block_id}_Fn_original", - ) - - torch._dynamo.graph_break() - if can_use_prune: - self.context_manager.add_pruned_step() - hidden_states, encoder_hidden_states = self.context_manager.apply_prune( - hidden_states, - encoder_hidden_states, - prefix=( - f"{self.cache_prefix}_{block_id}_Bn_residual" - if self.context_manager.is_cache_residual() - else f"{self.cache_prefix}_Bn_hidden_states" - ), - encoder_prefix=( - f"{self.cache_prefix}_{block_id}_Bn_encoder_residual" - if self.context_manager.is_encoder_cache_residual() - else f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states" - ), - ) - torch._dynamo.graph_break() - else: - # Normal steps: Compute the block and cache the residuals. - hidden_states = block( - hidden_states, - encoder_hidden_states, - *args, - **kwargs, - ) - hidden_states, encoder_hidden_states = self._process_block_outputs(hidden_states, encoder_hidden_states) - if not self._skip_prune(block_id): - hidden_states = hidden_states.contiguous() - hidden_states_residual = hidden_states - original_hidden_states - - if ( - encoder_hidden_states is not None - and original_encoder_hidden_states is not None - and isinstance(encoder_hidden_states, torch.Tensor) # FIX: Added Check - ): - encoder_hidden_states = encoder_hidden_states.contiguous() - encoder_hidden_states_residual = encoder_hidden_states - original_encoder_hidden_states - else: - encoder_hidden_states_residual = None - - self.context_manager.set_Fn_buffer( - original_hidden_states, - prefix=f"{self.cache_prefix}_{block_id}_Fn_original", - ) - if self.context_manager.is_cache_residual(): - self.context_manager.set_Bn_buffer( - hidden_states_residual, - prefix=f"{self.cache_prefix}_{block_id}_Bn_residual", - ) - else: - self.context_manager.set_Bn_buffer( - hidden_states, - prefix=f"{self.cache_prefix}_{block_id}_Bn_hidden_states", - ) - if encoder_hidden_states_residual is not None: - if self.context_manager.is_encoder_cache_residual(): - self.context_manager.set_Bn_encoder_buffer( - encoder_hidden_states_residual, - prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_residual", - ) - else: - self.context_manager.set_Bn_encoder_buffer( - encoder_hidden_states_residual, - prefix=f"{self.cache_prefix}_{block_id}_Bn_encoder_hidden_states", - ) - torch._dynamo.graph_break() - - 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. - """ - - @classmethod - def create_context( - cls, - block_adapter: BlockAdapter, - **context_kwargs, - ) -> tuple[list[str], list[dict[str, Any]]]: - # Override to use BagelCachedContextManager - - BlockAdapter.assert_normalized(block_adapter) - - if BlockAdapter.is_cached(block_adapter.pipe): - return block_adapter.pipe - - # Check context_kwargs - context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs) - - # Each Pipeline should have it's own context manager instance. - cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None) - assert cache_config is not None, "cache_config can not be None." - - # Apply cache on pipeline: wrap cache context - pipe_cls_name = block_adapter.pipe.__class__.__name__ - - # USE CUSTOM CONTEXT MANAGER - context_manager = BagelCachedContextManager( - name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}", - persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline), - ) - - flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs) - - block_adapter.pipe._context_manager = context_manager # instance level - - if not context_manager.persistent_context: - original_call = block_adapter.pipe.__class__.__call__ - - @functools.wraps(original_call) - def new_call(self, *args, **kwargs): - with ExitStack() as stack: - # cache context will be reset for each pipe inference - for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs): - stack.enter_context( - context_manager.enter_context( - context_manager.reset_context( - context_name, - **context_kwargs, - ), - ) - ) - outputs = original_call(self, *args, **kwargs) - cls.apply_stats_hooks(block_adapter) - return outputs - - block_adapter.pipe.__class__.__call__ = new_call - block_adapter.pipe.__class__._original_call = original_call - - else: - # Init persistent cache context for transformer - for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs): - context_manager.reset_context( - context_name, - **context_kwargs, - ) - - block_adapter.pipe.__class__._is_cached = True - - cls.apply_params_hooks(block_adapter, contexts_kwargs) - - return flatten_contexts, contexts_kwargs - - @classmethod - def collect_unified_blocks( - cls, - block_adapter: BlockAdapter, - contexts_kwargs: list[dict], - ) -> list[dict[str, torch.nn.ModuleList]]: - # Override to use BagelCachedBlocks - - BlockAdapter.assert_normalized(block_adapter) - - total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = [] - assert hasattr(block_adapter.pipe, "_context_manager") - # Skipping isinstance check for ContextManager._supported_managers to avoid import issues - - 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"] - - # Directly instantiate BagelCachedBlocks - unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList( - [ - BagelCachedBlocks( - # 0. Transformer blocks configuration - 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, - # 1. Cache/Prune context configuration - 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 - - -def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Bagel model (via OmniDiffusion pipeline). - - Args: - pipeline: The OmniDiffusion pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - # Build DBCacheConfig - db_cache_config = _build_db_cache_config(cache_config) - - # Build calibrator config if TaylorSeer is enabled - 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}") - - # Access the transformer: BagelPipeline -> Qwen2MoTForCausalLM -> Qwen2MoTModel - # BagelPipeline has self.language_model which is Qwen2MoTForCausalLM - # Qwen2MoTForCausalLM has self.model which is Qwen2MoTModel - transformer = pipeline.language_model.model - - logger.info( - f"Enabling cache-dit on Bagel transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit on the transformer - # Pattern_0 corresponds to (hidden_states, encoder_hidden_states) input, output - # Custom adapter for Bagel to handle NaiveCache correctly - # from vllm_omni.diffusion.cache.bagel_cache_adapter import BagelCachedAdapter # No longer needed - BagelCachedAdapter.apply( - BlockAdapter( - transformer=transformer, - blocks=transformer.layers, - forward_pattern=ForwardPattern.Pattern_0, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = pipeline.language_model.model - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - 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 - - -class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5): - """ - Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden - states during denoising. - """ - - @classmethod - def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool: - if kwargs.get("cache_dit_skip", False): - return False - - # Prefix/text forwards either omit image_gen_indicators or update the - # DynamicCache. Denoising forwards are gen-only and use update_cache=False. - if kwargs.get("update_cache", True): - return False - - exist_gen = kwargs.get("exist_gen") - exist_und = kwargs.get("exist_und") - if exist_gen is None or exist_und is None: - image_gen_indicators = kwargs.get("image_gen_indicators") - if image_gen_indicators is None: - return False - exist_gen = image_gen_indicators.any().item() - exist_und = (~image_gen_indicators).any().item() - - return exist_gen and not exist_und - - @staticmethod - def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: - kwargs = dict(kwargs) - kwargs.pop("cache_dit_skip", None) - return kwargs - - def forward(self, hidden_states: torch.Tensor, *args, **kwargs): - block_kwargs = self._strip_cache_only_kwargs(kwargs) - if not self._is_denoising_call(kwargs): - hidden_states, new_encoder_hidden_states = self.call_blocks( + 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, - **block_kwargs, - ) - return self._process_forward_outputs(hidden_states, new_encoder_hidden_states) - - return super().forward(hidden_states, *args, **block_kwargs) - - -class SensenovaCachedAdapter(CachedAdapter): - """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks.""" - - @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( - [ - SensenovaCachedBlocks( - # 0. Transformer blocks configuration - 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, - # 1. Cache/Prune context configuration - 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 - - -def enable_cache_for_sensenova_u1(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for SenseNova-U1 model. - - Args: - pipeline: The SenseNova-U1 pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - transformer = getattr(getattr(pipeline, "language_model", None), "model", None) - if transformer is None or not hasattr(transformer, "layers"): - raise ValueError("SenseNovaU1Pipeline cache-dit expects pipeline.language_model.model.layers.") - - # Build DBCacheConfig - 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}") - - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - logger.info( - f"Enabling cache-dit on SenseNova-U1 decoder layers: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - SensenovaCachedAdapter.apply( - BlockAdapter( - transformer=transformer, - blocks=transformer.layers, - forward_pattern=ForwardPattern.Pattern_3, - params_modifiers=[modifier], - check_forward_pattern=True, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = pipeline.language_model.model - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - 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_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Flux.2-dev pipeline. - - Args: - pipeline: The Flux2 pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) - - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") - - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - - logger.info( - f"Enabling cache-dit on Flux transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - # Enable cache-dit using BlockAdapter for transformer - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[ - pipeline.transformer.transformer_blocks, - pipeline.transformer.single_transformer_blocks, - ], - forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2], - params_modifiers=[modifier], + **kwargs, ) - ), - cache_config=db_cache_config, - ) - - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. - - Args: - pipeline: The Flux2 pipeline instance. - num_inference_steps: New number of inference steps. - """ - 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 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, ) - return refresh_cache_context + 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 enable_cache_for_glm_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for GlmImage pipeline. + 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 - Args: - pipeline: The GlmImage pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. - """ - # Build DBCacheConfig for transformer - db_cache_config = _build_db_cache_config(cache_config) + 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, + ) - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + return hidden_states, new_encoder_hidden_states - # Build ParamsModifier for transformer - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) - logger.info( - f"Enabling cache-dit on GlmImage transformer with BlockAdapter: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) +class Wan22S2VCachedAdapter(CachedAdapter): + """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection. - # Enable cache-dit using BlockAdapter for transformer - # Note: We don't use patch_functor here because it's designed for diffusers' GlmImage, - # and our vllm-omni implementation has a different forward signature. - # We use ForwardPattern.Pattern_0 because our block returns (hidden_states, encoder_hidden_states) - cache_dit.enable_cache( - ( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_0, - params_modifiers=[modifier], - patch_functor=None, - has_separate_cfg=True, - ) - ), - cache_config=db_cache_config, - ) + 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. + """ - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - """Refresh cache context for the transformer with new num_inference_steps. + @classmethod + def collect_unified_blocks( + cls, + block_adapter: BlockAdapter, + contexts_kwargs: list[dict], + ) -> list[dict[str, torch.nn.ModuleList]]: + BlockAdapter.assert_normalized(block_adapter) - Args: - pipeline: The GlmImage pipeline instance. - num_inference_steps: New number of inference steps. - """ - 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, - ) + total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = [] + assert hasattr(block_adapter.pipe, "_context_manager") - return refresh_cache_context + 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) -def enable_cache_for_dreamid_omni(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for DreamID-Omni fused pipeline. + return total_cached_blocks - Args: - pipeline: The DreamIDOmni pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called with a new ``num_inference_steps`` - to update the cache context for the pipeline. +class BagelCachedAdapter(CachedAdapter): + """ + Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks. """ - transformer = getattr(pipeline, "transformer", None) - fused_blocks = getattr(transformer, "fused_blocks", None) - if transformer is None or fused_blocks is None: - raise ValueError("DreamIDOmniPipeline cache-dit expects pipeline.transformer.fused_blocks.") - db_cache_config = _build_db_cache_config(cache_config) + @classmethod + def create_context( + cls, + block_adapter: BlockAdapter, + **context_kwargs, + ) -> tuple[list[str], list[dict[str, Any]]]: + # Override to use BagelCachedContextManager - calibrator = None - if cache_config.enable_taylorseer: - taylorseer_order = cache_config.taylorseer_order - calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order) - logger.info(f"TaylorSeer enabled with order={taylorseer_order}") + BlockAdapter.assert_normalized(block_adapter) - modifier = ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator, - ) + if BlockAdapter.is_cached(block_adapter.pipe): + return block_adapter.pipe - logger.info( - f"Enabling cache-dit on DreamID-Omni fused blocks: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) + # Check context_kwargs + context_kwargs = cls.check_context_kwargs(block_adapter, **context_kwargs) - cache_dit.enable_cache( - BlockAdapter( - transformer=transformer, - blocks=fused_blocks, - forward_pattern=ForwardPattern.Pattern_0, - params_modifiers=[modifier], - has_separate_cfg=True, - ), - cache_config=db_cache_config, - ) + # Each Pipeline should have it's own context manager instance. + cache_config: BasicCacheConfig = context_kwargs.get("cache_config", None) + assert cache_config is not None, "cache_config can not be None." - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - transformer = pipeline.transformer - if cache_config.scm_steps_mask_policy is None: - cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) - else: - cache_dit.refresh_context( - 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, - ) + # Apply cache on pipeline: wrap cache context + pipe_cls_name = block_adapter.pipe.__class__.__name__ - return refresh_cache_context + # USE CUSTOM CONTEXT MANAGER + context_manager = BagelCachedContextManager( + name=f"{pipe_cls_name}_{hash(id(block_adapter.pipe))}", + persistent_context=isinstance(block_adapter.pipe, FakeDiffusionPipeline), + ) + flatten_contexts, contexts_kwargs = cls.modify_context_params(block_adapter, **context_kwargs) -def enable_cache_for_ernie_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for ERNIE-Image pipeline. + block_adapter.pipe._context_manager = context_manager # instance level - ERNIE-Image blocks have signature: - forward(x, rotary_pos_emb, temb, attention_mask) -> x + if not context_manager.persistent_context: + original_call = block_adapter.pipe.__class__.__call__ - Where x is hidden_states (concatenated image + text tokens). - This matches Pattern_3 which expects: - - Input: hidden_states only - - Output: hidden_states only + @functools.wraps(original_call) + def new_call(self, *args, **kwargs): + with ExitStack() as stack: + # cache context will be reset for each pipe inference + for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs): + stack.enter_context( + context_manager.enter_context( + context_manager.reset_context( + context_name, + **context_kwargs, + ), + ) + ) + outputs = original_call(self, *args, **kwargs) + cls.apply_stats_hooks(block_adapter) + return outputs - Args: - pipeline: The ERNIE-Image pipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. - """ - db_cache_config = _build_db_cache_config(cache_config) + block_adapter.pipe.__class__.__call__ = new_call + block_adapter.pipe.__class__._original_call = original_call - 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}") + else: + # Init persistent cache context for transformer + for context_name, context_kwargs in zip(flatten_contexts, contexts_kwargs): + context_manager.reset_context( + context_name, + **context_kwargs, + ) - logger.info( - f"Enabling cache-dit on ERNIE-Image transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) + block_adapter.pipe.__class__._is_cached = True - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.layers, - forward_pattern=ForwardPattern.Pattern_3, - check_forward_pattern=False, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) + cls.apply_params_hooks(block_adapter, contexts_kwargs) + + return flatten_contexts, contexts_kwargs + + @classmethod + def collect_unified_blocks( + cls, + block_adapter: BlockAdapter, + contexts_kwargs: list[dict], + ) -> list[dict[str, torch.nn.ModuleList]]: + # Override to use BagelCachedBlocks - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - 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, - ) + BlockAdapter.assert_normalized(block_adapter) - return refresh_cache_context + total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = [] + assert hasattr(block_adapter.pipe, "_context_manager") + # Skipping isinstance check for ContextManager._supported_managers to avoid import issues + + 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"] + # Directly instantiate BagelCachedBlocks + unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList( + [ + BagelCachedBlocks( + # 0. Transformer blocks configuration + 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, + # 1. Cache/Prune context configuration + 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, + ) + ] + ) -def enable_cache_for_helios(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for Helios pipeline. + total_cached_blocks.append(unified_blocks_bind_context) - Helios extends Wan2.2 with multi-term memory patches and guidance cross-attention. - Its transformer blocks have the same single-output signature as Wan2.2 blocks - (returns hidden_states only), so we use ForwardPattern.Pattern_2. + return total_cached_blocks - Args: - pipeline: The HeliosPipeline instance. - cache_config: DiffusionCacheConfig instance with cache configuration. - Returns: - A refresh function that can be called to update cache context with new num_inference_steps. +class SensenovaCachedBlocks(CachedBlocks_Pattern_3_4_5): + """ + Custom CachedBlocks for SenseNova-U1 that only caches image-token hidden + states during denoising. """ - 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}") + @classmethod + def _is_denoising_call(cls, kwargs: dict[str, Any]) -> bool: + if kwargs.get("cache_dit_skip", False): + return False - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[pipeline.transformer.blocks], - forward_pattern=[ForwardPattern.Pattern_2], - 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, - ) + # Prefix/text forwards either omit image_gen_indicators or update the + # DynamicCache. Denoising forwards are gen-only and use update_cache=False. + if kwargs.get("update_cache", True): + return False - logger.info( - f"Enabling cache-dit on Helios transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) + exist_gen = kwargs.get("exist_gen") + exist_und = kwargs.get("exist_und") + if exist_gen is None or exist_und is None: + image_gen_indicators = kwargs.get("image_gen_indicators") + if image_gen_indicators is None: + return False + exist_gen = image_gen_indicators.any().item() + exist_und = (~image_gen_indicators).any().item() - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - 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 exist_gen and not exist_und - return refresh_cache_context + @staticmethod + def _strip_cache_only_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + kwargs = dict(kwargs) + kwargs.pop("cache_dit_skip", None) + return kwargs + def forward(self, hidden_states: torch.Tensor, *args, **kwargs): + block_kwargs = self._strip_cache_only_kwargs(kwargs) + if not self._is_denoising_call(kwargs): + hidden_states, new_encoder_hidden_states = self.call_blocks( + hidden_states, + *args, + **block_kwargs, + ) + return self._process_forward_outputs(hidden_states, new_encoder_hidden_states) -def enable_cache_for_hunyuan_video_15(pipeline: Any, cache_config: Any) -> Callable[[int], None]: - """Enable cache-dit for HunyuanVideo 1.5 pipeline. + return super().forward(hidden_states, *args, **block_kwargs) - HunyuanVideo 1.5 uses a single transformer with has_separate_cfg=True - (separate conditional/unconditional forward passes). The _sp_plan scatter - is applied at the transformer input boundary (empty-string key), so CacheDiT - sees already-sharded hidden_states throughout transformer_blocks. - """ - 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}") +class SensenovaCachedAdapter(CachedAdapter): + """Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks.""" - logger.info( - f"Enabling cache-dit on HunyuanVideo15 transformer: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) + @classmethod + def collect_unified_blocks( + cls, + block_adapter: BlockAdapter, + contexts_kwargs: list[dict], + ) -> list[dict[str, torch.nn.ModuleList]]: + BlockAdapter.assert_normalized(block_adapter) - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=pipeline.transformer.transformer_blocks, - forward_pattern=ForwardPattern.Pattern_0, - check_forward_pattern=True, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) + total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = [] + assert hasattr(block_adapter.pipe, "_context_manager") - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - 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, - ) + 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( + [ + SensenovaCachedBlocks( + # 0. Transformer blocks configuration + 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, + # 1. Cache/Prune context configuration + 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, + ) + ] + ) - return refresh_cache_context + total_cached_blocks.append(unified_blocks_bind_context) + + return total_cached_blocks -def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int], None]: +def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: """Enable cache-dit for Cosmos3. Cosmos3 has a dual-pathway architecture (UND + GEN) but only the GEN @@ -1959,41 +964,6 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int] Returns: A refresh function that can be called to update cache context with new num_inference_steps. """ - 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}") - - logger.info( - f"Enabling cache-dit on Cosmos3 gen_layers: " - f"Fn={db_cache_config.Fn_compute_blocks}, " - f"Bn={db_cache_config.Bn_compute_blocks}, " - f"W={db_cache_config.max_warmup_steps}, " - ) - - cache_dit.enable_cache( - BlockAdapter( - transformer=pipeline.transformer, - blocks=[pipeline.transformer.gen_layers], - # Cosmos3 GEN blocks return only hidden_states. Per-layer UND K/V - # conditioning uses the transformer's cache-dit fallback path. - forward_pattern=[ForwardPattern.Pattern_3], - params_modifiers=[ - ParamsModifier( - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ), - ], - check_forward_pattern=False, - has_separate_cfg=True, - ), - cache_config=db_cache_config, - calibrator_config=calibrator_config, - ) - # The T2I denoising loop skips the unconditional forward outside the # guidance interval as a speed optimization. cache-dit distinguishes the # conditional vs unconditional passes purely by transformer-forward parity @@ -2001,25 +971,40 @@ def enable_cache_for_cosmos3(pipeline: Any, cache_config: Any) -> Callable[[int] # step accounting. Still do both cond/uncond CFG steps when cache-dit is active. # CFG is instead neutralized via scale=1.0 outside the interval. pipeline._cache_dit_requires_paired_cfg = True + block_adapter = CacheDiTBackend.maybe_build_block_adapter(pipeline) + return enable_cache_for_dit(pipeline, cache_config, block_adapter) - def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: - 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_krea2(pipeline: Any, cache_config: Any) -> RefreshCacheContextFunc: + """Enable cache-dit for Krea 2. + + Krea 2 is a single-stream MMDiT: each ``Krea2TransformerBlock`` takes and returns only + ``hidden_states`` (text is fused into the token stream), so the blocks follow + ``ForwardPattern.Pattern_3``. + + ``has_separate_cfg`` is checkpoint-dependent, which is why this needs a custom enabler + rather than a static ``_cache_dit_adapter_config``: the distilled Turbo checkpoint runs + no-CFG (a single transformer forward per denoise step), while the Raw checkpoint runs CFG + as two separate forwards. cache-dit tells cond/uncond apart purely by transformer-forward + parity, so the flag must match the actual per-step forward count — ``True`` only for the + CFG (Raw) path. The pipeline exposes this via ``is_distilled`` (read from ``model_index.json``). + + Args: + pipeline: The Krea2Pipeline instance. + cache_config: DiffusionCacheConfig instance with cache configuration. + + Returns: + A refresh function that can be called to update cache context with new num_inference_steps. + """ + transformer = default_get_pipeline_transformer(pipeline) + block_adapter = BlockAdapter( + transformer=transformer, + blocks=[transformer.transformer_blocks], + forward_pattern=[ForwardPattern.Pattern_3], + has_separate_cfg=not pipeline.is_distilled, + check_forward_pattern=False, + ) + return enable_cache_for_dit(pipeline, cache_config, block_adapter) # Register custom cache-dit enablers after function definitions @@ -2027,30 +1012,12 @@ def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool { "Wan22Pipeline": enable_cache_for_wan22, "Wan22I2VPipeline": enable_cache_for_wan22, + "Wan22TI2VPipeline": enable_cache_for_wan22, + "Wan22VACEPipeline": 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, - "LongCatImagePipeline": enable_cache_for_longcat_image, - "LongCatImageEditPipeline": enable_cache_for_longcat_image, - "StableAudioPipeline": enable_cache_for_stable_audio_open, - "StableDiffusion3Pipeline": enable_cache_for_sd3, - "LTX2Pipeline": enable_cache_for_ltx2, - "LTX2ImageToVideoPipeline": enable_cache_for_ltx2, - "LTX2TwoStagesPipeline": enable_cache_for_ltx2, - "LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2, - "LTX23Pipeline": enable_cache_for_ltx2, - "LTX23ImageToVideoPipeline": enable_cache_for_ltx2, - "BagelPipeline": enable_cache_for_bagel, - "SenseNovaU1Pipeline": enable_cache_for_sensenova_u1, - "GlmImagePipeline": enable_cache_for_glm_image, - "Flux2Pipeline": enable_cache_for_flux2, - "DreamIDOmniPipeline": enable_cache_for_dreamid_omni, - "ErnieImagePipeline": enable_cache_for_ernie_image, - "HunyuanVideo15Pipeline": enable_cache_for_hunyuan_video_15, - "HunyuanVideo15I2VPipeline": enable_cache_for_hunyuan_video_15, - "HeliosPipeline": enable_cache_for_helios, "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, + "Cosmos3OmniPipeline": enable_cache_for_cosmos3, + "Krea2Pipeline": enable_cache_for_krea2, } ) @@ -2092,6 +1059,43 @@ def __init__(self, cache_config: Any = None): self._refresh_func: Callable[[Any, int, bool], None] | None = None self._last_num_inference_steps: int | None = None + @staticmethod + def maybe_build_block_adapter(pipeline) -> BlockAdapter | None: + """If a module defines `_cache_dit_adapter_config`, build the corresponding + block adapter. + """ + transformer = default_get_pipeline_transformer(pipeline) + + adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) + if adapter_cfg is None: + return None + + block_attrs, forward_pattern = zip(*(adapter_cfg.block_forward_patterns).items()) + missing_attrs = [block_attr for block_attr in block_attrs if not hasattr(transformer, block_attr)] + + if missing_attrs: + logger.warning("Missing Cache DiT block attributes: %s", missing_attrs) + + block_adapter = BlockAdapter( + transformer=transformer, + blocks=[getattr(transformer, block_attr) for block_attr in block_attrs], + forward_pattern=list(forward_pattern), + has_separate_cfg=adapter_cfg.has_separate_cfg, + check_forward_pattern=adapter_cfg.check_forward_pattern, + ) + return block_adapter + + @staticmethod + def maybe_get_cached_adapter_cls(pipeline) -> type[CachedAdapter] | None: + """If a module has a custom cached adapter type registered, e.g., SenseNova, retrieve it + from the transformer's CacheDiTAdapterConfig.""" + transformer = default_get_pipeline_transformer(pipeline) + + adapter_cfg: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) + if adapter_cfg is None: + return None + return adapter_cfg.cached_adapter_cls + def enable(self, pipeline: Any) -> None: """Enable cache-dit on the pipeline if configured. @@ -2110,8 +1114,12 @@ def enable(self, pipeline: Any) -> None: logger.info(f"Using custom cache-dit enabler for model: {pipeline_name}") self._refresh_func = CUSTOM_DIT_ENABLERS[pipeline_name](pipeline, self.config) else: - # For regular single-transformer models - self._refresh_func = enable_cache_for_dit(pipeline, self.config) + # Common case; either the model doesn't explicitly support dit cache yet, + # Or it defines its _cache_dit_adapter_config, which describes how we should + # create its block adapter. + block_adapter = self.maybe_build_block_adapter(pipeline) + adapter_cls = self.maybe_get_cached_adapter_cls(pipeline) + self._refresh_func = enable_cache_for_dit(pipeline, self.config, block_adapter, adapter_cls) self.enabled = True logger.info(f"Cache-dit enabled successfully on {pipeline_name}") diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py index e6c2e70e318..3b002abe0bf 100644 --- a/vllm_omni/diffusion/cache/selector.py +++ b/vllm_omni/diffusion/cache/selector.py @@ -3,6 +3,7 @@ from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend +from vllm_omni.diffusion.cache.stepcache import StepCacheBackend from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig @@ -14,14 +15,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface - mag_cache: Uses MagCacheBackend with enable()/refresh() interface + - step_cache: Uses StepCacheBackend for DreamZero velocity step skip Args: - cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None). + cache_backend: Cache backend name ("cache_dit", "tea_cache", + "mag_cache", "step_cache", or None). cache_config: Cache configuration (dict or DiffusionCacheConfig instance). Returns: - Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend) - if cache_backend is set, None otherwise. + Cache backend instance if cache_backend is set, None otherwise. Raises: ValueError: If cache_backend is unsupported. @@ -38,7 +40,10 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack return TeaCacheBackend(cache_config) elif cache_backend == "mag_cache": return MagCacheBackend(cache_config) + elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"): + return StepCacheBackend(cache_config) else: raise ValueError( - f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'" + f"Unsupported cache backend: {cache_backend}. " + "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'" ) diff --git a/vllm_omni/diffusion/cache/stepcache/__init__.py b/vllm_omni/diffusion/cache/stepcache/__init__.py new file mode 100644 index 00000000000..a7e1852b0d4 --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/__init__.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +StepCache: velocity-based step skipping for DreamZero-style DiT denoising. + +Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the +pipeline denoise loop when successive velocity predictions are highly aligned. + +Usage: + from vllm_omni import Omni + + omni = Omni( + model="...", + cache_backend="step_cache", + ) + + # Alternative: environment variable + # export DIFFUSION_CACHE_BACKEND=step_cache +""" + +from vllm_omni.diffusion.cache.stepcache.backend import ( + CUSTOM_STEP_CACHE_DIT_ENABLERS, + CUSTOM_STEPCACHE_ENABLERS, + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_DIT_CONFIG_ATTR, + STEP_CACHE_DIT_STATE_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheBackend, + StepCacheDiTBackend, + StepCacheDiTConfig, + StepCacheDiTState, + enable_dreamzero_step_cache_dit, + enable_dreamzero_stepcache, + get_step_cache_dit_config, + get_step_cache_dit_state, + get_stepcache_config, + get_stepcache_state, + is_step_cache_dit_active, + is_stepcache_active, +) +from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig +from vllm_omni.diffusion.cache.stepcache.state import StepCacheState + +__all__ = [ + "CUSTOM_STEP_CACHE_DIT_ENABLERS", + "CUSTOM_STEPCACHE_ENABLERS", + "STEP_CACHE_CONFIG_ATTR", + "STEP_CACHE_DIT_CONFIG_ATTR", + "STEP_CACHE_DIT_STATE_ATTR", + "STEP_CACHE_STATE_ATTR", + "StepCacheBackend", + "StepCacheConfig", + "StepCacheDiTBackend", + "StepCacheDiTConfig", + "StepCacheDiTState", + "StepCacheState", + "enable_dreamzero_step_cache_dit", + "enable_dreamzero_stepcache", + "get_step_cache_dit_config", + "get_step_cache_dit_state", + "get_stepcache_config", + "get_stepcache_state", + "is_step_cache_dit_active", + "is_stepcache_active", +] diff --git a/vllm_omni/diffusion/cache/stepcache/backend.py b/vllm_omni/diffusion/cache/stepcache/backend.py new file mode 100644 index 00000000000..517d2369a3c --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/backend.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +StepCache backend implementation. + +This module provides the stepcache backend that implements the CacheBackend +interface for DreamZero-style velocity-based step skipping in the denoise loop. +""" + +from __future__ import annotations + +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.stepcache.config import ( + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheConfig, +) +from vllm_omni.diffusion.cache.stepcache.state import StepCacheState +from vllm_omni.diffusion.data import DiffusionCacheConfig + +logger = init_logger(__name__) + +# Backward-compatible aliases for older imports and pipeline attribute names. +STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR +STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR +StepCacheDiTConfig = StepCacheConfig +StepCacheDiTState = StepCacheState + + +def is_stepcache_active(pipeline: Any) -> bool: + """Return True when stepcache is enabled on *pipeline*.""" + cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None) + return isinstance(cfg, StepCacheConfig) and cfg.enabled + + +def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None: + cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None) + return cfg if isinstance(cfg, StepCacheConfig) else None + + +def get_stepcache_state(pipeline: Any) -> StepCacheState | None: + state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None) + return state if isinstance(state, StepCacheState) else None + + +def is_step_cache_dit_active(pipeline: Any) -> bool: + return is_stepcache_active(pipeline) + + +def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None: + return get_stepcache_config(pipeline) + + +def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None: + return get_stepcache_state(pipeline) + + +def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """Enable stepcache for DreamZeroPipeline.""" + cache_config = StepCacheConfig.from_diffusion_cache_config(config) + setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config) + setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config)) + + logger.info( + "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)", + cache_config.enabled, + cache_config.sim_thresholds, + cache_config.skip_countdowns, + ) + + +enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache + +CUSTOM_STEPCACHE_ENABLERS = { + "DreamZeroPipeline": enable_dreamzero_stepcache, +} + +CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS + + +class StepCacheBackend(CacheBackend): + """ + Velocity cosine step-skipping cache backend for DreamZero. + + Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported + pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to + decide whether to run ``predict_noise``. + + Example: + >>> from vllm_omni.diffusion.data import DiffusionCacheConfig + >>> backend = StepCacheBackend(DiffusionCacheConfig()) + >>> backend.enable(pipeline) + >>> backend.refresh(pipeline, num_inference_steps=16) + """ + + def enable(self, pipeline: Any) -> None: + """Enable stepcache on a supported pipeline.""" + pipeline_type = pipeline.__class__.__name__ + + if pipeline_type in CUSTOM_STEPCACHE_ENABLERS: + logger.info("Using custom stepcache enabler for model: %s", pipeline_type) + CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config) + else: + raise ValueError( + f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}" + ) + + self.enabled = True + + def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh stepcache state for a new generation.""" + state = get_stepcache_state(pipeline) + if state is not None: + state.reset() + if verbose: + logger.debug( + "stepcache state refreshed (num_inference_steps=%d)", + num_inference_steps, + ) + elif verbose and is_stepcache_active(pipeline): + logger.warning("stepcache config active but state not found on pipeline") + + +StepCacheDiTBackend = StepCacheBackend diff --git a/vllm_omni/diffusion/cache/stepcache/config.py b/vllm_omni/diffusion/cache/stepcache/config.py new file mode 100644 index 00000000000..00a7984ae7d --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/config.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration for step-level DiT velocity caching.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from vllm_omni.diffusion.data import DiffusionCacheConfig + +# Pipeline attributes set by StepCacheBackend.enable(). +STEP_CACHE_CONFIG_ATTR = "_stepcache_config" +STEP_CACHE_STATE_ATTR = "_stepcache_state" + + +@dataclass(frozen=True) +class StepCacheConfig: + """Runtime config for velocity-based step skipping. + + Skips entire DiT forwards when successive video velocity predictions are + highly aligned (cosine similarity above configured thresholds). + + Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``. + """ + + enabled: bool = True + min_history_steps: int = 2 + max_history: int = 2 + sim_thresholds: tuple[float, ...] = (0.95, 0.93) + skip_countdowns: tuple[int, ...] = (4, 2) + + @classmethod + def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig: + enabled = config.step_cache_dit_enabled + + thresholds = tuple(config.velocity_sim_thresholds) + countdowns = tuple(config.velocity_skip_countdowns) + if len(thresholds) != len(countdowns): + raise ValueError( + "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; " + f"got {len(thresholds)} and {len(countdowns)}." + ) + + return cls( + enabled=enabled, + min_history_steps=config.step_cache_dit_min_history, + max_history=config.step_cache_dit_max_history, + sim_thresholds=thresholds, + skip_countdowns=countdowns, + ) diff --git a/vllm_omni/diffusion/cache/stepcache/state.py b/vllm_omni/diffusion/cache/stepcache/state.py new file mode 100644 index 00000000000..db786d4f210 --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/state.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Mutable state for step-level DiT velocity caching.""" + +from __future__ import annotations + +import torch + +from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig + + +class StepCacheState: + """Per-generation mutable state (skip countdown).""" + + def __init__(self, config: StepCacheConfig) -> None: + self.config = config + self.skip_countdown = 0 + + def reset(self) -> None: + self.skip_countdown = 0 + + def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool: + """Return True when the DiT forward should execute at this scheduler step.""" + if not self.config.enabled: + return True + + if len(prev_predictions) < self.config.min_history_steps: + return True + + if self.skip_countdown > 1: + self.skip_countdown -= 1 + return False + if self.skip_countdown == 1: + self.skip_countdown = 0 + return True + + v_last = prev_predictions[-1][0].flatten(1).float() + v_prev = prev_predictions[-2][0].flatten(1).float() + sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean() + + for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns): + if sim > threshold: + self.skip_countdown = countdown + return False + + return True + + def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None: + while len(prev_predictions) > self.config.max_history: + prev_predictions.pop(0) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 772dec78913..f6f9db88dc4 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -57,6 +57,22 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: ) +def enable_sensenova_u1_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """Enable TeaCache for SenseNova-U1 denoising forwards.""" + teacache_config = TeaCacheConfig( + transformer_type="SenseNovaU1ForCausalLM", + rel_l1_thresh=config.rel_l1_thresh, + coefficients=config.coefficients, + ) + transformer = pipeline.denoising_transformer + apply_teacache_hook(transformer, teacache_config) + + logger.info( + f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " + f"transformer_class={teacache_config.transformer_type}" + ) + + def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: """ Enable TeaCache for Flux2 Klein model. @@ -80,6 +96,7 @@ def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> "BagelPipeline": enable_bagel_teacache, "Flux2KleinPipeline": enable_flux2_klein_teacache, "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache, + "SenseNovaU1Pipeline": enable_sensenova_u1_teacache, } @@ -180,6 +197,8 @@ def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) # Extract transformer from pipeline transformer = pipeline.transformer + if not hasattr(transformer, "_hook_registry") and hasattr(pipeline, "denoising_transformer"): + transformer = pipeline.denoising_transformer if hasattr(transformer, "_hook_registry"): hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME) diff --git a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py index 60638cca876..49ba806d55b 100644 --- a/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py +++ b/vllm_omni/diffusion/cache/teacache/coefficient_estimator.py @@ -201,15 +201,17 @@ def __init__( def collect_from_prompt(self, prompt: str, **generate_kwargs): self.hook.start_collection() req = OmniDiffusionRequest( - prompts=[prompt], + prompt=prompt, request_id="teacache-coefficient-estimator", sampling_params=OmniDiffusionSamplingParams( num_inference_steps=generate_kwargs.get("num_inference_steps", 20), seed=generate_kwargs.get("seed", 42), ), ) + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + with torch.no_grad(): - self.pipeline.forward(req) + self.pipeline.forward(DiffusionRequestBatch(requests=[req])) trajectory = self.hook.stop_collection() if trajectory: self.collected_data.append(trajectory) diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 7efdd418e12..bcf2d3fe4a1 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -37,6 +37,14 @@ # Bagel transformer coefficients # Using Qwen's coefficients as reasonable default given shared architecture "Bagel": [1.33313129e06, -1.68644226e05, 7.95050740e03, -1.63747873e02, 1.26352397e00], + # SenseNova-U1 transformer coefficients + "SenseNovaU1ForCausalLM": [ + 9.07281930e04, + -2.17699186e04, + 1.83940990e03, + -6.30339273e01, + 7.61309272e-01, + ], # Z-Image transformer coefficients # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future "ZImageTransformer2DModel": [ diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 1f0c38aee58..b0bd58bc8d4 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -1187,6 +1187,68 @@ def postprocess(h: torch.Tensor) -> Any: ) +def extract_sensenova_u1_context( + module: nn.Module, + input_ids: torch.Tensor | None = None, + indexes: torch.Tensor | None = None, + attention_mask: dict[str, torch.Tensor] | torch.Tensor | None = None, + past_key_values: Any | None = None, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + embed_only: bool = False, + compute_logits: bool = True, + **kwargs: Any, +) -> CacheContext: + """Extract cache context for SenseNovaU1ForCausalLM denoising forwards.""" + from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import ( + SenseNovaU1CausalLMOutput, + ) + + layer_kwargs = dict(kwargs) + layer_kwargs.pop("cache_dit_skip", None) + image_gen_indicators = layer_kwargs.pop("image_gen_indicators", None) + exist_und = (~image_gen_indicators).any().item() + exist_gen = image_gen_indicators.any().item() + causal_mask_mapping = attention_mask + + first_layer = module.model.layers[0] + modulated_input = first_layer.input_layernorm_mot_gen(inputs_embeds) + + def run_transformer_blocks(): + h = inputs_embeds + for layer in module.model.layers: + h = layer( + h, + image_gen_indicators=image_gen_indicators, + exist_und=exist_und, + exist_gen=exist_gen, + indexes=indexes, + attention_mask=causal_mask_mapping, + past_key_values=past_key_values, + **layer_kwargs, + ) + return (h,) + + def postprocess(h: torch.Tensor) -> SenseNovaU1CausalLMOutput: + h = module.model.norm_mot_gen(h) + + logits = module.logits_processor(module.lm_head, h) if compute_logits else None + return SenseNovaU1CausalLMOutput( + logits=logits, + past_key_values=past_key_values if use_cache else None, + hidden_states=h, + ) + + return CacheContext( + modulated_input=modulated_input, + hidden_states=inputs_embeds, + encoder_hidden_states=None, + temb=modulated_input, + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + ) + + # Registry for model-specific extractors # Key: Transformer class name # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext @@ -1202,6 +1264,7 @@ def postprocess(h: torch.Tensor) -> Any: "Flux2Transformer2DModel": extract_flux2_context, "LongCatImageTransformer2DModel": extract_longcat_context, "FluxTransformer2DModel": extract_flux_context, + "SenseNovaU1ForCausalLM": extract_sensenova_u1_context, # Future models: # "CogVideoXTransformer3DModel": extract_cogvideox_context, } diff --git a/vllm_omni/diffusion/compile.py b/vllm_omni/diffusion/compile.py index 66a4675f3ce..45216febc0e 100644 --- a/vllm_omni/diffusion/compile.py +++ b/vllm_omni/diffusion/compile.py @@ -2,13 +2,37 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import Any +import torch import torch.nn as nn from vllm.logger import init_logger logger = init_logger(__name__) -def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: Any) -> nn.Module: +def _matches_repeated_block( + name: str, + module: nn.Module, + repeated_blocks: list[str], + repeated_block_attrs: list[str], +) -> bool: + class_name = module.__class__.__name__ + if class_name in repeated_blocks: + return True + + for attr in ("_fsdp_wrapped_module", "module", "_orig_mod"): + wrapped = getattr(module, attr, None) + if wrapped is not None and wrapped.__class__.__name__ in repeated_blocks: + return True + + parts = name.split(".") + return len(parts) >= 2 and parts[-2] in repeated_block_attrs and parts[-1].isdigit() + + +def regionally_compile( + model: nn.Module, + *compile_args: Any, + **compile_kwargs: Any, +) -> nn.Module: """ Apply regional compilation to a PyTorch model. @@ -27,15 +51,26 @@ def regionally_compile(model: nn.Module, *compile_args: Any, **compile_kwargs: A logger.warning("Regional compilation skipped because the model does not define `_repeated_blocks`.") return model + repeated_block_attrs = getattr(model, "_layerwise_offload_blocks_attrs", []) + # Check if we have modules with the specified class names has_compiled_region = False - for submod in model.modules(): - if submod.__class__.__name__ in repeated_blocks: - # Compile this submodule - submod.compile(*compile_args, **compile_kwargs) + compiled_region_count = 0 + for name, submod in model.named_modules(): + if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs): + # Compile the block compute while keeping nn.Module.__call__ hooks + # outside the compiled graph. + submod.forward = torch.compile(submod.forward, *compile_args, **compile_kwargs) has_compiled_region = True + compiled_region_count += 1 if not has_compiled_region: logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.") + else: + logger.info( + "Regional compilation applied to %d module(s) for repeated blocks %s.", + compiled_region_count, + repeated_blocks, + ) return model diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 00f3b0b5a1e..e131b833b17 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -6,6 +6,7 @@ import random from collections.abc import Callable, Mapping from dataclasses import dataclass, field, fields +from enum import Enum from typing import TYPE_CHECKING, Any import diffusers @@ -33,6 +34,62 @@ logger = init_logger(__name__) +def normalize_omni_diffusion_kwargs(kwargs: Mapping[str, Any]) -> dict[str, Any]: + """Normalize legacy diffusion kwargs before config construction.""" + normalized = dict(kwargs) + + # Backwards-compatibility: older callers may use a diffusion-specific + # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale". + if "static_lora_scale" in normalized: + if "lora_scale" not in normalized: + normalized["lora_scale"] = normalized["static_lora_scale"] + normalized.pop("static_lora_scale", None) + + # Backwards-compatibility: map "quantization" to "quantization_config" + # so callers using the old field name still work. + if "quantization" in normalized and normalized.get("quantization_config", None) is None: + normalized["quantization_config"] = normalized.pop("quantization") + else: + normalized.pop("quantization", None) + + # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype. + if normalized.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in normalized: + normalized["diffusion_kv_cache_dtype"] = normalized.pop("kv_cache_dtype") + else: + normalized.pop("kv_cache_dtype", None) + if normalized.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in normalized: + normalized["diffusion_kv_cache_skip_steps"] = normalized.pop("kv_cache_skip_steps") + else: + normalized.pop("kv_cache_skip_steps", None) + if normalized.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in normalized: + normalized["diffusion_kv_cache_skip_layers"] = normalized.pop("kv_cache_skip_layers") + else: + normalized.pop("kv_cache_skip_layers", None) + + # Handle "diffusion_attention_backend" shorthand: merge into + # diffusion_attention_config before field filtering. + diffusion_attn_backend = normalized.pop("diffusion_attention_backend", None) + if diffusion_attn_backend is not None: + existing = normalized.get("diffusion_attention_config") + normalized["diffusion_attention_config"] = parse_attention_config( + existing, + attention_backend=diffusion_attn_backend, + ) + + # Check environment variable as fallback for cache_backend. + # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND. + if "cache_backend" not in normalized: + cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") + normalized["cache_backend"] = cache_backend.lower() if cache_backend else "none" + + # Convert optional YAML null values to empty containers. + for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"): + if key in normalized and normalized[key] is None: + normalized[key] = {} + + return normalized + + def parse_kv_cache_skip_selector( selector: str | list[int] | tuple[int, ...] | set[int] | None, ) -> set[int] | None: @@ -127,6 +184,21 @@ class DiffusionParallelConfig: vae_patch_parallel_size: int = 1 """Number of ranks used for VAE patch/tile parallelism (decode/encode).""" + vae_parallel_mode: str = "tile" + """VAE parallel decode strategy. + + - "tile": Patch/tile parallel decode (default). Each rank decodes a subset + of spatial tiles and the results are stitched on rank 0. + - "spatial_shard_height": Spatially-sharded decode that splits decoder + feature maps along height and exchanges halo rows around spatial + convolutions. + - "spatial_shard_width": Same as "spatial_shard_height" but sharded along width. + + The "spatial_shard_*" modes are decode-only and currently require + ``vae_patch_parallel_size`` to match the DiT group size; otherwise the VAE + falls back to tile-parallel decode at runtime. + """ + use_hsdp: bool = False """Enable Hybrid Sharded Data Parallel (HSDP) for model weight sharding.""" @@ -158,6 +230,10 @@ def _validate_parallel_config(self) -> Self: f"CFG parallel size must be 1, 2, or 3, but got {self.cfg_parallel_size}" ) assert self.vae_patch_parallel_size > 0, "VAE patch parallel size must be > 0" + assert self.vae_parallel_mode in {"tile", "spatial_shard_height", "spatial_shard_width"}, ( + "vae_parallel_mode must be one of {'tile', 'spatial_shard_height', 'spatial_shard_width'}, " + f"but got {self.vae_parallel_mode!r}." + ) assert self.sequence_parallel_size == self.ulysses_degree * self.ring_degree, ( "Sequence parallel size must be equal to the product of ulysses degree and ring degree," f" but got {self.sequence_parallel_size} != {self.ulysses_degree} * {self.ring_degree}" @@ -302,6 +378,8 @@ class DiffusionCacheConfig: scm_steps_mask_policy, scm_steps_policy - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio, mag_ratios, mag_calibrate + - step_cache: step_cache_dit_enabled, velocity_sim_thresholds, + velocity_skip_countdowns, step_cache_dit_min_history Example: >>> # From dict (user-facing API) - partial config uses defaults for missing keys @@ -375,6 +453,13 @@ class DiffusionCacheConfig: # "repeat" refreshes every force_refresh_step_hint steps. force_refresh_step_policy: str = "once" + # step_cache parameters [step_cache only] — DreamZero velocity schedule + step_cache_dit_enabled: bool = True + velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93]) + velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2]) + step_cache_dit_min_history: int = 2 + step_cache_dit_max_history: int = 2 + # Additional parameters that may be passed but not explicitly defined _extra_params: dict[str, Any] = field(default_factory=dict, repr=False) @@ -423,6 +508,67 @@ def __getattr__(self, item: str) -> Any: raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{item}'") +def resolve_model_class_name(model: str | None, diffusion_load_format: str = "default") -> str | None: + """Resolve the diffusion pipeline class name from the model config. + + Read-only counterpart of ``OmniDiffusionConfig.enrich_config``, safe to call + client-side. Returns ``None`` if the pipeline can't be determined. + """ + from vllm.transformers_utils.config import get_hf_file_to_dict + + if not model: + return None + + is_lance_subfolder = os.path.basename(str(model).rstrip("/")) in {"Lance_3B", "Lance_3B_Video"} + + # Diffusers models: read _class_name from model_index.json. + try: + model_index = get_hf_file_to_dict("model_index.json", model) + except Exception: + model_index = None + if model_index is not None: + return model_index.get("_class_name") + if diffusion_load_format == "diffusers": + return "DiffusersAdapterPipeline" + + # Other models: map model_type / architecture from config.json. + try: + cfg = get_hf_file_to_dict("config.json", model) or {} + except Exception: + cfg = {} + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + return "BagelPipeline" + if ( + model_type == "lance" + or "LancePipeline" in architectures + or cfg.get("model_name") == "Lance" + or is_lance_subfolder + ): + return "LancePipeline" + if model_type == "neo_chat": + return "SenseNovaU1Pipeline" + if "BailingMM2NativeForConditionalGeneration" in architectures or model_type in ( + "bailingmm_moe_v2_lite", + "ming_flash_omni", + "ming_flash_omni_thinker", + ): + return "MingImagePipeline" + if model_type == "nextstep": + return "NextStep11Pipeline" + if model_type == "s2v": + return "WanS2VPipeline" + if model_type == "vla": + from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero + + return "DreamZeroPipeline" if _looks_like_dreamzero(model) else None + if len(architectures) == 1: + return architectures[0] + return None + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -468,6 +614,20 @@ class OmniDiffusionConfig: distributed_executor_backend: str = "mp" nccl_port: int | None = None + # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class`` + # (mirrors ``DiffusionExecutor.get_class``). Config files use a string: + # "default" -> DiffusionEngine, or an import path (set e.g. by a deploy + # config). Programmatic callers may pass a DiffusionEngine subclass + # directly — hence ``str | type`` (structured-config mirrors should expose + # the string form only). + engine_backend: str | type = "default" + + # Optional override for the diffusion model runner class (import path). + # Precedence in the worker: this override > the runner declared by the + # selected engine class (``default_diffusion_model_runner_cls``) > the + # platform default. Never mutated by engines. + diffusion_model_runner_cls: str | None = None + # HuggingFace specific parameters trust_remote_code: bool = False revision: str | None = None @@ -504,6 +664,9 @@ class OmniDiffusionConfig: # STA_mode: STA_Mode = STA_Mode.STA_INFERENCE skip_time_steps: int = 15 + # MoE kernel backend selection + moe_backend: str = "auto" + # Compilation enforce_eager: bool = False @@ -619,9 +782,18 @@ class OmniDiffusionConfig: # Step mode settings step_execution: bool = False + # Streaming mode settings + streaming_output: bool = False # Start (video) generation with initial prompt, but streaming output in chunks + # Maximum number of sequences to generate in a batch max_num_seqs: int = 1 + # Request-mode batch admission: wait briefly for compatible requests to + # accumulate in the scheduler waiting queue before the first schedule() of + # a wave. Improves fused forward batch sizes under bursty HTTP ingress. + # 0 disables admission (default; no added latency). + request_batch_max_wait_ms: float = 0.0 + # Supplementary model specific parameters extras: dict[str, Any] = Field(default_factory=dict) @@ -694,6 +866,9 @@ def settle_port(self, port: int, port_inc: int = 42, max_attempts: int = 100) -> def __post_init__(self): self.master_port = self._resolve_master_port() + self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0) + if self.request_batch_max_wait_ms < 0: + raise ValueError(f"request_batch_max_wait_ms must be non-negative, got {self.request_batch_max_wait_ms}.") if isinstance(self.profiler_config, dict): from vllm.config import ProfilerConfig @@ -903,7 +1078,12 @@ def enrich_config(self) -> None: ) else: tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) - self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) + if tf_config_dict is None: + tf_config_dict = get_hf_file_to_dict("unet/config.json", self.model) + if tf_config_dict is not None: + self.set_tf_model_config(TransformerConfig.from_dict(tf_config_dict)) + else: + self.set_tf_model_config(TransformerConfig()) else: raise FileNotFoundError("model_index.json not found") except (AttributeError, OSError, ValueError, FileNotFoundError): @@ -978,6 +1158,10 @@ def enrich_config(self) -> None: self.model_class_name = "WanS2VPipeline" self.tf_model_config = TransformerConfig() self.update_multimodal_support() + elif model_type == "Gr00tN1d7" or "Gr00tN1d7" in architectures: + self.model_class_name = "Gr00tN1d7Pipeline" + self.set_tf_model_config(TransformerConfig()) + self.update_multimodal_support() elif model_type == "vla": from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero @@ -1001,61 +1185,14 @@ def enrich_config(self) -> None: @classmethod def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": - # Backwards-compatibility: older callers may use a diffusion-specific - # "static_lora_scale" kwarg. Normalize it to the canonical "lora_scale" - # before constructing the dataclass to avoid TypeError on unknown fields. - if "static_lora_scale" in kwargs: - if "lora_scale" not in kwargs: - kwargs["lora_scale"] = kwargs["static_lora_scale"] - kwargs.pop("static_lora_scale", None) - - # Backwards-compatibility: map "quantization" to "quantization_config" - # so callers using the old field name still work. - if "quantization" in kwargs and kwargs.get("quantization_config", None) is None: - kwargs["quantization_config"] = kwargs.pop("quantization") - else: - kwargs.pop("quantization", None) - - # Renamed from kv_cache_* to avoid clashing with vLLM's --kv-cache-dtype. - if kwargs.get("diffusion_kv_cache_dtype") is None and "kv_cache_dtype" in kwargs: - kwargs["diffusion_kv_cache_dtype"] = kwargs.pop("kv_cache_dtype") - else: - kwargs.pop("kv_cache_dtype", None) - if kwargs.get("diffusion_kv_cache_skip_steps") is None and "kv_cache_skip_steps" in kwargs: - kwargs["diffusion_kv_cache_skip_steps"] = kwargs.pop("kv_cache_skip_steps") - else: - kwargs.pop("kv_cache_skip_steps", None) - if kwargs.get("diffusion_kv_cache_skip_layers") is None and "kv_cache_skip_layers" in kwargs: - kwargs["diffusion_kv_cache_skip_layers"] = kwargs.pop("kv_cache_skip_layers") - else: - kwargs.pop("kv_cache_skip_layers", None) - - # Handle "diffusion_attention_backend" shorthand: merge into - # diffusion_attention_config before field filtering. - diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None) - if diffusion_attn_backend is not None: - existing = kwargs.get("diffusion_attention_config") - kwargs["diffusion_attention_config"] = parse_attention_config( - existing, - attention_backend=diffusion_attn_backend, - ) - - # Check environment variable as fallback for cache_backend - # Support both old DIFFUSION_CACHE_ADAPTER and new DIFFUSION_CACHE_BACKEND for backwards compatibility - if "cache_backend" not in kwargs: - cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") - kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" - - # Convert optional YAML null values to empty containers. - for key in ("diffusers_load_kwargs", "diffusers_call_kwargs"): - if key in kwargs and kwargs[key] is None: - kwargs[key] = {} + kwargs = normalize_omni_diffusion_kwargs(kwargs) # Filter kwargs to only include valid fields valid_fields = {f.name for f in fields(cls)} filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields} - return cls(**filtered_kwargs) + instance = cls(**filtered_kwargs) + return instance @dataclass @@ -1064,8 +1201,10 @@ class DiffusionOutput: Final output (after pipeline completion) """ - # Fields may be replaced with SHM handle dicts by ipc.pack_diffusion_output_shm output: torch.Tensor | tuple[Any, ...] | dict[str, Any] | None = None + + # Legacy compatibility fields. New pipeline-specific payloads should be + # carried by output["payload"] instead. trajectory_timesteps: torch.Tensor | dict[str, Any] | None = None trajectory_latents: torch.Tensor | dict[str, Any] | None = None trajectory_log_probs: torch.Tensor | dict[str, Any] | None = None @@ -1078,23 +1217,23 @@ class DiffusionOutput: post_process_func: Callable[..., Any] | None = None - # Extra custom output data (e.g. latent trajectories, prompt embeds) - # passed through to OmniRequestOutput.custom_output - custom_output: dict[str, Any] = field(default_factory=dict) - # logged timings info, directly from Req.timings # timings: Optional["RequestTimings"] = None + # Streaming info (the defaults should make sense for non-streaming mode) + finished: bool = True + chunk_index: int = 0 + total_chunks: int = 1 + # logged duration of stages stage_durations: dict[str, float] = field(default_factory=dict) # memory usage info peak_memory_mb: float = 0.0 - # When True, move all tensor fields (including tensors inside - # ``custom_output``) to CPU at construction time. Useful when the output - # is shipped across process boundaries (e.g. step-execution mode) and the - # receiving side must not initialise a stray CUDA context. + # When True, move tensor fields to CPU at construction time. Useful when + # the output is shipped across process boundaries (e.g. step-execution + # mode) and the receiving side must not initialise a stray CUDA context. to_cpu: bool = False def __post_init__(self) -> None: @@ -1104,14 +1243,18 @@ def __post_init__(self) -> None: def _maybe_to_cpu(value: Any) -> Any: if isinstance(value, torch.Tensor): return value.detach().cpu() + if isinstance(value, dict): + return {key: _maybe_to_cpu(item) for key, item in value.items()} + if isinstance(value, list): + return [_maybe_to_cpu(item) for item in value] + if isinstance(value, tuple): + return tuple(_maybe_to_cpu(item) for item in value) return value self.output = _maybe_to_cpu(self.output) self.trajectory_timesteps = _maybe_to_cpu(self.trajectory_timesteps) self.trajectory_latents = _maybe_to_cpu(self.trajectory_latents) self.trajectory_log_probs = _maybe_to_cpu(self.trajectory_log_probs) - 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": @@ -1350,5 +1493,12 @@ class OmniWakeTask: tags: list[str] | None = None +class CuMemTag(str, Enum): + """Tags representing specific CuMem allocations for sleep/wake state tracking.""" + + WEIGHTS = "weights" + KV_CACHE = "kv_cache" + + # Special message broadcast via scheduler queues to signal worker shutdown. SHUTDOWN_MESSAGE = {"type": "shutdown"} diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index b3405e3bb5a..d6e7b35dec6 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -9,14 +9,15 @@ import queue import threading import time -from collections.abc import Iterable +from collections.abc import AsyncGenerator, Iterable from dataclasses import dataclass, field -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import PIL.Image import torch from vllm.logger import init_logger +from vllm.utils.import_utils import resolve_obj_by_qualname from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import ( @@ -25,9 +26,20 @@ OmniDiffusionConfig, ) from vllm_omni.diffusion.executor.abstract import DiffusionExecutor +from vllm_omni.diffusion.io_support import ( + get_dummy_run_num_frames, + image_color_format, + supports_audio_output, + supports_multimodal_input, +) +from vllm_omni.diffusion.output_formatter import ( + DiffusionStepTimings, + format_diffusion_outputs, + format_empty_diffusion_outputs, + normalize_diffusion_postprocess_output, +) from vllm_omni.diffusion.registry import ( DiffusionModelRegistry, - get_diffusion_action_post_process_func, get_diffusion_post_process_func, get_diffusion_pre_process_func, ) @@ -37,61 +49,61 @@ 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 + +if TYPE_CHECKING: + from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) +__all__ = [ + "DiffusionEngine", + "_RpcTask", + "_move_tensor_tree_to_cpu", + "get_dummy_run_num_frames", + "image_color_format", + "supports_audio_output", + "supports_multimodal_input", +] -@dataclass -class _RpcTask: - """A pending collective_rpc invocation queued for the busy loop.""" - method: str - args: tuple - kwargs: dict | None - deadline: float | None - unique_reply_rank: int | None - future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future) - - -def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]: - if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None: - signature = inspect.signature(pipe_cls.__call__) - support_image_input = "image" in signature.parameters - support_audio_input = ( - "audio" in signature.parameters or "audio_latents" in signature.parameters - ) # ref. LTX-2 format - return support_image_input, support_audio_input +def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool: + if func is None: + return False + parameters = inspect.signature(func).parameters + return parameter_name in parameters or any( + parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() + ) - supports_image_input = False - supports_audio_input = False - model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) - if model_cls is not None: - supports_image_input = bool(getattr(model_cls, "support_image_input", False)) - supports_audio_input = bool(getattr(model_cls, "support_audio_input", False)) - return supports_image_input, supports_audio_input +def _resolve_custom_pipeline_cls(custom_pipeline_args: dict[str, Any] | None) -> type | None: + if custom_pipeline_args is None: + return None + try: + pipeline_cls = custom_pipeline_args["pipeline_class"] + except KeyError as exc: + raise ValueError("custom_pipeline_args must include 'pipeline_class'.") from exc -def image_color_format(model_class_name: str) -> str: - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) - return getattr(model_cls, "color_format", "RGB") + if isinstance(pipeline_cls, type): + return pipeline_cls + if isinstance(pipeline_cls, str): + try: + return resolve_obj_by_qualname(pipeline_cls) + except (AttributeError, ImportError, ValueError) as exc: + raise ValueError(f"Failed to resolve custom diffusion pipeline class {pipeline_cls!r}.") from exc + raise TypeError( + f"custom_pipeline_args['pipeline_class'] must be a qualified name string or a class, " + f"got {type(pipeline_cls).__name__}" + ) -def supports_audio_output(model_class_name: str) -> bool: - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) +def supports_request_batch(od_config: OmniDiffusionConfig) -> bool: + model_cls = _resolve_custom_pipeline_cls(getattr(od_config, "custom_pipeline_args", None)) + if model_cls is None: + model_cls = DiffusionModelRegistry._try_load_model_cls(getattr(od_config, "model_class_name", None)) if model_cls is None: return False - return bool(getattr(model_cls, "support_audio_output", False)) - - -def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool: - if func is None: - return False - parameters = inspect.signature(func).parameters - return parameter_name in parameters or any( - parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values() - ) + return bool(getattr(model_cls, "supports_request_batch", False)) def _move_tensor_tree_to_cpu(value: object) -> object: @@ -106,46 +118,29 @@ def _move_tensor_tree_to_cpu(value: object) -> object: return value -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")) - return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1 - - -def get_extra_body_params(model_class_name: str) -> frozenset[str]: - """Return the set of extra_body keys accepted by a pipeline. - - Each pipeline can declare ``EXTRA_BODY_PARAMS: ClassVar[frozenset[str]]`` - to advertise which request-level parameters should be forwarded from - ``extra_body`` to ``OmniDiffusionSamplingParams.extra_args``. - Returns an empty frozenset when the pipeline does not declare any. - """ - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) - if model_cls is None: - return frozenset() - return frozenset(getattr(model_cls, "EXTRA_BODY_PARAMS", frozenset())) - - -def get_extra_output_params(model_class_name: str) -> frozenset[str]: - """Return the set of custom_output keys to expose in API response metrics. +@dataclass +class _RpcTask: + """A pending collective_rpc invocation queued for the busy loop.""" - Each pipeline can declare ``EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]]`` - to advertise which ``DiffusionOutput.custom_output`` keys should be - copied into the response ``metrics`` dict. - Returns an empty frozenset when the pipeline does not declare any. - """ - model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) - if model_cls is None: - return frozenset() - return frozenset(getattr(model_cls, "EXTRA_OUTPUT_PARAMS", frozenset())) + method: str + args: tuple + kwargs: dict | None + deadline: float | None + unique_reply_rank: int | None + future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future) class DiffusionEngine: """The diffusion engine for vLLM-Omni diffusion models.""" + #: Import path of the model runner this engine's workers should build, or + #: ``None`` for the platform default. Subclasses declare their runner here + #: (e.g. the AR-Diffusion engine), the worker resolves it through + #: :meth:`resolve_engine_class` — so engine->runner routing lives on the + #: engine class itself, and ``od_config.diffusion_model_runner_cls`` + #: remains a pure explicit user override (never mutated by engines). + default_diffusion_model_runner_cls: str | None = None + def __init__( self, od_config: OmniDiffusionConfig, @@ -159,29 +154,24 @@ def __init__( self.od_config = od_config self.post_process_func = get_diffusion_post_process_func(od_config) - self.action_post_process_func = get_diffusion_action_post_process_func(od_config) self.pre_process_func = get_diffusion_pre_process_func(od_config) # Cache whether the model-specific postprocess accepts request-level # sampling params so step() can support both legacy and extended hooks. self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params") - self._action_post_process_accepts_sampling_params = _func_accepts_parameter( - self.action_post_process_func, "sampling_params" - ) - self._action_post_process_accepts_custom_output = _func_accepts_parameter( - self.action_post_process_func, "custom_output" - ) + + self.step_execution = bool(getattr(od_config, "step_execution", False)) + if self.od_config.streaming_output and not self.step_execution: + logger.warning("streaming_output=True requires step_execution=True; enabling step execution.") + self.od_config.step_execution = True + self.step_execution = True executor_class = DiffusionExecutor.get_class(od_config) self.executor = executor_class(od_config) - self.step_execution = bool(getattr(od_config, "step_execution", False)) self.scheduler: SchedulerInterface = scheduler or ( StepScheduler() if self.step_execution else RequestScheduler() ) self.scheduler.initialize(od_config) - if self.scheduler.max_num_running_reqs > 1 and not self.step_execution: - max_num_seqs = self.scheduler.max_num_running_reqs - self.scheduler.max_num_running_reqs = 1 - logger.warning(f"Non-stepwise-execution does not support max-num-seqs={max_num_seqs}, set it to 1.") + self.supports_request_batch = False if self.step_execution else supports_request_batch(od_config) self.main_loop: asyncio.AbstractEventLoop | None = None self.stop_event: threading.Event | None = None self.worker_thread: threading.Thread | None = None @@ -194,11 +184,24 @@ def __init__( self._rpc_lock = threading.RLock() self._cv = threading.Condition(self._rpc_lock) self._out_queue: dict[str, asyncio.Future] = {} + self._out_queue_streaming: dict[str, asyncio.Queue[DiffusionOutput]] = {} self._closed = False self._shutdown_complete = False self.abort_queue: queue.Queue[str] = queue.Queue() self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue() - self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request + if self.step_execution: + self.execute_fn = self.executor.execute_step + elif self.supports_request_batch: + self.execute_fn = self.executor.execute_batch + else: + self.execute_fn = self.executor.execute_request + + if self.supports_request_batch: + logger.info( + "[RequestBatch] engine init max_num_seqs=%s max_wait_ms=%s", + getattr(od_config, "max_num_seqs", None), + getattr(od_config, "request_batch_max_wait_ms", None), + ) try: self._dummy_run() @@ -242,7 +245,37 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: exec_start_time = time.perf_counter() output = await self.async_add_req_and_wait_for_response(request) exec_total_time = time.perf_counter() - exec_start_time + return self.postprocess_output(request, output, diffusion_engine_start_time, preprocess_time, exec_total_time) + + async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[list[OmniRequestOutput], None]: + await self._check_and_start_background_loop() + + diffusion_engine_start_time = time.perf_counter() + + preprocess_time = 0.0 + if self.pre_process_func is not None: + preprocess_start_time = time.perf_counter() + request = self.pre_process_func(request) + preprocess_time = time.perf_counter() - preprocess_start_time + logger.debug("Pre-processing completed in %.4f seconds", preprocess_time) + exec_start_time = time.perf_counter() + generator = self.async_add_req_and_stream_response(request) + async for output in generator: + exec_total_time = time.perf_counter() - exec_start_time + yield self.postprocess_output( + request, output, diffusion_engine_start_time, preprocess_time, exec_total_time + ) + + def postprocess_output( + self, + request: OmniDiffusionRequest, + output: DiffusionOutput, + diffusion_engine_start_time: float, + preprocess_time: float, + exec_total_time: float, + ) -> list[OmniRequestOutput]: + """Convert a DiffusionOutput to a list of OmniRequestOutput, attaching profiling metrics.""" if output.aborted: raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.") if output.error: @@ -257,16 +290,7 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: if output.output is None: logger.warning("Output is None, returning empty OmniRequestOutput") - return [ - OmniRequestOutput.from_diffusion( - request_id=request.request_id, - images=[], - prompt=prompt, - metrics={}, - latents=None, - ) - for i, prompt in enumerate(request.prompts) - ] + return format_empty_diffusion_outputs(request, finished=output.finished) # When CPU offload is enabled, move output to CPU before # post-processing to avoid device OOM — model weights may still @@ -275,45 +299,18 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: if self.od_config.enable_cpu_offload: output_data = _move_tensor_tree_to_cpu(output_data) - custom_output = output.custom_output or {} - action_payload = None - action_only_output = bool(custom_output.get("action_only_output")) - postprocess_start_time = time.perf_counter() - if action_only_output: - outputs = [] - elif self.post_process_func is not None: + if self.post_process_func is not None: # Some video pipelines need request-level controls during # postprocess (for example worker-side frame interpolation). + postprocess_kwargs: dict[str, object] = {} if self._post_process_accepts_sampling_params: - outputs = self.post_process_func(output_data, sampling_params=request.sampling_params) - else: - outputs = self.post_process_func(output_data) + postprocess_kwargs["sampling_params"] = request.sampling_params + outputs = self.post_process_func(output_data, **postprocess_kwargs) else: outputs = output_data - audio_payload = None - model_audio_sample_rate = None - model_fps = None - if isinstance(outputs, dict): - audio_payload = outputs.get("audio") - action_payload = outputs.get("actions") - custom_output.update(outputs.get("custom_output") or {}) - model_audio_sample_rate = outputs.get("audio_sample_rate") - model_fps = outputs.get("fps") - outputs = outputs.get("video", outputs) - if action_payload is None: - action_payload = custom_output.get("actions") - action_post_process_func = getattr(self, "action_post_process_func", None) - if action_payload is None and action_post_process_func is not None: - raw_action_payload = custom_output.get("action", action_payload) - if raw_action_payload is not None: - action_kwargs: dict[str, Any] = {} - if getattr(self, "_action_post_process_accepts_custom_output", False): - action_kwargs["custom_output"] = custom_output - if getattr(self, "_action_post_process_accepts_sampling_params", False): - action_kwargs["sampling_params"] = request.sampling_params - action_payload = action_post_process_func(raw_action_payload, **action_kwargs) - custom_output["actions"] = action_payload + + postprocess_output = normalize_diffusion_postprocess_output(outputs) postprocess_time = time.perf_counter() - postprocess_start_time logger.debug("Post-processing completed in %.4f seconds", postprocess_time) @@ -327,184 +324,19 @@ async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]: step_total_ms, ) - # Convert to OmniRequestOutput format - # Ensure outputs is a list - if not isinstance(outputs, list): - outputs = [outputs] if outputs is not None else [] - - metrics = { - "preprocess_time_ms": preprocess_time * 1000, - "diffusion_engine_exec_time_ms": exec_total_time * 1000, - "diffusion_engine_total_time_ms": step_total_ms, - "image_num": int(request.sampling_params.num_outputs_per_prompt), - "resolution": int(request.sampling_params.resolution), - "postprocess_time_ms": postprocess_time * 1000, - } - - # Detect text output: when the pipeline returns a string (e.g., - # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it - # as a text-type response instead of an image. - is_text_output = isinstance(output_data, str) and custom_output.get("text_output") is not None - - # Handle single request or multiple requests - is_audio_output = supports_audio_output(self.od_config.model_class_name) - if is_audio_output and model_audio_sample_rate is None: - model_cls = DiffusionModelRegistry._try_load_model_cls(self.od_config.model_class_name) - model_audio_sample_rate = getattr(model_cls, "audio_sample_rate", None) - - def _audio_mm(payload: Any) -> dict[str, Any]: - mm: dict[str, Any] = {"audio": payload} - if model_audio_sample_rate is not None: - mm["audio_sample_rate"] = model_audio_sample_rate - return mm - - if len(request.prompts) == 1: - # Single request: return single OmniRequestOutput - prompt = request.prompts[0] - request_id = request.request_id - - if is_text_output: - return [ - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=[], - prompt=prompt, - metrics=metrics, - custom_output=custom_output, - final_output_type="text", - stage_durations=output.stage_durations, - peak_memory_mb=output.peak_memory_mb, - ), - ] - if is_audio_output: - request_audio_payload = outputs[0] if len(outputs) == 1 else outputs - return [ - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=[], - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - trajectory_latents=output.trajectory_latents, - trajectory_timesteps=output.trajectory_timesteps, - trajectory_log_probs=output.trajectory_log_probs, - trajectory_decoded=output.trajectory_decoded, - multimodal_output=_audio_mm(request_audio_payload), - final_output_type="audio", - stage_durations=output.stage_durations, - peak_memory_mb=output.peak_memory_mb, - ), - ] - else: - mm_output = {} - if audio_payload is not None: - mm_output["audio"] = audio_payload - if model_audio_sample_rate is not None: - mm_output["audio_sample_rate"] = model_audio_sample_rate - if model_fps is not None: - mm_output["fps"] = model_fps - if action_payload is not None: - mm_output["actions"] = action_payload - return [ - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=outputs, - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - trajectory_latents=output.trajectory_latents, - trajectory_timesteps=output.trajectory_timesteps, - trajectory_log_probs=output.trajectory_log_probs, - trajectory_decoded=output.trajectory_decoded, - custom_output=custom_output, - multimodal_output=mm_output, - stage_durations=output.stage_durations, - peak_memory_mb=output.peak_memory_mb, - ), - ] - else: - # Multiple requests: return list of OmniRequestOutput - # Split images based on num_outputs_per_prompt for each request - results = [] - output_idx = 0 - request_id = request.request_id - - for i, prompt in enumerate(request.prompts): - # Get images for this request - num_outputs = request.sampling_params.num_outputs_per_prompt - start_idx = output_idx - end_idx = start_idx + num_outputs - request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else [] - output_idx = end_idx - - if is_audio_output: - request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs - results.append( - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=[], - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - trajectory_latents=output.trajectory_latents, - trajectory_timesteps=output.trajectory_timesteps, - trajectory_log_probs=output.trajectory_log_probs, - trajectory_decoded=output.trajectory_decoded, - multimodal_output=_audio_mm(request_audio_payload), - final_output_type="audio", - stage_durations=output.stage_durations, - peak_memory_mb=output.peak_memory_mb, - ), - ) - else: - mm_output = {} - if audio_payload is not None: - sliced_audio = audio_payload - if isinstance(audio_payload, (list, tuple)): - sliced_audio = audio_payload[start_idx:end_idx] - if len(sliced_audio) == 1: - sliced_audio = sliced_audio[0] - elif hasattr(audio_payload, "shape") and getattr(audio_payload, "shape", None) is not None: - if len(audio_payload.shape) > 0 and audio_payload.shape[0] >= end_idx: - sliced_audio = audio_payload[start_idx:end_idx] - if num_outputs == 1: - sliced_audio = sliced_audio[0] - mm_output["audio"] = sliced_audio - if model_audio_sample_rate is not None: - mm_output["audio_sample_rate"] = model_audio_sample_rate - if model_fps is not None: - mm_output["fps"] = model_fps - if action_payload is not None: - sliced_actions = action_payload - if isinstance(action_payload, (list, tuple)): - sliced_actions = action_payload[start_idx:end_idx] - if len(sliced_actions) == 1: - sliced_actions = sliced_actions[0] - elif hasattr(action_payload, "shape") and getattr(action_payload, "shape", None) is not None: - if len(action_payload.shape) > 0 and action_payload.shape[0] >= end_idx: - sliced_actions = action_payload[start_idx:end_idx] - if num_outputs == 1: - sliced_actions = sliced_actions[0] - mm_output["actions"] = sliced_actions - results.append( - OmniRequestOutput.from_diffusion( - request_id=request_id, - images=request_outputs, - prompt=prompt, - metrics=metrics, - latents=output.trajectory_latents, - trajectory_latents=output.trajectory_latents, - trajectory_timesteps=output.trajectory_timesteps, - trajectory_log_probs=output.trajectory_log_probs, - trajectory_decoded=output.trajectory_decoded, - custom_output=custom_output, - multimodal_output=mm_output, - stage_durations=output.stage_durations, - peak_memory_mb=output.peak_memory_mb, - ), - ) - - return results + return format_diffusion_outputs( + request=request, + od_config=self.od_config, + diffusion_output=output, + output_data=output_data, + postprocess_output=postprocess_output, + timings=DiffusionStepTimings( + preprocess_time_s=preprocess_time, + exec_time_s=exec_total_time, + postprocess_time_s=postprocess_time, + total_time_ms=step_total_ms, + ), + ) def _busy_loop(self): while not self.stop_event.is_set(): @@ -527,14 +359,20 @@ def _busy_loop(self): # Only RPC / abort work pending; loop back to drain it. continue + if self.supports_request_batch: + self._wait_for_request_batch_admission_locked() + sched_output = self.scheduler.schedule() if sched_output.is_empty: - self._handle_finished_requests(sched_output.finished_req_ids, None) + if self.od_config.streaming_output: + self._handle_empty_streaming_requests(sched_output.finished_req_ids) + else: + self._handle_finished_requests(sched_output.finished_req_ids, None) continue try: - runner_output = self.execute_fn(sched_output) + runner_output: BaseRunnerOutput = self.execute_fn(sched_output) # pyright: ignore[reportAssignmentType] except Exception as exc: logger.error( "Execution failed for diffusion requests %s", sched_output.scheduled_request_ids, exc_info=True @@ -554,11 +392,73 @@ def _busy_loop(self): self._process_aborts_queue() self._process_rpc_queue() finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output) - self._handle_finished_requests(finished_req_ids, runner_output) + if self.od_config.streaming_output: + self._handle_step_streaming_runner_output( + finished_req_ids, + sched_output.scheduled_request_ids, + runner_output, + ) + else: + self._handle_finished_requests(finished_req_ids, runner_output) # Engine is stopping: fail any RPCs still queued so callers don't hang. self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down.")) + def _wait_for_request_batch_admission_locked(self) -> None: + """Wait for compatible requests to accumulate before scheduling a wave. + + Caller must hold ``self._cv``. + """ + if self.step_execution or not self.supports_request_batch: + return + + max_wait_s = self.od_config.request_batch_max_wait_ms / 1000.0 + if max_wait_s == 0: + return + + max_batch = self.scheduler.max_num_running_reqs + waiting = self.scheduler.num_waiting_requests() + running = self.scheduler.num_running_requests() + + if running > 0: + return + + start = time.monotonic() + deadline = start + max_wait_s + last_waiting = -1 + stable_since = start + # Require a short idle period with no queue growth so bursty HTTP + # ingress can land before the first schedule() of a wave. + stable_window_s = min(0.05, max_wait_s / 5.0) + + while not self.stop_event.is_set(): + waiting = self.scheduler.num_waiting_requests() + now = time.monotonic() + + if waiting >= max_batch: + break + if waiting > 0 and (now - stable_since) >= stable_window_s: + break + if now >= deadline: + break + + if waiting > last_waiting: + stable_since = now + last_waiting = waiting + + remaining = deadline - now + self._cv.wait(timeout=min(remaining, 0.002)) + + waited_ms = (time.monotonic() - start) * 1000.0 + final_waiting = self.scheduler.num_waiting_requests() + if final_waiting > 0: + logger.info( + "[RequestBatch] admission wait done waiting=%d max_batch=%d waited_ms=%.1f", + final_waiting, + max_batch, + waited_ms, + ) + def _process_rpc_queue(self) -> None: """Execute pending collective_rpc tasks from the busy-loop thread. @@ -629,28 +529,120 @@ def _handle_finished_requests( out = self._finalize_finished_request(rid, _output, missing_result_error) self._complete_future(fut, out) + def _handle_empty_streaming_requests( + self, + finished_ids: set[str], + missing_result_error: str = "Diffusion streaming request finished without execution output.", + ) -> None: + """Mirrors `_handle_finished_requests()` in non-streaming mode when used for empty scheduler output.""" + for rid in finished_ids: + out = self._finalize_finished_request(rid, None, missing_result_error=missing_result_error) + self._put_streaming_output_with_cv(rid, out) + + def _handle_step_streaming_runner_output( + self, + finished_req_ids: set[str], + scheduled_request_ids: list[str], + runner_output: BaseRunnerOutput, + ) -> None: + """ + Deliver partial step-execution outputs in streaming mode. + + Step execution returns one ``RunnerOutput`` per scheduled request per + engine tick. Most denoise steps have ``result=None``; chunk boundaries + return a ``DiffusionOutput`` that must be delivered even before the + scheduler marks the request finished. + """ + delivered_finished_req_ids: set[str] = set() + + # finished_ids may have some requests that are not scheduler in this round. + # First handle this-round requests. + for request_id in scheduled_request_ids: + req_output = runner_output.get_request_output(request_id) + if request_id in finished_req_ids: + # This entire request is finished (this is the last chunk) + out = self._finalize_finished_request( + request_id, + req_output, + missing_result_error="Diffusion streaming execution finished without a final output.", + ) + self._put_streaming_output_with_cv(request_id, out) + delivered_finished_req_ids.add(request_id) + elif req_output is not None and req_output.result is not None: + # This is a non-terminal chunk. So it is not in scheduler's finished_req_ids, but still need delivering. + self._put_streaming_output_with_cv(request_id, req_output.result) + + # Then handle other requests that are finished in this round. + for request_id in finished_req_ids - delivered_finished_req_ids: + out = self._finalize_finished_request( + request_id, + missing_result_error="Diffusion streaming request finished without execution output.", + ) + self._put_streaming_output_with_cv(request_id, out) + + @staticmethod + def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]: + """Resolve the engine class selected by ``config.engine_backend``. + + Mirrors ``DiffusionExecutor.get_class``: accepts ``"default"``, a + ``DiffusionEngine`` subclass, or an import-path string (e.g. a deploy + config's ``engine_backend``). Kept separate from :meth:`make_engine` so the + selection is testable without constructing an engine (which runs a dummy forward). + + Args: + config: The configuration for the diffusion engine. + + Returns: + The ``DiffusionEngine`` (sub)class to instantiate. + """ + backend = getattr(config, "engine_backend", "default") or "default" + + if isinstance(backend, type): + if not issubclass(backend, DiffusionEngine): + raise TypeError(f"engine_backend must be a DiffusionEngine subclass. Got {backend}.") + return backend + if backend == "default": + return DiffusionEngine + if isinstance(backend, str): + try: + engine_class = resolve_obj_by_qualname(backend) + except (ImportError, ValueError) as e: + raise ValueError( + f"Failed to load engine_backend '{backend}'. Ensure it is a valid python path. Error: {e}" + ) from e + if not issubclass(engine_class, DiffusionEngine): + raise TypeError(f"engine_backend must resolve to a DiffusionEngine subclass. Got {engine_class}.") + return engine_class + raise ValueError(f"Unknown engine_backend: {backend!r}") + @staticmethod def make_engine( config: OmniDiffusionConfig, scheduler: SchedulerInterface | None = None, ) -> DiffusionEngine: - """Factory method to create a DiffusionEngine instance. + """Factory method to create the engine selected by ``config.engine_backend``. Args: config: The configuration for the diffusion engine. Returns: - An instance of DiffusionEngine. + An instance of the resolved ``DiffusionEngine`` (sub)class. """ - return DiffusionEngine(config, scheduler=scheduler) + engine_class = DiffusionEngine.resolve_engine_class(config) + return engine_class(config, scheduler=scheduler) def add_request(self, request: OmniDiffusionRequest) -> str: with self._cv: if self._closed: raise RuntimeError("DiffusionEngine is closed.") - fut = self.main_loop.create_future() - request_id = self.scheduler.add_request(request) - self._out_queue[request_id] = fut + if not self.od_config.streaming_output: + fut = self.main_loop.create_future() + request_id = self.scheduler.add_request(request) + self._out_queue[request_id] = fut + else: + queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue() + request_id = self.scheduler.add_request(request) + self._out_queue_streaming[request_id] = queue self._cv.notify_all() return request_id @@ -666,12 +658,39 @@ async def get_result(self, request_id: str) -> DiffusionOutput: logger.error(f"Wait for response failed: {e}") raise + async def get_streaming_result(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]: + """Mirrors `get_result()` in non-streaming mode.""" + with self._cv: + queue = self._out_queue_streaming.get(request_id) + if queue is None: + raise RuntimeError(f"Request {request_id} not found in output queue.") + try: + while True: + output: DiffusionOutput = await queue.get() + yield output + if output.finished: + break + except Exception as e: + logger.error(f"Wait for response failed: {e}") + raise + finally: + # In streaming mode, an output queue is maintained until the terminal chunk is met. + # So unlike the non-streaming mode where output Future is popped in `_handle_finished_requests` (immediately + # after the request is returned), the streaming mode needs to pop the output queue here (one layer above). + with self._cv: + if self._out_queue_streaming.get(request_id) is queue: + self._out_queue_streaming.pop(request_id, None) + async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: # No lock needed: add_request is already protected by self._cv, and # all executor calls are serialized inside the busy loop. request_id = self.add_request(request) return await self.get_result(request_id) + def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]: + request_id = self.add_request(request) + return self.get_streaming_result(request_id) + def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: with self._rpc_lock: if self._closed: @@ -693,7 +712,7 @@ def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> Diffus # within _dummy_run, only one request will be scheduled request_id = sched_output.scheduled_request_ids[0] try: - runner_output = self.execute_fn(sched_output) + runner_output: BaseRunnerOutput = self.execute_fn(sched_output) # pyright: ignore[reportAssignmentType] except EngineDeadError: raise except Exception as exc: @@ -713,10 +732,10 @@ def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> Diffus if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1: raise ValueError("Sync func should receive one result at one time") if target_request_id in finished_req_ids: - runner_output = runner_output.get_request_output(target_request_id) + req_output = runner_output.get_request_output(target_request_id) return self._finalize_finished_request( target_request_id, - runner_output=runner_output, + runner_output=req_output, missing_result_error="Diffusion execution finished without a final output.", ) @@ -766,7 +785,7 @@ def _dummy_run(self): logger.info("Skipping dummy warmup run (num_frames=0)") return req = OmniDiffusionRequest( - prompts=[prompt], + prompt=prompt, request_id=DUMMY_DIFFUSION_REQUEST_ID, sampling_params=OmniDiffusionSamplingParams( height=height, @@ -911,8 +930,28 @@ def _set_result() -> None: else: _set_result() + def _put_streaming_queue_output( + self, + queue: asyncio.Queue[DiffusionOutput], + output: DiffusionOutput, + ) -> None: + """Append to streaming output queue in a safe event loop. Mirrors `_complete_future()` in non-streaming mode.""" + loop = self.main_loop + if loop is not None and loop.is_running(): + loop.call_soon_threadsafe(queue.put_nowait, output) + else: + queue.put_nowait(output) + + def _put_streaming_output_with_cv(self, request_id: str, output: DiffusionOutput) -> None: + with self._cv: + queue = self._out_queue_streaming.get(request_id) + if queue is None: + return + self._put_streaming_queue_output(queue, output) + def close(self) -> None: pending_futures: list[asyncio.Future] = [] + pending_streaming_queues: list[asyncio.Queue[DiffusionOutput]] = [] with self._cv: if self._closed and self._shutdown_complete: return @@ -921,12 +960,16 @@ def close(self) -> None: if self.stop_event is not None: self.stop_event.set() pending_futures = list(self._out_queue.values()) + pending_streaming_queues = list(self._out_queue_streaming.values()) self._out_queue.clear() + self._out_queue_streaming.clear() self._cv.notify_all() closed_output = DiffusionOutput(error="DiffusionEngine is closed.") for fut in pending_futures: self._complete_future(fut, closed_output) + for streaming_queue in pending_streaming_queues: + self._put_streaming_queue_output(streaming_queue, closed_output) worker_thread = self.worker_thread if worker_thread is not None: @@ -992,9 +1035,12 @@ def _finalize_finished_request( raise RuntimeError(f"Diffusion scheduler lost state for request {request_id}.") if state.status == DiffusionRequestStatus.FINISHED_ABORTED: + # Preserve runner-provided abort details when available. + if runner_output is not None and runner_output.result is not None and runner_output.result.aborted: + return runner_output.result return DiffusionOutput( aborted=True, - abort_message=f"Request {state.req.request_id} aborted.", + abort_message=f"Request {request_id} aborted.", ) if runner_output is not None and runner_output.result is not None: diff --git a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py index 35c9434d063..44c6efe6659 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py +++ b/vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py @@ -5,11 +5,13 @@ from typing import Any import torch +import torch.distributed as dist from diffusers.models.autoencoders import AutoencoderKLWan from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify from diffusers.models.autoencoders.vae import DecoderOutput from vllm.logger import init_logger +from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( DistributedOperator, DistributedVaeMixin, @@ -268,10 +270,54 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid dec = torch.clamp(dec, min=-1.0, max=1.0) return dec + def _spatial_shard_decode_split_dim(self) -> str | None: + mode = self.distributed_executor.parallel_mode + if mode == "spatial_shard_width": + return "width" + if mode == "spatial_shard_height": + return "height" + return None + + def _spatial_shard_decode_enabled(self, z: torch.Tensor) -> bool: + split_dim = self._spatial_shard_decode_split_dim() + if split_dim is None: + return False + if z.ndim != 5: + logger.warning("Wan VAE spatial sharded decode expects 5D latent input; falling back to tiled decode.") + return False + if not self.is_distributed_enabled(): + return False + + group = self.distributed_executor.group + world_size = dist.get_world_size(group=group) + requested_size = int(self.distributed_executor.parallel_size) + if requested_size != world_size: + logger.warning( + "Wan VAE spatial sharded decode currently requires vae_patch_parallel_size " + "to match the DIT group size; falling back to tiled decode. " + "requested=%s dit_group=%s split_dim=%s", + requested_size, + world_size, + split_dim, + ) + return False + return True + def tiled_decode(self, z: torch.Tensor, return_dict: bool = True): if not self.is_distributed_enabled(): return super().tiled_decode(z, return_dict=return_dict) + if self._spatial_shard_decode_enabled(z): + split_dim = self._spatial_shard_decode_split_dim() + logger.debug("Decode running with Wan VAE spatial_shard_%s mode", split_dim) + return wan_spatial_shard.spatial_shard_decode( + self, + z, + group=self.distributed_executor.group, + return_dict=return_dict, + split_dim=split_dim, + ) + logger.debug("Decode running with distributed executor") result = self.distributed_executor.execute( z, diff --git a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py index 8fb4ebeb353..d804740f90d 100644 --- a/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py +++ b/vllm_omni/diffusion/distributed/autoencoders/distributed_vae_executor.py @@ -49,9 +49,11 @@ def __init__(self): self.world_size = dist.get_world_size(self.group) self.rank = dist.get_rank(self.group) self.parallel_size = 1 + self.parallel_mode = "tile" - def set_parallel_size(self, parallel_size: int): + def set_parallel_size(self, parallel_size: int, mode: str = "tile"): self.parallel_size = parallel_size + self.parallel_mode = mode def gather_tensors(self, tensor: torch.Tensor): gather_list = [torch.empty_like(tensor) for _ in range(self.world_size)] @@ -170,8 +172,8 @@ class DistributedVaeMixin: def init_distributed(self): self.distributed_executor = DistributedVaeExecutor() - def set_parallel_size(self, parallel_size: int) -> None: - self.distributed_executor.set_parallel_size(parallel_size) + def set_parallel_size(self, parallel_size: int, mode: str = "tile") -> None: + self.distributed_executor.set_parallel_size(parallel_size, mode=mode) def is_distributed_enabled(self) -> bool: if ( diff --git a/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py new file mode 100644 index 00000000000..bda72854016 --- /dev/null +++ b/vllm_omni/diffusion/distributed/autoencoders/wan_spatial_shard.py @@ -0,0 +1,816 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# The halo-exchange spatial-parallel decode here is adapted from SGLang's +# spatial-parallel VAE decode +# https://github.com/sgl-project/sglang +# (python/sglang/multimodal_gen/runtime/layers/parallel_conv.py) +# which is in turn adapted from FastVideo (https://github.com/hao-ai-lab/FastVideo). +# This version generalizes the height-only sharding to shard along height or +# width and adds the Wan causal-conv ``feat_cache`` handling. +"""Spatially-sharded Wan VAE decode. + +The existing distributed Wan VAE path shards *tiles*. This module adds an +opt-in decode backend that shards decoder feature maps along height or width and +exchanges boundary rows/columns before spatial convolutions. It is +intentionally decode-only and keeps checkpoint loading unchanged by patching the +already-loaded decoder. +""" + +from __future__ import annotations + +import math +from contextlib import nullcontext +from contextvars import ContextVar +from dataclasses import dataclass +from types import MethodType +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.autoencoders.autoencoder_kl_wan import unpatchify +from diffusers.models.autoencoders.vae import DecoderOutput +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class SpatialShardContext: + input_extent: int + local_input_extent: int + split_dim: str + rank: int + world_size: int + + +_SPATIAL_SHARD_CONTEXT: ContextVar[SpatialShardContext | None] = ContextVar( + "wan_vae_spatial_shard_context", + default=None, +) + + +def _spatial_dim(split_dim: str) -> int: + if split_dim == "height": + return -2 + if split_dim == "width": + return -1 + raise ValueError(f"Unsupported Wan VAE split_dim={split_dim!r}; expected 'height' or 'width'.") + + +def _narrow_along_dim(x: torch.Tensor, dim: int, start: int, length: int) -> torch.Tensor: + if dim < 0: + dim += x.dim() + return x.narrow(dim, start, length) + + +def _global_rank(group: dist.ProcessGroup, group_rank: int) -> int: + try: + return dist.get_global_rank(group, group_rank) + except Exception: + return group_rank + + +def _rank_world(group: dist.ProcessGroup) -> tuple[int, int]: + return dist.get_rank(group), dist.get_world_size(group) + + +def _pad_along_dim(x: torch.Tensor, pad: int, dim: int, value: float = 0.0) -> torch.Tensor: + if pad <= 0: + return x + shape = list(x.shape) + shape[dim] = pad + padding = torch.full(shape, value, dtype=x.dtype, device=x.device) + return torch.cat([x, padding], dim=dim) + + +def _maybe_contiguous_for_shard_gather(x: torch.Tensor) -> torch.Tensor: + if ( + x.dim() == 5 + and hasattr(torch, "channels_last_3d") + and x.is_contiguous(memory_format=torch.channels_last_3d) + and not x.is_contiguous() + ): + return x.contiguous() + return x + + +def _halo_memory_format(reference: torch.Tensor) -> torch.memory_format: + if reference.dim() > 1 and reference.stride(1) == 1: + if reference.dim() == 5 and hasattr(torch, "channels_last_3d"): + return torch.channels_last_3d + if reference.dim() == 4: + return torch.channels_last + return torch.contiguous_format + + +def _current_full_extent(local_extent: int) -> int | None: + ctx = _SPATIAL_SHARD_CONTEXT.get() + if ctx is None: + return None + if ctx.local_input_extent <= 0: + return None + scale = local_extent / ctx.local_input_extent + rounded_scale = round(scale) + if not math.isclose(scale, rounded_scale, rel_tol=0.0, abs_tol=1e-6): + return None + return ctx.input_extent * rounded_scale + + +def _local_valid_extent(local_extent: int) -> int: + ctx = _SPATIAL_SHARD_CONTEXT.get() + full_extent = _current_full_extent(local_extent) + if ctx is None or full_extent is None: + return local_extent + start = ctx.rank * local_extent + return max(0, min(local_extent, full_extent - start)) + + +def _zero_invalid_extent(x: torch.Tensor, *, split_dim: str) -> torch.Tensor: + dim = _spatial_dim(split_dim) + dim_size = x.shape[dim] + valid_extent = _local_valid_extent(dim_size) + if valid_extent >= dim_size: + return x + x = x.clone() + invalid = _narrow_along_dim(x, dim, valid_extent, dim_size - valid_extent) + invalid.zero_() + return x + + +def split_for_parallel_decode( + x: torch.Tensor, + *, + upsample_count: int, + split_dim: str = "height", + group: dist.ProcessGroup | None = None, + rank: int | None = None, + world_size: int | None = None, +) -> tuple[torch.Tensor, int]: + """Shard latent/feature spatial extent and return expected full output extent.""" + if group is not None: + rank, world_size = _rank_world(group) + rank = 0 if rank is None else int(rank) + world_size = 1 if world_size is None else int(world_size) + if world_size < 1: + raise ValueError(f"Wan VAE world_size must be >= 1, got {world_size}.") + if not 0 <= rank < world_size: + raise ValueError(f"Wan VAE rank must satisfy 0 <= rank < world_size, got rank={rank}, world_size={world_size}.") + + dim = _spatial_dim(split_dim) + expected_extent = x.shape[dim] * (2**upsample_count) + if world_size <= 1: + return x, expected_extent + + pad = (world_size - (x.shape[dim] % world_size)) % world_size + if pad: + x = _pad_along_dim(x, pad, dim=dim) + chunk_size = x.shape[dim] // world_size + return _narrow_along_dim(x, dim, rank * chunk_size, chunk_size).contiguous(), expected_extent + + +def all_gather_along_dim( + x: torch.Tensor, + *, + group: dist.ProcessGroup, + dim: int, + dst: int | None = None, +) -> torch.Tensor: + rank, world_size = _rank_world(group) + if world_size <= 1: + return x + x = _maybe_contiguous_for_shard_gather(x) + gathered = [torch.empty_like(x) for _ in range(world_size)] + # NCCL has no rank-local gather, so every rank joins the collective; only ``dst`` + # keeps the assembled tensor while the rest drop their copies. + dist.all_gather(gathered, x.contiguous(), group=group) + if dst is not None and rank != dst: + return x.new_zeros(0) + return torch.cat(gathered, dim=dim) + + +def reshard_from_trimmed_extent( + x: torch.Tensor, + *, + local_extent: int, + split_dim: str, + group: dist.ProcessGroup, +) -> torch.Tensor: + rank, world_size = _rank_world(group) + if world_size <= 1: + return x + + dim = _spatial_dim(split_dim) + valid_extent = _local_valid_extent(local_extent) + start = rank * local_extent + local = _narrow_along_dim(x, dim, start, valid_extent).contiguous() + if valid_extent < local_extent: + local = _pad_along_dim(local, local_extent - valid_extent, dim=dim) + return local + + +def gather_and_trim_extent( + x: torch.Tensor, + *, + expected_extent: int | None, + split_dim: str, + group: dist.ProcessGroup, + dst: int | None = None, +) -> torch.Tensor: + dim = _spatial_dim(split_dim) + rank, _ = _rank_world(group) + out = all_gather_along_dim(x, group=group, dim=dim, dst=dst) + if dst is not None and rank != dst: + return out + if expected_extent is not None and out.shape[dim] != expected_extent: + out = _narrow_along_dim(out, dim, 0, expected_extent).contiguous() + return out + + +def _ensure_recv_buf(recv_buf: torch.Tensor | None, reference: torch.Tensor) -> torch.Tensor: + memory_format = _halo_memory_format(reference) + if ( + recv_buf is None + or recv_buf.shape != reference.shape + or recv_buf.dtype != reference.dtype + or recv_buf.device != reference.device + or not recv_buf.is_contiguous(memory_format=memory_format) + ): + return torch.empty( + reference.shape, + dtype=reference.dtype, + device=reference.device, + memory_format=memory_format, + ) + return recv_buf + + +def _halo_exchange_p2p( + *, + rank: int, + world_size: int, + group: dist.ProcessGroup, + top_row_ref: torch.Tensor, + bottom_row_ref: torch.Tensor, + recv_top_buf: torch.Tensor, + recv_bottom_buf: torch.Tensor, +) -> None: + p2p_ops = [] + if rank > 0: + prev_rank = _global_rank(group, rank - 1) + top_row = top_row_ref.contiguous(memory_format=_halo_memory_format(top_row_ref)) + p2p_ops.append(dist.P2POp(dist.irecv, recv_top_buf, prev_rank, group)) + p2p_ops.append(dist.P2POp(dist.isend, top_row, prev_rank, group)) + else: + recv_top_buf.zero_() + + if rank < world_size - 1: + next_rank = _global_rank(group, rank + 1) + bottom_row = bottom_row_ref.contiguous(memory_format=_halo_memory_format(bottom_row_ref)) + p2p_ops.append(dist.P2POp(dist.isend, bottom_row, next_rank, group)) + p2p_ops.append(dist.P2POp(dist.irecv, recv_bottom_buf, next_rank, group)) + else: + recv_bottom_buf.zero_() + + if p2p_ops: + reqs = dist.batch_isend_irecv(p2p_ops) + for req in reqs: + req.wait() + + +def halo_exchange( + x: torch.Tensor, + *, + group: dist.ProcessGroup, + halo_size: int, + split_dim: str = "height", + recv_top_buf: torch.Tensor | None = None, + recv_bottom_buf: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + if halo_size <= 0: + return x, recv_top_buf, recv_bottom_buf + + rank, world_size = _rank_world(group) + if world_size <= 1: + return x, recv_top_buf, recv_bottom_buf + + dim = _spatial_dim(split_dim) + top_row_ref = _narrow_along_dim(x, dim, 0, halo_size) + bottom_row_ref = _narrow_along_dim(x, dim, x.shape[dim] - halo_size, halo_size) + recv_top_buf = _ensure_recv_buf(recv_top_buf, top_row_ref) + recv_bottom_buf = _ensure_recv_buf(recv_bottom_buf, bottom_row_ref) + + _halo_exchange_p2p( + rank=rank, + world_size=world_size, + group=group, + top_row_ref=top_row_ref, + bottom_row_ref=bottom_row_ref, + recv_top_buf=recv_top_buf, + recv_bottom_buf=recv_bottom_buf, + ) + + return torch.cat([recv_top_buf, x, recv_bottom_buf], dim=dim), recv_top_buf, recv_bottom_buf + + +class WanDistZeroPad2d(nn.Module): + """Apply ZeroPad2d only at global split-dimension boundaries.""" + + def __init__( + self, + padding: tuple[int, int, int, int], + group: dist.ProcessGroup, + *, + split_dim: str = "height", + split_padding: tuple[int, int] | None = None, + ) -> None: + super().__init__() + self.padding = padding + self.split_dim = split_dim + default_split_padding = (padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]) + self.split_padding = split_padding or default_split_padding + self.group = group + + def forward(self, x: torch.Tensor) -> torch.Tensor: + rank, world_size = _rank_world(self.group) + left, right, top, bottom = self.padding + if world_size > 1: + if self.split_dim == "height": + top = top if rank == 0 else 0 + bottom = bottom if rank == world_size - 1 else 0 + else: + left = left if rank == 0 else 0 + right = right if rank == world_size - 1 else 0 + return F.pad(x, (left, right, top, bottom)) + + +class WanDistConv2d(nn.Conv2d): + def __init__( + self, + source: nn.Conv2d, + group: dist.ProcessGroup, + split_dim: str = "height", + split_padding: tuple[int, int] | None = None, + ): + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self.load_state_dict(source.state_dict()) + self.group = group + self.split_dim = split_dim + self.split_tensor_dim = _spatial_dim(split_dim) + kernel_extent = self.kernel_size[-2] if split_dim == "height" else self.kernel_size[-1] + self.halo_size = (kernel_extent - 1) // 2 + pad_h = source.padding[-2] if isinstance(source.padding, tuple) else source.padding + pad_w = source.padding[-1] if isinstance(source.padding, tuple) else source.padding + if split_dim == "height": + if split_padding is None: + split_padding = (pad_h, pad_h) + self._non_split_padding = (pad_w, pad_w, 0, 0) + self.kernel_extent = self.kernel_size[-2] + self.stride_extent = self.stride[-2] + else: + if split_padding is None: + split_padding = (pad_w, pad_w) + self._non_split_padding = (0, 0, pad_h, pad_h) + self.kernel_extent = self.kernel_size[-1] + self.stride_extent = self.stride[-1] + self.split_pad_left, self.split_pad_right = split_padding + self._halo_recv_top_buf: torch.Tensor | None = None + self._halo_recv_bottom_buf: torch.Tensor | None = None + self._trim_cache: dict[int, tuple[int, int, int]] = {} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.pad(x, self._non_split_padding) + x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( + x, + group=self.group, + halo_size=self.halo_size, + split_dim=self.split_dim, + recv_top_buf=self._halo_recv_top_buf, + recv_bottom_buf=self._halo_recv_bottom_buf, + ) + shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim]) + if shift: + x_padded = _narrow_along_dim( + x_padded, + self.split_tensor_dim, + shift, + x_padded.shape[self.split_tensor_dim] - shift, + ) + out = super().forward(x_padded) + out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim) + return _zero_invalid_extent(out, split_dim=self.split_dim) + + def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]: + trim_params = self._trim_cache.get(local_extent) + if trim_params is None: + rank, world_size = _rank_world(self.group) + trim_params = _compute_conv_trim_params( + local_extent=local_extent, + rank=rank, + world_size=world_size, + halo_size=self.halo_size, + pad_before=self.split_pad_left, + pad_after=self.split_pad_right, + kernel_extent=self.kernel_extent, + stride_extent=self.stride_extent, + ) + self._trim_cache[local_extent] = trim_params + return trim_params + + +class WanDistCausalConv3d(nn.Conv3d): + def __init__( + self, + source: nn.Conv3d, + group: dist.ProcessGroup, + split_dim: str = "height", + ): + super().__init__( + source.in_channels, + source.out_channels, + source.kernel_size, + stride=source.stride, + padding=0, + dilation=source.dilation, + groups=source.groups, + bias=source.bias is not None, + padding_mode=source.padding_mode, + device=source.weight.device, + dtype=source.weight.dtype, + ) + self.load_state_dict(source.state_dict()) + self.group = group + self.split_dim = split_dim + self.split_tensor_dim = _spatial_dim(split_dim) + source_padding = getattr(source, "_padding", None) + if source_padding is None: + p_t, p_h, p_w = source.padding + source_padding = (p_w, p_w, p_h, p_h, 2 * p_t, 0) + self._source_padding = tuple(source_padding) + if split_dim == "height": + self.split_pad_left = int(self._source_padding[2]) + self.split_pad_right = int(self._source_padding[3]) + self.kernel_extent = self.kernel_size[-2] + self.stride_extent = self.stride[-2] + padding = ( + self._source_padding[0], + self._source_padding[1], + 0, + 0, + self._source_padding[4], + self._source_padding[5], + ) + else: + self.split_pad_left = int(self._source_padding[0]) + self.split_pad_right = int(self._source_padding[1]) + self.kernel_extent = self.kernel_size[-1] + self.stride_extent = self.stride[-1] + padding = ( + 0, + 0, + self._source_padding[2], + self._source_padding[3], + self._source_padding[4], + self._source_padding[5], + ) + self.halo_size = (self.kernel_extent - 1) // 2 + self._padding = padding if self.halo_size > 0 else self._source_padding + self._halo_recv_top_buf: torch.Tensor | None = None + self._halo_recv_bottom_buf: torch.Tensor | None = None + self._trim_cache: dict[int, tuple[int, int, int]] = {} + + def forward(self, x: torch.Tensor, cache_x: torch.Tensor | None = None) -> torch.Tensor: + padding = list(self._padding) + if cache_x is not None and padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + + x = F.pad(x, padding) + x_padded, self._halo_recv_top_buf, self._halo_recv_bottom_buf = halo_exchange( + x, + group=self.group, + halo_size=self.halo_size, + split_dim=self.split_dim, + recv_top_buf=self._halo_recv_top_buf, + recv_bottom_buf=self._halo_recv_bottom_buf, + ) + shift, start, upper_bound = self._get_trim_params(x.shape[self.split_tensor_dim]) + if shift: + x_padded = _narrow_along_dim( + x_padded, + self.split_tensor_dim, + shift, + x_padded.shape[self.split_tensor_dim] - shift, + ) + out = super().forward(x_padded) + out = _trim_local_conv_output(out, self.halo_size, start, upper_bound, split_dim=self.split_dim) + return _zero_invalid_extent(out, split_dim=self.split_dim) + + def _get_trim_params(self, local_extent: int) -> tuple[int, int, int]: + trim_params = self._trim_cache.get(local_extent) + if trim_params is None: + rank, world_size = _rank_world(self.group) + trim_params = _compute_conv_trim_params( + local_extent=local_extent, + rank=rank, + world_size=world_size, + halo_size=self.halo_size, + pad_before=self.split_pad_left, + pad_after=self.split_pad_right, + kernel_extent=self.kernel_extent, + stride_extent=self.stride_extent, + ) + self._trim_cache[local_extent] = trim_params + return trim_params + + +def _compute_conv_trim_params( + *, + local_extent: int, + rank: int, + world_size: int, + halo_size: int, + pad_before: int, + pad_after: int, + kernel_extent: int, + stride_extent: int, +) -> tuple[int, int, int]: + global_start = rank * local_extent + shift = 0 + if halo_size > 0 and stride_extent > 1: + shift = (global_start - halo_size + pad_before) % stride_extent + if shift: + global_start += shift + + global_extent = local_extent * world_size + min_i = math.ceil(((-pad_before) - (global_start - halo_size)) / stride_extent) + max_i = math.floor( + ((global_extent - 1 + pad_after) - (kernel_extent - 1) - (global_start - halo_size)) / stride_extent + ) + return shift, max(min_i, 0), max_i + 1 + + +def _trim_local_conv_output( + out: torch.Tensor, + halo_size: int, + start: int, + upper_bound: int, + *, + split_dim: str, +) -> torch.Tensor: + if halo_size <= 0: + return out + dim = _spatial_dim(split_dim) + end = min(upper_bound, out.shape[dim]) + if start != 0 or end != out.shape[dim]: + out = _narrow_along_dim(out, dim, start, end - start) + return out + + +def _patch_attention_block(module: nn.Module, group: dist.ProcessGroup, split_dim: str) -> None: + if getattr(module, "_vllm_omni_spatial_shard_attention", False): + return + orig_forward = module.forward + + def _forward(self: nn.Module, x: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: + _, world_size = _rank_world(group) + if world_size <= 1: + return orig_forward(x, *args, **kwargs) + dim = _spatial_dim(split_dim) + local_extent = x.shape[dim] + gathered = all_gather_along_dim(x, group=group, dim=dim).contiguous() + full_extent = _current_full_extent(local_extent) + if full_extent is not None: + gathered = _narrow_along_dim(gathered, dim, 0, full_extent).contiguous() + out = orig_forward(gathered, *args, **kwargs) + return reshard_from_trimmed_extent(out, local_extent=local_extent, split_dim=split_dim, group=group) + + module.forward = MethodType(_forward, module) + module._vllm_omni_spatial_shard_attention = True # type: ignore[attr-defined] + + +def _replace_child( + parent: nn.Module, + name: str, + child: nn.Module, + group: dist.ProcessGroup, + split_dim: str, +) -> None: + if child.__class__.__name__ == "WanCausalConv3d": + setattr( + parent, + name, + WanDistCausalConv3d( + child, + group, + split_dim=split_dim, + ), + ) + return + if isinstance(child, nn.ZeroPad2d): + padding = tuple(int(p) for p in child.padding) + module_padding = padding + if parent.__class__.__name__ == "Sequential": + # Let the following WanDistConv2d account for global after-edge + # padding; this module handles the non-split dimension and before edge. + if split_dim == "height": + module_padding = (padding[0], padding[1], padding[2], 0) + else: + module_padding = (padding[0], 0, padding[2], padding[3]) + setattr( + parent, + name, + WanDistZeroPad2d( + module_padding, + group, + split_dim=split_dim, + split_padding=(padding[2], padding[3]) if split_dim == "height" else (padding[0], padding[1]), + ), + ) + return + if isinstance(child, nn.Conv2d): + split_padding = None + if name == "1" and parent.__class__.__name__ == "Sequential": + # WanResample downsample uses ZeroPad2d((0, 1, 0, 1)) before a + # stride-2 conv with padding=0. Only the last rank should see the + # bottom/right global padding, which is approximated by split_padding. + prev = getattr(parent, "0", None) + if isinstance(prev, WanDistZeroPad2d): + split_padding = prev.split_padding + elif isinstance(prev, nn.ZeroPad2d): + pad = prev.padding + split_padding = (int(pad[2]), int(pad[3])) if split_dim == "height" else (int(pad[0]), int(pad[1])) + setattr( + parent, + name, + WanDistConv2d( + child, + group, + split_dim=split_dim, + split_padding=split_padding, + ), + ) + + +def _patch_decoder_modules( + module: nn.Module, + group: dist.ProcessGroup, + split_dim: str, + inside_attention: bool = False, +) -> None: + if module.__class__.__name__ == "WanAttentionBlock": + _patch_attention_block(module, group, split_dim) + inside_attention = True + + for name, child in list(module.named_children()): + if child.__class__.__name__ == "WanCausalConv3d": + _replace_child(module, name, child, group, split_dim) + continue + if isinstance(child, nn.Conv2d) and not inside_attention: + _replace_child(module, name, child, group, split_dim) + continue + _patch_decoder_modules(child, group, split_dim, inside_attention=inside_attention) + + +def _decoder_upsample_count(decoder: nn.Module) -> int: + count = 0 + for block in getattr(decoder, "up_blocks", []): + if getattr(block, "upsampler", None) is not None or getattr(block, "upsamplers", None) is not None: + count += 1 + return count + + +def install_wan_spatial_shard_decode(vae: Any, group: dist.ProcessGroup, split_dim: str = "height") -> None: + """Patch ``vae.decoder`` once for spatially-sharded decode. + + This mutates the already-loaded decoder in place by swapping its spatial + convolutions/padding for halo-exchanging variants and wrapping + ``decoder.forward``. The patch is permanent for the lifetime of the VAE + instance and is applied only once (subsequent calls are no-ops). A given + instance is bound to a single ``split_dim``; switching between + ``"height"`` and ``"width"`` requires a fresh VAE instance and raises here + otherwise. + + Only group-relative rank 0 assembles the final decoded frame, mirroring the + distributed tiled-decode ``broadcast_result=False`` contract; the other ranks + take part in the collectives but return an empty placeholder. + """ + _spatial_dim(split_dim) + if getattr(vae, "_vllm_omni_wan_spatial_shard_installed", False): + installed_split_dim = getattr(vae, "_vllm_omni_wan_spatial_shard_split_dim", "height") + if installed_split_dim != split_dim: + raise ValueError( + "Wan spatial-shard VAE decoder was already patched for " + f"{installed_split_dim!r} split; create a fresh VAE instance to use {split_dim!r} split." + ) + return + decoder = getattr(vae, "decoder", None) + if decoder is None: + raise ValueError("Wan spatial-shard VAE decode requires a decoder module.") + + _patch_decoder_modules(decoder, group, split_dim) + upsample_count = _decoder_upsample_count(decoder) + orig_forward = decoder.forward + + def _forward( + self: nn.Module, + x: torch.Tensor, + feat_cache: list[torch.Tensor] | None = None, + feat_idx: list[int] | None = None, + first_chunk: bool = False, + ) -> torch.Tensor: + if feat_idx is None: + feat_idx = [0] + tensor_dim = _spatial_dim(split_dim) + input_extent = x.shape[tensor_dim] + x, expected_extent = split_for_parallel_decode( + x, + upsample_count=upsample_count, + split_dim=split_dim, + group=group, + ) + rank, world_size = _rank_world(group) + token = _SPATIAL_SHARD_CONTEXT.set( + SpatialShardContext( + input_extent=input_extent, + local_input_extent=x.shape[tensor_dim], + split_dim=split_dim, + rank=rank, + world_size=world_size, + ) + ) + try: + out = orig_forward(x, feat_cache=feat_cache, feat_idx=feat_idx, first_chunk=first_chunk) + finally: + _SPATIAL_SHARD_CONTEXT.reset(token) + return gather_and_trim_extent(out, expected_extent=expected_extent, split_dim=split_dim, group=group, dst=0) + + decoder.forward = MethodType(_forward, decoder) + vae._vllm_omni_wan_spatial_shard_installed = True + vae._vllm_omni_wan_spatial_shard_split_dim = split_dim + logger.info("Installed Wan VAE %s-sharded decode.", split_dim) + + +def spatial_shard_decode( + vae: Any, + z: torch.Tensor, + *, + group: dist.ProcessGroup, + return_dict: bool = True, + split_dim: str = "height", +) -> DecoderOutput | tuple[torch.Tensor]: + install_wan_spatial_shard_decode(vae, group, split_dim=split_dim) + + if z.shape[2] == 0: + raise ValueError("Wan spatial-shard VAE decode expects at least one latent frame.") + + # Non-rank-0 ranks must still run the decoder every chunk to stay in lockstep with + # the halo/all-gather collectives; they just skip keeping/assembling the output. + rank, world_size = _rank_world(group) + produce_output = world_size <= 1 or rank == 0 + + vae.clear_cache() + try: + context = vae._execution_context() if hasattr(vae, "_execution_context") else nullcontext() + with context: + x = vae.post_quant_conv(z) + decoded_chunks = [] + for i in range(z.shape[2]): + vae._conv_idx = [0] + chunk = vae.decoder( + x[:, :, i : i + 1, :, :], + feat_cache=vae._feat_map, + feat_idx=vae._conv_idx, + first_chunk=(i == 0), + ) + if produce_output: + decoded_chunks.append(chunk) + + if produce_output: + out = torch.cat(decoded_chunks, dim=2) + if vae.config.patch_size is not None: + out = unpatchify(out, patch_size=vae.config.patch_size) + out = torch.clamp(out, min=-1.0, max=1.0) + else: + out = z.new_zeros(0) + finally: + vae.clear_cache() + + if not return_dict: + return (out,) + return DecoderOutput(sample=out) diff --git a/vllm_omni/diffusion/distributed/cfg_parallel.py b/vllm_omni/diffusion/distributed/cfg_parallel.py index 6810a447b44..55130e0cf55 100644 --- a/vllm_omni/diffusion/distributed/cfg_parallel.py +++ b/vllm_omni/diffusion/distributed/cfg_parallel.py @@ -81,6 +81,7 @@ def predict_noise_maybe_with_cfg( negative_kwargs: dict[str, Any] | None, cfg_normalize: bool = True, output_slice: int | None = None, + kwargs: dict[str, Any] | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Predict noise with optional classifier-free guidance. @@ -92,6 +93,7 @@ def predict_noise_maybe_with_cfg( negative_kwargs: Kwargs for negative/unconditional prediction cfg_normalize: Whether to normalize CFG output (default: True) output_slice: If set, slice each output to [:, :output_slice] for image editing + kwargs: Optional extra context for custom combine implementations. Returns: Predicted noise tensor or tuple of tensors. @@ -111,8 +113,12 @@ def predict_noise_maybe_with_cfg( cfg_rank = get_classifier_free_guidance_rank() # Each rank computes one branch - kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs - local_pred = _wrap(self.predict_noise(**kwargs)) + if cfg_rank == 0: + logger.debug("CFG Parallel: Rank 0 computing positive branch") + local_pred = _wrap(self.predict_noise(**positive_kwargs)) + else: + logger.debug("CFG Parallel: Rank %d computing negative branch", cfg_rank) + local_pred = _wrap(self.predict_noise(**negative_kwargs)) if output_slice is not None: local_pred = _slice_pred(local_pred, output_slice) @@ -128,6 +134,7 @@ def predict_noise_maybe_with_cfg( negative_noise_pred, true_cfg_scale, cfg_normalize, + **({} if kwargs is None else {"kwargs": kwargs}), ) else: # Sequential CFG: compute both positive and negative @@ -143,6 +150,7 @@ def predict_noise_maybe_with_cfg( negative_noise_pred, true_cfg_scale, cfg_normalize, + **({} if kwargs is None else {"kwargs": kwargs}), ) else: # No CFG: only compute positive/conditional prediction @@ -173,6 +181,7 @@ def combine_cfg_noise( negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], true_cfg_scale: float, cfg_normalize: bool = False, + kwargs: dict[str, Any] | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ Combine conditional and unconditional noise predictions with CFG. @@ -196,6 +205,7 @@ def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, scale, nor negative_noise_pred: Negative/unconditional prediction(s) — Tensor or tuple true_cfg_scale: CFG scale factor cfg_normalize: Whether to normalize the combined prediction (default: False) + kwargs: Optional extra context for custom combine implementations. Returns: Combined noise prediction(s) — same type as inputs @@ -276,7 +286,7 @@ def _predict_multi_branch_parallel( branches_kwargs: list[dict[str, Any]], n_branches: int, cfg_world_size: int, - true_cfg_scale: float, + true_cfg_scale: float | dict[str, float], cfg_normalize: bool, output_slice: int | None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: @@ -453,6 +463,31 @@ def diffuse(self, video_latents, audio_latents, timesteps_video, timesteps_audio """ raise NotImplementedError("Subclasses must implement diffuse") + def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool) -> bool: + """ + Check if CFG parallel configuration is valid. + + Args: + true_cfg_scale: The classifier-free guidance scale value + has_neg_prompt: Whether a negative prompt is provided + + Returns: + True if CFG parallel configuration is valid, False otherwise + """ + if get_classifier_free_guidance_world_size() == 1: + return True + + if true_cfg_scale <= 1: + logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.") + return False + + if not has_neg_prompt: + logger.warning( + "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings." + ) + return False + return True + def scheduler_step( self, noise_pred: torch.Tensor | tuple[torch.Tensor, ...], diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 5c5b092ef21..0a67fd941a9 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -589,11 +589,7 @@ def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: if dst is None: dst = self.group_next_rank - torch.distributed.send( - tensor, - self.ranks[dst], - group=(self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group), - ) + torch.distributed.send(tensor, self.ranks[dst], group=self.device_group) def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> torch.Tensor: """Receives a tensor from the src rank.""" @@ -602,11 +598,7 @@ def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> src = self.group_prev_rank tensor = torch.empty(size, dtype=dtype, device=self.device) - torch.distributed.recv( - tensor, - self.ranks[src], - (self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group), - ) + torch.distributed.recv(tensor, self.ranks[src], self.device_group) return tensor def destroy(self): @@ -710,6 +702,22 @@ def __init__( self.skip_device_group = skip_device_group assert self.skip_device_group is not None + def send(self, tensor: torch.Tensor, dst: int | None = None) -> None: + if dst is None: + dst = self.group_next_rank + + group = self.device_groups[self.rank_in_group % 2] if self.world_size == 2 else self.device_group + torch.distributed.send(tensor, self.ranks[dst], group=group) + + def recv(self, size: torch.Size, dtype: torch.dtype, src: int | None = None) -> torch.Tensor: + if src is None: + src = self.group_prev_rank + + tensor = torch.empty(size, dtype=dtype, device=self.device) + group = self.device_groups[(self.rank_in_group + 1) % 2] if self.world_size == 2 else self.device_group + torch.distributed.recv(tensor, self.ranks[src], group) + return tensor + def reset_buffer(self): self.recv_tasks_queue = [] self.receiving_tasks = [] diff --git a/vllm_omni/diffusion/distributed/hsdp.py b/vllm_omni/diffusion/distributed/hsdp.py index 69506c92834..cae4356f43a 100644 --- a/vllm_omni/diffusion/distributed/hsdp.py +++ b/vllm_omni/diffusion/distributed/hsdp.py @@ -26,6 +26,15 @@ logger = init_logger(__name__) +def _unshardable_parameters(model: nn.Module) -> set[nn.Parameter]: + """Return packed/integer or scalar parameters unsupported by FSDP2.""" + return { + param + for param in model.parameters() + if param.ndim == 0 or not (param.is_floating_point() or param.is_complex()) + } + + @dataclass class HSDPInferenceConfig: """Configuration for HSDP inference. @@ -88,6 +97,7 @@ def _create_hsdp_mesh( def apply_hsdp_to_model( model: nn.Module, hsdp_config: HSDPInferenceConfig, + target_device: torch.device | None = None, ) -> nn.Module: """ Apply HSDP sharding to a model that already has weights loaded. @@ -98,6 +108,10 @@ def apply_hsdp_to_model( Args: model: Model instance with weights already loaded hsdp_config: HSDP configuration with HSDP mesh dimensions + target_device: Worker's execution device. When the model declares + _hsdp_ignored_modules, those modules are excluded from FSDP's + mesh-driven device placement, so the caller must specify where to + put them. Optional only when there are no ignored modules. Returns: HSDP-wrapped model ready for inference @@ -135,8 +149,13 @@ def apply_hsdp_to_model( fs_rank, ) + # When the model contains FP8 parameters (online quantization), let FSDP + # keep the original storage dtype on all-gather instead of casting to + # hsdp_config.param_dtype (typically bfloat16). FP8 GEMM kernels expect + # FP8 inputs; an implicit FP8 -> bf16 cast would silently break them. + has_fp8_params = any(p.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for p in model.parameters()) mp_policy = MixedPrecisionPolicy( - param_dtype=hsdp_config.param_dtype, + param_dtype=None if has_fp8_params else hsdp_config.param_dtype, reduce_dtype=hsdp_config.reduce_dtype, output_dtype=hsdp_config.output_dtype, cast_forward_inputs=False, @@ -158,6 +177,66 @@ def apply_hsdp_to_model( if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0: raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined") + # Collect parameters of any modules the model wants excluded from FSDP sharding. + # See _hsdp_ignored_modules on each model class. These params keep their + # original dtype/storage and are not subject to MixedPrecisionPolicy on the root + # FSDP wrap. Useful for small auxiliary modules (e.g., timestep embedders) that + # need to stay in a higher-precision dtype than the bulk of the model. + ignored_module_names = getattr(model, "_hsdp_ignored_modules", []) or [] + ignored_params: set[nn.Parameter] = set() + # FSDP wraps move sharded params to mesh devices automatically, but + # ignored modules stay wherever load_weights left them (CPU under HSDP). + # The caller (diffusers_loader / similar) is responsible for telling us + # which device the worker should run on, matching the convention used + # for VAEs / encoders / resident_modules. + if ignored_module_names and target_device is None: + raise ValueError( + f"Model {type(model).__name__} declares _hsdp_ignored_modules=" + f"{ignored_module_names} but apply_hsdp_to_model was called " + "without target_device. The caller must pass target_device so " + "the ignored modules can be placed on the worker's execution device." + ) + for mod_name in ignored_module_names: + sub_mod = getattr(model, mod_name, None) + if sub_mod is None: + logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name) + continue + sub_mod.to(target_device) + ignored_params.update(sub_mod.parameters()) + if ignored_params: + logger.info( + "HSDP excluding %d parameter tensors from sharding (modules: %s, moved to %s)", + len(ignored_params), + ignored_module_names, + target_device, + ) + + # Serialized low-bit checkpoints may store packed weights in integer + # Parameters (for example, ModelOpt NVFP4 uses uint8) and global scales as + # scalar Parameters. FSDP2 cannot represent non-floating or zero-dimensional + # sharded parameters. Keep those tensors resident and replicated; eligible + # scales, biases, and other parameters remain HSDP-sharded. + unshardable_params = _unshardable_parameters(model) + if unshardable_params: + if target_device is None: + raise ValueError( + f"Model {type(model).__name__} has parameters that HSDP must ignore, " + "but apply_hsdp_to_model was called without target_device." + ) + for param in unshardable_params: + if param.device != target_device: + param.data = param.data.to(target_device) + ignored_params.update(unshardable_params) + logger.info( + "HSDP excluding %d unshardable parameter tensors from sharding " + "(non_floating=%d, scalar=%d, dtypes=%s, moved to %s)", + len(unshardable_params), + sum(not (param.is_floating_point() or param.is_complex()) for param in unshardable_params), + sum(param.ndim == 0 for param in unshardable_params), + sorted({str(param.dtype) for param in unshardable_params}), + target_device, + ) + # Apply HSDP sharding, this will automatically handle weight distribution shard_model( model, @@ -165,6 +244,7 @@ def apply_hsdp_to_model( mp_policy=mp_policy, mesh=device_mesh, hsdp_shard_conditions=hsdp_shard_conditions, + ignored_params=ignored_params if ignored_params else None, ) for param in model.parameters(): @@ -181,8 +261,17 @@ def shard_model( mp_policy: MixedPrecisionPolicy | None = None, mesh: DeviceMesh | None = None, hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]], + ignored_params: set[nn.Parameter] | None = None, ) -> None: - """Apply HSDP sharding to model modules based on shard conditions.""" + """Apply HSDP sharding to model modules based on shard conditions. + + ignored_params (if provided) are excluded from the root fully_shard + wrap, so they are not collected into the root flat-parameter, are not + subject to MixedPrecisionPolicy, and retain their original dtype. + Each per-submodule wrap receives the subset of ignored_params that it owns. + This is required for packed integer parameters inside sharded transformer + blocks; the root wrap receives the full set for all remaining parameters. + """ hsdp_kwargs: dict[str, Any] = { "reshard_after_forward": reshard_after_forward, "mesh": mesh, @@ -192,11 +281,23 @@ def shard_model( num_sharded = 0 for name, module in reversed(list(model.named_modules())): if any(cond(name, module) for cond in hsdp_shard_conditions): - fully_shard(module, **hsdp_kwargs) + module_kwargs = dict(hsdp_kwargs) + if ignored_params: + module_ignored_params = ignored_params.intersection(module.parameters()) + if module_ignored_params: + module_kwargs["ignored_params"] = module_ignored_params + fully_shard(module, **module_kwargs) num_sharded += 1 if num_sharded == 0: raise ValueError("No modules were sharded") - fully_shard(model, **hsdp_kwargs) - logger.info("Sharded %d modules + root", num_sharded) + root_kwargs = dict(hsdp_kwargs) + if ignored_params: + root_kwargs["ignored_params"] = ignored_params + fully_shard(model, **root_kwargs) + logger.info( + "Sharded %d modules + root (ignored_params=%d)", + num_sharded, + len(ignored_params) if ignored_params else 0, + ) diff --git a/vllm_omni/diffusion/distributed/parallel_state.py b/vllm_omni/diffusion/distributed/parallel_state.py index 163f9dd45b0..7b38d377563 100644 --- a/vllm_omni/diffusion/distributed/parallel_state.py +++ b/vllm_omni/diffusion/distributed/parallel_state.py @@ -35,7 +35,6 @@ from vllm.logger import init_logger from vllm_omni.diffusion import envs -from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.platforms import current_omni_platform @@ -61,6 +60,11 @@ _FS: GroupCoordinator | None = None # Fully Sharded (HSDP shard dimension) _DIT: GroupCoordinator | None = None +# Rank-layout metadata for expert parallelism. This is not a process group; +# it is reused by platform-specific runtimes that must build companion groups +# with the same rank layout as vLLM EP. +_EXPERT_PARALLEL_GROUP_RANKS: list[list[int]] | None = None + def generate_masked_orthogonal_rank_groups( world_size: int, parallel_size: list[int], mask: list[bool] @@ -190,7 +194,6 @@ def __init__( self.fs = fs self.rank_offset = rank_offset self.world_size = tp * sp * pp * cfg * dp - self.ep = tp * sp * cfg * dp # EP level exclude PP self.name_to_size = { "tp": self.tp, @@ -253,15 +256,6 @@ def get_ranks(self, token, independent_ranks: bool = False): ranks.append(group) return ranks - if token == "ep": - ranks = [] - num_pp_stages = self.pp - for i in range(num_pp_stages): - start = i * self.ep + self.rank_offset - end = start + self.ep - ranks.append(list(range(start, end))) - return ranks - mask = self.get_mask(self.order, token) ranks = generate_masked_orthogonal_rank_groups(self.world_size, self.ordered_size, mask) if self.rank_offset > 0: @@ -309,6 +303,12 @@ def get_ring_parallel_rank(): return get_sp_group().ring_rank +def get_expert_parallel_group_ranks() -> list[list[int]]: + assert vllm_parallel_state._EP is not None, "expert parallel group is not initialized" + assert _EXPERT_PARALLEL_GROUP_RANKS is not None, "expert parallel group ranks are not initialized" + return _EXPERT_PARALLEL_GROUP_RANKS + + # PP def get_pp_group() -> PipelineGroupCoordinator: assert _PP is not None, "pipeline model parallel group is not initialized" @@ -512,6 +512,21 @@ def init_model_parallel_group( ) +def init_vllm_model_parallel_group( + group_ranks: list[list[int]], + local_rank: int, + backend: str, + group_name: str, +) -> vllm_parallel_state.GroupCoordinator: + return vllm_parallel_state.init_model_parallel_group( + group_ranks=group_ranks, + local_rank=local_rank, + backend=backend, + group_name=group_name, + use_device_communicator=True, + ) + + def init_dit_group( dit_parallel_size: int, backend: str, @@ -782,6 +797,13 @@ def initialize_model_parallel( fs=fully_shard_degree, order="tp-sp-pp-cfg-dp", ) + use_moe_parallel_mapping = False + if enable_expert_parallel: + od_config = get_forward_context().omni_diffusion_config + use_moe_parallel_mapping = bool(od_config and od_config.is_moe) + if not use_moe_parallel_mapping: + raise RuntimeError("Expert parallelism enabled for a non-MoE model") + sp_group_ranks = rank_generator.get_ranks("sp") global _DP assert _DP is None, "data parallel group is already initialized" @@ -828,16 +850,46 @@ def initialize_model_parallel( ulysses_group=ulysses_pg, ring_group=ring_pg, ) + if use_moe_parallel_mapping: + # Diffusion normally uses its own SP group. Map it to vLLM PCP only for + # expert-parallel runtimes that rely on vLLM FusedMoE group semantics. + # vLLM 0.24 MoE kernels require GroupCoordinator.device_communicator + # and reduce_scatter(), which the diffusion SP coordinator intentionally + # does not own. Keep the rank layout but build a vLLM coordinator. + vllm_parallel_state._PCP = init_vllm_model_parallel_group( + group_ranks=sp_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + group_name="pcp", + ) assert vllm_parallel_state._TP is None, "Tensor parallel group is already initialized" - vllm_parallel_state._TP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("tp"), - local_rank=get_world_group().local_rank, - backend=backend, - parallel_mode="tensor", - ) + tp_group_ranks = rank_generator.get_ranks("tp") + if use_moe_parallel_mapping: + vllm_parallel_state._TP = init_vllm_model_parallel_group( + group_ranks=tp_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + group_name="tp", + ) + else: + vllm_parallel_state._TP = init_model_parallel_group( + group_ranks=tp_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="tensor", + ) + if use_moe_parallel_mapping: + # CFG is a diffusion-specific replica dimension. Fold it into vLLM DP + # only when constructing the vLLM EP layout for expert-parallel paths. + vllm_parallel_state._DP = init_vllm_model_parallel_group( + group_ranks=rank_generator.get_ranks("cfg-dp"), + local_rank=get_world_group().local_rank, + backend=backend, + group_name="dp", + ) - global _FS + global _FS, _EXPERT_PARALLEL_GROUP_RANKS assert _FS is None, "fully shard group is already initialized" _FS = init_model_parallel_group( group_ranks=rank_generator.get_ranks("fs", independent_ranks=True), @@ -846,34 +898,40 @@ def initialize_model_parallel( parallel_mode="fully_shard", ) - if enable_expert_parallel: - od_config: OmniDiffusionConfig | None = get_forward_context().omni_diffusion_config - if od_config and od_config.is_moe: - vllm_parallel_state._EP = init_model_parallel_group( - group_ranks=rank_generator.get_ranks("ep"), - local_rank=get_world_group().local_rank, - backend=backend, - parallel_mode="expert", - ) - else: - raise RuntimeError("Expert parallelism enabled for a non-MoE model ") + _EXPERT_PARALLEL_GROUP_RANKS = None + if use_moe_parallel_mapping: + ep_group_ranks = rank_generator.get_ranks("tp-sp-cfg-dp") + vllm_parallel_state._EP = init_vllm_model_parallel_group( + group_ranks=ep_group_ranks, + local_rank=get_world_group().local_rank, + backend=backend, + group_name="ep", + ) + _EXPERT_PARALLEL_GROUP_RANKS = ep_group_ranks init_dit_group(dit_parallel_size, backend) def destroy_model_parallel(): """Set the groups to none and destroy them.""" - global _DP + global _DP, _CFG, _SP, _PP, _FS, _EXPERT_PARALLEL_GROUP_RANKS + + if vllm_parallel_state._DP and vllm_parallel_state._DP is not _DP: + vllm_parallel_state._DP.destroy() + vllm_parallel_state._DP = None + if _DP: _DP.destroy() _DP = None - global _CFG if _CFG: _CFG.destroy() _CFG = None - global _SP + if vllm_parallel_state._PCP and vllm_parallel_state._PCP is not _SP: + vllm_parallel_state._PCP.destroy() + vllm_parallel_state._PCP = None + if _SP: _SP.destroy() _SP = None @@ -885,13 +943,12 @@ def destroy_model_parallel(): if vllm_parallel_state._EP: vllm_parallel_state._EP.destroy() vllm_parallel_state._EP = None + _EXPERT_PARALLEL_GROUP_RANKS = None - global _PP if _PP: _PP.destroy() _PP = None - global _FS if _FS: _FS.destroy() _FS = None diff --git a/vllm_omni/diffusion/distributed/utils.py b/vllm_omni/diffusion/distributed/utils.py index 851328e6233..fd4310e3f5f 100644 --- a/vllm_omni/diffusion/distributed/utils.py +++ b/vllm_omni/diffusion/distributed/utils.py @@ -14,3 +14,26 @@ def get_local_device() -> torch.device: return current_omni_platform.get_torch_device(current_omni_platform.current_device()) local_rank = int(os.environ.get("LOCAL_RANK", 0)) return current_omni_platform.get_torch_device(local_rank) + + +def build_local_sp_padding_mask( + batch_size: int, + local_seq_len: int, + sp_original_seq_len: int | None, + sp_padding_size: int, + sequence_parallel_rank: int, + device, +): + """Build a per-rank SP padding mask that matches the local shard shape.""" + if sp_original_seq_len is None or sp_padding_size <= 0: + return None + + shard_start = sequence_parallel_rank * local_seq_len + shard_end = shard_start + local_seq_len + valid_tokens = max(0, min(sp_original_seq_len, shard_end) - shard_start) + + if valid_tokens >= local_seq_len: + return None + + token_mask = torch.arange(local_seq_len, device=device) < valid_tokens + return token_mask.unsqueeze(0).expand(batch_size, -1) diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index 370010d1b21..8bfb5923aa7 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/vllm_omni/diffusion/executor/abstract.py @@ -73,6 +73,11 @@ def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRun """Execute request-mode work from a scheduler output.""" pass + @abstractmethod + def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + """Execute request-mode work as a single batched RPC.""" + pass + @abstractmethod def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: """Execute step-mode work from a scheduler output.""" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index d78377e50cd..2e491ff6971 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -308,36 +308,68 @@ def register_failure_callback( self._failure_callbacks.append(callback) def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: - """Adapt request-mode scheduler output to worker execute_model RPC.""" - from vllm_omni.diffusion.worker.utils import RunnerOutput + """Adapt request-mode scheduler output to worker execute_model RPCs. + + Returns a BatchRunnerOutput with one RunnerOutput per scheduled request. + """ + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput self._ensure_open() - if scheduler_output.num_scheduled_reqs != 1: - raise ValueError( - f"Request mode currently supports batch_size=1, " - f"but got {scheduler_output.num_scheduled_reqs} scheduled requests." - ) + runner_outputs: list[RunnerOutput] = [] + + for new_req in scheduler_output.scheduled_new_reqs: + req = new_req.req + try: + result = self.collective_rpc( + "execute_model", + args=(req, self.od_config, scheduler_output.kv_prefetch_jobs), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if not isinstance(result, DiffusionOutput): + raise RuntimeError(f"Unexpected response type: {type(result)!r}") + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=result, + ) + ) + except Exception as exc: + runner_outputs.append( + RunnerOutput( + request_id=new_req.request_id, + step_index=None, + finished=True, + result=DiffusionOutput(error=str(exc)), + ) + ) - new_req = scheduler_output.scheduled_new_reqs[0] + return BatchRunnerOutput.from_list(runner_outputs) + + def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: + """Execute request-mode work through a single batched worker RPC. + + The worker builds DiffusionRequestBatch from scheduler output and returns + BatchRunnerOutput with one RunnerOutput per scheduled request. + """ + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput + + self._ensure_open() result = self.collective_rpc( - "execute_model", - args=(new_req.req, self.od_config), + "execute_model_batch", + args=(scheduler_output, self.od_config), unique_reply_rank=0, exec_all_ranks=True, ) - if not isinstance(result, DiffusionOutput): - raise RuntimeError(f"Unexpected response type for execute_request: {type(result)!r}") - - return RunnerOutput( - request_id=new_req.request_id, - step_index=None, - finished=True, - result=result, - ) + if not isinstance(result, BatchRunnerOutput): + raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}") + return result def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: """Forward step-mode scheduler output to worker execute_stepwise RPC.""" - from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, RunnerOutput + from vllm_omni.diffusion.worker.utils import BaseRunnerOutput self._ensure_open() result = self.collective_rpc( @@ -349,18 +381,7 @@ def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunner if isinstance(result, BaseRunnerOutput): return result - # TODO: Remove this fallback; DiffusionOutput cannot faithfully represent - # failed multi-request step batches. - if isinstance(result, DiffusionOutput): - request_id = scheduler_output.scheduled_request_ids[0] if scheduler_output.scheduled_request_ids else "" - return RunnerOutput( - request_id=request_id, - step_index=None, - finished=True, - result=result, - ) - else: - raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}") + raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}") def collective_rpc( self, diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index f2b6927d2fd..ea230ed9bed 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -4,6 +4,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING +import torch import vllm.ir from vllm.config import VllmConfig @@ -88,6 +89,40 @@ def is_forward_context_available() -> bool: return _forward_context is not None +def build_local_sp_padding_mask( + batch_size: int, + local_seq_len: int, + device, +): + """Build a per-rank SP padding mask that matches the local shard shape. + + Auto-padding is applied before sequence-parallel sharding, so attention on each + rank must receive a mask for its local shard, not for the global padded sequence. + """ + if not is_forward_context_available(): + return None + + ctx = get_forward_context() + if ctx.sp_original_seq_len is None or ctx.sp_padding_size <= 0: + return None + + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + ) + from vllm_omni.diffusion.distributed.utils import ( + build_local_sp_padding_mask as build_local_sp_padding_mask_for_rank, + ) + + return build_local_sp_padding_mask_for_rank( + batch_size=batch_size, + local_seq_len=local_seq_len, + sp_original_seq_len=ctx.sp_original_seq_len, + sp_padding_size=ctx.sp_padding_size, + sequence_parallel_rank=get_sequence_parallel_rank(), + device=device, + ) + + def get_ulysses_mode(*, default: str = "strict") -> str: """Resolve the Ulysses-SP mode from the current ForwardContext. diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index 137f2b887e4..8b3d90f0097 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -12,8 +12,6 @@ from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any -import torch -from PIL import Image from vllm.logger import init_logger from vllm.v1.engine.exceptions import EngineDeadError @@ -128,18 +126,24 @@ async def _dispatch_request( ) -> None: try: request = OmniDiffusionRequest( - prompts=[prompt], + prompt=prompt, sampling_params=sampling_params, request_id=request_id, kv_sender_info=kv_sender_info, ) - results = await self._engine.step(request) - result = results[0] - if not result.request_id: - result.request_id = request_id - - self._output_queue.put_nowait(result) + if self.od_config.streaming_output: + async for results in self._engine.step_streaming(request): + result = results[0] + if not result.request_id: + result.request_id = request_id + self._output_queue.put_nowait(result) + else: + results = await self._engine.step(request) + result = results[0] + if not result.request_id: + result.request_id = request_id + self._output_queue.put_nowait(result) except DiffusionRequestAbortedError as e: logger.info("request_id: %s aborted: %s", request_id, str(e)) except Exception as e: @@ -155,113 +159,6 @@ async def _dispatch_request( finally: self._tasks.pop(request_id, None) - async def add_batch_request_async( - self, - request_id: str, - prompts: list[OmniPromptType], - sampling_params: OmniDiffusionSamplingParams, - kv_sender_info: dict[int, dict[str, Any]] | None = None, - ) -> None: - logger.debug( - "[InlineStageDiffusionClient] stage-%s [rep-%s] add batch request: %s (%d prompts)", - self.stage_id, - self.replica_id, - request_id, - len(prompts), - ) - task = asyncio.create_task( - self._dispatch_batch( - request_id, - prompts, - sampling_params, - kv_sender_info, - ) - ) - self._tasks[request_id] = task - - async def _dispatch_batch( - self, - request_id: str, - prompts: list[Any], - sampling_params: OmniDiffusionSamplingParams, - kv_sender_info: dict[str, Any] | None = None, - ) -> None: - try: - request = OmniDiffusionRequest( - prompts=prompts, - sampling_params=sampling_params, - request_id=request_id, - kv_sender_info=kv_sender_info, - ) - - results = await self._engine.step(request) - - all_images: list = [] - merged_mm: dict[str, Any] = {} - merged_metrics: dict[str, Any] = {} - merged_durations: dict[str, float] = {} - merged_custom: dict[str, Any] = {} - peak_mem = 0.0 - latents = None - trajectory_latents: list[torch.Tensor] | None = None - trajectory_timesteps: list[torch.Tensor] | None = None - trajectory_log_probs: torch.Tensor | None = None - trajectory_decoded: list[Image.Image] | None = None - final_output_type = "image" - - for r in results: - all_images.extend(r.images) - merged_mm.update(r._multimodal_output) - merged_metrics.update(r.metrics) - merged_durations.update(r.stage_durations) - merged_custom.update(r._custom_output) - peak_mem = max(peak_mem, r.peak_memory_mb) - if latents is None and r.latents is not None: - latents = r.latents - if trajectory_latents is None: - trajectory_latents = r.trajectory_latents - if trajectory_timesteps is None: - trajectory_timesteps = r.trajectory_timesteps - if trajectory_log_probs is None: - trajectory_log_probs = r.trajectory_log_probs - if trajectory_decoded is None: - trajectory_decoded = r.trajectory_decoded - if r.final_output_type != "image": - final_output_type = r.final_output_type - - result = OmniRequestOutput.from_diffusion( - request_id=request_id, - images=all_images, - prompt=prompts[0] if len(prompts) == 1 else None, - metrics=merged_metrics, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - trajectory_log_probs=trajectory_log_probs, - trajectory_decoded=trajectory_decoded, - custom_output=merged_custom or None, - multimodal_output=merged_mm or None, - final_output_type=final_output_type, - stage_durations=merged_durations, - peak_memory_mb=peak_mem, - ) - - self._output_queue.put_nowait(result) - except DiffusionRequestAbortedError as e: - 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) - status_code, error_type = client_error_metadata(e) - error_output = OmniRequestOutput.from_error( - request_id=request_id, - error_message=str(e), - status_code=status_code, - error_type=error_type, - ) - self._output_queue.put_nowait(error_output) - finally: - self._tasks.pop(request_id, None) - def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: try: return self._output_queue.get_nowait() diff --git a/vllm_omni/diffusion/io_support.py b/vllm_omni/diffusion/io_support.py new file mode 100644 index 00000000000..8d1405a6cf5 --- /dev/null +++ b/vllm_omni/diffusion/io_support.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import inspect + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.registry import DiffusionModelRegistry + + +def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]: + if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None: + signature = inspect.signature(pipe_cls.__call__) + support_image_input = "image" in signature.parameters + support_audio_input = ( + "audio" in signature.parameters or "audio_latents" in signature.parameters + ) # ref. LTX-2 format + return support_image_input, support_audio_input + + supports_image_input = False + supports_audio_input = False + + model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) + if model_cls is not None: + supports_image_input = bool(getattr(model_cls, "support_image_input", False)) + supports_audio_input = bool(getattr(model_cls, "support_audio_input", False)) + return supports_image_input, supports_audio_input + + +def image_color_format(model_class_name: str) -> str: + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) + return getattr(model_cls, "color_format", "RGB") + + +def supports_audio_output(model_class_name: str) -> bool: + model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) + if model_cls is None: + return False + return bool(getattr(model_cls, "support_audio_output", False)) + + +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")) + return 2 if supports_audio_input or supports_audio_output(model_class_name) else 1 diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index 2268b45ab1f..cd83e718c19 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/vllm_omni/diffusion/ipc.py @@ -138,9 +138,10 @@ def _is_rpc_result_envelope(output: object) -> bool: def pack_diffusion_output_shm(output: object) -> object: """Replace large tensors in diffusion worker outputs with SHM handles. - Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one - in ``.result`` (for example ``RunnerOutput``), or an RPC result envelope - carrying the diffusion output in ``["result"]``. + Supports a bare ``DiffusionOutput``, a wrapper object carrying one in + ``.result`` (for example ``RunnerOutput``), an RPC result envelope carrying + the diffusion output in ``["result"]``, or a batch wrapper carrying + ``RunnerOutput`` objects in ``.runner_outputs``. """ if isinstance(output, DiffusionOutput): return _pack_diffusion_fields(output) @@ -154,6 +155,11 @@ def pack_diffusion_output_shm(output: object) -> object: result = getattr(output, "result", None) if isinstance(result, DiffusionOutput): output.result = _pack_diffusion_fields(result) + + runner_outputs = getattr(output, "runner_outputs", None) + if isinstance(runner_outputs, list): + for runner_output in runner_outputs: + pack_diffusion_output_shm(runner_output) return output @@ -179,4 +185,9 @@ def unpack_diffusion_output_shm(output: object) -> object: result = getattr(output, "result", None) if isinstance(result, DiffusionOutput): output.result = _unpack_diffusion_fields(result) + + runner_outputs = getattr(output, "runner_outputs", None) + if isinstance(runner_outputs, list): + for runner_output in runner_outputs: + unpack_diffusion_output_shm(runner_output) return output diff --git a/vllm_omni/diffusion/layers/mot/__init__.py b/vllm_omni/diffusion/layers/mot/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/diffusion/layers/mot/mot_layernorm.py b/vllm_omni/diffusion/layers/mot/mot_layernorm.py new file mode 100644 index 00000000000..0c6bd5b2ee8 --- /dev/null +++ b/vllm_omni/diffusion/layers/mot/mot_layernorm.py @@ -0,0 +1,96 @@ +# ruff: noqa: N803, E741 +"""Mixture-of-Tokens (MoT) RMS Normalization layer. + +Holds two sets of weights (text / gen) and routes tokens to the +appropriate weight based on indices. When text_indices is None the +layer degrades to a standard RMSNorm using self.weight (und mode). +""" + +import torch +import torch.nn as nn +from vllm import ir + +from vllm_omni.diffusion.layers.custom_op import CustomOp + + +class MoTRMSNorm(CustomOp): + """Mixture-of-Tokens RMS Normalization. + + In *und* mode (``text_indices is None``), every token is normalised + with ``self.weight`` – exactly like a vanilla RMSNorm. + + In *gen* mode, text tokens are normalised with ``self.weight`` and + gen tokens are normalised with ``self.gen_weight``, using a single + fused Triton kernel that avoids the gather / scatter overhead. + """ + + def __init__( + self, + hidden_size: int, + head_norm: bool = False, + eps: float = 1e-6, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.variance_epsilon = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.gen_weight = nn.Parameter(torch.ones(hidden_size)) + self.head_norm = head_norm + + # ------------------------------------------------------------------ + # Native (pure-PyTorch) fallback + # ------------------------------------------------------------------ + def forward_native( + self, + x: torch.Tensor, + text_indices: torch.Tensor | None = None, + vae_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + if text_indices is None: + return self._rms_norm_native(x, self.weight) + + output = torch.empty_like(x) + output[text_indices] = self._rms_norm_native(x[text_indices], self.weight) + output[vae_indices] = self._rms_norm_native(x[vae_indices], self.gen_weight) + return output + + # ------------------------------------------------------------------ + # CUDA fast-path + # ------------------------------------------------------------------ + def forward_cuda( + self, + x: torch.Tensor, + text_indices: torch.Tensor | None = None, + vae_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + if text_indices is None: + # und mode – delegate to vllm's highly-optimised CUDA kernel + return ir.ops.rms_norm(x, self.weight.data, self.variance_epsilon) + + # gen mode – fused MoT Triton kernel + from vllm_omni.diffusion.layers.mot.ops.mot_rms_norm import ( + mot_rms_norm, + ) + + return mot_rms_norm( + x, + self.weight.data, + self.gen_weight.data, + text_indices, + vae_indices, + head_norm=self.head_norm, + eps=self.variance_epsilon, + ) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + def _rms_norm_native(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + orig_dtype = x.dtype + x = x.float() + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + self.variance_epsilon) + return (x * weight.float()).to(orig_dtype) + + def extra_repr(self) -> str: + return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}" diff --git a/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py new file mode 100644 index 00000000000..ce83d884247 --- /dev/null +++ b/vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py @@ -0,0 +1,349 @@ +from __future__ import annotations + +import torch +from torch.nn.parameter import Parameter +from vllm.distributed import ( + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + WEIGHT_LOADER_V2_SUPPORTED, + QKVParallelLinear, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform + +from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm + +logger = init_logger(__name__) + + +class MoTQKVParallelLinear(QKVParallelLinear): + """QKVParallelLinear with Mixture-of-Tokens routing. + + Text weights: stored directly on self (self.weight, self.weight_scale, ...), + created through the standard QKVParallelLinear.__init__ process. + + VAE weights: stored in the permanent submodule self.gen_exp + (self.gen_exp.weight, ...), + created via quant_method.create_weights(self.gen_exp, ...). + gen_exp.quant_method points to the same quant_method, so that + the vLLM framework’s process_weights_after_loading + can automatically detect and process it. + + Forward behavior: + - und mode (text_indices is None): fully reuse super().forward() + - gen mode: call the MoT fused GEMM kernel + """ + + def __init__( + self, + hidden_size: int, + head_size: int, + total_num_heads: int, + total_num_kv_heads: int | None = None, + bias: bool = True, + vae_bias: bool = False, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + *, + return_bias: bool = True, + disable_tp: bool = False, + v_head_size: int | None = None, + ): + # ---- 1) Parent class creates text weights ---- + # super().__init__ will do the following: + # quant_method.create_weights(self, ...) → self.weight, self.weight_scale, ... + # QKVParallelLinear hardcodes gather_output=False + super().__init__( + hidden_size, + head_size, + total_num_heads, + total_num_kv_heads, + bias, + skip_bias_add, + params_dtype, + quant_config, + prefix, + return_bias=return_bias, + disable_tp=disable_tp, + v_head_size=v_head_size, + ) + + # ---- 2) Create vae weights (permanent submodule) ---- + if self.quant_method is None: + raise ValueError( + f"quant_method must not be None for MoTQKVParallelLinear (prefix={prefix!r}). " + "Ensure a valid QuantizationConfig is provided or the default UnquantizedLinearMethod is set." + ) + + # NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace + # container to hold the secondary VAE weights. + # This is a design choice for dynamically attaching `quant_method` to the bare module. + # This is slightly unconventional but necessary. It allows vLLM's framework + # (specifically `process_weights_after_loading`) to automatically discover + # and process these secondary weights without requiring a dedicated, + # boilerplate Module subclass just for the VAE pathway. + self.gen_exp = torch.nn.Module() + + # Use the same weight_loader as text + vae_weight_loader = ( + self.weight_loader_v2 + if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ) + self.quant_method.create_weights( + layer=self.gen_exp, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=vae_weight_loader, + ) + + # Make gen_exp discoverable by vLLM framework's process_weights_after_loading + self.gen_exp.quant_method = self.quant_method + + # ---- 3) vae bias ---- + if vae_bias: + self.gen_exp.bias = Parameter(torch.empty(self.output_size_per_partition, dtype=self.params_dtype)) + set_weight_attrs( + self.gen_exp.bias, + {"output_dim": 0, "weight_loader": self.weight_loader}, + ) + else: + self.gen_exp.register_parameter("bias", None) + + self.update_param_tp_status() + + # ================================================================== + # Forward + # ================================================================== + def forward( + self, + input_: torch.Tensor, + text_indices: torch.Tensor | None = None, + vae_indices: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: + # ---- und mode: fully reuse parent class (only text path) ---- + if text_indices is None: + return super().forward(input_) + + # ---- gen mode: fuse MoT GEMM ---- + output_parallel = self._mot_gemm_dispatch(input_, text_indices, vae_indices) + + # QKVParallelLinear hardcodes gather_output=False, this branch never enters; + # retained for future subclass changes + if self.gather_output and self.tp_size > 1: + output = tensor_model_parallel_all_gather(output_parallel) + else: + output = output_parallel + + if not self.return_bias: + return output + + if self.skip_bias_add and text_indices is not None and (self.bias is not None or self.gen_exp.bias is not None): + merged_bias = torch.zeros( + output.size(0), + self.output_size_per_partition, + dtype=output.dtype, + device=output.device, + ) + if self.bias is not None: + merged_bias[text_indices] = self.bias + if self.gen_exp.bias is not None: + merged_bias[vae_indices] = self.gen_exp.bias + return output, merged_bias + + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + # ================================================================== + # MoT GEMM dispatcher + # ================================================================== + def _mot_gemm_dispatch( + self, + x: torch.Tensor, + text_indices: torch.Tensor, + vae_indices: torch.Tensor, + ) -> torch.Tensor: + # Any other backend (ROCm, XPU, TPU, CPU) uses the safe vllm fallback. + if not current_platform.is_cuda(): + return self._mot_fallback(x, text_indices, vae_indices) + + """Dispatch to different MoT kernel paths based on weight dtype / quant attributes.""" + bias_text = self.bias if not self.skip_bias_add else None + bias_vae = self.gen_exp.bias if not self.skip_bias_add else None + + w_text = self.weight + w_vae = self.gen_exp.weight + assert w_text.dtype.is_floating_point == w_vae.dtype.is_floating_point, ( + "weight of text expert and image expert should be the same dtype." + ) + + # w_text.dtype.itemsize >= 2 means bytes_per_element >= 2 (16bits or 32bits) + if w_text.dtype.is_floating_point and w_text.dtype.itemsize >= 2: + return self._mot_gemm_unquantized( + x, + text_indices, + vae_indices, + bias_text, + bias_vae, + ) + elif w_text.dtype == torch.float8_e4m3fn: + return self._mot_gemm_fp8_w8a8( + x, + text_indices, + vae_indices, + bias_text, + bias_vae, + ) + elif w_text.dtype == torch.int8: + return self._mot_gemm_weight_only( + x, + text_indices, + vae_indices, + bias_text, + bias_vae, + ) + else: + return self._mot_fallback( + x, + text_indices, + vae_indices, + ) + + # ================================================================== + # Implementation of each quantization path + # ================================================================== + def _mot_gemm_unquantized(self, x, text_idx, vae_idx, bias_t, bias_v): + """BF16/FP16 path.""" + N = self.output_size_per_partition + C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device) + invoke_mot_gemm( + A=x, + B_text=self.weight.data.t(), + B_vae=self.gen_exp.weight.data.t(), + C=C, + bias_text=bias_t, + bias_vae=bias_v, + text_indices=text_idx, + vae_indices=vae_idx, + A_scale=None, + B_text_scale=None, + B_vae_scale=None, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + A_per_channel_quant=False, + B_per_channel_quant=False, + ) + return C + + def _mot_gemm_fp8_w8a8(self, x, text_idx, vae_idx, bias_t, bias_v): + """FP8 W8A8 path. + + 1) Activation quantization: reuse vllm's ops.scaled_fp8_quant + 2) MoT GEMM: text/vae use different fp8 weights and weight_scale + 3) De-quantization: done by MoT kernel's epilogue + """ + from vllm import _custom_ops as ops + + x_2d = x.view(-1, x.shape[-1]) + input_scale = getattr(self, "input_scale", None) + x_fp8, x_scale = ops.scaled_fp8_quant( + x_2d, + input_scale, + use_per_token_if_dynamic=True, + ) + + N = self.output_size_per_partition + C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device) + + invoke_mot_gemm( + A=x_fp8, + B_text=self.weight.data, + B_vae=self.gen_exp.weight.data, + C=C, + bias_text=bias_t, + bias_vae=bias_v, + text_indices=text_idx, + vae_indices=vae_idx, + A_scale=x_scale, + B_text_scale=self.weight_scale.data, + B_vae_scale=self.gen_exp.weight_scale.data, + use_fp8_w8a8=True, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + A_per_channel_quant=True, + B_per_channel_quant=False, + ) + return C + + def _mot_gemm_weight_only(self, x, text_idx, vae_idx, bias_t, bias_v): + """Weight-Only W8A16 path. + + Activation values kept as bf16/fp16, weights are int8 + per-channel scale. + De-quantization done by MoT kernel's epilogue. + """ + N = self.output_size_per_partition + C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device) + invoke_mot_gemm( + A=x, + B_text=self.weight.data.t(), + B_vae=self.gen_exp.weight.data.t(), + C=C, + bias_text=bias_t, + bias_vae=bias_v, + text_indices=text_idx, + vae_indices=vae_idx, + A_scale=None, + B_text_scale=self.weight_scale.data, + B_vae_scale=self.gen_exp.weight_scale.data, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=True, + use_int4_w4a16=False, + A_per_channel_quant=False, + B_per_channel_quant=True, + ) + return C + + def _mot_fallback(self, x, text_idx, vae_idx): + """Fallback: fall back to gather/scatter + quant_method.apply. + + For unsupported quantization types, call standard forward for text/vae tokens. + """ + if self.quant_method is None: + raise ValueError( + "quant_method must not be None in MoTQKVParallelLinear._mot_fallback. " + "This indicates an initialization error." + ) + + bias_text = self.bias if not self.skip_bias_add else None + bias_vae = self.gen_exp.bias if not self.skip_bias_add else None + + output = torch.zeros( + x.size(0), + self.output_size_per_partition, + dtype=x.dtype, + device=x.device, + ) + output_text = self.quant_method.apply(self, x[text_idx], bias_text) + + output_vae = self.quant_method.apply( + self.gen_exp, + x[vae_idx], + bias_vae, + ) + output[text_idx] = output_text + output[vae_idx] = output_vae + return output diff --git a/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py new file mode 100644 index 00000000000..2811b110f45 --- /dev/null +++ b/vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py @@ -0,0 +1,361 @@ +from __future__ import annotations + +import torch +from torch.nn.parameter import Parameter +from vllm.distributed import ( + split_tensor_along_last_dim, + tensor_model_parallel_all_reduce, +) +from vllm.model_executor.layers.linear import ( + WEIGHT_LOADER_V2_SUPPORTED, + RowParallelLinear, +) +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform + +from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm + + +class MoTRowParallelLinear(RowParallelLinear): + """RowParallelLinear with Mixture-of-Tokens routing. + + text weights: directly on self (self.weight, self.weight_scale, ...), + created by RowParallelLinear.__init__ standard process. + vae weights: on permanent submodule self.gen_exp (self.gen_exp.weight, ...), + created by quant_method.create_weights(self.gen_exp, ...). + gen_exp.quant_method points to the same quant_method, enabling + vLLM framework's process_weights_after_loading to automatically + discover and process it. + + Forward behavior: + - und mode (text_indices is None): fully reuse super().forward() + - gen mode: call MoT fused GEMM kernel, then execute TP all-reduce + """ + + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + vae_bias: bool = False, + input_is_parallel: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + reduce_results: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + *, + return_bias: bool = True, + disable_tp: bool = False, + ): + # ---- Step 1: Parent class creates text weights ---- + # super().__init__ internally calls: + # quant_method.create_weights(self, ...) → self.weight, self.weight_scale, ... + super().__init__( + input_size, + output_size, + bias, + input_is_parallel, + skip_bias_add, + params_dtype, + reduce_results, + quant_config, + prefix, + return_bias=return_bias, + disable_tp=disable_tp, + ) + + # ---- Step 2: Create vae weights (permanent submodule) ---- + if self.quant_method is None: + raise ValueError( + f"quant_method must not be None for MoTRowParallelLinear (prefix={prefix!r}). " + "Ensure a valid QuantizationConfig is provided or the default UnquantizedLinearMethod is set." + ) + + # NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace + # container to hold the secondary VAE weights. + # This is a design choice for dynamically attaching `quant_method` to the bare module. + # This is slightly unconventional but necessary. It allows vLLM's framework + # (specifically `process_weights_after_loading`) to automatically discover + # and process these secondary weights without requiring a dedicated, + # boilerplate Module subclass just for the VAE pathway. + self.gen_exp = torch.nn.Module() + + # Select weight_loader consistent with text + vae_weight_loader = ( + self.weight_loader_v2 + if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED + else self.weight_loader + ) + self.quant_method.create_weights( + layer=self.gen_exp, + input_size_per_partition=self.input_size_per_partition, + output_partition_sizes=self.output_partition_sizes, + input_size=self.input_size, + output_size=self.output_size, + params_dtype=self.params_dtype, + weight_loader=vae_weight_loader, + ) + + # Enable gen_exp to be automatically discovered by vLLM framework's process_weights_after_loading + self.gen_exp.quant_method = self.quant_method + + # ---- Step 3: vae bias ---- + # RowParallelLinear's bias is full output_size (not sharded) + if vae_bias: + self.gen_exp.bias = Parameter(torch.empty(self.output_size, dtype=self.params_dtype)) + set_weight_attrs( + self.gen_exp.bias, + {"output_dim": 0, "weight_loader": self.weight_loader}, + ) + else: + self.gen_exp.register_parameter("bias", None) + + self.update_param_tp_status() + + # ================================================================== + # Forward + # ================================================================== + def forward( + self, + input_: torch.Tensor, + text_indices: torch.Tensor | None = None, + vae_indices: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]: + # ---- und mode: fully reuse parent class (only text path) ---- + if text_indices is None: + return super().forward(input_) + + # ---- gen mode ---- + # Handle input_is_parallel (consistent with parent class logic) + if self.input_is_parallel: + input_parallel = input_ + else: + split_input = split_tensor_along_last_dim(input_, num_partitions=self.tp_size) + input_parallel = split_input[self.tp_rank].contiguous() + + # Fused MoT GEMM + output_parallel = self._mot_gemm_dispatch(input_parallel, text_indices, vae_indices) + + # ---- TP communication: all-reduce (consistent with parent class) ---- + if self.reduce_results and self.tp_size > 1: + output = tensor_model_parallel_all_reduce(output_parallel) + else: + output = output_parallel + + if not self.return_bias: + return output + + if self.skip_bias_add and text_indices is not None and (self.bias is not None or self.gen_exp.bias is not None): + # Construct per-token mixed bias + merged_bias = torch.zeros( + output.size(0), + self.output_size_per_partition, + dtype=output.dtype, + device=output.device, + ) + if self.bias is not None: + merged_bias[text_indices] = self.bias + if self.gen_exp.bias is not None: + merged_bias[vae_indices] = self.gen_exp.bias + return output, merged_bias + + output_bias = self.bias if self.skip_bias_add else None + return output, output_bias + + # ================================================================== + # MoT GEMM Dispatcher + # ================================================================== + def _mot_gemm_dispatch( + self, + x: torch.Tensor, + text_indices: torch.Tensor, + vae_indices: torch.Tensor, + ) -> torch.Tensor: + """Dispatch to different MoT kernel paths based on weight dtype / quant attributes.""" + + # Any other backend (ROCm, XPU, TPU, CPU) uses the safe vllm fallback. + if not current_platform.is_cuda(): + return self._mot_fallback(x, text_indices, vae_indices) + + # RowParallelLinear: bias only fused into GEMM at rank 0, + # other ranks pass None to avoid duplicate accumulation + bias_text = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias + bias_vae = None if (self.tp_rank > 0 or self.skip_bias_add) else self.gen_exp.bias + + # Determine quantization type by weight dtype, avoid isinstance coupling to specific quant_method + # TODO: Currently does not support online quantization fp8, does not support quantization types involving int4 + w_text = self.weight + w_vae = self.gen_exp.weight + assert w_text.dtype.is_floating_point == w_vae.dtype.is_floating_point, ( + "weight of text expert and image expert should be the same dtype." + ) + + # w_text.dtype.itemsize >= 2 means bytes_per_element >= 2 (16bits or 32bits) + if w_text.dtype.is_floating_point and w_text.dtype.itemsize >= 2: + # ---- Path 0: BF16 / FP16 (unquantized) ---- + return self._mot_gemm_unquantized( + x, + text_indices, + vae_indices, + bias_text, + bias_vae, + ) + elif w_text.dtype == torch.float8_e4m3fn: + # ---- Path 1: FP8 W8A8 ---- + return self._mot_gemm_fp8_w8a8( + x, + text_indices, + vae_indices, + bias_text, + bias_vae, + ) + elif w_text.dtype == torch.int8: + # ---- Path 2: Weight-Only INT8 W8A16 ---- + return self._mot_gemm_weight_only( + x, + text_indices, + vae_indices, + bias_text, + bias_vae, + ) + else: + # ---- Fallback: gather/scatter + quant_method.apply ---- + return self._mot_fallback( + x, + text_indices, + vae_indices, + ) + + # ================================================================== + # Specific implementations for each quantization path + # ================================================================== + def _mot_gemm_unquantized(self, x, text_idx, vae_idx, bias_t, bias_v): + """BF16/FP16 path.""" + N = self.output_size_per_partition + C = torch.empty(x.size(0), N, dtype=x.dtype, device=x.device) + invoke_mot_gemm( + A=x, + B_text=self.weight.data.t(), + B_vae=self.gen_exp.weight.data.t(), + C=C, + bias_text=bias_t, + bias_vae=bias_v, + text_indices=text_idx, + vae_indices=vae_idx, + A_scale=None, + B_text_scale=None, + B_vae_scale=None, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + A_per_channel_quant=False, + B_per_channel_quant=False, + ) + return C + + def _mot_gemm_fp8_w8a8(self, x, text_idx, vae_idx, bias_t, bias_v): + """FP8 W8A8 path. + + 1) Activation quantization: reuse vllm's ops.scaled_fp8_quant + 2) MoT GEMM: text/vae each use different fp8 weights and weight_scale + 3) Dequantization: completed by MoT kernel internal epilogue + """ + from vllm import _custom_ops as ops + + x_2d = x.view(-1, x.shape[-1]) + input_scale = getattr(self, "input_scale", None) + x_fp8, x_scale = ops.scaled_fp8_quant( + x_2d, + input_scale, + use_per_token_if_dynamic=True, + ) + + N = self.output_size_per_partition + C = torch.empty(x.size(0), N, dtype=x.dtype, device=x.device) + + # weight has been transposed to (K, N) in process_weights_after_loading + invoke_mot_gemm( + A=x_fp8, + B_text=self.weight.data, + B_vae=self.gen_exp.weight.data, + C=C, + bias_text=bias_t, + bias_vae=bias_v, + text_indices=text_idx, + vae_indices=vae_idx, + A_scale=x_scale, + B_text_scale=self.weight_scale.data, + B_vae_scale=self.gen_exp.weight_scale.data, + use_fp8_w8a8=True, + use_int8_w8a8=False, + use_int8_w8a16=False, + use_int4_w4a16=False, + A_per_channel_quant=True, + B_per_channel_quant=False, + ) + return C + + def _mot_gemm_weight_only(self, x, text_idx, vae_idx, bias_t, bias_v): + """Weight-Only W8A16 path. + + Activations remain bf16/fp16, weights are int8 + per-channel scale. + MoT kernel internally performs immediate dequantization. + """ + N = self.output_size_per_partition + C = torch.empty(x.size(0), N, dtype=x.dtype, device=x.device) + invoke_mot_gemm( + A=x, + B_text=self.weight.data.t(), + B_vae=self.gen_exp.weight.data.t(), + C=C, + bias_text=bias_t, + bias_vae=bias_v, + text_indices=text_idx, + vae_indices=vae_idx, + A_scale=None, + B_text_scale=self.weight_scale.data, + B_vae_scale=self.gen_exp.weight_scale.data, + use_fp8_w8a8=False, + use_int8_w8a8=False, + use_int8_w8a16=True, + use_int4_w4a16=False, + A_per_channel_quant=False, + B_per_channel_quant=True, + ) + return C + + def _mot_fallback(self, x, text_idx, vae_idx): + """Fallback: degrade to gather/scatter + quant_method.apply. + + For unsupported quantization types, call standard forward separately for text/vae tokens. + """ + if self.quant_method is None: + raise ValueError( + "quant_method must not be None in MoTRowParallelLinear._mot_fallback. " + "This indicates an initialization error." + ) + + # RowParallelLinear: bias only fused at rank 0 + bias_text = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias + bias_vae = None if (self.tp_rank > 0 or self.skip_bias_add) else self.gen_exp.bias + + output = torch.empty( + x.size(0), + self.output_size_per_partition, + dtype=x.dtype, + device=x.device, + ) + # text tokens → standard quant_method (operate on weights on self) + output[text_idx] = self.quant_method.apply(self, x[text_idx], bias_text) + # vae tokens → same quant_method (operate on weights on gen_exp) + output[vae_idx] = self.quant_method.apply( + self.gen_exp, + x[vae_idx], + bias_vae, + ) + return output diff --git a/vllm_omni/diffusion/layers/mot/ops/__init__.py b/vllm_omni/diffusion/layers/mot/ops/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py b/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py new file mode 100644 index 00000000000..f1bffbbac13 --- /dev/null +++ b/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py @@ -0,0 +1,908 @@ +# ruff: noqa: N803, E741 +import functools +import json +import logging +import os +import pathlib +import re +from typing import Any + +import torch +from vllm.triton_utils import tl, triton + +logger = logging.getLogger(__name__) + +# ===================================================================== +# MoT GEMM Config Loading (3-tier: env → built-in → default) +# +# Usage pattern (mirrors vLLM fused_moe): +# 1. Layer invokes invoke_mot_gemm(...) without extra config context. +# 2. invoke_mot_gemm lazily calls +# get_mot_configs(K, N, dtype_str) +# which is @lru_cache'd and lazily loads the JSON on first hit. +# 3. If get_mot_configs returns None, fall back to +# get_mot_default_config(M, N, K, ...) +# ===================================================================== + +_CONFIGS_DIR = pathlib.Path(__file__).resolve().parent.parent / "configs" +_ENV_CONFIG_FOLDER = "VLLM_TUNED_CONFIG_FOLDER" + + +def get_device_name() -> str: + """Sanitized GPU device name, matching mot_linear_benchmarks.py output.""" + raw = torch.cuda.get_device_name(0) + name = re.sub(r"[^a-zA-Z0-9]", "", raw.replace(" ", "")) + for prefix in ("NVIDIA", "AMD"): + if name.startswith(prefix): + name = name[len(prefix) :] + # --- Device Aliasing Patch --- + alias_map = { + "A800": "A100", + "H800": "H100", + } + # map A800/H800 to A100/H100 for chinese market + for key, target in alias_map.items(): + if key in name: + name = name.replace(key, target) + break + return name + + +def build_config_filename(device_name: str, dtype_str: str) -> str: + return f"device_name={device_name},dtype={dtype_str}.json" + + +def _try_load_json(filepath: str) -> dict | None: + if os.path.isfile(filepath): + with open(filepath) as f: + return json.load(f) + return None + + +@functools.lru_cache +def _load_mot_config_file(dtype_str: str) -> dict | None: + """Load and cache the full MoT config JSON (one file per device/dtype). + + Search order: + 1. ``$VLLM_TUNED_CONFIG_FOLDER/device_name=...,dtype=....json`` + 2. ``vllm_omni/.../mot/configs/device_name=...,dtype=....json`` + 3. Return ``None`` (caller falls back to ``get_mot_default_config``). + """ + device_name = get_device_name() + filename = build_config_filename(device_name, dtype_str) + + config_file_paths: list[str] = [] + env_dir = os.environ.get(_ENV_CONFIG_FOLDER) + if env_dir: + config_file_paths.append(str(pathlib.Path(env_dir) / filename)) + config_file_paths.append(str(_CONFIGS_DIR / filename)) + + for path in config_file_paths: + data = _try_load_json(path) + if data is not None: + logger.info("MoT config loaded from %s", path) + return data + + logger.warning( + f"\n{'=' * 80}\n" + f" ⚠️ [WARNING] No tuned MoT config found.\n" + f" Searched paths: {', '.join(config_file_paths)}\n" + f" Using conservative default configs which are NOT optimal.\n" + f" Run `python benchmarks/kernels/mot_linear_benchmarks.py --tune` \n" + f" to generate hardware-specific optimal configs.\n" + f"{'=' * 80}\n" + ) + return None + + +@functools.lru_cache +def get_mot_configs( + K: int, + N: int, + dtype_str: str | None = None, +) -> dict[int, dict[str, int]] | None: + """Return ``{M: tile_config}`` for a given (K, N) shape, or ``None``. + + The return value maps an irregular grid of batch sizes (M) to Triton + tile configurations. The caller should pick the entry whose M is + closest to the actual batch size. + + Config file is selected by ``device_name + dtype``. + """ + file_data = _load_mot_config_file(dtype_str or "w16a16") + if file_data is None: + return None + + shape_entry = file_data.get(f"{K}_{N}") + if shape_entry is None: + logger.warning( + f"\n{'=' * 80}\n" + f" ⚠️ [WARNING] MoT config file found, but NO tuned entry for shape K={K}, N={N}.\n" + f" Using conservative default configs which are NOT optimal for this specific shape.\n" + f" Run `python benchmarks/kernels/mot_linear_benchmarks.py --tune` \n" + f" to generate hardware-specific optimal configs.\n" + f"{'=' * 80}\n" + ) + return None + + return {int(k): dict(v) for k, v in shape_entry.items() if k != "_comment"} + + +def get_mot_default_config( + M: int, + N: int, + K: int, + dtype: str | None = None, + block_quant_shape: list[int] | None = None, +) -> dict[str, int]: + """Conservative fallback config guaranteed to compile on all hardware. + + Trades peak performance for universal compatibility (T4 / V100 / A100 / + H100, CUDA & ROCm). + """ + # FP8 block-wise quantization requires strict alignment + if dtype == "fp8_w8a8" and block_quant_shape is not None: + return { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": block_quant_shape[0], + "BLOCK_SIZE_K": block_quant_shape[1], + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + } + + # Very small M (tail batches, final feature concat, etc.) + # NOTE If M is not a multiple of the block size (e.g., M=3 with BLOCK_SIZE_M=16), + # the kernel will allocate a full tile. Correctness is guaranteed by boundary masks. + # We tolerate this minor compute/SRAM padding to avoid generating excessive + # Triton compiled variants for extremely small sequence lengths. + if M <= 64: + return { + "BLOCK_SIZE_M": 16 if M <= 16 else 32, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + } + + # Standard fallback for typical image-generation M (2048 / 4096 / …) + # SRAM usage: (64*32 + 64*32) * 2 * 2 = 16 KB — safe on decade-old GPUs + return { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_warps": 4, + "num_stages": 2, + } + + +def get_best_mot_config(M: int, N: int, K: int, dtype_str: str | None = None) -> tuple[int, dict[str, int]]: + configs = get_mot_configs(K, N, dtype_str) + if configs: + loaded_m_key = min(configs.keys(), key=lambda x: abs(x - M)) + return loaded_m_key, configs[loaded_m_key] + else: + return -1, get_mot_default_config(M, N, K, dtype=dtype_str) + + +# ================================================================= +# Part 1: The Router (Routing Component) +# Responsibilities: Handle PID mapping, Text/VAE distribution, indirect index loading, pointer calculation +# ================================================================= +@triton.jit +def _get_mot_pointers( + # System Inputs + pid, + # Matrix Pointers + a_ptr, + b_text_ptr, + b_vae_ptr, + bias_text_ptr, + bias_vae_ptr, + scale_a_ptr, + scale_b_text_ptr, + scale_b_vae_ptr, + # Indices & Meta + text_indices_ptr, + vae_indices_ptr, + M_text, + M_vae, + N, + # Strides (need to select based on Text/VAE) + stride_bk_text, + stride_bn_text, + stride_bk_vae, + stride_bn_vae, + # Block Config + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + # 1. Calculate Text/VAE task boundaries + num_pid_m_text = tl.cdiv(M_text, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + # 2. MoT Routing + # Initialize VAE variables first and overwrite for Text + + # VAE Path + cur_pid = pid - (num_pid_m_text * num_pid_n) + # Select VAE Pointers + cur_b_ptr = b_vae_ptr + cur_bias_ptr = bias_vae_ptr + cur_scale_b_ptr = scale_b_vae_ptr + cur_indices_ptr = vae_indices_ptr + # Select VAE Strides / Limits + cur_stride_bk = stride_bk_vae + cur_stride_bn = stride_bn_vae + M_limit = M_vae + + # Text Path + if pid < num_pid_m_text * num_pid_n: + cur_pid = pid + # Select Text Pointers + cur_b_ptr = b_text_ptr + cur_bias_ptr = bias_text_ptr + cur_scale_b_ptr = scale_b_text_ptr + cur_indices_ptr = text_indices_ptr + # Select Text Strides / Limits + cur_stride_bk = stride_bk_text + cur_stride_bn = stride_bn_text + M_limit = M_text + + # 3. Calculate Grid coordinates(grouping) + cur_num_pid_m = tl.cdiv(M_limit, BLOCK_SIZE_M) + + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = cur_pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m_adj = tl.minimum(cur_num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (cur_pid % group_size_m_adj) + pid_n = (cur_pid % num_pid_in_group) // group_size_m_adj + + # 4. Load indirect indices (Indirect Indexing for A) + # Calculate the M range covered by current Block [0, BLOCK_SIZE_M] + offs_m_idx = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + m_mask = offs_m_idx < M_limit + # Load real row indices from the index array + # for very big K (eg. ffn down proj*TP=1,K=16k), if M is also huge(M>130K) + # loading int32 indices may result in integer overflow when we compute offs_m + real_row_idxs = tl.load(cur_indices_ptr + offs_m_idx, mask=m_mask, other=0).to(tl.int64) + + # 5. Calculate N-dimension Offsets + offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64) + n_mask = offs_n < N + + return ( + pid_m, + pid_n, # Grid Coordinates + real_row_idxs, + m_mask, # M-dim info (Indirect) + offs_n, + n_mask, # N-dim info + M_limit, # Boundary + cur_b_ptr, + cur_bias_ptr, # Selected Pointers + cur_scale_b_ptr, # Selected Scale Pointer + cur_stride_bk, + cur_stride_bn, # Selected Strides + ) + + +# ================================================================= +# Part 2: Compute Cores +# Responsibilities: Execute specific Loop structures based on QUANT_TYPE +# ================================================================= + + +# Core A: Standard GEMM (for BF16/FP16 and W8A8) +# Feature: No dequantization inside Loop, Scale is applied after Loop ends +@triton.jit +def _core_standard_gemm( + # Pointers + a_ptr, + b_ptr, + # Offsets & Masks + real_row_idxs, + m_mask, + offs_n, + n_mask, + offs_k, + # Strides + stride_am, + stride_ak, + stride_bk, + stride_bn, + # Quant-Related + scale_a_ptr, + scale_b_ptr, + stride_scale_a, + stride_scale_b, + # Loop Info + K, + BLOCK_SIZE_K: tl.constexpr, + # Configs + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + ACCUMULATOR_DTYPE: tl.constexpr, + IS_W8A8: tl.constexpr, + # Accelerate Configs + EVEN_K: tl.constexpr, + EVEN_N: tl.constexpr, + STRIDE_AK_IS_1: tl.constexpr, + STRIDE_BK_IS_1: tl.constexpr, + STRIDE_BN_IS_1: tl.constexpr, +): + # 1. Stride optimizations (Bypassing compiler limitations for unit strides) + _stride_ak = 1 if STRIDE_AK_IS_1 else stride_ak + _stride_bk = 1 if STRIDE_BK_IS_1 else stride_bk + _stride_bn = 1 if STRIDE_BN_IS_1 else stride_bn + + # Pointer initialization (A uses indirect indexing, B uses standard striding) + a_ptrs = a_ptr + (stride_am * real_row_idxs[:, None] + _stride_ak * offs_k[None, :]) + b_ptrs = b_ptr + (_stride_bk * offs_k[:, None] + _stride_bn * offs_n[None, :]) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=ACCUMULATOR_DTYPE) + + # 2. Unified Main Loop + # Triton evaluates `constexpr` conditions at compile time, ensuring zero runtime overhead + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + # Compute K-mask only if needed + if not EVEN_K: + mask_k = offs_k < K - k * BLOCK_SIZE_K + + # Load A + if EVEN_K: + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) + else: + a = tl.load(a_ptrs, mask=m_mask[:, None] & mask_k[None, :], other=0.0) + + # Load B + if EVEN_K and EVEN_N: + b = tl.load(b_ptrs) + elif not EVEN_K and EVEN_N: + b = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0) + elif EVEN_K and not EVEN_N: + b = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0) + else: + b = tl.load(b_ptrs, mask=mask_k[:, None] & n_mask[None, :], other=0.0) + + # Compute & Advance + accumulator = tl.dot(a, b, accumulator, out_dtype=ACCUMULATOR_DTYPE) + a_ptrs += BLOCK_SIZE_K * _stride_ak + b_ptrs += BLOCK_SIZE_K * _stride_bk + + # 3. Epilogue (only needed for W8A8) + if IS_W8A8: + accumulator = accumulator.to(tl.float32) + # Load Scale A + scale_a_ptrs = scale_a_ptr + real_row_idxs * stride_scale_a + sa = tl.load(scale_a_ptrs, mask=m_mask, other=1.0) + accumulator = accumulator * sa[:, None] + + # Load Scale B + scale_b_ptrs = scale_b_ptr + offs_n * stride_scale_b + sb = tl.load(scale_b_ptrs, mask=n_mask, other=1.0) + accumulator = accumulator * sb[None, :] + + return accumulator + + +# Core B: Weight Only GEMM (for W4A16 / W8A16) +# Feature: Dequantization inside Loop (Dequantize-on-the-fly) +# Currently only supports W8A16 +@triton.jit +def _core_weight_only_gemm( + # Pointers + a_ptr, + b_ptr, + # Offsets & Masks + real_row_idxs, + m_mask, + offs_n, + n_mask, + offs_k, + # Strides + stride_am, + stride_ak, + stride_bk, + stride_bn, + # Quant-Related + scale_b_ptr, + stride_scale_b, # 1(per-channel) or 0(per-tensor) + # Loop Info + K, + BLOCK_SIZE_K: tl.constexpr, + # Configs + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + ACCUMULATOR_DTYPE: tl.constexpr, + COMPUTE_DTYPE: tl.constexpr, # bf16 or fp16 + WEIGHT_BITS: tl.constexpr, # 4 or 8 + # Accelerate Configs + EVEN_K: tl.constexpr, + EVEN_N: tl.constexpr, + STRIDE_AK_IS_1: tl.constexpr, + STRIDE_BK_IS_1: tl.constexpr, + STRIDE_BN_IS_1: tl.constexpr, +): + # Compile-time Check + tl.static_assert(WEIGHT_BITS == 8, "For weight-only, we only support W8A16 at this point") + + # 1. Stride optimizations + _stride_ak = 1 if STRIDE_AK_IS_1 else stride_ak + _stride_bk = 1 if STRIDE_BK_IS_1 else stride_bk + _stride_bn = 1 if STRIDE_BN_IS_1 else stride_bn + + a_ptrs = a_ptr + (stride_am * real_row_idxs[:, None] + _stride_ak * offs_k[None, :]) + b_ptrs = b_ptr + (_stride_bk * offs_k[:, None] + _stride_bn * offs_n[None, :]) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=ACCUMULATOR_DTYPE) + + # 2. Unified Main Loop + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + if not EVEN_K: + mask_k = offs_k < K - k * BLOCK_SIZE_K + + # Load A + if EVEN_K: + a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) + else: + a = tl.load(a_ptrs, mask=m_mask[:, None] & mask_k[None, :], other=0.0) + + # Load B (Int) + if EVEN_K and EVEN_N: + b_int = tl.load(b_ptrs) + elif not EVEN_K and EVEN_N: + b_int = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0) + elif EVEN_K and not EVEN_N: + b_int = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0) + else: + b_int = tl.load(b_ptrs, mask=mask_k[:, None] & n_mask[None, :], other=0.0) + + # --- Dequantize Logic (Type Cast Only) --- + # No multiplication here, correct as long as per-tensor/per-channel scaling (scale_b.shape=(N,)) + b_compute = b_int.to(COMPUTE_DTYPE) + + # Compute & Advance + accumulator = tl.dot(a, b_compute, accumulator, out_dtype=ACCUMULATOR_DTYPE) + + a_ptrs += BLOCK_SIZE_K * _stride_ak + b_ptrs += BLOCK_SIZE_K * _stride_bk + + # 3. Epilogue: Apply Scale B safely outside the loop + # Load Scale B + scale_b_ptrs = scale_b_ptr + offs_n * stride_scale_b + sb = tl.load(scale_b_ptrs, mask=n_mask, other=1.0) # shape=(N，) + + accumulator = accumulator * sb[None, :] + return accumulator + + +# ================================================================= +# Part 3: Unified Entry Kernel +# Responsibilities: Call Router, statically dispatch Core, store results +# ================================================================= +@triton.jit +def mot_unified_gemm_kernel( + # Inputs + a_ptr, + b_text_ptr, + b_vae_ptr, + c_ptr, + bias_text_ptr, + bias_vae_ptr, + text_indices_ptr, + vae_indices_ptr, + # Dimensions & Strides + M_text, + M_vae, + N, + K, + stride_am, + stride_ak, + stride_bk_text, + stride_bn_text, + stride_bk_vae, + stride_bn_vae, + stride_cm, + stride_cn, + # Scales (pass 0 if None) + scale_a_ptr, + scale_b_text_ptr, + scale_b_vae_ptr, + stride_scale_a, + stride_scale_b, # 1(per-channel) or 0(per-tensor) + # Metas + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + # accelerate config + EVEN_K: tl.constexpr, + EVEN_N: tl.constexpr, + STRIDE_AK_IS_1: tl.constexpr, + STRIDE_BK_IS_1: tl.constexpr, + STRIDE_BN_IS_1: tl.constexpr, + # Quant-related dtypes + ACCUMULATOR_DTYPE: tl.constexpr, + COMPUTE_DTYPE: tl.constexpr, + OUTPUT_DTYPE: tl.constexpr, + # Quant Control + # 0=None, 1=W8A8, 2=W8A16, 3=W4A16 + QUANT_TYPE: tl.constexpr, + HAS_BIAS: tl.constexpr = False, +): + pid = tl.program_id(axis=0) + + # ----------------------------------------------------------- + # 1. Routing Phase (General) + # ----------------------------------------------------------- + ( + pid_m, + pid_n, + real_row_idxs, + m_mask, + offs_n, + n_mask, + M_limit, + cur_b_ptr, + cur_bias_ptr, + cur_scale_b_ptr, + cur_stride_bk, + cur_stride_bn, + ) = _get_mot_pointers( + pid, + a_ptr, + b_text_ptr, + b_vae_ptr, + bias_text_ptr, + bias_vae_ptr, + scale_a_ptr, + scale_b_text_ptr, + scale_b_vae_ptr, + text_indices_ptr, + vae_indices_ptr, + M_text, + M_vae, + N, + stride_bk_text, + stride_bn_text, + stride_bk_vae, + stride_bn_vae, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + GROUP_SIZE_M, + ) + + offs_k = tl.arange(0, BLOCK_SIZE_K) + + # ----------------------------------------------------------- + # 2. Compute Phase (Static Dispatch) + # ----------------------------------------------------------- + if QUANT_TYPE == 0: # FP16 / BF16 Standard + c = _core_standard_gemm( + a_ptr, + cur_b_ptr, + real_row_idxs, + m_mask, + offs_n, + n_mask, + offs_k, + stride_am, + stride_ak, + cur_stride_bk, + cur_stride_bn, + 0, + 0, + 0, + 0, + K, + BLOCK_SIZE_K, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + ACCUMULATOR_DTYPE, + IS_W8A8=False, + EVEN_K=EVEN_K, + EVEN_N=EVEN_N, + STRIDE_AK_IS_1=STRIDE_AK_IS_1, + STRIDE_BK_IS_1=STRIDE_BK_IS_1, + STRIDE_BN_IS_1=STRIDE_BN_IS_1, + ) + elif QUANT_TYPE == 1: # W8A8 + c = _core_standard_gemm( + a_ptr, + cur_b_ptr, + real_row_idxs, + m_mask, + offs_n, + n_mask, + offs_k, + stride_am, + stride_ak, + cur_stride_bk, + cur_stride_bn, + scale_a_ptr, + cur_scale_b_ptr, + stride_scale_a, + stride_scale_b, + K, + BLOCK_SIZE_K, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + ACCUMULATOR_DTYPE, + IS_W8A8=True, + EVEN_K=EVEN_K, + EVEN_N=EVEN_N, + STRIDE_AK_IS_1=STRIDE_AK_IS_1, + STRIDE_BK_IS_1=STRIDE_BK_IS_1, + STRIDE_BN_IS_1=STRIDE_BN_IS_1, + ) + elif QUANT_TYPE == 2 or QUANT_TYPE == 3: # Weight Only + bits = 8 if QUANT_TYPE == 2 else 4 + c = _core_weight_only_gemm( + a_ptr, + cur_b_ptr, + real_row_idxs, + m_mask, + offs_n, + n_mask, + offs_k, + stride_am, + stride_ak, + cur_stride_bk, + cur_stride_bn, + cur_scale_b_ptr, + stride_scale_b, + K, + BLOCK_SIZE_K, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + ACCUMULATOR_DTYPE, + COMPUTE_DTYPE, + WEIGHT_BITS=bits, + EVEN_K=EVEN_K, + EVEN_N=EVEN_N, + STRIDE_AK_IS_1=STRIDE_AK_IS_1, + STRIDE_BK_IS_1=STRIDE_BK_IS_1, + STRIDE_BN_IS_1=STRIDE_BN_IS_1, + ) + + # ----------------------------------------------------------- + # 3. Store Phase (General) + # ----------------------------------------------------------- + + # Bias Add + if HAS_BIAS: + bias = tl.load(cur_bias_ptr + offs_n, mask=n_mask, other=0.0) + c = c + bias[None, :] # reshape to (1, BLOCK_SIZE_N) + + # Cast C into the output dtype + c = c.to(OUTPUT_DTYPE) + + # Store + offs_cm = real_row_idxs + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_n[None, :] + + # Need to be careful with the Mask here: real_row_idxs may contain 0 + # (out-of-bounds padding), but c_mask needs the real boundary + store_mask = m_mask[:, None] & n_mask[None, :] + tl.store(c_ptrs, c, mask=store_mask) + + +# Define is_weak_contiguous +def is_weak_contiguous(x: torch.Tensor): + strides = x.stride() + sizes = x.shape + is_not_transpose = strides[0] == 1 and (strides[1] >= max(1, sizes[0])) + is_transpose = strides[1] == 1 and (strides[0] >= max(1, sizes[1])) + return is_transpose or is_not_transpose + + +def invoke_mot_gemm( + # Inputs + A: torch.Tensor, + B_text: torch.Tensor, + B_vae: torch.Tensor, + C: torch.Tensor, + bias_text: torch.Tensor | None, + bias_vae: torch.Tensor | None, + # Indices + text_indices: torch.Tensor, + vae_indices: torch.Tensor, + # Quant Scales (None if disabled) + A_scale: torch.Tensor | None, + B_text_scale: torch.Tensor | None, + B_vae_scale: torch.Tensor | None, + # Quant Flags + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + use_int4_w4a16: bool, + # Quant Type + A_per_channel_quant: bool, + B_per_channel_quant: bool, + # Config + config: dict[str, Any] | None = None, +): + # ------ 1. Basic Assertions ------ + M = A.size(0) + K = B_text.size(0) + N = B_text.size(1) + + if config is None: + if use_fp8_w8a8: + _dtype_str: str | None = "fp8_w8a8" + elif use_int8_w8a16: + _dtype_str = "int8_w8a16" + else: + _dtype_str = None + + loaded_m_key, config = get_best_mot_config(M, N, K, _dtype_str) + + assert len(A.shape) == 2 and len(C.shape) == 2, ( + "The input tensor and output tensor should be flattened to (batch_size*seq_len, hidden_dim)" + ) + + assert K == A.size(1), "the weights' first dimension should matchinputtensor's last dimension (hidden_dim)" + + assert K == B_vae.size(0) and N == B_vae.size(1), ( + "the weights dimension for text expert andimage expert should be the same" + ) + + assert C.size(0) == M and C.size(1) == N, "the output tensor shape is not correct" + + M_text = text_indices.size(0) + M_vae = vae_indices.size(0) + assert M_text + M_vae == M, "the length sum of text and image indices should match input tensor's first dimension" + if bias_text is not None: + assert bias_text.dtype == C.dtype, "the bias tensor dtype should match the output tensor dtype" + if bias_vae is not None: + assert bias_vae.dtype == C.dtype, "the bias tensor dtype should match the output tensor dtype" + + assert is_weak_contiguous(A) + assert is_weak_contiguous(B_text) + assert is_weak_contiguous(B_vae) + + # --- 2. Quantization Logic Translation --- + + def triton_dtype(torch_dtype): + if torch_dtype == torch.float8_e4m3fn: + return getattr(tl, "float8e4m3fn", tl.float8e4nv) + elif torch_dtype == torch.float8_e5m2: + return tl.float8e5 + return { + torch.float16: tl.float16, + torch.bfloat16: tl.bfloat16, + torch.float32: tl.float32, + torch.int8: tl.int8, + torch.float8_e4m3fn: tl.float8e4nv, + torch.float8_e5m2: tl.float8e5, + }[torch_dtype] + + # Determine QUANT_TYPE + # 0=None, 1=W8A8, 2=W8A16, 3=W4A16 + quant_type = 0 + ACCUMULATOR_DTYPE = tl.float32 + COMPUTE_DTYPE = triton_dtype(A.dtype) + OUTPUT_DTYPE = triton_dtype(C.dtype) + + if use_int8_w8a8 or use_fp8_w8a8: + quant_type = 1 + assert A_scale is not None, "W8A8 requires A_scale" + assert B_text_scale is not None and B_vae_scale is not None, "W8A8 requires B_text_scale and B_vae_scale" + if use_int8_w8a8: + ACCUMULATOR_DTYPE = tl.int32 + assert ( + A.dtype == torch.int8 + and B_text.dtype == torch.int8 + and B_vae.dtype == torch.int8 + and C.dtype in [torch.float16, torch.bfloat16] + ), "if you want to use INT8_W8A8, A should be INT8, B should be INT8, C should be FP16/BF16" + else: + ACCUMULATOR_DTYPE = tl.float32 + assert ( + A.dtype in [torch.float8_e4m3fn, torch.float8_e5m2] + and B_text.dtype in [torch.float8_e4m3fn, torch.float8_e5m2] + and B_vae.dtype in [torch.float8_e4m3fn, torch.float8_e5m2] + and C.dtype in [torch.float16, torch.bfloat16] + ), "if you want to use FP8_W8A8, A should be FP8, B should be FP8, C should be FP16/BF16" + elif use_int8_w8a16: + quant_type = 2 + assert B_text_scale is not None and B_vae_scale is not None, "W8A16 requires B_text_scale and B_vae_scale" + ACCUMULATOR_DTYPE = tl.float32 + assert ( + A.dtype in [torch.float16, torch.bfloat16] + and B_text.dtype == torch.int8 + and B_vae.dtype == torch.int8 + and C.dtype in [torch.float16, torch.bfloat16] + ), "if you want to use INT8_W8A16, A should be FP16/BF16, B should be INT8, C should be FP16/BF16" + + elif use_int4_w4a16: + raise NotImplementedError("For weight-only, we only support W8A16 at this point") + # quant_type = 3 + # ACCUMULATOR_DTYPE=tl.float32 + + # accelerate config + EVEN_K = K % config["BLOCK_SIZE_K"] == 0 + EVEN_N = N % config["BLOCK_SIZE_N"] == 0 + STRIDE_AK_IS_1 = A.stride(1) == 1 + STRIDE_BK_IS_1 = (B_text.stride(0) == 1) and (B_vae.stride(0) == 1) + STRIDE_BN_IS_1 = (B_text.stride(1) == 1) and (B_vae.stride(1) == 1) + + # bias check + assert (bias_text is None) == (bias_vae is None), ( + "Bias must be provided for both Text and VAE simultaneously, or neither." + ) + has_bias = bias_text is not None + + # --- 3. Grid Calculation --- + def grid(META): + return ( + (triton.cdiv(M_text, META["BLOCK_SIZE_M"]) + triton.cdiv(M_vae, META["BLOCK_SIZE_M"])) + * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + # --- 4. Launch --- + run_config = config.copy() + run_config.update( + { + "QUANT_TYPE": quant_type, + "ACCUMULATOR_DTYPE": ACCUMULATOR_DTYPE, + "COMPUTE_DTYPE": COMPUTE_DTYPE, + "OUTPUT_DTYPE": OUTPUT_DTYPE, + "HAS_BIAS": has_bias, + "EVEN_K": EVEN_K, + "EVEN_N": EVEN_N, + "STRIDE_AK_IS_1": STRIDE_AK_IS_1, + "STRIDE_BK_IS_1": STRIDE_BK_IS_1, + "STRIDE_BN_IS_1": STRIDE_BN_IS_1, + } + ) + + # Pointers (Handle None -> 0) + p_a_scale = A_scale if A_scale is not None else 0 + p_b_text_scale = B_text_scale if B_text_scale is not None else 0 + p_b_vae_scale = B_vae_scale if B_vae_scale is not None else 0 + p_bias_text = bias_text if bias_text is not None else 0 + p_bias_vae = bias_vae if bias_vae is not None else 0 + + # Quantization granularity + stride_scale_a = 1 if A_per_channel_quant else 0 + stride_scale_b = 1 if B_per_channel_quant else 0 + + mot_unified_gemm_kernel[grid]( + # Inputs + A, + B_text, + B_vae, + C, + p_bias_text, + p_bias_vae, + text_indices, + vae_indices, + # Dimensions + M_text, + M_vae, + N, + K, + # Strides + A.stride(0), + A.stride(1), + B_text.stride(0), + B_text.stride(1), + B_vae.stride(0), + B_vae.stride(1), + C.stride(0), + C.stride(1), + # Scales + p_a_scale, + p_b_text_scale, + p_b_vae_scale, + stride_scale_a, + stride_scale_b, + # Config + **run_config, + ) diff --git a/vllm_omni/diffusion/layers/mot/ops/mot_rms_norm.py b/vllm_omni/diffusion/layers/mot/ops/mot_rms_norm.py new file mode 100644 index 00000000000..bbf97916475 --- /dev/null +++ b/vllm_omni/diffusion/layers/mot/ops/mot_rms_norm.py @@ -0,0 +1,283 @@ +# ruff: noqa: N803, E741 +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _mot_rms_norm_kernel( + input_ptr, + text_weight_ptr, + vae_weight_ptr, + text_indices_ptr, # MoT Routing Info + vae_indices_ptr, # MoT Routing Info + M_text, # MoT Routing Info + output_ptr, + input_row_stride, + output_row_stride, + n_cols, + eps, + BLOCK_SIZE: tl.constexpr, +): + """ + input shape: (batch_size*seq_len, hidden_size) + RMSNorm: y = x / sqrt(mean(x^2) + eps) * weight + """ + # Step 0: MoT Routing + pid = tl.program_id(0).to(tl.int64) + + # dummy init (must be scalar int64[] to match tl.load return type) + row_idx = tl.cast(0, tl.int64) + weight_ptr = text_weight_ptr + + if pid < M_text: + # --- Text Path --- + row_idx = tl.load(text_indices_ptr + pid) + weight_ptr = text_weight_ptr + else: + # --- VAE Path --- + vae_pid = pid - M_text + row_idx = tl.load(vae_indices_ptr + vae_pid) + weight_ptr = vae_weight_ptr + + row_start_ptr = input_ptr + row_idx * input_row_stride + output_row_start_ptr = output_ptr + row_idx * output_row_stride + + # Step 1: Compute sum of squares in float32 to avoid overflow + sum_sq = tl.zeros([1], dtype=tl.float32) + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + # Convert to float32 for accumulation to prevent overflow + vals_f32 = vals.to(tl.float32) + sq_vals = vals_f32 * vals_f32 + sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0)) + + # Step 2: Compute RMS (root mean square) in float32 + mean_sq = sum_sq / n_cols + rms = tl.sqrt(mean_sq + eps) + inv_rms = 1.0 / rms + + # Step 3: Normalize and apply weight + for col_offset in range(0, n_cols, BLOCK_SIZE): + col_idx = col_offset + tl.arange(0, BLOCK_SIZE) + mask = col_idx < n_cols + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0) + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0) + # Compute in float32 then convert back to input dtype + vals_f32 = vals.to(tl.float32) + weight_f32 = weight.to(tl.float32) + output_f32 = vals_f32 * inv_rms * weight_f32 + output = output_f32.to(vals.dtype) + tl.store(output_row_start_ptr + col_idx, output, mask=mask) + + +@triton.jit +def _mot_rms_norm_qk_kernel( + input_ptr, + text_weight_ptr, + vae_weight_ptr, + text_indices_ptr, + vae_indices_ptr, + M_text, + output_ptr, + stride_tok, + stride_head, + out_stride_tok, + out_stride_head, + num_heads, + head_dim, + eps, + BLOCK_DIM: tl.constexpr, # for head_dim=128,should be 128 + SHARED_WEIGHT: tl.constexpr, # weights are shared across heads +): + pid_tok = tl.program_id(0).to(tl.int64) + + # MoT Routing + if pid_tok < M_text: + row_idx = tl.load(text_indices_ptr + pid_tok) + weight_ptr = text_weight_ptr + else: + vae_pid = pid_tok - M_text + row_idx = tl.load(vae_indices_ptr + vae_pid) + weight_ptr = vae_weight_ptr + + dim_offsets = tl.arange(0, BLOCK_DIM) + mask = dim_offsets < head_dim + + # input.shape=(4098,28,128)=(seq_len,num_heads,head_dim) + # weight.shape=(28,128)=(num_heads,head_dim) or (128,)=(head_dim) + tok_in_ptr = input_ptr + row_idx * stride_tok + tok_out_ptr = output_ptr + row_idx * out_stride_tok + + # if weight.shape=(128,) + if SHARED_WEIGHT: + weight = tl.load(weight_ptr + dim_offsets, mask=mask, other=1.0) + weight_f32 = weight.to(tl.float32) + + # one program loop over all heads + for h in range(num_heads): + head_in_ptr = tok_in_ptr + h * stride_head + head_out_ptr = tok_out_ptr + h * out_stride_head + + # if weight.shape=(28,128) + if not SHARED_WEIGHT: + weight_offset = h * head_dim + dim_offsets + weight = tl.load(weight_ptr + weight_offset, mask=mask, other=1.0) + weight_f32 = weight.to(tl.float32) + + vals = tl.load(head_in_ptr + dim_offsets, mask=mask, other=0.0) + vals_f32 = vals.to(tl.float32) + + sq_vals = vals_f32 * vals_f32 + sum_sq = tl.sum(tl.where(mask, sq_vals, 0.0)) + mean_sq = sum_sq / head_dim + rms = tl.sqrt(mean_sq + eps) + inv_rms = 1.0 / rms + + output_f32 = vals_f32 * inv_rms * weight_f32 + tl.store(head_out_ptr + dim_offsets, output_f32.to(vals.dtype), mask=mask) + + +def mot_rms_norm( + input: torch.Tensor, + text_weight: torch.Tensor, + vae_weight: torch.Tensor, + text_indices: torch.Tensor, + vae_indices: torch.Tensor, + head_norm: bool = False, + eps: float = 1e-6, + block_size: int | None = None, +) -> torch.Tensor: + """ + Compute RMS normalization using Triton kernel. + + RMS Norm normalizes the input by the root mean square and scales by weight: + output = input / sqrt(mean(input^2) + eps) * weight + + Args: + input: Input tensor of shape (batch_size*seq_len, hidden_size) or (batch_size,seq_len, hidden_size) + text_weight: Weight for text tokens tensor of shape (hidden_size,) + vae_weight: Weight for vae tokens tensor of shape (hidden_size,) + text_indices: indices of text tokens, (batch_size*2,) + vae_indices: indices of vae tokens, (batch_size*(seq_len-2),) + eps: Small constant for numerical stability + + Returns: + Tensor with RMS normalization applied along the last dimension + """ + + assert input.shape[-1] == text_weight.shape[-1], ( + f"Input last dimension ({input.shape[-1]}) must match Text weight dimension ({text_weight.shape[-1]})" + ) + assert input.shape[-1] == vae_weight.shape[-1], ( + f"Input last dimension ({input.shape[-1]}) must match VAE weight dimension ({vae_weight.shape[-1]})" + ) + + original_shape = input.shape + text_indices = text_indices.reshape(-1) + vae_indices = vae_indices.reshape(-1) + M_text = text_indices.shape[0] + M_vae = vae_indices.shape[0] + num_tokens = M_text + M_vae + + if not head_norm: + input_2d = input.reshape(-1, input.shape[-1]) + assert input_2d.shape[0] == num_tokens, ( + f"batch_size={input_2d.shape[0]}, len(text_indices)={M_text}, len(vae_indices)={M_vae}" + f"for layer norm, batched_token_length should match the sum of indices_length" + ) + input_2d = input_2d.contiguous() + text_indices = text_indices.contiguous() + vae_indices = vae_indices.contiguous() + + text_weight = text_weight.contiguous() + vae_weight = vae_weight.contiguous() + + n_rows, n_cols = input_2d.shape + + output = torch.empty_like(input_2d) + + if block_size is None: + block_size = triton.next_power_of_2(n_cols) + block_size = min(block_size, 4096) + + num_warps = 4 + if block_size >= 2048: + num_warps = 8 + elif block_size >= 1024: + num_warps = 4 + else: + num_warps = 2 + + grid = (n_rows,) + _mot_rms_norm_kernel[grid]( + input_2d, + text_weight, + vae_weight, + text_indices, + vae_indices, + M_text, + output, + input_2d.stride(0), + output.stride(0), + n_cols, + eps, + BLOCK_SIZE=block_size, + num_warps=num_warps, + ) + return output.reshape(original_shape) + else: + # qk norm scenarios: + # input.shape=(batch_size, seq_len,head_num, head_dim) or + # input.shape=(batch_size*seq_len,head_num, head_dim) + assert len(original_shape) > 2, ( + "If head_norm=True,input shape be 3D or 3D, last 2 dimensions should be head_num and head_dim" + ) + num_heads = input.shape[-2] + head_dim = input.shape[-1] + + is_shared_weight = text_weight.dim() == 1 + if not is_shared_weight: + assert num_heads == text_weight.shape[0] and num_heads == vae_weight.shape[0], ( + "when weights are not shared across heads, the first dimension of " + "weights should be num of heads and match input.shape[-2]" + ) + + # reshape to 3D + input_3d = input.view(-1, num_heads, head_dim) + + input_3d = input_3d.contiguous() + text_indices = text_indices.contiguous() + vae_indices = vae_indices.contiguous() + text_weight = text_weight.contiguous() + vae_weight = vae_weight.contiguous() + + output_3d = torch.empty_like(input_3d) + + block_dim = triton.next_power_of_2(head_dim) + num_warps = 4 if block_dim > 128 else 2 + num_stages = 3 + + _mot_rms_norm_qk_kernel[(num_tokens,)]( + input_3d, + text_weight, + vae_weight, + text_indices, + vae_indices, + M_text, + output_3d, + input_3d.stride(0), + input_3d.stride(1), + output_3d.stride(0), + output_3d.stride(1), + num_heads, + head_dim, + eps, + BLOCK_DIM=block_dim, + SHARED_WEIGHT=is_shared_weight, + num_warps=num_warps, + num_stages=num_stages, + ) + return output_3d.view(original_shape) diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 4cb714465f5..6a045fca2bf 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -188,6 +188,12 @@ def forward_native( cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: + # All batch elements share the same rotary position encoding. + # Strip the batch dim so the underlying op broadcasts over the batch, + # consistent with forward_cuda / forward_hip / apply_rotary_emb_mindiesd. + if cos.dim() == 3: + cos = cos[0] + sin = sin[0] return apply_rotary_emb_torch( x, cos, diff --git a/vllm_omni/diffusion/lora/manager.py b/vllm_omni/diffusion/lora/manager.py index 8ba30b79e1f..2fcf43b6e31 100644 --- a/vllm_omni/diffusion/lora/manager.py +++ b/vllm_omni/diffusion/lora/manager.py @@ -368,7 +368,17 @@ def _matches_target(module_name: str) -> bool: # Default denoising components plus any a pipeline opts into via # ``_lora_components`` (opt-in; other pipelines unchanged). - default_components = ("transformer", "transformer_2", "dit", "bagel") + # + # NOTE: SDXL-style pipelines expose the denoiser as ``unet``. + # Without scanning this component, adapters can load/activate while + # effectively applying to zero layers, producing base-identical output. + default_components = ( + "transformer", + "transformer_2", + "dit", + "bagel", + "unet", + ) extra_components = tuple(getattr(self.pipeline, "_lora_components", ()) or ()) for component_name in (*default_components, *extra_components): if not hasattr(self.pipeline, component_name): diff --git a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py index 04fc8ea32f4..68ec98d5ae7 100644 --- a/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py +++ b/vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py @@ -11,7 +11,12 @@ logger = init_logger(__name__) -MODEL_OPT_SCALE_SUFFIXES = (".input_scale", ".weight_scale", ".weight_scale_2", ".weight_scale_inv") +MODEL_OPT_SCALE_SUFFIXES = ( + ".input_scale", + ".weight_scale", + ".weight_scale_2", + ".weight_scale_inv", +) DEFAULT_PACKED_MODULES_MAPPING = { "to_qkv": ("to_q", "to_k", "to_v"), "add_kv_proj": ("add_q_proj", "add_k_proj", "add_v_proj"), @@ -32,7 +37,7 @@ @dataclass class _AdaptState: scale_tensors: dict[str, torch.Tensor] = field(default_factory=dict) - pending_weights: dict[str, list[tuple[str, torch.Tensor, torch.dtype]]] = field(default_factory=dict) + pending_weights: dict[str, list[tuple[str, str, torch.Tensor, torch.dtype]]] = field(default_factory=dict) skipped_scales: int = 0 dequantized_weights: int = 0 @@ -41,6 +46,7 @@ class ModelOptFp8CheckpointAdapter: def __init__(self, model: nn.Module, source: object): self._loadable_tensors = self._get_model_loadable_tensors(model) self._weights_mapper = self._get_weights_mapper(model) + self._checkpoint_key_mapper = getattr(model, "remap_checkpoint_key", None) self._source_label = getattr(source, "prefix", "") or getattr(source, "subfolder", None) or "model" @classmethod @@ -110,14 +116,20 @@ def _get_weights_mapper(cls, model: nn.Module) -> WeightsMapper: orig_to_new_prefix=orig_to_new_prefix, ) - def _resolve_target_name(self, name: str) -> str | None: + def _resolve_target_and_output_names(self, name: str) -> tuple[str | None, str]: if name in self._loadable_tensors: - return name + return name, name + + if callable(self._checkpoint_key_mapper): + candidate = self._checkpoint_key_mapper(name) + if candidate in self._loadable_tensors: + return candidate, candidate for candidate in self._weights_mapper.apply_list([name]): if candidate != name and candidate in self._loadable_tensors: - return candidate - return None + # Preserve shard names so downstream packed-weight loaders can route them. + return candidate, name + return None, name @staticmethod def _reshape_weight_scale(scale: torch.Tensor, weight_shape: torch.Size) -> torch.Tensor: @@ -160,13 +172,19 @@ def _flush_pending_weights( scale_name: str, state: _AdaptState, ) -> Generator[tuple[str, torch.Tensor], None, None]: - for weight_name, weight_tensor, target_dtype in state.pending_weights.pop(scale_name, []): - yield weight_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype) + for ( + weight_name, + output_name, + weight_tensor, + target_dtype, + ) in state.pending_weights.pop(scale_name, []): + yield output_name, self._dequantize_weight(weight_name, weight_tensor, state, target_dtype) state.dequantized_weights += 1 def _handle_scale_tensor( self, name: str, + output_name: str, tensor: torch.Tensor, target_name: str | None, state: _AdaptState, @@ -175,7 +193,7 @@ def _handle_scale_tensor( if target_name is None: state.skipped_scales += 1 else: - yield name, tensor + yield output_name, tensor yield from self._flush_pending_weights(name, state) def _target_dtype_for_dequantization( @@ -194,6 +212,7 @@ def _target_dtype_for_dequantization( def _maybe_dequantize_or_defer_weight( self, name: str, + output_name: str, tensor: torch.Tensor, target_dtype: torch.dtype, state: _AdaptState, @@ -203,7 +222,7 @@ def _maybe_dequantize_or_defer_weight( raise ValueError(f"Missing ModelOpt FP8 weight_scale name for weight {name!r}") if scale_name not in state.scale_tensors: - state.pending_weights.setdefault(scale_name, []).append((name, tensor, target_dtype)) + state.pending_weights.setdefault(scale_name, []).append((name, output_name, tensor, target_dtype)) return None state.dequantized_weights += 1 @@ -235,23 +254,25 @@ def adapt( state = _AdaptState() for name, tensor in weights: - target_name = self._resolve_target_name(name) + target_name, output_name = self._resolve_target_and_output_names(name) if self._is_scale(name): - yield from self._handle_scale_tensor(name, tensor, target_name, state) + yield from self._handle_scale_tensor(name, output_name, tensor, target_name, state) continue target_dtype = self._target_dtype_for_dequantization(tensor, target_name) if target_dtype is not None: - tensor = self._maybe_dequantize_or_defer_weight(name, tensor, target_dtype, state) + tensor = self._maybe_dequantize_or_defer_weight(name, output_name, tensor, target_dtype, state) if tensor is None: continue - yield name, tensor + yield output_name, tensor self._check_pending_weights(state) self._log_adaptation_summary(state) class ModelOptNvFp4CheckpointAdapter(ModelOptFp8CheckpointAdapter): + _PRE_QUANT_SCALE_SUFFIX = ".pre_quant_scale" + @staticmethod def _is_checkpoint_quant_config(quant_config: object | None) -> bool: return ( @@ -261,6 +282,22 @@ def _is_checkpoint_quant_config(quant_config: object | None) -> bool: and bool(getattr(quant_config, "is_checkpoint_nvfp4_serialized", False)) ) + def adapt( + self, + weights: Iterable[tuple[str, torch.Tensor]], + ) -> Generator[tuple[str, torch.Tensor], None, None]: + def validated_weights() -> Generator[tuple[str, torch.Tensor], None, None]: + for name, tensor in weights: + if name.endswith(self._PRE_QUANT_SCALE_SUFFIX): + raise ValueError( + f"ModelOpt NVFP4 checkpoint tensor {name!r} is unsupported: " + "vLLM 0.25.0 does not consume pre_quant_scale. Export the checkpoint " + "with pre-quant scales folded into the weights." + ) + yield name, tensor + + yield from super().adapt(validated_weights()) + class ModelOptMixedPrecisionCheckpointAdapter(ModelOptFp8CheckpointAdapter): @staticmethod diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index e5ef5e3191d..8ee172e6bd8 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -11,14 +11,12 @@ from typing import cast import torch -from huggingface_hub import hf_hub_download from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import ( - download_gguf, download_safetensors_index_file_from_hf, download_weights_from_hf, filter_duplicate_safetensors_files, @@ -37,11 +35,43 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import ( get_checkpoint_adapter, ) -from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery from vllm_omni.diffusion.registry import initialize_model + +# download_gguf was removed from upstream vLLM (commit 6635279d8). +# Inlined from the last upstream version before the GGUF plugin migration. +def download_gguf( + repo_id: str, + quant_type: str, + cache_dir: str | None = None, + revision: str | None = None, + ignore_patterns: str | list[str] | None = None, +) -> str: + allow_patterns = [ + f"*-{quant_type}.gguf", + f"*-{quant_type}-*.gguf", + f"*/*-{quant_type}.gguf", + f"*/*-{quant_type}-*.gguf", + ] + folder = download_weights_from_hf( + model_name_or_path=repo_id, + cache_dir=cache_dir, + allow_patterns=allow_patterns, + revision=revision, + ignore_patterns=ignore_patterns, + ) + local_files: list[str] = [] + for pattern in allow_patterns: + glob_pattern = os.path.join(folder, pattern) + local_files.extend(glob.glob(glob_pattern)) + if not local_files: + raise ValueError(f"Downloaded GGUF files not found in {folder} for quant_type {quant_type}") + local_files.sort(key=lambda x: (x.count("-"), x)) + return local_files[0] + + logger = init_logger(__name__) @@ -349,15 +379,11 @@ def load_model( # DiffusionPipeline.from_pretrained() internally; it does # not use our native customized pipeline classes. cast(DiffusersAdapterPipeline, model).load_weights() - 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) - - # Process weights after loading for quantization (e.g., FP8 online quantization) - # This is needed for vLLM's quantization methods that need to transform weights - self._process_weights_after_loading(model, target_device) + # HSDP processes quantized weights before wrapping parameters as + # DTensors. The non-HSDP path can process them here as usual. + self._process_weights_after_loading(model, target_device) if offload_after_quant: model.to("cpu") @@ -365,16 +391,77 @@ def load_model( return model.eval() + @staticmethod + def _has_online_quant(model: nn.Module) -> bool: + """Whether any layer uses an online-quant method that defers weight + materialization onto the ``meta`` device (upstream vLLM + ``uses_meta_device=True``, e.g. online FP8).""" + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if getattr(quant_method, "uses_meta_device", False): + return True + return False + def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None: """Process weights after loading for quantization methods. This handles vLLM's quantization methods that need to process weights after loading (e.g., FP8 online quantization from BF16/FP16 weights). """ + # Newer upstream vLLM online-quant methods (uses_meta_device=True) create + # weights on the ``meta`` device and materialize them just-in-time as each + # layer's weights finish loading (via the layerwise online-process loader). + # Any "straggler" layers whose weights were not fully materialized during + # load (padded / partially-loaded layers) remain on ``meta``. Upstream's + # base_loader calls finalize_layerwise_processing() to materialize them; + # the diffusion loader must mirror that, otherwise the module.to() below + # raises "Cannot copy out of meta tensor; no data!". This whole meta-device + # handling is gated on online quant actually being in use, so that the + # proven code path for everything else (in particular FSDP/HSDP-sharded + # params, whose per-parameter .data cannot be cross-device reassigned) is + # left untouched. Import lazily so older vLLM (no meta-device quant) is + # unaffected. + has_online_quant = self._has_online_quant(model) + if has_online_quant: + from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_processing, + ) + + # model_config is only dereferenced by finalize for vLLM Attention / + # MLAAttention layers; diffusion DiT models use their own attention and + # have none, so passing None is safe here. + finalize_layerwise_processing(model, model_config=None) + for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) - if quant_method is not None and isinstance(quant_method, QuantizeMethodBase): - # Move module to target device for processing if needed + if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): + continue + + if has_online_quant: + # Online quant may leave straggler params on the ``meta`` device. + # Move only real (non-meta) params onto the target device for + # processing and restore them afterward, mirroring upstream vLLM's + # device_loading_context — a blanket module.to(target_device) would + # raise NotImplementedError on meta params. Online quant initializes + # on the accelerator, so params are normally already on the target + # device and this loop is a no-op move; the point is to skip meta. + original_devices: dict[str, torch.device] = {} + for name, param in module.named_parameters(): + if param.device.type != "meta" and param.device != target_device: + original_devices[name] = param.device + param.data = param.data.to(target_device) + + quant_method.process_weights_after_loading(module) + + # Restore pre-existing params to their original device; leave any + # newly created (e.g. quantized) params on the target device. + for name, param in module.named_parameters(): + if name in original_devices: + param.data = param.data.to(original_devices[name]) + else: + # No meta params possible here. Preserve the original FSDP/HSDP-aware + # whole-module move (module.to()), which correctly handles sharded + # DTensor params that per-parameter .data reassignment cannot. module_device = next(module.parameters(), None) if module_device is not None: module_device = module_device.device @@ -422,7 +509,7 @@ def load_weights(self, model: nn.Module) -> None: def _is_expected_quantized_weight(name: str) -> bool: """Return True if *name* is a quantization-specific parameter. - Quantization methods (GPTQ, AWQ, FP8, GGUF, Autoround, etc.) create extra + Quantization methods (GPTQ, AWQ, FP8, Autoround, etc.) create extra parameters that have no counterpart in an unquantized checkpoint. These are expected to be absent and should not trigger a load error. """ @@ -435,8 +522,6 @@ def _is_expected_quantized_weight(name: str) -> bool: ".weight_scale", ".weight_scale_inv", ".input_scale", - # GGUF - ".qweight_type", # INT8 (weight_scale already covered above) ) return name.endswith(_QUANTIZED_WEIGHT_SUFFIXES) @@ -468,116 +553,6 @@ def _check_unloaded_weights( if unexpected_missing: raise ValueError(f"Following weights were not initialized from checkpoint: {unexpected_missing}") - 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) - - 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 - - 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 - - if isinstance(self.quant_config, dict): - return False - - return bool(maybe_gguf_model) - - def _is_transformer_source(self, source: "ComponentSource") -> bool: - if source.subfolder == "transformer": - return True - return source.prefix.startswith("transformer.") - - def _get_model_loadable_names(self, model: nn.Module) -> set[str]: - # Avoid model.state_dict() here because GGUF uses UninitializedParameter - # which raises during detach(). Collect names directly. - return {name for name, _ in model.named_parameters()} | {name for name, _ in model.named_buffers()} - - def _resolve_gguf_model_path(self, gguf_model: str, revision: str | None) -> str: - if os.path.isfile(gguf_model): - return gguf_model - # repo_id/filename.gguf - if "/" in gguf_model and gguf_model.endswith(".gguf"): - repo_id, filename = gguf_model.rsplit("/", 1) - return hf_hub_download( - repo_id=repo_id, - filename=filename, - revision=revision, - cache_dir=self.load_config.download_dir, - ) - # repo_id:quant_type - if "/" in gguf_model and ":" in gguf_model: - repo_id, quant_type = gguf_model.rsplit(":", 1) - return download_gguf( - repo_id, - quant_type, - cache_dir=self.load_config.download_dir, - revision=revision, - ignore_patterns=self.load_config.ignore_patterns, - ) - raise ValueError( - f"Unrecognized GGUF reference: {gguf_model!r} (expected local file, " - "<repo_id>/<filename>.gguf, or <repo_id>:<quant_type>)" - ) - - def _get_gguf_weights_iterator( - self, - source: "ComponentSource", - model: nn.Module, - ) -> Generator[tuple[str, torch.Tensor], None, 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, 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) -> 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)) - - # GGUF checkpoints can be transformer-only or partially quantized. - # Only fall back to HF if this source still has missing loadable weights. - loadable_names = loadable_names or self._get_model_loadable_names(model) - has_missing_for_source = any( - name.startswith(source.prefix) and name not in loaded for name in loadable_names - ) - if not has_missing_for_source: - continue - - hf_iter = self._get_weights_iterator(source) - hf_iter = ( - (name, tensor) for (name, tensor) in hf_iter if name in loadable_names and name not in loaded - ) - loaded |= model.load_weights(hf_iter) - else: - loaded |= model.load_weights(self._get_weights_iterator(source)) - - weights_to_load = self._get_expected_parameter_names(model) - weights_not_loaded = weights_to_load - loaded - if weights_not_loaded: - raise ValueError(f"Following weights were not initialized from checkpoint: {weights_not_loaded}") - return loaded - def _init_from_load_format( self, load_format: str, @@ -602,10 +577,17 @@ def _init_from_load_format( 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": + # HSDP normally defers GPU placement to apply_hsdp_to_model to keep peak + # load-time memory on CPU. Online quantization (e.g. fp8) runs CUDA-only + # kernels inside load_weights via the layerwise loader, so when a quant + # config is set we initialize on the accelerator like the non-HSDP path; + # apply_hsdp_to_model shards GPU-resident params equally well. + hsdp_defer_to_cpu = is_hsdp and self.quant_config is None + if not hsdp_defer_to_cpu and target_device.type != "cpu": model.to(target_device) else: - device_ctx = target_device if not is_hsdp else contextlib.nullcontext() + hsdp_defer_to_cpu = is_hsdp and self.quant_config is None + device_ctx = contextlib.nullcontext() if hsdp_defer_to_cpu else target_device with device_ctx: if load_format == "default": model = initialize_model(self.od_config) @@ -646,26 +628,50 @@ def _load_model_with_hsdp( 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) - transformers_to_shard = [] - transformer = getattr(model, "transformer", None) - if transformer is None: - raise ValueError("Model has no transformer attribute for HSDP") - transformers_to_shard.append(("transformer", transformer)) + # Quantization methods must finish while parameters are ordinary local + # tensors. Some post-load transforms use operations (for example, + # torch.unique in ModelOpt NVFP4) that do not support DTensor inputs. + self._process_weights_after_loading(model, target_device) + + # Discover pipeline components (DiT, encoders, VAEs) via + # ModuleDiscovery, which consults SupportsComponentDiscovery + # when available and falls back to well-known attribute names. + # This supports nested pipelines (e.g. LTX2TwoStagesPipeline + # where the transformer lives at "pipe.transformer"). + discovered_modules = ModuleDiscovery.discover(model) + + # Shard only the outermost DiTs. A pipeline may list a DiT and one of its + # submodules as separate DiTs (e.g. Cosmos3's transformer and the nested + # transformer.language_model) for offload's independent rings; for HSDP an + # inner DiT is already covered by its ancestor's _hsdp_shard_conditions, so + # sharding it again would double-wrap blocks and require the inner stack to + # declare its own conditions. + outer_dit_names, outer_dits = discovered_modules.outermost_dits() + + # Online FP8 quantization (Fp8OnlineLinearMethod) leaves layer weights + # as non-contiguous transpose views (qweight.t()) so the Cutlass kernel + # gets a column-major B. FSDP2 fully_shard rejects non-contiguous params. + # Rewrite affected layers in-place to row-major contiguous storage and + # shift the .t() to GEMM-call time. Layers using other quant methods or + # already-contiguous weights are left untouched. + if self.quant_config is not None: + from vllm_omni.diffusion.quantization.hsdp_fp8 import ( + prepare_fp8_layers_for_fsdp, + ) - # Check for transformer_2 (MoE two-stage models like Wan2.2-I2V) - transformer_2 = getattr(model, "transformer_2", None) - if transformer_2 is not None: - transformers_to_shard.append(("transformer_2", transformer_2)) + for trans in outer_dits: + prepare_fp8_layers_for_fsdp(trans) - # Apply HSDP sharding to all transformers - for name, trans in transformers_to_shard: + if not outer_dits: + raise ValueError("No DiT modules discovered for HSDP sharding") + + # Apply HSDP sharding to each outermost DiT transformer + for name, trans in zip(outer_dit_names, outer_dits): logger.debug("Applying HSDP to %s", name) - apply_hsdp_to_model(trans, hsdp_config) + apply_hsdp_to_model(trans, hsdp_config, target_device=target_device) - # # HSDP only shards transformer modules. All other runtime modules must - # # be placed on the execution device explicitly after sharding. - discovered_modules = ModuleDiscovery.discover(model) + # HSDP only shards transformer modules. All other runtime modules must + # be placed on the execution device explicitly after sharding. modules_to_move: list[nn.Module] = [] if discovered_modules.vaes is not None: modules_to_move.extend(discovered_modules.vaes) diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py b/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py deleted file mode 100644 index 3266ba06c6a..00000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/__init__.py +++ /dev/null @@ -1,47 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -from .base import GGUFAdapter -from .flux2_klein import Flux2KleinGGUFAdapter -from .qwen_image import QwenImageGGUFAdapter -from .z_image import ZImageGGUFAdapter - -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - from vllm_omni.diffusion.model_loader.diffusers_loader import ( - DiffusersPipelineLoader, - ) - - -def get_gguf_adapter( - gguf_file: str, - model: torch.nn.Module, - source: DiffusersPipelineLoader.ComponentSource, - od_config: OmniDiffusionConfig, -) -> GGUFAdapter: - adapter_classes = (QwenImageGGUFAdapter, ZImageGGUFAdapter, Flux2KleinGGUFAdapter) - for adapter_cls in adapter_classes: - if adapter_cls.is_compatible(od_config, model, source): - return adapter_cls(gguf_file, model, source, od_config) - model_type = None - if od_config.tf_model_config is not None: - model_type = od_config.tf_model_config.get("model_type") - supported = ", ".join(cls.__name__ for cls in adapter_classes) - raise ValueError( - "No GGUF adapter matched diffusion model " - f"(model_class_name={od_config.model_class_name!r}, model_type={model_type!r}). " - f"Supported adapters: {supported}." - ) - - -__all__ = [ - "GGUFAdapter", - "Flux2KleinGGUFAdapter", - "QwenImageGGUFAdapter", - "ZImageGGUFAdapter", - "get_gguf_adapter", -] diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/base.py b/vllm_omni/diffusion/model_loader/gguf_adapters/base.py deleted file mode 100644 index 3bc38f16d96..00000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/base.py +++ /dev/null @@ -1,96 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from abc import ABC, abstractmethod -from collections.abc import Generator -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any - -import gguf -import numpy as np -import torch - -if TYPE_CHECKING: - from vllm_omni.diffusion.data import OmniDiffusionConfig - from vllm_omni.diffusion.model_loader.diffusers_loader import ( - DiffusersPipelineLoader, - ) - - -@dataclass -class MappedTensor: - name: str - tensor: Any - tensor_type: Any - row_slice: slice | None = None - swap_scale_shift: bool = False - - -class GGUFAdapter(ABC): - """Base class for model-specific GGUF adapters.""" - - def __init__( - self, - gguf_file: str, - model: torch.nn.Module, - source: DiffusersPipelineLoader.ComponentSource, - od_config: OmniDiffusionConfig, - ) -> None: - self.gguf_file = gguf_file - self.model = model - self.source = source - self.od_config = od_config - - @staticmethod - def is_compatible( - od_config: OmniDiffusionConfig, - model: torch.nn.Module, - source: DiffusersPipelineLoader.ComponentSource, - ) -> bool: - return False - - @abstractmethod - def weights_iterator(self) -> Generator[tuple[str, torch.Tensor], None, None]: - raise NotImplementedError - - -# FIXME(Isotr0py): Sync implemnentation with upstream vLLM? -def gguf_quant_weights_iterator(gguf_file: str) -> Generator[tuple[str, torch.Tensor]]: - """ - Iterate over the quant weights in the model gguf files and convert - them to torch tensors. - Be careful of the order of yielding weight types and weights data, - we have to yield all weight types first before yielding any weights. - Otherwise it would cause issue when loading weights with for packed - layer with different quant types. - """ - - reader = gguf.GGUFReader(gguf_file) - - for tensor in reader.tensors: - weight_type = tensor.tensor_type - name = tensor.name - - if weight_type.name not in ("F32", "F16"): - weight_type_name = name.replace("weight", "qweight_type") - weight_type = torch.tensor(weight_type) - yield weight_type_name, weight_type - - for tensor in reader.tensors: - weight = tensor.data - weight_type = tensor.tensor_type - name = tensor.name - if weight_type.name not in ("F32", "F16"): - name = name.replace("weight", "qweight") - if weight_type.name == "BF16" and tensor.data.dtype == np.uint8: - # BF16 is currently the only "quantization" type that isn't - # actually quantized but is read as a raw byte tensor. - # Reinterpret as `torch.bfloat16` tensor. - weight = weight.view(np.uint16) - if reader.byte_order == "S": - # GGUF endianness != system endianness - weight = weight.byteswap() - param = torch.tensor(weight).view(torch.bfloat16) - else: - param = torch.tensor(weight) - yield name, param diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py b/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py deleted file mode 100644 index 4ee6439b282..00000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py +++ /dev/null @@ -1,98 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -import torch -from vllm.model_executor.models.utils import WeightsMapper - -from .base import GGUFAdapter, gguf_quant_weights_iterator - -FLUX2_TRANSFORMER_KEYS_RENAME_DICT = { - "single_blocks.": "single_transformer_blocks.", - # Image and text input projections - "img_in": "x_embedder", - "txt_in": "context_embedder", - # Timestep and guidance embeddings - "time_in.in_layer": "time_guidance_embed.timestep_embedder.linear_1", - "time_in.out_layer": "time_guidance_embed.timestep_embedder.linear_2", - "guidance_in.in_layer": "time_guidance_embed.guidance_embedder.linear_1", - "guidance_in.out_layer": "time_guidance_embed.guidance_embedder.linear_2", - # Modulation parameters - "double_stream_modulation_img.lin": "double_stream_modulation_img.linear", - "double_stream_modulation_txt.lin": "double_stream_modulation_txt.linear", - "single_stream_modulation.lin": "single_stream_modulation.linear", - # Final output layer - # "final_layer.adaLN_modulation.1": "norm_out.linear", # Handle separately since we need to swap mod params - "final_layer.linear": "proj_out", -} - -FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP = { - "final_layer.adaLN_modulation.1": "norm_out.linear", -} - -FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP = { - "double_blocks.": "transformer_blocks.", - # Handle fused QKV projections separately as we need to break into Q, K, V projections - "img_attn.norm.query_norm": "attn.norm_q", - "img_attn.norm.key_norm": "attn.norm_k", - "img_attn.proj": "attn.to_out.0", - "img_mlp.0": "ff.linear_in", - "img_mlp.2": "ff.linear_out", - "txt_attn.norm.query_norm": "attn.norm_added_q", - "txt_attn.norm.key_norm": "attn.norm_added_k", - "txt_attn.proj": "attn.to_add_out", - "txt_mlp.0": "ff_context.linear_in", - "txt_mlp.2": "ff_context.linear_out", - # Additional for fuse qkv - "img_attn.qkv": "attn.to_qkv", - "txt_attn.qkv": "attn.add_kv_proj", -} - -FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP = { - "linear1": "attn.to_qkv_mlp_proj", - "linear2": "attn.to_out", - "norm.query_norm": "attn.norm_q", - "norm.key_norm": "attn.norm_k", -} - - -class Flux2KleinGGUFAdapter(GGUFAdapter): - """GGUF adapter for Flux2-Klein models with qkv splitting and adaLN swap.""" - - @staticmethod - def is_compatible(od_config, model: torch.nn.Module, source) -> bool: - model_class = od_config.model_class_name or "" - if model_class.startswith("Flux2"): - return True - cfg = od_config.tf_model_config - if cfg is not None: - model_type = str(cfg.get("model_type", "")).lower() - if model_type.startswith("flux"): - return True - return False - - gguf_to_hf_mapper = WeightsMapper( - # double_stream_modulation - orig_to_new_prefix=FLUX2_TRANSFORMER_KEYS_RENAME_DICT | FLUX2_TRANSFORMER_ADA_LAYER_NORM_KEY_MAP, - orig_to_new_substr=FLUX2_TRANSFORMER_DOUBLE_BLOCK_KEY_MAP | FLUX2_TRANSFORMER_SINGLE_BLOCK_KEY_MAP, - ) - - def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: - def custom_weights_adapter( - weights: Iterable[tuple[str, torch.Tensor]], - ) -> Iterable[tuple[str, torch.Tensor]]: - for name, weight in weights: - # Handle the special case for adaLN modulation parameters that require swapping shift and scale - if name.endswith(".scale"): - name = name.replace(".scale", ".weight") - if name == "norm_out.linear.weight": - shift, scale = weight.chunk(2, dim=0) - weight = torch.cat([scale, shift], dim=0) - yield name, weight - else: - yield name, weight - - weights = gguf_quant_weights_iterator(self.gguf_file) - weights = self.gguf_to_hf_mapper.apply(weights) - yield from custom_weights_adapter(weights) diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py deleted file mode 100644 index cc8ccea5cbd..00000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -import torch - -from .base import GGUFAdapter, gguf_quant_weights_iterator - - -class QwenImageGGUFAdapter(GGUFAdapter): - """GGUF adapter for the Qwen-Image transformer family.""" - - @staticmethod - def is_compatible(od_config, model: torch.nn.Module, source) -> bool: - model_class = od_config.model_class_name or "" - if model_class.startswith("QwenImage"): - return True - cfg = od_config.tf_model_config - if cfg is not None: - model_type = str(cfg.get("model_type", "")).lower() - if model_type.startswith("qwen_image"): - return True - return False - - def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: - yield from gguf_quant_weights_iterator(self.gguf_file) diff --git a/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py b/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py deleted file mode 100644 index 7d896335590..00000000000 --- a/vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -import torch -from vllm.model_executor.models.utils import WeightsMapper - -from .base import GGUFAdapter, gguf_quant_weights_iterator - -Z_IMAGE_KEYS_RENAME_DICT = { - "final_layer.": "all_final_layer.2-1.", - "x_embedder.": "all_x_embedder.2-1.", - ".attention.qkv": ".attention.to_qkv", - ".attention.k_norm": ".attention.norm_k", - ".attention.q_norm": ".attention.norm_q", - ".attention.out": ".attention.to_out.0", - "model.diffusion_model.": "", -} - - -class ZImageGGUFAdapter(GGUFAdapter): - """GGUF adapter for Z-Image models with QKV/FFN shard support.""" - - @staticmethod - def is_compatible(od_config, model: torch.nn.Module, source) -> bool: - model_class = od_config.model_class_name or "" - if model_class.startswith("ZImage"): - return True - cfg = od_config.tf_model_config - if cfg is not None: - model_type = str(cfg.get("model_type", "")).lower() - if model_type in {"z_image", "zimage", "z-image"}: - return True - return False - - gguf_to_hf_mapper = WeightsMapper( - orig_to_new_substr=Z_IMAGE_KEYS_RENAME_DICT, - ) - - def weights_iterator(self) -> Iterable[tuple[str, torch.Tensor]]: - weights = gguf_quant_weights_iterator(self.gguf_file) - yield from self.gguf_to_hf_mapper.apply(weights) diff --git a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py index b2035b7a8d0..395e06dc860 100644 --- a/vllm_omni/diffusion/models/audiox/pipeline_audiox.py +++ b/vllm_omni/diffusion/models/audiox/pipeline_audiox.py @@ -27,7 +27,7 @@ from vllm_omni.diffusion.models.audiox.audiox_transformer import MMDiffusionTransformer from vllm_omni.diffusion.models.interface import SupportAudioOutput from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.transformers_utils.processors import audiox as _audiox_transforms _VIDEO_ONLY_TASKS = _audiox_transforms.VIDEO_ONLY_TASKS @@ -60,6 +60,15 @@ def get_audiox_post_process_func(od_config: OmniDiffusionConfig): """Convert the pipeline's float audio tensor to a CPU numpy array for serving.""" def post_process_func(audio: torch.Tensor) -> Any: + if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): + payload = dict(audio["payload"]) + audio_payload = payload.get("audio") + if isinstance(audio_payload, torch.Tensor): + payload["audio"] = audio_payload.detach().cpu().float().numpy() + return { + "payload": payload, + "metadata": audio.get("metadata") or {}, + } if isinstance(audio, torch.Tensor): return audio.detach().cpu().float().numpy() return audio @@ -359,6 +368,7 @@ def __call__(self, sigma: torch.Tensor, sigma_next: torch.Tensor) -> torch.Tenso class AudioXPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMixin): + supports_request_batch = False support_audio_output: ClassVar[bool] = True audio_sample_rate: ClassVar[int] = 44100 audio_channels: ClassVar[int] = 2 @@ -811,7 +821,7 @@ def _build_conditioning_batch( ) return batch - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: if req.prompts is None or len(req.prompts) == 0: raise ValueError("AudioXPipeline requires at least one prompt.") normalized_prompts = _normalize_prompts(list(req.prompts)) @@ -913,10 +923,11 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: target_samples = int(seconds_total * self._sample_rate) audio = audio[..., :target_samples] + stage_durations = self.stage_durations if getattr(self, "enable_diffusion_pipeline_profiler", False) else None return DiffusionOutput( - output=audio, - custom_output={"audiox_task": task_norm}, - stage_durations=self.stage_durations - if getattr(self, "enable_diffusion_pipeline_profiler", False) - else None, + output={ + "payload": {"audio": audio}, + "metadata": {"audio": {"audiox_task": task_norm}}, + }, + stage_durations=stage_durations, ) diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index 1cb20e5c990..30f4a041b72 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -7,12 +7,13 @@ # Original file was released under Apache-2.0, with the full license text # available at https://github.com/huggingface/transformers/blob/main/LICENSE. -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from dataclasses import dataclass import numpy as np import torch import torch.distributed as dist +from cache_dit import ForwardPattern from torch import nn from transformers.models.qwen2.configuration_qwen2 import Qwen2Config from transformers.models.qwen2.modeling_qwen2 import ( @@ -20,11 +21,11 @@ ) from transformers.utils import ModelOutput from vllm.distributed import get_tensor_model_parallel_world_size -from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, - QKVParallelLinear, RowParallelLinear, ) from vllm.model_executor.layers.quantization.base_config import ( @@ -36,18 +37,25 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata as DiffusionAttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention +from vllm_omni.diffusion.cache.cache_dit_backend import BagelCachedAdapter, CacheDiTAdapterConfig from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, - get_classifier_free_guidance_rank, get_classifier_free_guidance_world_size, get_sequence_parallel_rank, get_sp_group, ) from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available +from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm +from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import MoTQKVParallelLinear +from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import MoTRowParallelLinear from vllm_omni.diffusion.layers.rope import RotaryEmbedding +from vllm_omni.diffusion.utils.kv_utils import left_pad_stack from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding +logger = init_logger(__name__) + def patchify(imgs, p): """ @@ -203,6 +211,13 @@ def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Te class BagelMLP(nn.Module): + """FFN with Mixture-of-Tokens routing via MoT parallel linear layers. + + gate_proj + up_proj are fused into a single MoTMergedColumnParallelLinear. + down_proj uses MoTRowParallelLinear. Both layers hold text weights on self + and vae weights on self.gen_exp, routing by text_indices / vae_indices. + """ + def __init__( self, hidden_size: int, @@ -212,29 +227,35 @@ def __init__( prefix: str = "", ) -> None: super().__init__() + + self.intermediate_size = intermediate_size + self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size, intermediate_size], + input_size=hidden_size, + output_sizes=[intermediate_size, intermediate_size], bias=False, + gather_output=False, quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - input_is_parallel=True, + input_size=intermediate_size, + output_size=hidden_size, bias=False, + input_is_parallel=True, quant_config=quant_config, prefix=f"{prefix}.down_proj", ) if hidden_act != "silu": raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.") - self.act_fn = nn.SiLU() + self.act_fn = SiluAndMul() - def forward(self, x): + def forward( + self, + x: torch.Tensor, + ) -> torch.Tensor: gate_up, _ = self.gate_up_proj(x) - gate, up = gate_up.chunk(2, dim=-1) - x = self.act_fn(gate) * up + x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x @@ -306,6 +327,10 @@ class NaiveCache: def __init__(self, num_layers): self.key_cache = {k: None for k in range(num_layers)} self.value_cache = {k: None for k in range(num_layers)} + # Track kv_lens; we need this because we pack the forward passes + # for CFG into a single forward call and the kv length may be different, + # e.g., due to 0 kvs for text_cfg path and nonzero for others + self.key_values_lens: list[int] | None = None @property def num_layers(self): @@ -318,6 +343,76 @@ def seq_lens(self): else: return 0 + @classmethod + def from_object(cls, obj) -> "NaiveCache": + """Convert a duck-typed cache (e.g., SimpleNamespace from KV transfer) + to NaiveCache; in the future, we should find a better way to handle this, + e.g., a model agnostic abstraction for key cache transfer instead of having + this cache live in bagel. + + NOTE: If a NaiveCache is provided, the object is just returned. Otherwise, + we enumerate over the key/value cache values and map layer indices to the + corresponding tensors. + """ + if isinstance(obj, cls): + return obj + cache = cls(len(obj.key_cache)) + for i, (k, v) in enumerate(zip(obj.key_cache, obj.value_cache, strict=True)): + cache.key_cache[i] = k + cache.value_cache[i] = v + return cache + + @staticmethod + def merge(caches: Sequence["NaiveCache"]) -> "NaiveCache": + """Merge per-branch NaiveCaches into one for batched attention; + this lets us do the forward passes for CFG in one batched pass, + although it's worth noting that this is currently only used + for single request. We need this so that gen mode knows the + respective kv lengths, and can split things back out as needed. + """ + num_layers = caches[0].num_layers + merged = NaiveCache(num_layers) + lens = [c.seq_lens for c in caches] + merged.key_values_lens = lens + + nonempty = [c for c in caches if c.key_cache[0] is not None] + if not nonempty: + return merged + + for layer in range(num_layers): + merged.key_cache[layer] = torch.cat([c.key_cache[layer] for c in nonempty], dim=0) + merged.value_cache[layer] = torch.cat([c.value_cache[layer] for c in nonempty], dim=0) + + return merged + + @staticmethod + def split_with_zeros( + tensor: torch.Tensor, + lengths: Sequence[int], + ) -> list[torch.Tensor | None]: + """Split tensor by lengths, which may include 0 entries, e.g., for splitting cfg + branches out, since text_cfg may have 0 kv length. + + 0 lengths will be replaced with None in the returned list. + """ + # Ensure that the lengths are all nonzero and sum to the first dim of our tensor + if not all(isinstance(ln, int) and ln >= 0 for ln in lengths): + raise ValueError("split lengths must be greater than or equal to zero") + + expected = sum(ln for ln in lengths if ln > 0) + if tensor.shape[0] != expected: + raise ValueError(f"tensor dim 0 ({tensor.shape[0]}) != sum of nonzero lengths ({expected})") + + result: list[torch.Tensor | None] = [] + offset = 0 + for ln in lengths: + if ln > 0: + result.append(tensor[offset : offset + ln]) + offset += ln + else: + result.append(None) + return result + @dataclass class BaseNavitOutputWithPast(ModelOutput): @@ -328,12 +423,9 @@ class BaseNavitOutputWithPast(ModelOutput): class PackedAttentionMoT(nn.Module): """Packed attention with Mixture-of-Tokens routing for understanding/generation. - Uses vLLM's QKVParallelLinear and RowParallelLinear for tensor parallelism - support, following the same pattern as vLLM's Qwen2Attention. - - The q/k/v projections are stacked into a single QKVParallelLinear: - - qkv_proj : stacks q_proj + k_proj + v_proj (understanding + gen text) - - qkv_proj_moe_gen : stacks q_proj_moe_gen + k_proj_moe_gen + v_proj_moe_gen (gen vae) + Uses MoTQKVParallelLinear and MoTRowParallelLinear for tensor parallelism. + Text and vae weights are held within the same MoT layer (text on self, + vae on self.gen_exp). Token routing is driven by text_indices / vae_indices. """ def __init__( @@ -363,47 +455,26 @@ def __init__( self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim - # Understanding mode projections (stacked q/k/v) - self.qkv_proj = QKVParallelLinear( + self.qkv_proj = MoTQKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=True, + vae_bias=True, quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - self.hidden_size, + self.o_proj = MoTRowParallelLinear( + input_size=self.total_num_heads * self.head_dim, + output_size=self.hidden_size, bias=False, quant_config=quant_config, prefix=f"{prefix}.o_proj", ) - # Generation mode MoE projections (stacked q/k/v) - self.qkv_proj_moe_gen = QKVParallelLinear( - self.hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj_moe_gen", - ) - self.o_proj_moe_gen = RowParallelLinear( - self.total_num_heads * self.head_dim, - self.hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj_moe_gen", - ) - - # QK normalization - self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.q_norm_moe_gen = RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm_moe_gen = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.q_norm = MoTRMSNorm(self.head_dim, head_norm=True, eps=config.rms_norm_eps) + self.k_norm = MoTRMSNorm(self.head_dim, head_norm=True, eps=config.rms_norm_eps) self.rotary_op = RotaryEmbedding(is_neox_style=True) @@ -431,6 +502,7 @@ def _is_sp_active(self) -> bool: def _forward_gen( self, packed_query_sequence: torch.Tensor, + query_lens: torch.Tensor, packed_query_position_embeddings: torch.Tensor, past_key_values: NaiveCache | None, packed_vae_token_indexes: torch.Tensor, @@ -455,51 +527,43 @@ def _forward_gen( Currently we shouldn't need it in the model, and it would be ideal to handle packing/batching etc in a more model agnostic way. """ - packed_query_sequence = packed_query_sequence.to(torch.bfloat16) - packed_text_query_sequence = packed_query_sequence[packed_text_indexes] - packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes] + text_indices = packed_text_indexes + vae_indices = packed_vae_token_indexes - # Project text tokens through base qkv - text_qkv, _ = self.qkv_proj(packed_text_query_sequence) - text_q, text_k, text_v = text_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + cache_k = cache_v = None + packed_query_sequence = packed_query_sequence.to(torch.bfloat16) - # Project vae tokens through moe_gen qkv - vae_qkv, _ = self.qkv_proj_moe_gen(packed_vae_query_sequence) - vae_q, vae_k, vae_v = vae_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + # MoT QKV projection routes text/vae tokens to the matching weights. + qkv, _ = self.qkv_proj(packed_query_sequence, text_indices, vae_indices) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) # Reshape to (tokens, heads, head_dim) - text_q = text_q.view(-1, self.num_heads, self.head_dim) - text_k = text_k.view(-1, self.num_kv_heads, self.head_dim) - text_v = text_v.view(-1, self.num_kv_heads, self.head_dim) - vae_q = vae_q.view(-1, self.num_heads, self.head_dim) - vae_k = vae_k.view(-1, self.num_kv_heads, self.head_dim) - vae_v = vae_v.view(-1, self.num_kv_heads, self.head_dim) - - # Apply QK norms - text_q = self.q_norm(text_q.to(torch.float32)) - text_k = self.k_norm(text_k.to(torch.float32)) - vae_q = self.q_norm_moe_gen(vae_q.to(torch.float32)) - vae_k = self.k_norm_moe_gen(vae_k.to(torch.float32)) - - # Apply RoPE - need to build per-token cos/sin for text and vae separately - # packed_query_position_embeddings are ordered as the packed sequence - cos_full, sin_full = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings] - text_cos = cos_full[packed_text_indexes] - text_sin = sin_full[packed_text_indexes] - vae_cos = cos_full[packed_vae_token_indexes] - vae_sin = sin_full[packed_vae_token_indexes] - - text_q = self.rotary_op(text_q.to(text_cos.dtype).unsqueeze(0), text_cos, text_sin).squeeze(0) - text_k = self.rotary_op(text_k.to(text_cos.dtype).unsqueeze(0), text_cos, text_sin).squeeze(0) - vae_q = self.rotary_op(vae_q.to(vae_cos.dtype).unsqueeze(0), vae_cos, vae_sin).squeeze(0) - vae_k = self.rotary_op(vae_k.to(vae_cos.dtype).unsqueeze(0), vae_cos, vae_sin).squeeze(0) - - text_q = text_q.to(torch.bfloat16) - text_k = text_k.to(torch.bfloat16) - text_v = text_v.to(torch.bfloat16) - vae_q = vae_q.to(torch.bfloat16) - vae_k = vae_k.to(torch.bfloat16) - vae_v = vae_v.to(torch.bfloat16) + q = q.view(-1, self.num_heads, self.head_dim) + k = k.view(-1, self.num_kv_heads, self.head_dim) + v = v.view(-1, self.num_kv_heads, self.head_dim) + + # MoT QK norms route text/vae tokens to weight/gen_weight internally. + q = self.q_norm(q.to(torch.float32), text_indices, vae_indices) + k = self.k_norm(k.to(torch.float32), text_indices, vae_indices) + + cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings] + q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0) + k = self.rotary_op(k.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0) + + q = q.to(torch.bfloat16) + k = k.to(torch.bfloat16) + v = v.to(torch.bfloat16) + + text_q = q[text_indices] + text_k = k[text_indices] + text_v = v[text_indices] + vae_q = q[vae_indices] + vae_k = k[vae_indices] + vae_v = v[vae_indices] + + num_branches = len(query_lens) + text_per_branch = text_q.shape[0] // num_branches + vae_per_branch = vae_q.shape[0] // num_branches # Build joint K/V: [kv_cache, text_markers] (replicated across SP ranks) if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None: @@ -527,19 +591,54 @@ def _forward_gen( ), ) else: - q = torch.cat([text_q, vae_q], dim=0).unsqueeze(0) - k = torch.cat([ctx_k, vae_k], dim=0).unsqueeze(0) - v = torch.cat([ctx_v, vae_v], dim=0).unsqueeze(0) - attn_out = self.attn_noncausal(q, k, v) - - text_len = text_q.shape[0] - attn_out = attn_out.squeeze(0) - text_attn = attn_out[:text_len].reshape(text_len, self.q_size) - vae_attn = attn_out[text_len:].reshape(-1, self.q_size) - - # Apply output projections - text_out, _ = self.o_proj(text_attn) - vae_out, _ = self.o_proj_moe_gen(vae_attn) + text_q_parts = text_q.split([text_per_branch] * num_branches) + vae_q_parts = vae_q.split([vae_per_branch] * num_branches) + text_k_parts = text_k.split([text_per_branch] * num_branches) + vae_k_parts = vae_k.split([vae_per_branch] * num_branches) + text_v_parts = text_v.split([text_per_branch] * num_branches) + vae_v_parts = vae_v.split([vae_per_branch] * num_branches) + + # Query lengths should not be variable since we + # just split above, so we just concat + stack to 4D + q_4d = torch.stack([torch.cat([t, v]) for t, v in zip(text_q_parts, vae_q_parts)]) + + if cache_k is not None and cache_v is not None: + kv_lens = getattr(past_key_values, "key_values_lens", None) + if kv_lens is None: + per_branch = cache_k.shape[0] // num_branches + kv_lens = [per_branch] * num_branches + + ck_per_branch = NaiveCache.split_with_zeros(cache_k, kv_lens) + cv_per_branch = NaiveCache.split_with_zeros(cache_v, kv_lens) + k_branches = [ + torch.cat([t for t in (ck_per_branch[i], text_k_parts[i], vae_k_parts[i]) if t is not None]) + for i in range(num_branches) + ] + v_branches = [ + torch.cat([t for t in (cv_per_branch[i], text_v_parts[i], vae_v_parts[i]) if t is not None]) + for i in range(num_branches) + ] + k_4d, mask = left_pad_stack(k_branches) + v_4d, _ = left_pad_stack(v_branches) + metadata = DiffusionAttentionMetadata(attn_mask=mask) if mask is not None else None + else: + k_4d = torch.stack([torch.cat([t, v]) for t, v in zip(text_k_parts, vae_k_parts)]) + v_4d = torch.stack([torch.cat([t, v]) for t, v in zip(text_v_parts, vae_v_parts)]) + metadata = None + attn_out = self.attn_noncausal(q_4d, k_4d, v_4d, metadata) + + attn_out = attn_out.reshape(num_branches, -1, self.q_size) + text_attn = attn_out[:, :text_per_branch].reshape(-1, self.q_size) + vae_attn = attn_out[:, text_per_branch:].reshape(-1, self.q_size) + text_len = text_attn.shape[0] + + # MoT output projection over local [text, vae] order. + local_packed = torch.cat([text_attn, vae_attn], dim=0) + local_text_idx = torch.arange(text_len, device=local_packed.device) + local_vae_idx = torch.arange(text_len, text_len + vae_attn.shape[0], device=local_packed.device) + local_out, _ = self.o_proj(local_packed, local_text_idx, local_vae_idx) + text_out = local_out[:text_len] + vae_out = local_out[text_len:] # Merge back into packed format total_len = packed_query_sequence.shape[0] @@ -675,6 +774,7 @@ def forward( raise ValueError("Generation model for Bagel requires non-causal attention") return self._forward_gen( packed_query_sequence=packed_query_sequence, + query_lens=query_lens, packed_query_position_embeddings=packed_query_position_embeddings, past_key_values=past_key_values, packed_vae_token_indexes=packed_vae_token_indexes, @@ -702,12 +802,15 @@ def __init__( prefix: str = "", ): super().__init__() + self.layer_idx = layer_idx self.hidden_size = config.hidden_size self.self_attn = attn_module( config, layer_idx, parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.self_attn" ) + self.input_layernorm = MoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.mlp = BagelMLP( config.hidden_size, config.intermediate_size, @@ -722,10 +825,8 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.mlp_moe_gen", ) - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.input_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.post_attention_layernorm = MoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward( self, @@ -743,18 +844,12 @@ def forward( ) -> BaseNavitOutputWithPast: if packed_query_sequence is None: packed_query_sequence = hidden_states + + text_indices = packed_text_indexes if mode == "gen" else None + vae_indices = packed_vae_token_indexes if mode == "gen" else None + residual = packed_query_sequence - if mode == "und": - packed_query_sequence = self.input_layernorm(packed_query_sequence) - elif mode == "gen": - packed_query_sequence_ = torch.zeros_like(packed_query_sequence) - packed_query_sequence_[packed_text_indexes] = self.input_layernorm( - packed_query_sequence[packed_text_indexes] - ) - packed_query_sequence_[packed_vae_token_indexes] = self.input_layernorm_moe_gen( - packed_query_sequence[packed_vae_token_indexes] - ) - packed_query_sequence = packed_query_sequence_ + packed_query_sequence = self.input_layernorm(packed_query_sequence, text_indices, vae_indices) # Self Attention packed_query_sequence, past_key_values = self.self_attn( @@ -776,13 +871,11 @@ def forward( packed_query_sequence = self.post_attention_layernorm(packed_query_sequence) packed_query_sequence = self.mlp(packed_query_sequence) elif mode == "gen": - packed_text_query_sequence = packed_query_sequence[packed_text_indexes] - packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes] - packed_text_query_sequence = self.post_attention_layernorm(packed_text_query_sequence).to(torch.bfloat16) - packed_vae_query_sequence = self.post_attention_layernorm_moe_gen(packed_vae_query_sequence).to( + packed_normed = self.post_attention_layernorm(packed_query_sequence, text_indices, vae_indices).to( torch.bfloat16 ) - + packed_text_query_sequence = packed_normed[packed_text_indexes] + packed_vae_query_sequence = packed_normed[packed_vae_token_indexes] packed_query_sequence_ = torch.zeros_like(packed_query_sequence).to(torch.bfloat16) packed_query_sequence_[packed_text_indexes] = self.mlp(packed_text_query_sequence) packed_query_sequence_[packed_vae_token_indexes] = self.mlp_moe_gen(packed_vae_query_sequence) @@ -794,6 +887,13 @@ def forward( class Qwen2MoTModel(Qwen2PreTrainedModel): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_0, + }, + cached_adapter_cls=BagelCachedAdapter, + ) + _layerwise_offload_blocks_attrs = ["layers"] @staticmethod @@ -829,9 +929,7 @@ def __init__( ] ) - self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - if self.use_moe: - self.norm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm = MoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = BagelRotaryEmbedding(config=config) # Initialize weights and apply final processing @@ -895,18 +993,9 @@ def forward( **extra_inputs, ) - if self.use_moe: - if mode == "und": - packed_query_sequence = self.norm(packed_query_sequence) - elif mode == "gen": - packed_query_sequence_ = torch.zeros_like(packed_query_sequence) - packed_query_sequence_[packed_text_indexes] = self.norm(packed_query_sequence[packed_text_indexes]) - packed_query_sequence_[packed_vae_token_indexes] = self.norm_moe_gen( - packed_query_sequence[packed_vae_token_indexes] - ) - packed_query_sequence = packed_query_sequence_ - else: - packed_query_sequence = self.norm(packed_query_sequence) + text_indices = packed_text_indexes if self.use_moe and mode == "gen" else None + vae_indices = packed_vae_token_indexes if self.use_moe and mode == "gen" else None + packed_query_sequence = self.norm(packed_query_sequence, text_indices, vae_indices) return BaseNavitOutputWithPast( packed_query_sequence=packed_query_sequence, @@ -983,57 +1072,87 @@ def forward( return outputs def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights for vLLM parallel layers. + """Load weights for MoT parallel layers. - Handles stacked parameter remapping for QKVParallelLinear: - - q_proj, k_proj, v_proj -> qkv_proj (shard ids: q, k, v) - - q_proj_moe_gen, k_proj_moe_gen, v_proj_moe_gen -> qkv_proj_moe_gen - Other parallel layers (gate_proj, up_proj, down_proj, embed_tokens, etc.) - keep HF checkpoint names and use weight_loader for TP sharding. + Stacked parameter remapping (checkpoint name → model parameter): + - q/k/v_proj → qkv_proj (text, shard q/k/v) + - q/k/v_proj_moe_gen → qkv_proj.gen_exp (gen, shard q/k/v) + + Direct remapping (no shard dimension): + - o_proj_moe_gen → o_proj.gen_exp + - {norm}_moe_gen.weight → {norm}.gen_weight (all MoTRMSNorm layers) + + Text norm weights (input_layernorm.weight, q_norm.weight, etc.) and + other names (embed_tokens, lm_head) pass through unchanged. """ stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - # More specific _moe_gen patterns FIRST to avoid substring - # ambiguity (`.q_proj` is a substring of `.q_proj_moe_gen`). - (".qkv_proj_moe_gen", ".q_proj_moe_gen", "q"), - (".qkv_proj_moe_gen", ".k_proj_moe_gen", "k"), - (".qkv_proj_moe_gen", ".v_proj_moe_gen", "v"), + # (param_name, weight_name, shard_id) + # _moe_gen patterns MUST come first — `.q_proj` is a substring + # of `.q_proj_moe_gen`, so the more specific pattern must match first. + (".qkv_proj.gen_exp", ".q_proj_moe_gen", "q"), + (".qkv_proj.gen_exp", ".k_proj_moe_gen", "k"), + (".qkv_proj.gen_exp", ".v_proj_moe_gen", "v"), (".qkv_proj", ".q_proj", "q"), (".qkv_proj", ".k_proj", "k"), (".qkv_proj", ".v_proj", "v"), - # MLP gate/up projections — fused into MergedColumnParallelLinear. - # HF checkpoints store separate gate_proj / up_proj weights; - # these entries remap them to the fused gate_up_proj parameter. - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 1), + (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.gate_proj", 0), + (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.up_proj", 1), + (".mlp.gate_up_proj", ".mlp.gate_proj", 0), + (".mlp.gate_up_proj", ".mlp.up_proj", 1), ] - self.stacked_params_mapping = stacked_params_mapping + + direct_remap = [ + (".o_proj_moe_gen.", ".o_proj.gen_exp."), + # Norm _moe_gen.weight → {norm_name}.gen_weight + (".input_layernorm_moe_gen.", ".input_layernorm.gen_"), + (".post_attention_layernorm_moe_gen.", ".post_attention_layernorm.gen_"), + (".q_norm_moe_gen.", ".q_norm.gen_"), + (".k_norm_moe_gen.", ".k_norm.gen_"), + (".norm_moe_gen.", ".norm.gen_"), + ] + params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + def handle_weight(name, loaded_weight, shard_id=None): + param = params_dict.get(name) + if param is None: + logger.warning_once("Skipping weight %r: no matching parameter found in model.", name) + return + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is not None: + weight_loader(param, loaded_weight, shard_id) + else: + weight_loader(param, loaded_weight) + loaded_params.add(name) + for name, loaded_weight in weights: - loaded = False + # match direct remap + handled = False + for old_substr, new_substr in direct_remap: + if old_substr in name: + name = name.replace(old_substr, new_substr) + handle_weight(name, loaded_weight) + handled = True + break + + if handled: + continue + + # match stacked params mapping for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue - stacked_name = name.replace(weight_name, param_name) - param = params_dict.get(stacked_name) - if param is None: - break - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight, shard_id) - name = stacked_name - loaded = True + name = name.replace(weight_name, param_name) + handle_weight(name, loaded_weight, shard_id) + handled = True break - if not loaded: - param = params_dict.get(name) - if param is None: - continue - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) + if handled: + continue - loaded_params.add(name) + # no-name-match cases are handled here + handle_weight(name, loaded_weight) return loaded_params @@ -1128,10 +1247,17 @@ def get_flattened_position_ids_extrapolate(img_h, img_w, patch_size, max_num_pat return pos_ids -class Bagel(nn.Module): +class Bagel(CFGParallelMixin, nn.Module): config_class = BagelConfig base_model_prefix = "bagel" + # Flow-matching denoise schedule convention. Official BAGEL samples + # ``num_timesteps`` points over [1, 0] and drops the terminal t=0, yielding + # ``num_timesteps - 1`` Euler steps. Lance samples one extra point + # (``num_timesteps + 1``) for ``num_timesteps`` steps; ``LanceBagel`` flips + # this on. See https://github.com/vllm-project/vllm-omni/issues/4470. + _denoise_schedule_extra_step: bool = False + def __init__( self, language_model, @@ -1693,11 +1819,11 @@ def generate_image( frame_condition_token_indexes = frame_condition_token_indexes.to(x_t.device).long() pinned_x_t = x_t[frame_condition_token_indexes].clone() - # Use num_timesteps + 1 sample points so we get `num_timesteps` denoise - # steps after dropping the terminal t=0 (which has no dt). Upstream - # Lance / BAGEL both use this convention; without the +1 we silently - # run one fewer denoise iteration than the user asked for. - timesteps = torch.linspace(1, 0, num_timesteps + 1, device=x_t.device) + # Build the flow-matching schedule. BAGEL drops the terminal t=0 for + # ``num_timesteps - 1`` Euler steps; Lance keeps it for ``num_timesteps``. + # ``_denoise_schedule_extra_step`` (overridden by ``LanceBagel``) selects which. + num_sample_points = num_timesteps + 1 if self._denoise_schedule_extra_step else num_timesteps + timesteps = torch.linspace(1, 0, num_sample_points, device=x_t.device) timesteps = timestep_shift * timesteps / (1 + (timestep_shift - 1) * timesteps) dts = timesteps[:-1] - timesteps[1:] timesteps = timesteps[:-1] @@ -1854,15 +1980,15 @@ def generate_image( # Each CFG branch runs its own LLM forward; we just need the # per-branch packed_position_ids and past_key_values for # ``Bagel.forward`` to dispatch through. - cfg_branches: dict | None = None + cfg_branch_pids: list[torch.Tensor] | None = None + cfg_branch_caches: list[NaiveCache] | None = None if use_cfg_text: - branches_pid = [packed_position_ids, cfg_text_packed_position_ids] - branches_cache = [past_key_values, cfg_text_past_key_values] + cfg_branch_pids = [packed_position_ids, cfg_text_packed_position_ids] + cfg_branch_caches = [past_key_values, cfg_text_past_key_values] if use_cfg_img: - branches_pid.append(cfg_img_packed_position_ids) - branches_cache.append(cfg_img_past_key_values) - cfg_branches = {"pids": branches_pid, "caches": branches_cache} + cfg_branch_pids.append(cfg_img_packed_position_ids) + cfg_branch_caches.append(cfg_img_past_key_values) if return_trajectory_latents and len(timesteps) > 0: trajectory_latents.append(x_t.clone()) @@ -1895,7 +2021,8 @@ def generate_image( cfg_renorm_type=cfg_renorm_type, cfg_text_scale=cfg_text_scale_, cfg_img_scale=cfg_img_scale_, - cfg_branches=cfg_branches, + cfg_branch_pids=cfg_branch_pids, + cfg_branch_caches=cfg_branch_caches, ) if scheduler is not None: @@ -1951,7 +2078,6 @@ def _generate_image_parallel( Rank 2: img_cfg branch (no image condition), only when cfg_img_scale > 1.0 """ cfg_group = get_cfg_group() - cfg_rank = get_classifier_free_guidance_rank() cfg_world_size = get_classifier_free_guidance_world_size() use_cfg_img = cfg_img_scale > 1.0 @@ -1974,22 +2100,6 @@ def _generate_image_parallel( x_t = x_t.contiguous() cfg_group.broadcast(x_t, src=0) - # Select this rank's branch inputs - if cfg_rank == 0: - # Gen branch: use main inputs directly - branch_position_ids = packed_position_ids - branch_past_key_values = past_key_values - elif cfg_rank == 1: - # Text CFG branch - branch_position_ids = cfg_text_packed_position_ids - branch_past_key_values = cfg_text_past_key_values - elif cfg_rank == 2: - # Image CFG branch - branch_position_ids = cfg_img_packed_position_ids - branch_past_key_values = cfg_img_past_key_values - else: - raise RuntimeError(f"Unexpected cfg_rank={cfg_rank} for Bagel 3-branch CFG parallel") - trajectory_latents: list[torch.Tensor] | None = [] if return_trajectory_latents else None trajectory_timesteps: list[torch.Tensor] | None = [] if return_trajectory_latents else None trajectory_log_probs: list[torch.Tensor] | None = ( @@ -2010,43 +2120,46 @@ def _generate_image_parallel( timestep[frame_condition_token_indexes] = 0.0 use_cfg_this_step = t > cfg_interval[0] and t <= cfg_interval[1] and cfg_text_scale > 1.0 - if use_cfg_this_step: - # CFG interval: each rank computes its own branch - local_v_t = self.forward_single_branch( - x_t=x_t, - timestep=timestep, - packed_vae_token_indexes=packed_vae_token_indexes, - packed_vae_position_ids=packed_vae_position_ids, - packed_text_ids=packed_text_ids, - packed_text_indexes=packed_text_indexes, - packed_position_ids=branch_position_ids, - packed_seqlens=packed_seqlens, - past_key_values=branch_past_key_values, + # Per-branch kwargs. Branch 0 (gen, full conditioning) is also the + # branch used by all ranks when do_true_cfg is False (outside the + # CFG interval) — CFGParallelMixin only runs branches_kwargs[0] then, + # mirroring the previous "all ranks compute gen inputs, no comm" path. + common = dict( + x_t=x_t, + timestep=timestep, + packed_vae_token_indexes=packed_vae_token_indexes, + packed_vae_position_ids=packed_vae_position_ids, + packed_text_ids=packed_text_ids, + packed_text_indexes=packed_text_indexes, + packed_seqlens=packed_seqlens, + ) + branches_kwargs = [ + dict(**common, packed_position_ids=packed_position_ids, past_key_values=past_key_values), + dict( + **common, packed_position_ids=cfg_text_packed_position_ids, past_key_values=cfg_text_past_key_values + ), + ] + if use_cfg_img: + branches_kwargs.append( + dict( + **common, + packed_position_ids=cfg_img_packed_position_ids, + past_key_values=cfg_img_past_key_values, + ) ) - gathered = cfg_group.all_gather(local_v_t, separate_tensors=True) - v_t = self._combine_cfg( - gathered[0], - gathered[1], - gathered[2] if (use_cfg_img and len(gathered) > 2) else None, - cfg_text_scale, - cfg_img_scale, - cfg_renorm_type, - cfg_renorm_min, - ) - else: - # Outside CFG interval: all ranks compute with gen inputs, no comm - v_t = self.forward_single_branch( - x_t=x_t, - timestep=timestep, - packed_vae_token_indexes=packed_vae_token_indexes, - packed_vae_position_ids=packed_vae_position_ids, - packed_text_ids=packed_text_ids, - packed_text_indexes=packed_text_indexes, - packed_position_ids=packed_position_ids, - packed_seqlens=packed_seqlens, - past_key_values=past_key_values, - ) + # Each rank computes its assigned branch, then all_gather + combine + # happen inside the mixin (identical result on every rank). + v_t = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=use_cfg_this_step, + true_cfg_scale={ + "cfg_text_scale": cfg_text_scale, + "cfg_img_scale": cfg_img_scale, + "cfg_renorm_type": cfg_renorm_type, + "cfg_renorm_min": cfg_renorm_min, + }, + branches_kwargs=branches_kwargs, + ) if scheduler is not None: out = scheduler.step(v_t.to(x_t.device), timesteps[i], x_t, dts[i], **_sched_kw) @@ -2115,6 +2228,45 @@ def _combine_cfg( return v_t + # ── CFGParallelMixin hooks ── + # Bagel mounts CFGParallelMixin (see class declaration) to reuse the shared + # N-branch CFG dispatch/all_gather logic in predict_noise_with_multi_branch_cfg. + # Only two hooks need model-specific behaviour: + # * predict_noise: one branch == one Bagel forward (per-branch KV cache). + # * combine_multi_branch_cfg_noise: Bagel's renorm-aware 3-branch combine. + + def predict_noise(self, **kwargs) -> torch.Tensor: + """Single-branch velocity prediction for CFGParallelMixin. + + Each CFG branch differs only by ``packed_position_ids`` and + ``past_key_values`` (carried in ``kwargs``); the heavy lifting is the + per-branch ``forward_single_branch`` pass. + """ + return self.forward_single_branch(**kwargs) + + def combine_multi_branch_cfg_noise( + self, + predictions: list[torch.Tensor], + true_cfg_scale: dict[str, float], + cfg_normalize: bool = False, + ) -> torch.Tensor: + """Combine gen/text/img branch velocities via Bagel's renorm CFG. + + ``predictions[0]`` is the gen branch, ``[1]`` the text-CFG branch, and + ``[2]`` (when present) the image-CFG branch. ``cfg_normalize`` is unused + because renormalization is folded into ``_combine_cfg`` itself. + """ + cfg_img_v_t = predictions[2] if len(predictions) > 2 else None + return self._combine_cfg( + predictions[0], + predictions[1], + cfg_img_v_t, + true_cfg_scale["cfg_text_scale"], + true_cfg_scale["cfg_img_scale"], + true_cfg_scale["cfg_renorm_type"], + true_cfg_scale["cfg_renorm_min"], + ) + def forward_single_branch( self, x_t: torch.Tensor, @@ -2239,7 +2391,8 @@ def forward( cfg_renorm_type: str = "global", cfg_text_scale: float = 1.0, cfg_img_scale: float = 1.0, - cfg_branches: dict | None = None, + cfg_branch_pids: list[torch.Tensor] | None = None, + cfg_branch_caches: list[NaiveCache] | None = None, ): # Build query sequence (identical for all CFG branches) packed_text_embedding = self.language_model.forward( @@ -2267,31 +2420,37 @@ def forward( cfg_text_v_t = None cfg_img_v_t = None - if use_cfg and cfg_branches is not None: - # Sequential per-branch CFG forwards (matches upstream lance.py). - # The previous batched path concatenated cond + cfg into one LLM - # forward, but the block-diagonal attention mask was lost when - # PR #3728 dropped flash_attn_varlen, so branches leaked into - # each other. Running each branch through its own forward is - # numerically identical to upstream. - def _run_branch(branch_pkv, branch_pids): - out = self.language_model.forward( - packed_query_sequence=packed_sequence, - query_lens=packed_seqlens, - packed_query_position_ids=branch_pids, - past_key_values=branch_pkv, - update_past_key_values=False, - is_causal=False, - **extra_inputs, - ) - return self.llm2vae(out.packed_query_sequence)[packed_vae_token_indexes] - - branches_pids = cfg_branches["pids"] - branches_caches = cfg_branches["caches"] - v_t = _run_branch(branches_caches[0], branches_pids[0]) - cfg_text_v_t = _run_branch(branches_caches[1], branches_pids[1]) - if cfg_img_scale > 1.0 and len(branches_caches) > 2: - cfg_img_v_t = _run_branch(branches_caches[2], branches_pids[2]) + if use_cfg and cfg_branch_pids is not None and cfg_branch_caches is not None: + num_branches = len(cfg_branch_pids) + seq_len = int(packed_seqlens.sum()) + + batched_sequence = packed_sequence.repeat(num_branches, 1) + batched_vae_indexes = torch.cat([packed_vae_token_indexes + i * seq_len for i in range(num_branches)]) + batched_position_ids = torch.cat(cfg_branch_pids) + batched_seqlens = packed_seqlens.repeat(num_branches) + merged_cache = NaiveCache.merge(cfg_branch_caches) + + if self.use_moe: + batched_text_indices = torch.cat([packed_text_indexes + i * seq_len for i in range(num_branches)]) + extra_inputs["packed_vae_token_indexes"] = batched_vae_indexes + extra_inputs["packed_text_indexes"] = batched_text_indices + + output = self.language_model.forward( + packed_query_sequence=batched_sequence, + query_lens=batched_seqlens, + packed_query_position_ids=batched_position_ids, + past_key_values=merged_cache, + update_past_key_values=False, + is_causal=False, + **extra_inputs, + ) + + all_vae_v_t = self.llm2vae(output.packed_query_sequence)[batched_vae_indexes] + vae_per_branch = packed_vae_token_indexes.shape[0] + branch_v_ts = all_vae_v_t.split(vae_per_branch) + v_t = branch_v_ts[0] + cfg_text_v_t = branch_v_ts[1] + cfg_img_v_t = branch_v_ts[2] if len(branch_v_ts) > 2 else None else: # Single forward (no CFG or outside cfg_interval). output = self.language_model.forward( diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 56a5b0dcc40..9786e2b473b 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -29,7 +29,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from .autoencoder import AutoEncoder, AutoEncoderParams, DistributedAutoEncoder @@ -341,7 +341,7 @@ def _regen_init_noise_on_device(self, gen_input: dict, seed: int | None) -> None ) @torch.inference_mode() - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: if len(req.prompts) > 1: logger.warning( """This model only supports a single prompt, not a batched request.""", @@ -395,6 +395,7 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: injected_kv = req.sampling_params.past_key_values if injected_kv is not None: logger.info("Using injected KV Cache (direct)") + injected_kv = NaiveCache.from_object(injected_kv) gen_context["past_key_values"] = injected_kv seq_len = injected_kv.key_cache[0].shape[0] gen_context["kv_lens"] = [seq_len] @@ -431,6 +432,7 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: ) if cfg_text_kv is not None: + cfg_text_kv = NaiveCache.from_object(cfg_text_kv) cfg_text_seq_len = cfg_text_kv.key_cache[0].shape[0] cfg_text_context["past_key_values"] = cfg_text_kv cfg_text_context["kv_lens"] = [cfg_text_seq_len] @@ -458,6 +460,7 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: else: cfg_img_context["ropes"] = [cfg_img_seq_len] else: + cfg_img_kv = NaiveCache.from_object(cfg_img_kv) cfg_img_seq_len = cfg_img_kv.key_cache[0].shape[0] cfg_img_context["past_key_values"] = cfg_img_kv cfg_img_context["kv_lens"] = [cfg_img_seq_len] @@ -553,6 +556,8 @@ def vae_transforms(img): for k, v in gen_input_img.items(): if torch.is_tensor(v): gen_input_img[k] = v.to(self.device) + for k in ("packed_indexes", "packed_key_value_indexes", "key_values_lens"): + gen_input_img.pop(k, None) with torch.autocast( device_type=self.device.type, enabled=self.device.type != "cpu", @@ -604,7 +609,7 @@ def vae_transforms(img): # cfg_text_context: update with negative prompt (no text condition). # When empty, keep cfg_text_context as-is (kv_lens=0) to match - # original BAGEL; _merge_naive_caches handles None KV entries. + # original BAGEL. neg_prompt = extra_args.get("negative_prompt", "") if neg_prompt: neg_input, neg_newlens, neg_rope = self.bagel.prepare_prompts( @@ -734,8 +739,10 @@ def vae_transforms(img): text_output = text_output.split("<|im_start|>")[-1] return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -842,22 +849,28 @@ def vae_transforms(img): if trajectory_log_probs: trajectory_log_probs_stacked = torch.stack(trajectory_log_probs) - custom = {} + payload = {"image": img} + metadata = {} if think_text is not None: - custom["think_text"] = think_text - # Mirror the PIL image into ``custom_output`` so callers reading via - # the orchestrator IPC boundary (which strips the bare ``output`` - # field) can still recover the result. ``video_frames`` already - # uses this pattern. - custom["image"] = img + metadata["text"] = {"think_text": think_text} + trajectory_payload = {} + if trajectory_latents_stacked is not None: + trajectory_payload["latents"] = trajectory_latents_stacked + if trajectory_timesteps_stacked is not None: + trajectory_payload["timesteps"] = trajectory_timesteps_stacked + if trajectory_log_probs_stacked is not None: + trajectory_payload["log_probs"] = trajectory_log_probs_stacked + if trajectory_decoded is not None: + trajectory_payload["decoded"] = trajectory_decoded + if trajectory_payload: + payload["trajectory"] = trajectory_payload + metadata["trajectory"] = {"type": "denoising"} return DiffusionOutput( - output=img, - trajectory_latents=trajectory_latents_stacked, - trajectory_timesteps=trajectory_timesteps_stacked, - trajectory_log_probs=trajectory_log_probs_stacked, - trajectory_decoded=trajectory_decoded, - custom_output=custom, + output={ + "payload": payload, + "metadata": metadata, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -869,18 +882,34 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: tp_aware_params = {name for name, p in self.named_parameters() if hasattr(p, "weight_loader")} # Expand allowed/tp_aware_params with stacked param source names. - # QKVParallelLinear merges q_proj+k_proj+v_proj into qkv_proj; the - # checkpoint stores the original separate names. We must recognise - # those names so _filtered_weights does not drop them. + # The model fuses several checkpoint projections into merged layers: + # QKV: q/k/v_proj → qkv_proj, q/k/v_proj_moe_gen → qkv_proj.gen_exp + # and remaps non-stacked weights like: + # {norm}_moe_gen.weight → {norm}.gen_weight (MoTRMSNorm layers) + # We expand allowed names so _filtered_weights does not drop them. _stacked_expansions = [ + # text QKV (".qkv_proj", ".q_proj"), (".qkv_proj", ".k_proj"), (".qkv_proj", ".v_proj"), - (".qkv_proj_moe_gen", ".q_proj_moe_gen"), - (".qkv_proj_moe_gen", ".k_proj_moe_gen"), - (".qkv_proj_moe_gen", ".v_proj_moe_gen"), - (".gate_up_proj", ".gate_proj"), - (".gate_up_proj", ".up_proj"), + # gen QKV + (".qkv_proj.gen_exp", ".q_proj_moe_gen"), + (".qkv_proj.gen_exp", ".k_proj_moe_gen"), + (".qkv_proj.gen_exp", ".v_proj_moe_gen"), + # gen o_proj (non-stacked, but still remapped) + (".o_proj.gen_exp", ".o_proj_moe_gen"), + # text FFN gate+up + (".mlp.gate_up_proj", ".mlp.gate_proj"), + (".mlp.gate_up_proj", ".mlp.up_proj"), + # gen FFN gate+up + (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.gate_proj"), + (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.up_proj"), + # MoTRMSNorm gen_weight ← checkpoint _moe_gen.weight + (".input_layernorm.gen_", ".input_layernorm_moe_gen."), + (".post_attention_layernorm.gen_", ".post_attention_layernorm_moe_gen."), + (".q_norm.gen_", ".q_norm_moe_gen."), + (".k_norm.gen_", ".k_norm_moe_gen."), + (".norm.gen_", ".norm_moe_gen."), ] stacked_source_names: set[str] = set() for name in list(allowed): diff --git a/vllm_omni/diffusion/models/cosmos3/__init__.py b/vllm_omni/diffusion/models/cosmos3/__init__.py index 43be0c37aa7..af4eeaacfab 100644 --- a/vllm_omni/diffusion/models/cosmos3/__init__.py +++ b/vllm_omni/diffusion/models/cosmos3/__init__.py @@ -3,16 +3,16 @@ from .pipeline_cosmos3 import ( Cosmos3OmniDiffusersPipeline, - get_cosmos3_action_post_process_func, get_cosmos3_post_process_func, get_cosmos3_pre_process_func, ) from .transformer_cosmos3 import Cosmos3VFMTransformer +from .transformer_cosmos3_edge import Cosmos3EdgeVFMTransformer __all__ = [ "Cosmos3OmniDiffusersPipeline", - "get_cosmos3_action_post_process_func", "get_cosmos3_post_process_func", "get_cosmos3_pre_process_func", "Cosmos3VFMTransformer", + "Cosmos3EdgeVFMTransformer", ] diff --git a/vllm_omni/diffusion/models/cosmos3/action.py b/vllm_omni/diffusion/models/cosmos3/action.py index 80a1bc7a9b8..c5fe41a91df 100644 --- a/vllm_omni/diffusion/models/cosmos3/action.py +++ b/vllm_omni/diffusion/models/cosmos3/action.py @@ -1,6 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Action-token helpers for Cosmos3 UVA/action generation.""" +"""Action-token helpers for Cosmos3 action generation. + +These helpers cover action modes that use action tokens as an auxiliary output +stream: ``policy``, ``forward_dynamics``, and ``inverse_dynamics``. The pipeline +returns predicted actions through the diffusion output payload/metadata +envelope rather than as a video frame stream. +""" from __future__ import annotations diff --git a/vllm_omni/diffusion/models/cosmos3/guardrails.py b/vllm_omni/diffusion/models/cosmos3/guardrails.py index e249d2826e3..3c6166e85ea 100644 --- a/vllm_omni/diffusion/models/cosmos3/guardrails.py +++ b/vllm_omni/diffusion/models/cosmos3/guardrails.py @@ -21,6 +21,7 @@ from vllm_omni.diffusion.models.progress_bar import _is_rank_zero from vllm_omni.errors import GuardrailViolationError +from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: from vllm_omni.diffusion.data import OmniDiffusionConfig @@ -28,6 +29,28 @@ logger = init_logger(__name__) +# --------------------------------------------------------------------------- +# NPU/XPU compatibility: cosmos_guardrail.cosmos_utils.load_model calls +# ``torch.load(path, weights_only=True)`` without ``map_location``. +# On NPU/XPU, ``torch.cuda.is_available()`` is False, so deserializing a +# CUDA checkpoint with ``weights_only=True`` raises an UnpicklingError. +# Work around by patching ``torch.load`` to default ``map_location="cpu"`` +# when CUDA is unavailable and no explicit map_location was given. +# --------------------------------------------------------------------------- +_original_torch_load = torch.load + + +def _patched_torch_load(*args, **kwargs): + if ( + "map_location" not in kwargs + and not torch.cuda.is_available() + and (current_omni_platform.is_npu() or current_omni_platform.is_xpu()) + ): + kwargs["map_location"] = "cpu" + return _original_torch_load(*args, **kwargs) + + +torch.load = _patched_torch_load # type: ignore[assignment] try: from cosmos_guardrail import CosmosSafetyChecker @@ -75,7 +98,7 @@ def text_guardrail(prompt: str) -> None: def _build_video_guardrail(checker: Any, offload_to_cpu: bool) -> VideoGuardrailFn: video_models = _nn_models(checker.video_guardrail) - compute_device = "cuda" + compute_device = current_omni_platform.device_type def video_guardrail(frames: np.ndarray) -> np.ndarray: if offload_to_cpu: @@ -112,9 +135,9 @@ def _init_default_guardrails(offload_to_cpu: bool = False) -> None: checker = CosmosSafetyChecker() # Place text models on their resting device permanently. Video models - # idle on CPU when offload is on and move to GPU per-call (handled in + # idle on CPU when offload is on and move to compute device per-call (handled in # the video guardrail closure). - idle_device = "cpu" if offload_to_cpu else "cuda" + idle_device = "cpu" if offload_to_cpu else current_omni_platform.device_type for m in _nn_models(checker.text_guardrail): m.to(idle_device) for m in _nn_models(checker.video_guardrail): diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index b0b1f7559c1..41a3e80fcbf 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -1,33 +1,45 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Cosmos3 text/image/video-to-video and text-to-image pipeline for vllm-omni. - -Single pipeline class supports T2V, I2V, V2V, and T2I; the mode is selected at -runtime by: - -* ``prompt["modalities"]`` contains ``"image"``: **T2I** (text-to-image). -* ``prompt["modalities"]`` contains ``"video"`` or is omitted: **T2V** - (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`) - +"""Cosmos3 text/image/video/sound/action pipeline for vllm-omni. + +One pipeline class serves the Cosmos3 family modes. Output modality is selected +mainly by ``prompt["modalities"]``: + +* ``"image"`` selects T2I (text-to-image) and forces a single visual frame. +* ``"video"`` or omitted modalities select video generation. +* ``"audio"`` is accepted for compatibility but does not request sound by + itself; sound is enabled with ``generate_sound`` or ``sound_gen``. + +Video generation is further specialized by inputs and extra args: + +* no image/video input: T2V (text-to-video). +* ``multi_modal_data["image"]``: I2V (image-to-video). +* ``multi_modal_data["video"]`` with no action/transfer mode: V2V + (video-to-video). +* transfer hints (``edge``, ``blur``, ``depth``, ``seg``, or ``wsm``): control + transfer video generation. +* ``action_mode``: action-capable video generation. RoboLab/OpenPI observation + payloads in ``extra_args["robot_obs"]`` or ``extra_args["observation"]`` + bypass normal video output and return an action-only payload/metadata + envelope. + +Generated sound is video-only, cannot be combined with action or transfer, and +is produced from sound latents rather than from ``multi_modal_data["audio"]``. """ from __future__ import annotations +import json import math import os import time -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import fields from typing import Any, ClassVar import numpy as np import PIL.Image import torch -from diffusers import UniPCMultistepScheduler from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor from torch import nn @@ -46,10 +58,19 @@ from vllm_omni.diffusion.models.interface import ( ReferenceVideoDecodeSpec, SupportImageInput, + SupportsComponentDiscovery, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero +from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.entrypoints.openai.video_api_utils import positive_float from vllm_omni.inputs.data import OmniDiffusionSamplingParams from .action import ( @@ -64,8 +85,23 @@ normalize_action_mode, pad_action_to_dim, resolve_domain_id, + vision_condition_indexes, +) +from .transfer import ( + Cosmos3TransferConfig, + has_transfer_hints, + load_or_compute_control_frames, + media_hw, + media_to_uint8_cthw, + normalized_video_to_uint8_cthw, + pad_temporal_frames, + resize_center_crop_uint8_cthw, + resolve_transfer_config, + transfer_max_frames_from_extra_args, + uint8_cthw_to_normalized_5d, ) -from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen +from .transformer_cosmos3 import Cosmos3VFMTransformer, _tf_config_get, resolve_sound_gen +from .transformer_cosmos3_edge import COSMOS3_EDGE_BACKBONE_TYPE, Cosmos3EdgeVFMTransformer from .utils import ( COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION, COSMOS3_VAE_TEMPORAL_COMPRESSION, @@ -105,6 +141,7 @@ logger = init_logger(__name__) + COSMOS3_DEFAULT_CONDITION_PIXEL_FRAMES = ( max(COSMOS3_DEFAULT_CONDITION_FRAME_INDEXES_VISION) * COSMOS3_VAE_TEMPORAL_COMPRESSION + 1 ) @@ -115,14 +152,16 @@ COSMOS3_INVERSE_DURATION_TEMPLATE = "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution." COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{width} resolution." -COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a given prompt." -COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a given prompt." +# NOTE: Intentional typo in "give" instead of "given" to match training setup. +COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a give prompt." +COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a give prompt." COSMOS3_T2V_DEFAULT_HEIGHT = 720 COSMOS3_T2V_DEFAULT_WIDTH = 1280 COSMOS3_T2V_DEFAULT_NUM_FRAMES = 189 COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS = 35 COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE = 6.0 +COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT = 10.0 COSMOS3_T2I_DEFAULT_HEIGHT = 1024 COSMOS3_T2I_DEFAULT_WIDTH = 1024 @@ -131,28 +170,98 @@ COSMOS3_T2I_DEFAULT_FLOW_SHIFT = 3.0 COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL: tuple[float, float] = (400.0, 1000.0) +COSMOS3_EDGE_T2V_DEFAULT_HEIGHT = 480 +COSMOS3_EDGE_T2V_DEFAULT_WIDTH = 832 +COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE = 5.0 +COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT = 3.0 +COSMOS3_EDGE_T2I_DEFAULT_HEIGHT = 640 +COSMOS3_EDGE_T2I_DEFAULT_WIDTH = 640 + +COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS = "FlowMatchEulerDiscreteScheduler" + # Truncation cap on the prompt token count (shared by T2I and T2V). Prompts # are tokenized to their natural length (no padding); this only bounds the # UND pathway / GEN cross-attention cost for pathologically long prompts. COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH = 4096 +def _ceil_video_num_frames(num_frames: int, temporal_compression_factor: int) -> int: + """Round a video length up to the causal VAE's ``factor * k + 1`` grid.""" + if temporal_compression_factor <= 0: + raise ValueError(f"Cosmos3 temporal_compression_factor must be positive, got {temporal_compression_factor}.") + if num_frames <= 1: + return num_frames + latent_frames = math.ceil((num_frames - 1) / temporal_compression_factor) + return latent_frames * temporal_compression_factor + 1 + + +def _format_json_object_prompt( + prompt: str, + *, + num_frames: int, + frame_rate: float, + height: int, + width: int, + aspect_ratio: str | None, +) -> str | None: + """Match Imaginaire4 metadata injection for JSON-object prompts.""" + try: + prompt_obj = json.loads(prompt) + except (json.JSONDecodeError, TypeError, ValueError): + return None + if not isinstance(prompt_obj, dict): + return None + + metadata: dict[str, Any] = {} + if num_frames > 1: + duration_seconds = int(num_frames / frame_rate) if frame_rate > 0 else 0 + metadata.update({"duration": f"{duration_seconds}s", "fps": float(frame_rate)}) + else: + prompt_obj.pop("duration", None) + prompt_obj.pop("fps", None) + metadata["resolution"] = {"H": int(height), "W": int(width)} + if aspect_ratio is not None: + metadata["aspect_ratio"] = aspect_ratio + + prompt_obj.update(metadata) + return json.dumps(prompt_obj) + + +def resolve_cosmos3_transformer_cls(model_config: Any) -> type[Cosmos3VFMTransformer]: + """Select the Cosmos3 transformer implementation from transformer/config.json.""" + backbone_type = _tf_config_get(model_config, "backbone_type", None) + if backbone_type is None: + return Cosmos3VFMTransformer + if backbone_type == COSMOS3_EDGE_BACKBONE_TYPE: + return Cosmos3EdgeVFMTransformer + raise ValueError(f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}.") + + # --------------------------------------------------------------------------- # Post-process function (registered in registry.py) # --------------------------------------------------------------------------- def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig): - """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``. + """Build the request preprocessor for Cosmos3 image/video inputs. + + For plain T2V (no image or video in ``multi_modal_data``), the request is + returned unchanged after the optional guardrail check. For I2V, the + conditioning image is loaded, aspect-resized, center-cropped, and stored as + ``additional_information.preprocessed_image``. For V2V, source frames are + cropped to the target size and stored as + ``additional_information.preprocessed_video``. + + Action modes reuse image/video preprocessing but use action-specific resize + and padding rules. Transfer requests store + ``additional_information.preprocessed_transfer_video`` for optional input + video conditioning. Cosmos3 sound generation is not driven by + ``multi_modal_data["audio"]``; it is enabled later from sampling params. """ from .guardrails import check_text_safety, ensure_initialized, is_guardrails_enabled video_processor = VideoProcessor(vae_scale_factor=16) + is_edge_model = ( + _tf_config_get(getattr(od_config, "tf_model_config", None), "backbone_type", None) == COSMOS3_EDGE_BACKBONE_TYPE + ) # Eager-load guardrail models at pipeline build time when the server-level # gate is on. Per-request overrides only decide whether the loaded models # are *invoked* — they cannot turn checks on without a server-side preload. @@ -166,6 +275,14 @@ def _extra_args(request: OmniDiffusionRequest) -> dict[str, Any]: def _request_action_mode(request: OmniDiffusionRequest) -> str | None: return normalize_action_mode(_extra_args(request).get("action_mode")) + def _set_transfer_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]: + extra = _extra_args(request) + resolution = extra.get("resolution", extra.get("image_size", 720)) + target_w, target_h = find_closest_target_size(image.height, image.width, resolution) + request.sampling_params.height = target_h + request.sampling_params.width = target_w + return int(target_h), int(target_w) + def _set_action_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]: sp = request.sampling_params if sp.height is not None and sp.width is not None: @@ -248,7 +365,43 @@ def _preprocess_condition_image(image: PIL.Image.Image, target_h: int, target_w: 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_value(video: Any, key: str) -> Any: + if isinstance(video, Mapping): + return video.get(key) + return getattr(video, key, None) + + def _video_payload_fps(video: Any) -> float | None: + for key in ("fps", "frame_rate", "source_fps", "input_fps", "avg_fps", "average_fps"): + fps = positive_float(_video_payload_value(video, key)) + if fps is not None: + return fps + for key in ("metadata", "info"): + metadata = _video_payload_value(video, key) + if metadata is None or metadata is video: + continue + fps = _video_payload_fps(metadata) + if fps is not None: + return fps + if isinstance(video, Mapping): + for key in ("frames", "data", "video"): + nested = video.get(key) + if nested is None or nested is video: + continue + fps = _video_payload_fps(nested) + if fps is not None: + return fps + return None + + def _unwrap_video_payload(video: Any) -> Any: + if isinstance(video, Mapping): + for key in ("frames", "data", "video"): + nested = video.get(key) + if nested is not None: + return nested + return video + def _video_payload_to_frames(video: Any) -> list[Any]: + video = _unwrap_video_payload(video) if isinstance(video, list): return video if isinstance(video, torch.Tensor): @@ -295,69 +448,92 @@ def _preprocess_condition_video( def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: action_mode = _request_action_mode(request) + prompt = request.prompt if is_guardrails_enabled(od_config, request.sampling_params): - for prompt in request.prompts: - text = prompt if isinstance(prompt, str) else prompt.get("prompt", "") - check_text_safety(text) - - for i, prompt in enumerate(request.prompts): - if isinstance(prompt, str): - continue - 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 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.") + text = prompt if isinstance(prompt, str) else prompt.get("prompt", "") + check_text_safety(text) + + if isinstance(prompt, str): + return request + 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 raw_video is None: + return request + 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"] = {} + 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.") + raw_video_frames: list[Any] | None = None + transfer_input_fps: float | None = None + if raw_video is not None: + transfer_input_fps = _video_payload_fps(raw_video) + 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: - 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) + if raw_image is None: + 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) + extra = _extra_args(request) + transfer_requested = action_mode is None and has_transfer_hints(extra) - # Auto-calculate H/W from aspect ratio (720p max area) - if request.sampling_params.height is None or request.sampling_params.width is None: - if action_mode is not None: - _set_action_size_from_image(request, image) - else: - max_area = 720 * 1280 - aspect_ratio = image.height / image.width - mod_value = 16 - height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value - width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width - - target_w = request.sampling_params.width - target_h = request.sampling_params.height + # Resolve missing H/W. + if transfer_requested: + _set_transfer_size_from_image(request, image) + elif request.sampling_params.height is None or request.sampling_params.width is None: if action_mode is not None: - prompt["additional_information"]["preprocessed_image"] = _preprocess_action_image( - image, - int(target_h), - int(target_w), + _set_action_size_from_image(request, image) + elif is_edge_model: + if request.sampling_params.height is None: + request.sampling_params.height = COSMOS3_EDGE_T2V_DEFAULT_HEIGHT + if request.sampling_params.width is None: + request.sampling_params.width = COSMOS3_EDGE_T2V_DEFAULT_WIDTH + else: + max_area = 720 * 1280 + aspect_ratio = image.height / image.width + mod_value = 16 + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + target_w = request.sampling_params.width + target_h = request.sampling_params.height + if action_mode is not None: + prompt["additional_information"]["preprocessed_image"] = _preprocess_action_image( + image, + int(target_h), + int(target_w), + ) + elif raw_video is None: + prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image( + image, + int(target_h), + int(target_w), + ) + else: + assert raw_video_frames is not None + if transfer_requested: + if transfer_input_fps is not None: + prompt["additional_information"]["transfer_input_fps"] = transfer_input_fps + transfer_frames = media_to_uint8_cthw( + raw_video_frames, + height=int(target_h), + width=int(target_w), + max_frames=transfer_max_frames_from_extra_args(extra), ) - elif raw_video is None: - prompt["additional_information"]["preprocessed_image"] = _preprocess_condition_image( - image, - int(target_h), - int(target_w), + prompt["additional_information"]["preprocessed_transfer_video"] = uint8_cthw_to_normalized_5d( + transfer_frames, + dtype=torch.float32, ) 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", @@ -378,13 +554,13 @@ def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: 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_frames, - int(target_h), - int(target_w), - ) - request.prompts[i] = prompt + if action_mode is not None and raw_video_frames is not None: + prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video( + raw_video_frames, + int(target_h), + int(target_w), + ) + request.prompt = prompt return request @@ -392,6 +568,13 @@ def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: def get_cosmos3_post_process_func(od_config: OmniDiffusionConfig): + """Build the postprocessor for Cosmos3 image, video, and video+audio output. + + The pipeline returns image payloads as ``{"image": tensor}`` and video + payloads as ``{"video": tensor}``. Sound-enabled video returns the same + video payload plus ``audio`` and ``audio_sample_rate``. Image output with + audio is rejected because Cosmos3 sound generation is video-only. + """ from .guardrails import check_video_safety, is_guardrails_enabled video_processor = VideoProcessor(vae_scale_factor=16) @@ -426,6 +609,37 @@ def post_process_func( if output_type == "latent": return output + def _postprocess_action(action: Any, metadata: dict[str, Any]) -> Any: + internal_metadata = metadata.get("internal") + inputs = ( + internal_metadata.get("robolab_action_postprocess") if isinstance(internal_metadata, dict) else None + ) + if isinstance(inputs, RoboLabActionPostprocessInputs): + return postprocess_robolab_action(action, inputs) + return action + + pending_action = None + pending_action_metadata: dict[str, Any] = {} + envelope_public_metadata: dict[str, Any] = {} + if isinstance(output, dict) and isinstance(output.get("payload"), dict): + envelope_payload = dict(output.get("payload") or {}) + metadata = output.get("metadata") or {} + envelope_metadata = metadata if isinstance(metadata, dict) else {} + envelope_public_metadata = {key: value for key, value in envelope_metadata.items() if key != "internal"} + action = envelope_payload.pop("actions", None) + if action is not None: + pending_action = _postprocess_action(action, envelope_metadata) + pending_action_metadata = envelope_public_metadata + if not envelope_payload: + return { + "payload": { + "video": [], + "actions": pending_action, + }, + "metadata": pending_action_metadata, + } + output = envelope_payload + audio = None audio_sample_rate = None if isinstance(output, dict): @@ -464,11 +678,31 @@ def post_process_func( # check_video_safety expects a 5D tensor; re-add T axis. checked = check_video_safety(image.unsqueeze(2)) image = checked.squeeze(2) - return video_processor.postprocess(image, output_type="pil") - if is_guardrails_enabled(od_config, sampling_params): + processed_image = video_processor.postprocess(image, output_type="pil") + if envelope_public_metadata: + return { + "payload": {"image": processed_image}, + "metadata": envelope_public_metadata, + } + return processed_image + guardrails_enabled = is_guardrails_enabled(od_config, sampling_params) + if guardrails_enabled: video = check_video_safety(video) processed_video = video_processor.postprocess_video(video, output_type=output_type) if audio is None: + if pending_action is not None: + return { + "payload": { + "video": processed_video, + "actions": pending_action, + }, + "metadata": pending_action_metadata, + } + if envelope_public_metadata: + return { + "payload": {"video": processed_video}, + "metadata": envelope_public_metadata, + } return processed_video if isinstance(audio, torch.Tensor): audio = audio.detach().cpu() @@ -479,26 +713,38 @@ def post_process_func( } if audio_sample_rate is not None: result["audio_sample_rate"] = int(audio_sample_rate) + if pending_action is not None: + return { + "payload": { + "video": result["video"], + "audio": result["audio"], + "actions": pending_action, + }, + "metadata": { + **pending_action_metadata, + "video": {"fps": result["fps"]}, + "audio": {"sample_rate": result.get("audio_sample_rate")}, + }, + } + if envelope_public_metadata: + return { + "payload": { + "video": result["video"], + "audio": result["audio"], + }, + "metadata": { + **envelope_public_metadata, + "video": {"fps": result["fps"], **envelope_public_metadata.get("video", {})} + if isinstance(envelope_public_metadata.get("video"), dict) + else {"fps": result["fps"]}, + "audio": {"sample_rate": result.get("audio_sample_rate")}, + }, + } return result return post_process_func -def get_cosmos3_action_post_process_func(od_config: OmniDiffusionConfig): - del od_config - - def action_post_process_func(action: Any, custom_output: dict[str, Any] | None = None, sampling_params=None): - del sampling_params - inputs = custom_output.get("robolab_action_postprocess") if isinstance(custom_output, dict) else None - if isinstance(inputs, RoboLabActionPostprocessInputs): - processed_action = postprocess_robolab_action(action, inputs) - custom_output.pop("robolab_action_postprocess", None) - return processed_action - return action - - return action_post_process_func - - def get_cosmos3_ir_op_priority_func(od_config: OmniDiffusionConfig): del od_config @@ -518,16 +764,22 @@ def ir_op_priority_func(ir_op_priority, vllm_config=None): # Pipeline # --------------------------------------------------------------------------- class Cosmos3OmniDiffusersPipeline( - nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin + nn.Module, + CFGParallelMixin, + SupportImageInput, + SupportsComponentDiscovery, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, ): - """Cosmos3 text/image/video-to-video / text-to-image pipeline. + """Cosmos3 text/image/video/sound/action 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) + - Generation pathway: cross-attention on visual latents and optional + transfer-control, action, and sound latents (runs each step) - Supports T2V, I2V, V2V, and T2I from the same class. Mode is selected at - runtime: + Supports T2V, I2V, V2V, T2I, transfer, sound-enabled video, and action + generation from the same class. Mode is selected at runtime: * **T2I** when ``prompt["modalities"]`` contains ``"image"``. Latent T-dim is forced to 1, T2I-specific scheduler defaults are applied (50 steps, @@ -544,11 +796,27 @@ class Cosmos3OmniDiffusersPipeline( ``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. + * **Transfer** when ``edge``, ``blur``, ``depth``, ``seg``, or ``wsm`` hints + are supplied. Transfer is video-output only and cannot be combined with + sound or action generation. + * **Sound-enabled video** when ``generate_sound`` or ``sound_gen`` is true. + Sound is generated from sound latents, not from ``multi_modal_data['audio']``; + T2I, transfer, and action+sound are rejected. + * **Action generation** when ``action_mode`` is provided. ``policy`` and + ``forward_dynamics`` require an image or video input; ``inverse_dynamics`` + requires video input. Action predictions are returned in the diffusion + output payload/metadata envelope. + RoboLab/OpenPI observations in ``extra_args['robot_obs']`` or + ``extra_args['observation']`` return an action-only envelope. * **T2V** otherwise (default video generation). """ support_image_input: ClassVar[bool] = True color_format: ClassVar[str] = "RGB" + _dit_modules: ClassVar[list[str]] = ["transformer.language_model", "transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [] @classmethod def reference_video_decode_spec( @@ -558,6 +826,12 @@ def reference_video_decode_spec( extra_args: dict[str, Any] | None = None, ) -> ReferenceVideoDecodeSpec: extra_args = extra_args if isinstance(extra_args, dict) else {} + if has_transfer_hints(extra_args): + max_frames = transfer_max_frames_from_extra_args(extra_args) + if num_frames is not None: + max_frames = min(max_frames, int(num_frames)) + return ReferenceVideoDecodeSpec(max_frames=max_frames, keep="first") + action_mode = normalize_action_mode(extra_args.get("action_mode")) if action_mode is not None: if num_frames is not None: @@ -586,12 +860,6 @@ def __init__( prefix: str = "", ) -> None: super().__init__() - if od_config.enable_cpu_offload: - raise ValueError( - "Cosmos3 has no separate text encoder, so CPU offloading " - "(transformer↔encoder swapping) is not supported. " - "Use --enable-layerwise-offload instead." - ) self.od_config = od_config self.device = get_local_device() self.dtype = od_config.dtype @@ -632,25 +900,53 @@ def __init__( sound_latent_fps = self._sound_tokenizer.latent_fps # --- Transformer (weights loaded later via weights_sources) --- - self.transformer = Cosmos3VFMTransformer( + transformer_cls = resolve_cosmos3_transformer_cls(od_config.tf_model_config) + self.transformer = transformer_cls( od_config=od_config, temporal_compression_factor=self.vae_scale_factor_temporal, sound_gen=sound_gen, sound_dim=sound_dim, sound_latent_fps=sound_latent_fps, ) + self.is_edge_model = transformer_cls is Cosmos3EdgeVFMTransformer # --- Scheduler --- - # Load from checkpoint to preserve solver_order, timestep_spacing, - # beta_schedule, sigma bounds, flow_shift, etc. Only override - # flow_shift when explicitly requested by the user. - self.scheduler = UniPCMultistepScheduler.from_pretrained( + # Distilled model differs from regular one only by scheduler, + # distilled one uses FlowMatchEulerDiscreteScheduler, while + # regular should use FlowUniPCMultistepScheduler + + scheduler_config = FlowUniPCMultistepScheduler.load_config( model_path, subfolder="scheduler", local_files_only=local_files_only, ) - if od_config.flow_shift is not None: - self.scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config, flow_shift=od_config.flow_shift) + + scheduler_class_name = scheduler_config.get("_class_name") + + self.is_distilled_model = False + if scheduler_class_name == COSMOS3_DISTILLED_CHECKPOINT_SCHEDULER_CLASS: + fixed_step_config = scheduler_config.get("fixed_step_sampler_config") + if not isinstance(fixed_step_config, dict) or fixed_step_config.get("sample_type") != "sde": + raise ValueError("Cosmos3 distilled scheduler requires fixed_step_sampler_config.sample_type=sde.") + t_list = fixed_step_config.get("t_list") + if not isinstance(t_list, list) or not t_list: + raise ValueError("Cosmos3 distilled scheduler requires a non-empty fixed_step_sampler_config.t_list.") + self.scheduler = FlowMatchEulerDiscreteScheduler.from_config( + scheduler_config, + stochastic_sampling=True, + ) + self._scheduler_init_t_list = list(t_list) + self.is_distilled_model = True + else: + # Preserve compatible solver settings from the checkpoint, but keep + # the base shift neutral. The concrete request shift is applied when + # FlowUniPC builds its timesteps. + self.scheduler = FlowUniPCMultistepScheduler.from_config( + scheduler_config, + shift=1.0, + use_dynamic_shifting=False, + prediction_type="flow_prediction", + ) self._cpu_scheduler_state() # --- Video processor for post-decode --- @@ -668,15 +964,20 @@ def __init__( ), ] - # Snapshot the loaded scheduler config so we can rebuild the - # scheduler at request time when a per-request flow_shift override - # is supplied (T2I uses shift=3.0; T2V/I2V use the engine default). - self._base_scheduler_config = self.scheduler.config - self._engine_init_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0) + # An engine-level override becomes the video default; otherwise use + # the matching regular/Edge default. Per-request values still take + # precedence. + default_video_flow_shift = ( + COSMOS3_EDGE_VIDEO_DEFAULT_FLOW_SHIFT if self.is_edge_model else COSMOS3_VIDEO_DEFAULT_FLOW_SHIFT + ) + self._engine_init_flow_shift = float( + od_config.flow_shift if od_config.flow_shift is not None else default_video_flow_shift + ) self._current_flow_shift = self._engine_init_flow_shift self._guidance_scale = None self._num_timesteps = None + self._cosmos3_branch_caches: dict[str, tuple[Any, Any]] | None = None self._robolab_transform = None # Set True by ``enable_cache_for_cosmos3`` when cache-dit is enabled on @@ -689,6 +990,32 @@ def __init__( enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + def enable_omni_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool = True, + use_hsdp: bool = False, + ) -> None: + """Enable Cosmos3 component-level model offload. + + Cosmos3 has a nested reasoner/generator transformer instead of separate + text-encoder and DiT pipeline components, so the transformer owns the + mutual-exclusion swaps. The VAE stays resident on GPU like the generic + model-level offloader. + """ + self.vae.to(device, non_blocking=True) + if isinstance(self._sound_tokenizer, nn.Module): + self._sound_tokenizer.to(device) + self.transformer.enable_model_cpu_offload( + device=device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + ) + + def disable_omni_model_cpu_offload(self) -> None: + self.transformer.disable_model_cpu_offload() + # -- Weight loading -------------------------------------------------------- @staticmethod @@ -766,6 +1093,7 @@ def _remap_ckpt_key(key: str) -> str | None: "self_attn.to_out.": f"{und_lp}.self_attn.to_out.", "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.", "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.", + "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.", # GEN attention "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.", "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.", @@ -795,6 +1123,9 @@ def _remap_ckpt_key(key: str) -> str | None: return None + # Checkpoint adapters use this hook before model-specific weight loading. + remap_checkpoint_key = _remap_ckpt_key + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Stream-remap checkpoint weights and load via AutoWeightsLoader. @@ -809,6 +1140,10 @@ def _remapped_weights() -> Iterable[tuple[str, torch.Tensor]]: total = kept = 0 for name, tensor in weights: total += 1 + if name in allowed or name in tp_aware: + kept += 1 + yield name, tensor + continue remapped = self._remap_ckpt_key(name) if remapped is not None and (remapped in allowed or remapped in tp_aware): kept += 1 @@ -824,6 +1159,7 @@ def _remapped_weights() -> Iterable[tuple[str, torch.Tensor]]: loaded = loader.load_weights(_remapped_weights()) self.transformer.post_load_weights() self.transformer.eval() + self.transformer.validate_loaded_weights(loaded) if getattr(self.transformer, "sound_gen", False): sound_markers = ("audio_proj_in.", "audio_proj_out.", "audio_modality_embed") missing = [marker.rstrip(".") for marker in sound_markers if not any(marker in name for name in loaded)] @@ -850,7 +1186,63 @@ def predict_noise(self, **kwargs) -> torch.Tensor | tuple[torch.Tensor, ...]: The transformer returns the raw prediction: video-only as a tensor, or a tuple in video, action, sound order for multimodal generation. """ - return self.transformer(**kwargs) + cache_key = kwargs.pop("_cosmos3_cache_key", None) + if cache_key is None: + return self.transformer(**kwargs) + + branch_caches = self._cosmos3_branch_caches + if branch_caches is None: + return self.transformer(**kwargs) + + cache_key = str(cache_key) + self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get(cache_key, (None, None)) + prediction = self.transformer(**kwargs) + branch_caches[cache_key] = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) + return prediction + + def combine_multi_branch_cfg_noise( + self, + predictions: list[torch.Tensor | tuple[torch.Tensor, ...]], + true_cfg_scale: float | dict[str, float], + cfg_normalize: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + if not isinstance(true_cfg_scale, dict) or true_cfg_scale.get("mode") != "cosmos3_transfer": + return super().combine_multi_branch_cfg_noise(predictions, true_cfg_scale, cfg_normalize) + mode = str(true_cfg_scale.get("branch_mode", "control_and_text")) + guidance_scale = float(true_cfg_scale.get("guidance_scale", 1.0)) + control_guidance = float(true_cfg_scale.get("control_guidance", 1.0)) + + if mode == "control_only": + if len(predictions) != 2: + raise ValueError(f"Cosmos3 transfer control-only CFG expects 2 branches, got {len(predictions)}.") + cond_full, cond_no_control = predictions + if isinstance(cond_full, tuple) or isinstance(cond_no_control, tuple): + raise ValueError("Cosmos3 transfer control-only CFG expects video-only tensor predictions.") + cfg_reference = cond_full + combined = cond_no_control + control_guidance * (cond_full - cond_no_control) + elif mode == "text_only": + if len(predictions) != 2: + raise ValueError(f"Cosmos3 transfer text-only CFG expects 2 branches, got {len(predictions)}.") + cond_full, uncond_full = predictions + if isinstance(cond_full, tuple) or isinstance(uncond_full, tuple): + raise ValueError("Cosmos3 transfer text-only CFG expects video-only tensor predictions.") + cfg_reference = cond_full + combined = uncond_full + guidance_scale * (cond_full - uncond_full) + elif mode == "control_and_text": + if len(predictions) != 3: + raise ValueError(f"Cosmos3 transfer control+text CFG expects 3 branches, got {len(predictions)}.") + cond_full, cond_no_control, uncond_full = predictions + if isinstance(cond_full, tuple) or isinstance(cond_no_control, tuple) or isinstance(uncond_full, tuple): + raise ValueError("Cosmos3 transfer control+text CFG expects video-only tensor predictions.") + cfg_reference = cond_full + control_cond = cond_no_control + control_guidance * (cond_full - cond_no_control) + combined = uncond_full + guidance_scale * (control_cond - uncond_full) + else: + raise ValueError(f"Unknown Cosmos3 transfer CFG branch_mode={mode!r}.") + + if cfg_normalize: + combined = self.cfg_normalize_function(cfg_reference, combined) + return combined @staticmethod def _cfg_parallel_active() -> bool: @@ -1073,6 +1465,8 @@ def _forward_robolab_policy( inputs: RoboLabPolicyInputs, pipeline_start: float, ) -> DiffusionOutput: + if self.is_distilled_model: + raise self._distilled_unsupported_error("RoboLab/action policy requests are unsupported.") if not getattr(self.transformer, "action_gen", False): raise ValueError( "Cosmos3 RoboLab policy serving was requested, but the transformer " @@ -1183,21 +1577,32 @@ def _forward_robolab_policy( guidance_interval=None, raw_action_dim=raw_action_dim, scheduler=scheduler, + generator=generator, ) if _is_rank_zero(): logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) action = action_latents[:, :, :raw_action_dim].detach().cpu() - custom_action_output: dict[str, Any] = { - "action": action, - "raw_action_dim": raw_action_dim, - "action_mode": action_mode, - "domain_id": domain_id, - "action_only_output": True, - "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs), + action_output: dict[str, Any] = { + "payload": { + "actions": action, + }, + "metadata": { + "actions": { + "raw_action_dim": raw_action_dim, + "action_mode": action_mode, + "domain_id": domain_id, + }, + "common": { + "action_only_output": True, + }, + "internal": { + "robolab_action_postprocess": make_robolab_action_postprocess_inputs(inputs), + }, + }, } - return DiffusionOutput(output={}, custom_output=custom_action_output) + return DiffusionOutput(output=action_output) @staticmethod def _truthy(value) -> bool: @@ -1239,8 +1644,15 @@ def _get_sound_tokenizer(self): return self._sound_tokenizer @staticmethod - def _is_t2i_request(req: OmniDiffusionRequest) -> bool: - """Detect text-to-image mode from request-level prompt modalities.""" + def _is_t2i_request(req: DiffusionRequestBatch) -> bool: + """Return whether request-level modalities select image output. + + Only ``"image"`` switches Cosmos3 into T2I. ``"video"`` and omitted + modalities select video output. ``"text"`` and ``"audio"`` are accepted + compatibility values for callers that share prompt schemas, but they do + not select text or audio output in this pipeline. ``"image"`` and + ``"video"`` cannot be requested together. + """ if not req.prompts: return False first_prompt = req.prompts[0] @@ -1257,22 +1669,34 @@ def _is_t2i_request(req: OmniDiffusionRequest) -> bool: raise ValueError(f"Incorrect modality value in {modalities}, expected one of {accepted_modalities}.") return "image" in modalities + def _set_timesteps(self, num_inference_steps: int, device: str | torch.device, shift: float) -> None: + if self.is_distilled_model: + self.scheduler.set_timesteps(sigmas=self._scheduler_init_t_list, device=device) + else: + self.scheduler.set_timesteps( + num_inference_steps, + device=device, + shift=shift, + ) + def _set_flow_shift(self, target_shift: float) -> None: - """Set the UniPC ``flow_shift`` to a concrete target value. - - The scheduler is rebuilt from the saved base config if - the target differs from the current shift. Tracking - ``self._current_flow_shift`` explicitly is required because the - previous mode may have rebuilt the scheduler - we cannot rely on - ``self.scheduler.config.flow_shift`` reflecting the last requested - target if a rebuild was skipped via the equality check. - """ - target = float(target_shift) - if target == float(self._current_flow_shift): - return - self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, flow_shift=target) - self._cpu_scheduler_state() - self._current_flow_shift = target + """Select the shift applied by FlowUniPC when timesteps are built.""" + self._current_flow_shift = float(target_shift) + + @staticmethod + def _resolve_seed(sp: OmniDiffusionSamplingParams, generator: torch.Generator | None, default: int = 42) -> int: + if sp.seed is not None: + return int(sp.seed) + if isinstance(generator, torch.Generator): + return int(generator.initial_seed()) + return int(default) + + def _resolve_guidance_scale(self, sp: OmniDiffusionSamplingParams, default: float) -> float: + if self.is_distilled_model: + return 1.0 + if sp.guidance_scale_provided: + return float(sp.guidance_scale) + return float(default) def _cpu_scheduler_state(self) -> None: # We need to move scheduler tensors to CPU, as unipc from diffusers assumes they are on CPU. @@ -1293,6 +1717,53 @@ def do_classifier_free_guidance(self): def num_timesteps(self): return self._num_timesteps + @staticmethod + def _distilled_unsupported_error(detail: str) -> ValueError: + return ValueError( + f"Cosmos3 distilled checkpoints support only text-to-image and image-to-video generation; {detail}" + ) + + def _validate_distilled_generation_mode( + self, + *, + is_t2i: bool, + image_tensor: torch.Tensor | None, + action_enabled: bool, + transfer_config: Cosmos3TransferConfig | None, + is_v2v: bool, + sound_enabled: bool, + ) -> None: + if not self.is_distilled_model: + return + if action_enabled: + raise self._distilled_unsupported_error("action requests are unsupported.") + if transfer_config is not None: + raise self._distilled_unsupported_error("transfer requests are unsupported.") + if is_v2v: + raise self._distilled_unsupported_error("video-to-video requests are unsupported.") + if sound_enabled: + raise self._distilled_unsupported_error("sound generation is unsupported.") + if not is_t2i and image_tensor is None: + raise self._distilled_unsupported_error("text-to-video requests are unsupported.") + if is_t2i and image_tensor is not None: + raise self._distilled_unsupported_error("image-conditioned image generation is unsupported.") + + def _validate_edge_generation_mode( + self, + *, + transfer_config: Cosmos3TransferConfig | None, + is_v2v: bool, + sound_enabled: bool, + ) -> None: + if not self.is_edge_model: + return + if sound_enabled: + raise ValueError("Cosmos3 Edge checkpoints do not support sound generation.") + if transfer_config is not None: + raise ValueError("Cosmos3 Edge checkpoints do not support transfer inference.") + if is_v2v: + raise ValueError("Cosmos3 Edge checkpoints do not support video-to-video generation.") + # -- Prompt formatting ----------------------------------------------------- @staticmethod @@ -1309,6 +1780,10 @@ def _apply_metadata_templates( """ Append duration and resolution metadata to a prompt. """ + prompt = prompt.strip() + if duration_template is None and resolution_template is None: + return prompt + parts: list[str] = [] head = prompt.rstrip(".").strip() if head: @@ -1454,12 +1929,21 @@ def _prepare_sound_latents( self, target_audio_samples: int, generator: torch.Generator, + *, + sp_video_shape: tuple[int, int, int] | None = None, + sp_num_vision_items: int = 1, ) -> tuple[torch.Tensor, int]: sound_tokenizer = self._get_sound_tokenizer() hop_size = int( getattr(sound_tokenizer, "hop_size", None) or getattr(sound_tokenizer, "temporal_compression_factor") ) latent_frames = max(1, math.ceil(max(1, int(target_audio_samples)) / hop_size)) + if sp_video_shape is not None: + latent_frames = self.transformer.sound_latent_frames_for_sequence_parallel( + video_shape=sp_video_shape, + sound_frames=latent_frames, + num_vision_items=sp_num_vision_items, + ) sound_dim = int(getattr(sound_tokenizer, "latent_ch", 64)) transformer_sound_dim = int(getattr(self.transformer, "sound_dim", sound_dim)) if sound_dim != transformer_sound_dim: @@ -1493,18 +1977,36 @@ def _resolve_sound_target_samples( # -- VAE decode ---------------------------------------------------------- + def _get_latents_mean_std(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + cached = getattr(self, "_latents_mean_std", None) + if cached is not None: + latents_mean, latents_std = cached + if latents_mean.device == device and latents_mean.dtype == dtype: + return latents_mean, latents_std + + latents_mean = torch.as_tensor(self.vae.config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + latents_std = torch.as_tensor(self.vae.config.latents_std, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + self._latents_mean_std = (latents_mean, latents_std) + return latents_mean, latents_std + + def _to_vae_device(self, tensor: torch.Tensor, *, pin_cpu: bool = False) -> torch.Tensor: + if tensor.device == self.device and tensor.dtype == self.vae.dtype: + return tensor + + non_blocking = False + if tensor.device.type == "cpu" and self.device.type == "cuda": + if pin_cpu and not tensor.is_pinned(): + tensor = tensor.pin_memory() + non_blocking = tensor.is_pinned() + + return tensor.to(device=self.device, dtype=self.vae.dtype, non_blocking=non_blocking) + def _decode_latents(self, latents: torch.Tensor) -> torch.Tensor: latents = latents.to(self.vae.dtype) if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - if not hasattr(self, "_latents_mean"): - self._latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype) - ) - self._latents_std = ( - torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(self.device, self.vae.dtype) - ) - latents = (latents * self._latents_std) + self._latents_mean + latents_mean, latents_std = self._get_latents_mean_std(latents.device, latents.dtype) + latents = (latents * latents_std) + latents_mean else: scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) latents = latents / scaling_factor @@ -1557,15 +2059,26 @@ def _format_and_tokenize_prompts( res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE else: res_tmpl = None - prompt = self._apply_metadata_templates( + json_prompt = _format_json_object_prompt( prompt, - num_frames, - frame_rate, - height, - width, - duration_template=dur_tmpl, - resolution_template=res_tmpl, + num_frames=num_frames, + frame_rate=frame_rate, + height=height, + width=width, + aspect_ratio=self._get_sp_param(sp, "aspect_ratio", None), ) + if json_prompt is not None: + prompt = json_prompt + else: + prompt = self._apply_metadata_templates( + prompt, + num_frames, + frame_rate, + height, + width, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + ) if _is_rank_zero(): logger.info("Final prompt: '%s'", prompt) @@ -1602,40 +2115,28 @@ def _format_and_tokenize_prompts( # -- I2V latent preparation --------------------------------------------- - def _encode_conditioning_video( - self, - image_tensor: torch.Tensor, - num_frames: int, - height: int, - width: int, - ) -> torch.Tensor: - """VAE-encode a conditioning image as a full-length video. - - The WAN VAE has temporal compression (factor 4), so encoding a single - frame produces degenerate temporal features. We fill the entire - pixel-space video with the conditioning image (repeating it across all - frames) so the temporal encoder sees plausible content everywhere. - The caller keeps only the conditioned latent frame(s) and replaces - the rest with noise. - """ - # image_tensor: [1, 3, H, W] -> [1, 3, num_frames, H, W] - video = image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() - video = video.to(device=self.device, dtype=self.vae.dtype) - - latent = self.vae.encode(video).latent_dist.mode() - - # Normalize (inverse of _decode_latents denormalization) + def _normalize_vae_latent(self, latent: torch.Tensor) -> torch.Tensor: if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) - ) - latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) + latents_mean, latents_std = self._get_latents_mean_std(latent.device, latent.dtype) latent = (latent - latents_mean) / latents_std else: scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) latent = latent * scaling_factor - return latent.to(self.dtype) + return latent + + def _encode_conditioning_image_latent(self, image_tensor: torch.Tensor) -> torch.Tensor: + """VAE-encode the first I2V conditioning frame. + + I2V only consumes the first conditioning latent frame. Encoding the + input image as a one-frame video keeps that latent while avoiding VAE + work for repeated frames that are later replaced by noise. + """ + image_tensor = self._to_vae_device(image_tensor, pin_cpu=True) + video = image_tensor.unsqueeze(2) + latent = self.vae.encode(video).latent_dist.mode() + latent = self._normalize_vae_latent(latent) + return latent[:, :, 0:1, :, :].to(self.dtype) def _latent_hw_from_image_size(self, image_size: Any | None) -> tuple[int, int] | None: if image_size is None: @@ -1667,19 +2168,9 @@ def _encode_video_tensor(self, video_tensor: torch.Tensor, image_size: Any | Non if 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)}.") - video = video_tensor.to(device=self.device, dtype=self.vae.dtype) + video = self._to_vae_device(video_tensor) latent = self.vae.encode(video).latent_dist.mode() - - if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): - latents_mean = ( - torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) - ) - latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype) - latent = (latent - latents_mean) / latents_std - else: - scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) - latent = latent * scaling_factor - + latent = self._normalize_vae_latent(latent) latent = self._crop_latent_to_image_size(latent, image_size) return latent.to(self.dtype) @@ -1710,13 +2201,12 @@ def _prepare_latents_i2v( dtype=self.dtype, ) - cond_latent = self._encode_conditioning_video(image_tensor, num_frames, height, width) - image_latent = cond_latent[:, :, 0:1, :, :] + image_latent = self._encode_conditioning_image_latent(image_tensor) + latents = noise + latents[:, :, 0:1, :, :] = image_latent - condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype) - condition_mask[:, :, 0, :, :] = 1.0 - latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise - velocity_mask = 1.0 - condition_mask + velocity_mask = torch.ones(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype) + velocity_mask[:, :, 0, :, :] = 0.0 return latents, velocity_mask, image_latent def _prepare_latents_v2v( @@ -1793,8 +2283,21 @@ def _prepare_latents_action_video( generator: torch.Generator, image_size: Any | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Prepare video latents for action modes with mode-specific conditioning.""" + """Prepare video latents for action modes with mode-specific conditioning. + + Policy and forward-dynamics modes condition only latent frame zero. + The Wan VAE is temporally causal, so encoding only the first pixel + frame preserves that latent while avoiding unused future-frame work. + Inverse dynamics conditions every latent and keeps the full encode. + """ 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 action 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 latent_hw = self._latent_hw_from_image_size(image_size) @@ -1810,12 +2313,24 @@ def _prepare_latents_action_video( device=self.device, dtype=self.dtype, ) - cond_latent = self._encode_video_tensor(video_tensor, image_size=image_size) - if cond_latent.shape[2:] != noise.shape[2:]: + condition_indexes = vision_condition_indexes(mode, num_frames, self.vae_scale_factor_temporal) + condition_video = video_tensor[:, :, :1] if condition_indexes == [0] else video_tensor + cond_prefix_latent = self._encode_video_tensor(condition_video, image_size=image_size) + expected_prefix = (1, C, max(condition_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 action video latent shape mismatch: " - f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}." + f"encoded={tuple(cond_prefix_latent.shape)}, expected at least {expected_prefix}." ) + + condition_latents = torch.zeros_like(noise) + for index in condition_indexes: + condition_latents[:, :, index : index + 1] = cond_prefix_latent[:, :, index : index + 1] condition_mask = build_vision_condition_mask( mode, num_frames, @@ -1823,9 +2338,9 @@ def _prepare_latents_action_video( device=self.device, dtype=self.dtype, ) - latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise + latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise velocity_mask = 1.0 - condition_mask - return latents, velocity_mask, cond_latent + return latents, velocity_mask, condition_latents def _prepare_action_latents( self, @@ -1924,6 +2439,7 @@ def diffuse( guidance_interval: tuple[float, float] | None = None, raw_action_dim: int | None = None, scheduler: Any | None = None, + generator: torch.Generator | None = None, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Denoising loop with 3-mode CFG support (parallel, sequential, none). @@ -2026,17 +2542,32 @@ def _step( ) -> torch.Tensor | tuple[torch.Tensor, ...]: video_pred, action_pred, sound_pred = _split_noise_pred(noise_pred) if velocity_mask is not None: - video_pred = video_pred * velocity_mask + if image_latent is not None and condition_latents is None: + video_pred[:, :, 0:1, :, :] = 0 + else: + video_pred = video_pred * velocity_mask if action_pred is not None and action_velocity_mask is not None: action_pred = action_pred * action_velocity_mask if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]: action_pred[..., raw_action_dim:] = 0 if action_latents is None and sound_latents is None: - latents = step_scheduler.step(video_pred, t, latents, return_dict=False)[0] + latents = step_scheduler.step( + video_pred, + t, + latents, + generator=generator, + return_dict=False, + )[0] else: packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred) packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents) - packed_next = step_scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0] + packed_next = step_scheduler.step( + packed_noise, + t, + packed_latents, + generator=generator, + return_dict=False, + )[0] unpacked = _unpack_joint(packed_next, shapes, numels) latents = unpacked[0] idx = 1 @@ -2176,20 +2707,476 @@ def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None: outputs.append(sound_latents) return outputs[0] if len(outputs) == 1 else tuple(outputs) + @staticmethod + def _get_transfer_num_chunks( + total_frames: int, + frames_per_chunk: int, + conditional_frames: int, + ) -> tuple[int, int]: + if frames_per_chunk <= 0: + raise ValueError("Cosmos3 transfer frames_per_chunk must be positive.") + if total_frames <= frames_per_chunk: + return 1, frames_per_chunk + stride = frames_per_chunk - conditional_frames + if stride <= 0: + raise ValueError("Cosmos3 transfer num_conditional_frames must be smaller than num_video_frames_per_chunk.") + remaining = total_frames - frames_per_chunk + extra_chunks = remaining // stride + (1 if remaining % stride else 0) + return 1 + extra_chunks, stride + + def _prepare_transfer_latents( + self, + target_video: torch.Tensor, + current_conditional_frames: int, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + condition_latents = self._encode_video_tensor(target_video) + noise = randn_tensor( + condition_latents.shape, + generator=generator, + device=self.device, + dtype=self.dtype, + ) + condition_mask = torch.zeros( + 1, + 1, + condition_latents.shape[2], + 1, + 1, + device=self.device, + dtype=self.dtype, + ) + if current_conditional_frames > 0: + latent_frames = (current_conditional_frames - 1) // self.vae_scale_factor_temporal + 1 + condition_mask[:, :, :latent_frames] = 1.0 + latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, condition_mask * condition_latents + + def _transfer_bucket_size( + self, + sp: OmniDiffusionSamplingParams, + source_hw: tuple[int, int] | None, + ) -> tuple[int, int]: + resolution = self._get_sp_param(sp, "resolution", self._get_sp_param(sp, "image_size", 720)) + source_h, source_w = source_hw or (COSMOS3_T2V_DEFAULT_HEIGHT, COSMOS3_T2V_DEFAULT_WIDTH) + target_w, target_h = find_closest_target_size(int(source_h), int(source_w), resolution) + return int(target_h), int(target_w) + + @staticmethod + def _first_transfer_control_hw(transfer_config: Cosmos3TransferConfig) -> tuple[int, int] | None: + for hint in transfer_config.ordered_hints: + if hint.control is not None: + detected = media_hw(hint.control) + if detected is not None: + return detected + if hint.control_path is not None: + detected = media_hw(hint.control_path) + if detected is not None: + return detected + return None + + def diffuse_transfer( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + cond_ids: torch.Tensor, + cond_mask: torch.Tensor, + uncond_ids: torch.Tensor, + uncond_mask: torch.Tensor, + guidance_scale: float, + control_guidance: float, + control_guidance_interval: tuple[float, float] | None, + control_latents: list[torch.Tensor], + shared_kwargs: dict[str, Any], + *, + velocity_mask: torch.Tensor, + condition_latents: torch.Tensor, + guidance_interval: tuple[float, float] | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + def _active_at(t: torch.Tensor, interval: tuple[float, float] | None) -> bool: + if interval is None: + return True + t_scalar = float(t.item()) if torch.is_tensor(t) else float(t) + lo, hi = interval + return lo <= t_scalar <= hi + + self.transformer.reset_cache() + self._cosmos3_branch_caches = {} + try: + for t in self.progress_bar(timesteps): + timestep = t.unsqueeze(0) + step_guidance = guidance_scale if _active_at(t, guidance_interval) else 1.0 + step_control = control_guidance if _active_at(t, control_guidance_interval) else 1.0 + needs_text_cfg = step_guidance > 1.0 + needs_control_cfg = step_control != 1.0 + + cond_full_kwargs = dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + control_latents=control_latents, + _cosmos3_cache_key="transfer_cond_full", + **shared_kwargs, + ) + if needs_control_cfg and needs_text_cfg: + branches_kwargs = [ + cond_full_kwargs, + dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + control_latents=None, + _cosmos3_cache_key="transfer_cond_no_control", + **shared_kwargs, + ), + dict( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + control_latents=control_latents, + _cosmos3_cache_key="transfer_uncond_full", + **shared_kwargs, + ), + ] + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "mode": "cosmos3_transfer", + "branch_mode": "control_and_text", + "guidance_scale": step_guidance, + "control_guidance": step_control, + }, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + elif needs_control_cfg: + branches_kwargs = [ + cond_full_kwargs, + dict( + hidden_states=latents, + timestep=timestep, + text_ids=cond_ids, + text_mask=cond_mask, + control_latents=None, + _cosmos3_cache_key="transfer_cond_no_control", + **shared_kwargs, + ), + ] + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "mode": "cosmos3_transfer", + "branch_mode": "control_only", + "control_guidance": step_control, + }, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + elif needs_text_cfg: + branches_kwargs = [ + cond_full_kwargs, + dict( + hidden_states=latents, + timestep=timestep, + text_ids=uncond_ids, + text_mask=uncond_mask, + control_latents=control_latents, + _cosmos3_cache_key="transfer_uncond_full", + **shared_kwargs, + ), + ] + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "mode": "cosmos3_transfer", + "branch_mode": "text_only", + "guidance_scale": step_guidance, + }, + branches_kwargs=branches_kwargs, + cfg_normalize=False, + ) + else: + noise_pred = self.predict_noise(**cond_full_kwargs) + if isinstance(noise_pred, tuple): + raise ValueError("Cosmos3 transfer diffusion expects video-only tensor predictions.") + noise_pred = noise_pred * velocity_mask + latents = self.scheduler.step( + noise_pred, + t, + latents, + generator=generator, + return_dict=False, + )[0] + latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents + finally: + self._cosmos3_branch_caches = None + self.transformer.reset_cache() + return latents + + def _forward_transfer( + self, + *, + prompt: str, + negative_prompt: str, + sp: OmniDiffusionSamplingParams, + transfer_config: Cosmos3TransferConfig, + transfer_video_tensor: torch.Tensor | None, + transfer_input_fps: float | None, + ) -> DiffusionOutput: + if self.is_distilled_model: + raise self._distilled_unsupported_error("transfer requests are unsupported.") + input_frames = None + if transfer_video_tensor is not None: + input_frames = normalized_video_to_uint8_cthw(transfer_video_tensor) + source_hw = (int(input_frames.shape[-2]), int(input_frames.shape[-1])) + else: + source_hw = self._first_transfer_control_hw(transfer_config) + height, width = self._transfer_bucket_size(sp, source_hw) + + if input_frames is not None: + if tuple(input_frames.shape[-2:]) != (height, width): + input_frames = resize_center_crop_uint8_cthw(input_frames, height, width) + input_frames = input_frames[:, : transfer_config.max_frames] + + per_hint_frames: dict[str, torch.Tensor] = {} + for hint in transfer_config.ordered_hints: + frames = load_or_compute_control_frames( + hint, + height=height, + width=width, + max_frames=transfer_config.max_frames, + input_frames=input_frames, + ) + if frames.shape[1] < 1: + raise ValueError(f"Cosmos3 transfer hint '{hint.key}' produced no frames.") + per_hint_frames[hint.key] = frames + if not per_hint_frames: + raise ValueError("Cosmos3 transfer requires at least one control hint.") + + total_frames = next(iter(per_hint_frames.values())).shape[1] + if transfer_config.num_frames is not None: + total_frames = min(total_frames, int(transfer_config.num_frames)) + total_frames = max(1, total_frames) + per_hint_frames = {key: pad_temporal_frames(frames, total_frames) for key, frames in per_hint_frames.items()} + if input_frames is not None: + input_frames = pad_temporal_frames(input_frames, total_frames) + + temporal_compression = self.vae_scale_factor_temporal + chunk_frames = 1 if total_frames == 1 else transfer_config.num_video_frames_per_chunk + chunk_frames = math.ceil((chunk_frames - 1) / temporal_compression) * temporal_compression + 1 + num_chunks, stride = self._get_transfer_num_chunks( + total_frames, + chunk_frames, + transfer_config.num_conditional_frames, + ) + padded_frames = max(total_frames, chunk_frames) + per_hint_frames = {key: pad_temporal_frames(frames, padded_frames) for key, frames in per_hint_frames.items()} + if input_frames is not None: + input_frames = pad_temporal_frames(input_frames, padded_frames) + + configured_frame_rate = positive_float(transfer_config.fps) + input_frame_rate = positive_float(transfer_input_fps) + sampling_frame_rate = ( + positive_float(self._get_sp_param(sp, "resolved_frame_rate")) + or positive_float(self._get_sp_param(sp, "frame_rate")) + or positive_float(self._get_sp_param(sp, "fps")) + ) + is_wsm_only = len(transfer_config.hints) == 1 and "wsm" in transfer_config.hints + if is_wsm_only: + frame_rate = configured_frame_rate or input_frame_rate or sampling_frame_rate or 24.0 + else: + frame_rate = input_frame_rate or configured_frame_rate or sampling_frame_rate or 24.0 + num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS + guidance_scale = ( + float(transfer_config.guidance_scale) + if transfer_config.guidance_scale is not None + else self._resolve_guidance_scale(sp, COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE) + ) + flow_shift_target = float( + transfer_config.flow_shift + if transfer_config.flow_shift is not None + else self._get_sp_param(sp, "flow_shift", COSMOS3_V2V_DEFAULT_FLOW_SHIFT) + ) + max_sequence_length = ( + self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH) + or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH + ) + use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False)) + + self._guidance_scale = guidance_scale + self._num_timesteps = num_inference_steps + self._set_flow_shift(flow_shift_target) + + generator = sp.generator + seed = self._resolve_seed(sp, generator) + if generator is None: + generator = torch.Generator(device=self.device).manual_seed(seed) + + cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( + prompt, + negative_prompt, + chunk_frames, + frame_rate, + height, + width, + max_sequence_length, + sp, + use_system_prompt, + is_t2i=False, + ) + + output_chunks: list[torch.Tensor] = [] + control_chunks_per_hint: dict[str, list[torch.Tensor]] = {key: [] for key in per_hint_frames} + previous_output: torch.Tensor | None = None + + for chunk_id in range(num_chunks): + start_frame = chunk_id * stride + end_frame = min(start_frame + chunk_frames, total_frames) + control_norms = { + key: uint8_cthw_to_normalized_5d( + pad_temporal_frames(frames[:, start_frame:end_frame], chunk_frames), + dtype=self.dtype, + ) + for key, frames in per_hint_frames.items() + } + target_norm = torch.zeros_like(next(iter(control_norms.values()))) + current_conditional_frames = 0 + + if chunk_id == 0 and transfer_config.num_first_chunk_conditional_frames > 0: + if input_frames is None: + raise ValueError("Cosmos3 transfer num_first_chunk_conditional_frames > 0 requires a video input.") + current_conditional_frames = min( + transfer_config.num_first_chunk_conditional_frames, + input_frames.shape[1], + chunk_frames, + ) + if current_conditional_frames > 0: + input_cond = uint8_cthw_to_normalized_5d( + input_frames[:, :current_conditional_frames], + dtype=self.dtype, + ) + target_norm[:, :, :current_conditional_frames] = input_cond + if current_conditional_frames < chunk_frames: + fill = target_norm[:, :, current_conditional_frames - 1 : current_conditional_frames] + target_norm[:, :, current_conditional_frames:] = fill.expand( + -1, + -1, + chunk_frames - current_conditional_frames, + -1, + -1, + ) + elif chunk_id > 0 and previous_output is not None: + current_conditional_frames = min( + transfer_config.num_conditional_frames, + previous_output.shape[2], + chunk_frames, + ) + if current_conditional_frames > 0: + target_norm[:, :, :current_conditional_frames] = previous_output[ + :, :, -current_conditional_frames: + ].to(target_norm) + if current_conditional_frames < chunk_frames: + fill = target_norm[:, :, current_conditional_frames - 1 : current_conditional_frames] + target_norm[:, :, current_conditional_frames:] = fill.expand( + -1, + -1, + chunk_frames - current_conditional_frames, + -1, + -1, + ) + + control_latents = [self._encode_video_tensor(video) for video in control_norms.values()] + latents, velocity_mask, condition_latents = self._prepare_transfer_latents( + target_norm, + current_conditional_frames, + generator, + ) + video_shape = (latents.shape[2], latents.shape[3], latents.shape[4]) + shared_kwargs = dict( + video_shape=video_shape, + fps=frame_rate, + noisy_frame_mask=velocity_mask, + transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions, + ) + + self._set_timesteps( + num_inference_steps, + device=self.device, + shift=self._current_flow_shift, + ) + latents = self.diffuse_transfer( + latents=latents, + timesteps=self.scheduler.timesteps, + cond_ids=cond_ids, + cond_mask=cond_mask, + uncond_ids=uncond_ids, + uncond_mask=uncond_mask, + guidance_scale=guidance_scale, + control_guidance=transfer_config.control_guidance, + control_guidance_interval=transfer_config.control_guidance_interval, + control_latents=control_latents, + shared_kwargs=shared_kwargs, + velocity_mask=velocity_mask, + condition_latents=condition_latents, + generator=generator, + ) + output_video = self._decode_latents(latents).clamp(-1, 1) + previous_output = output_video + + if chunk_id == 0: + output_chunks.append(output_video) + for key, control in control_norms.items(): + control_chunks_per_hint[key].append(control) + else: + output_chunks.append(output_video[:, :, current_conditional_frames:]) + for key, control in control_norms.items(): + control_chunks_per_hint[key].append(control[:, :, current_conditional_frames:]) + + full_output = torch.cat(output_chunks, dim=2)[:, :, :total_frames] + full_controls = { + key: torch.cat(chunks, dim=2)[:, :, :total_frames] for key, chunks in control_chunks_per_hint.items() + } + + if transfer_config.show_control_condition: + all_controls = torch.cat([full_controls[key] for key in per_hint_frames], dim=-1) + all_controls = all_controls.to(full_output) + full_output = torch.cat([all_controls, full_output], dim=-1) + if transfer_config.show_input and input_frames is not None: + normalized_input = uint8_cthw_to_normalized_5d(input_frames[:, :total_frames], dtype=torch.float32) + full_output = torch.cat([normalized_input.to(full_output), full_output], dim=-1) + + return DiffusionOutput( + output={ + "payload": {"video": full_output}, + "metadata": { + "transfer": { + "controls": full_controls, + "hints": list(per_hint_frames), + }, + "video": {"fps": frame_rate}, + }, + }, + ) + # -- Forward (main generation entry point) ------------------------------- def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, ) -> DiffusionOutput: pipeline_start = time.time() # --- Parse request --- + prompt_data = req.prompts[0] if req.prompts else "" if len(req.prompts) > 1: raise ValueError("Cosmos3OmniDiffusersPipeline currently supports a single prompt per request.") sp = req.sampling_params - prompt_data = req.prompts[0] robolab_inputs = self._build_robolab_policy_inputs(sp, prompt_data, getattr(req, "request_id", None)) if robolab_inputs is not None: return self._forward_robolab_policy(sp, robolab_inputs, pipeline_start) @@ -2199,18 +3186,44 @@ def forward( negative_prompt = None image_tensor = None video_tensor = None + transfer_video_tensor = None + transfer_input_fps = 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") video_tensor = additional_info.get("preprocessed_video") + transfer_video_tensor = additional_info.get("preprocessed_transfer_video") + transfer_input_fps = positive_float(additional_info.get("transfer_input_fps")) 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 + transfer_config = resolve_transfer_config(sp, prompt_data) action_video_tensor = video_tensor if action_enabled else None + is_v2v = video_tensor is not None and not is_t2i and not action_enabled + self._validate_distilled_generation_mode( + is_t2i=is_t2i, + image_tensor=image_tensor, + action_enabled=action_enabled, + transfer_config=transfer_config, + is_v2v=is_v2v, + sound_enabled=sound_enabled, + ) + self._validate_edge_generation_mode( + transfer_config=transfer_config, + is_v2v=is_v2v, + sound_enabled=sound_enabled, + ) + if transfer_config is not None: + if is_t2i: + raise ValueError("Cosmos3 transfer inference is supported only for video outputs.") + if action_enabled: + raise ValueError("Cosmos3 transfer inference cannot be combined with action generation.") + if sound_enabled: + raise ValueError("Cosmos3 transfer inference cannot be combined with sound generation.") if action_enabled and is_t2i: raise ValueError("Cosmos3 action generation is supported only for video outputs.") if action_enabled and sound_enabled: @@ -2234,35 +3247,53 @@ def forward( ) if negative_prompt is None: negative_prompt = "" + if transfer_config is not None: + if image_tensor is not None: + raise ValueError("Cosmos3 transfer inference accepts video inputs or control_path values, not images.") + if transfer_video_tensor is None and video_tensor is not None: + transfer_video_tensor = video_tensor + return self._forward_transfer( + prompt=prompt, + negative_prompt=negative_prompt, + sp=sp, + transfer_config=transfer_config, + transfer_video_tensor=transfer_video_tensor, + transfer_input_fps=transfer_input_fps, + ) 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: - # T2I: 1024x1024, 50 steps, shift=3.0, guidance_interval=[400, 1000] - # T2V: 720x1280, 35 steps, shift=engine-init, no interval + # T2I and T2V share the same model and forward path; their defaults are: + # T2I: regular 1024x1024, Edge 640x640; 50 steps, shift=3.0, + # guidance_interval=[400, 1000] + # T2V: 189 frames and 35 steps; regular guidance=6, shift=10; + # Edge guidance=5, shift=3; + # no guidance interval if is_t2i: - height = sp.height or COSMOS3_T2I_DEFAULT_HEIGHT - width = sp.width or COSMOS3_T2I_DEFAULT_WIDTH + height = sp.height or ( + COSMOS3_EDGE_T2I_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2I_DEFAULT_HEIGHT + ) + width = sp.width or (COSMOS3_EDGE_T2I_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2I_DEFAULT_WIDTH) num_frames = 1 num_inference_steps = sp.num_inference_steps or COSMOS3_T2I_DEFAULT_NUM_INFERENCE_STEPS - guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE + guidance_scale = self._resolve_guidance_scale(sp, COSMOS3_T2I_DEFAULT_GUIDANCE_SCALE) default_flow_shift = COSMOS3_T2I_DEFAULT_FLOW_SHIFT default_guidance_interval: tuple[float, float] | None = COSMOS3_T2I_DEFAULT_GUIDANCE_INTERVAL batch_size = max(1, int(sp.num_outputs_per_prompt or 1)) else: - height = sp.height or COSMOS3_T2V_DEFAULT_HEIGHT - width = sp.width or COSMOS3_T2V_DEFAULT_WIDTH + height = sp.height or ( + COSMOS3_EDGE_T2V_DEFAULT_HEIGHT if self.is_edge_model else COSMOS3_T2V_DEFAULT_HEIGHT + ) + width = sp.width or (COSMOS3_EDGE_T2V_DEFAULT_WIDTH if self.is_edge_model else COSMOS3_T2V_DEFAULT_WIDTH) + default_guidance_scale = ( + COSMOS3_EDGE_T2V_DEFAULT_GUIDANCE_SCALE if self.is_edge_model else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE + ) num_frames = sp.num_frames or COSMOS3_T2V_DEFAULT_NUM_FRAMES num_inference_steps = sp.num_inference_steps or COSMOS3_T2V_DEFAULT_NUM_INFERENCE_STEPS - guidance_scale = sp.guidance_scale if sp.guidance_scale else COSMOS3_T2V_DEFAULT_GUIDANCE_SCALE - # 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 = COSMOS3_V2V_DEFAULT_FLOW_SHIFT if is_v2v else self._engine_init_flow_shift + guidance_scale = self._resolve_guidance_scale(sp, default_guidance_scale) + default_flow_shift = self._engine_init_flow_shift default_guidance_interval = None batch_size = 1 # Existing video pipeline assumes B=1. @@ -2285,9 +3316,23 @@ def forward( f"or action_chunk_size + 1; got num_frames={num_frames}, action_chunk_size={action_chunk_size}." ) num_inference_steps = sp.num_inference_steps or 30 - guidance_scale = sp.guidance_scale if sp.guidance_scale is not None else 1.0 + guidance_scale = self._resolve_guidance_scale(sp, 1.0) default_flow_shift = 5.0 + if not is_t2i and not action_enabled: + requested_num_frames = int(num_frames) + num_frames = _ceil_video_num_frames( + requested_num_frames, + self.vae_scale_factor_temporal, + ) + if num_frames != requested_num_frames and _is_rank_zero(): + logger.info( + "Rounded Cosmos3 num_frames from %d to %d for temporal compression factor %d.", + requested_num_frames, + num_frames, + self.vae_scale_factor_temporal, + ) + domain_id = None if action_enabled: domain_id = resolve_domain_id( @@ -2306,7 +3351,7 @@ def forward( self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH) or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH ) - use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False)) + use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", is_v2v)) if action_enabled and action_video_tensor is None: extra_action_video = self._get_sp_param(sp, "action_video", None) @@ -2328,13 +3373,12 @@ def forward( self._guidance_scale = guidance_scale self._num_timesteps = num_inference_steps - # Always resolve to a concrete target shift for this request, then - # update the shared Diffusers scheduler. + # Always resolve to a concrete target shift for this request. self._set_flow_shift(flow_shift_target) generator = sp.generator + seed = self._resolve_seed(sp, generator) if generator is None: - seed = sp.seed if sp.seed is not None else 42 generator = torch.Generator(device=self.device).manual_seed(seed) # --- Format prompts & tokenize (B=1; reused across loop iterations @@ -2446,17 +3490,21 @@ def forward( image_latent = None condition_latents = None + T_latent = latents.shape[2] + H_latent = latents.shape[3] + W_latent = latents.shape[4] + video_shape = (T_latent, H_latent, W_latent) + sound_latents = None target_audio_samples = None sound_sample_rate = None if sound_enabled: target_audio_samples, _, sound_sample_rate = self._resolve_sound_target_samples(sp, num_frames, frame_rate) - sound_latents, _ = self._prepare_sound_latents(target_audio_samples, generator) - - T_latent = latents.shape[2] - H_latent = latents.shape[3] - W_latent = latents.shape[4] - video_shape = (T_latent, H_latent, W_latent) + sound_latents, _ = self._prepare_sound_latents( + target_audio_samples, + generator, + sp_video_shape=video_shape, + ) # --- Denoising loop --- shared_kwargs = dict(video_shape=video_shape, fps=frame_rate) @@ -2471,7 +3519,11 @@ def forward( ) def _run_diffusion(start_latents): - self.scheduler.set_timesteps(num_inference_steps, device=self.device) + self._set_timesteps( + num_inference_steps, + device=self.device, + shift=self._current_flow_shift, + ) scheduler = self.scheduler return self.diffuse( latents=start_latents, @@ -2492,6 +3544,7 @@ def _run_diffusion(start_latents): guidance_interval=guidance_interval, raw_action_dim=raw_action_dim, scheduler=scheduler, + generator=generator, ) if is_t2i and batch_size > 1: @@ -2525,26 +3578,39 @@ def _run_diffusion(start_latents): video = self._decode_latents(latents) if _is_rank_zero(): logger.info("Video decoded in %.2fs", time.time() - decode_start) - logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) + if not sound_enabled: + logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) if sound_enabled: if sound_latents is None or target_audio_samples is None or sound_sample_rate is None: raise ValueError("Cosmos3 sound generation finished without sound latents.") if _is_rank_zero(): logger.info("Decoding sound...") + sound_decode_start = time.time() audio = self._decode_sound_latents(sound_latents, target_audio_samples) + if _is_rank_zero(): + logger.info("Sound tokenizer decoded in %.2fs", time.time() - sound_decode_start) + logger.info("Total pipeline time: %.2fs", time.time() - pipeline_start) return DiffusionOutput(output={"video": video, "audio": audio, "audio_sample_rate": sound_sample_rate}) if action_enabled: if action_latents is None or raw_action_dim is None or domain_id is None: raise ValueError("Cosmos3 action generation finished without action latents.") action = action_latents[:, :, :raw_action_dim].detach().cpu() - custom_action_output: dict[str, Any] = { - "action": action, - "raw_action_dim": raw_action_dim, - "action_mode": action_mode, - "domain_id": domain_id, - } - return DiffusionOutput(output={"video": video}, custom_output=custom_action_output) + return DiffusionOutput( + output={ + "payload": { + "video": video, + "actions": action, + }, + "metadata": { + "actions": { + "raw_action_dim": raw_action_dim, + "action_mode": action_mode, + "domain_id": domain_id, + }, + }, + }, + ) return DiffusionOutput(output={"image": video} if is_t2i else {"video": video}) diff --git a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py index 66937b33c6b..5d6ff8dd968 100644 --- a/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py +++ b/vllm_omni/diffusion/models/cosmos3/sound_tokenizer.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Cosmos3 sound tokenizer integration.""" +"""Cosmos3 sound tokenizer integration. + +The tokenizer decodes model-generated sound latents for video+audio output. It +does not encode or condition on ``multi_modal_data["audio"]``; request-side +sound generation is enabled by ``generate_sound`` or ``sound_gen``. +""" from __future__ import annotations diff --git a/vllm_omni/diffusion/models/cosmos3/transfer.py b/vllm_omni/diffusion/models/cosmos3/transfer.py new file mode 100644 index 00000000000..42cbea559e7 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/transfer.py @@ -0,0 +1,628 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 transfer inference helpers. + +The reference Cosmos Framework transfer path accepts one or more control hints +and packs them before the noisy target video. This module keeps the vLLM Omni +pipeline-specific pieces small: request parsing, media normalization, and +optional on-the-fly edge/blur control generation. + +Transfer is a video-output mode. It cannot be combined with T2I, sound +generation, or action generation. +""" + +from __future__ import annotations + +import importlib +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import PIL.Image +import torch +import torch.nn.functional as F + +TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm") +TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { + "num_video_frames_per_chunk": 93, + "num_conditional_frames": 1, + "max_frames": 5000, + "show_control_condition": False, + "show_input": False, + "num_first_chunk_conditional_frames": 0, + "share_vision_temporal_positions": True, +} +TRANSFER_DEFAULTS: dict[str, dict[str, Any]] = { + "edge": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "blur": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "depth": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "seg": {"guidance_scale": 3.0, "control_guidance": 2.0, "flow_shift": 10.0}, + "wsm": { + "guidance_scale": 1.0, + "control_guidance": 3.0, + "flow_shift": 10.0, + "num_frames": 101, + "fps": 10, + "num_video_frames_per_chunk": 101, + }, +} +EDGE_PRESETS: dict[str, tuple[int, int]] = { + "none": (20, 50), + "very_low": (20, 50), + "low": (50, 100), + "medium": (100, 200), + "high": (200, 300), + "very_high": (300, 400), +} +BLUR_DOWNUP_PRESETS: dict[str, int] = { + "none": 1, + "very_low": 4, + "low": 4, + "medium": 10, + "high": 16, + "very_high": 16, +} +BLUR_PRE_BLUR_DOWNSCALE_PRESETS: dict[str, int] = { + "none": 1, + "very_low": 1, + "low": 4, + "medium": 2, + "high": 1, + "very_high": 4, +} +BILATERAL_REFERENCE_RESOLUTION = 720 +BILATERAL_D = 30 +BILATERAL_SIGMA_COLOR = 150 +BILATERAL_SIGMA_SPACE = 100 +BILATERAL_ITERATIONS = 1 +IMAGE_EXTENSIONS = {".bmp", ".gif", ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp"} + + +@dataclass +class Cosmos3TransferHint: + key: str + control_path: str | None = None + control: Any | None = None + preset_edge_threshold: str = "medium" + preset_blur_strength: str = "medium" + + +@dataclass +class Cosmos3TransferConfig: + hints: dict[str, Cosmos3TransferHint] = field(default_factory=dict) + guidance_scale: float | None = None + control_guidance: float = 1.0 + control_guidance_interval: tuple[float, float] | None = None + flow_shift: float | None = None + num_video_frames_per_chunk: int = 93 + num_conditional_frames: int = 1 + max_frames: int = 5000 + show_control_condition: bool = False + show_input: bool = False + num_first_chunk_conditional_frames: int = 0 + share_vision_temporal_positions: bool = True + num_frames: int | None = None + fps: float | None = None + + @property + def ordered_hints(self) -> list[Cosmos3TransferHint]: + return [self.hints[key] for key in TRANSFER_HINT_KEYS if key in self.hints] + + +def _extra_args(sp: Any) -> Mapping[str, Any]: + extra = getattr(sp, "extra_args", None) + return extra if isinstance(extra, Mapping) else {} + + +def _prompt_value(prompt_data: Any, key: str) -> Any: + if not isinstance(prompt_data, Mapping): + return None + if prompt_data.get(key) is not None: + return prompt_data[key] + additional = prompt_data.get("additional_information") + if isinstance(additional, Mapping): + return additional.get(key) + return None + + +def _param(extra: Mapping[str, Any], sp: Any, prompt_data: Any, key: str, default: Any = None) -> Any: + if extra.get(key) is not None: + return extra[key] + value = getattr(sp, key, None) + if value is not None: + return value + value = _prompt_value(prompt_data, key) + return default if value is None else value + + +def _param_any(extra: Mapping[str, Any], sp: Any, prompt_data: Any, keys: tuple[str, ...], default: Any = None) -> Any: + for key in keys: + value = _param(extra, sp, prompt_data, key, None) + if value is not None: + return value + return default + + +def _is_user_field(extra: Mapping[str, Any], sp: Any, prompt_data: Any, *keys: str) -> bool: + for key in keys: + if key in extra and extra[key] is not None: + return True + if getattr(sp, key, None) is not None: + return True + if _prompt_value(prompt_data, key) is not None: + return True + return False + + +# fps, frame_rate, and resolved_frame_rate are aliases for the same request-level +# control. The serving layer leaves all of them ``None`` on the sampling params when +# the user did not provide fps, so a non-``None`` value here unambiguously means the +# user set it. +_FPS_KEYS: tuple[str, ...] = ("resolved_frame_rate", "frame_rate", "fps") + + +def _transfer_fps_param(extra: Mapping[str, Any], sp: Any, prompt_data: Any) -> Any: + # extra_args is a model-specific override and wins over the generic sampling params. + for key in _FPS_KEYS: + if extra.get(key) is not None: + return extra[key] + return _param_any(extra, sp, prompt_data, _FPS_KEYS) + + +def _as_bool(value: Any, default: bool = False) -> bool: + if value is None: + return default + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _as_interval(value: Any) -> tuple[float, float] | None: + if value is None: + return None + if isinstance(value, str): + value = [item.strip() for item in value.split(",") if item.strip()] + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError("Cosmos3 transfer control_guidance_interval must contain exactly two values.") + lo, hi = float(value[0]), float(value[1]) + if lo > hi: + raise ValueError(f"Cosmos3 transfer control_guidance_interval must be ordered as [lo, hi], got {(lo, hi)}.") + return lo, hi + + +def has_transfer_hints(extra_args: Mapping[str, Any] | None) -> bool: + if not isinstance(extra_args, Mapping): + return False + return any(extra_args.get(key) is not None for key in TRANSFER_HINT_KEYS) + + +def transfer_max_frames_from_extra_args(extra_args: Mapping[str, Any] | None) -> int: + if not isinstance(extra_args, Mapping): + return int(TRANSFER_SAMPLE_DEFAULTS["max_frames"]) + value = extra_args.get("max_frames", TRANSFER_SAMPLE_DEFAULTS["max_frames"]) + return max(1, int(value)) + + +def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3TransferConfig | None: + extra = _extra_args(sp) + hints: dict[str, Cosmos3TransferHint] = {} + for key in TRANSFER_HINT_KEYS: + raw = _param(extra, sp, prompt_data, key, None) + if raw is None: + continue + if raw is True: + raw = {} + elif isinstance(raw, str | Path): + raw = {"control_path": str(raw)} + if not isinstance(raw, Mapping): + raise TypeError( + f"Cosmos3 transfer hint '{key}' must be an object, path string, or true; got {type(raw)!r}." + ) + hints[key] = Cosmos3TransferHint( + key=key, + control_path=str(raw["control_path"]) if raw.get("control_path") is not None else None, + control=raw.get("control"), + preset_edge_threshold=str(raw.get("preset_edge_threshold") or "medium").lower(), + preset_blur_strength=str(raw.get("preset_blur_strength") or "medium").lower(), + ) + + if not hints: + transfer_only = ( + "control_guidance", + "control_guidance_interval", + "num_video_frames_per_chunk", + "num_conditional_frames", + "max_frames", + "show_control_condition", + "show_input", + "num_first_chunk_conditional_frames", + "share_vision_temporal_positions", + ) + if any(_is_user_field(extra, sp, prompt_data, key) for key in transfer_only): + raise ValueError("Cosmos3 transfer options were provided, but no transfer hint was selected.") + return None + + fps_value = _transfer_fps_param(extra, sp, prompt_data) + config = Cosmos3TransferConfig( + hints=hints, + guidance_scale=( + float(_param_any(extra, sp, prompt_data, ("guidance_scale", "guidance"))) + if _param_any(extra, sp, prompt_data, ("guidance_scale", "guidance")) is not None + else None + ), + control_guidance=float(_param(extra, sp, prompt_data, "control_guidance", 1.0)), + control_guidance_interval=_as_interval(_param(extra, sp, prompt_data, "control_guidance_interval", None)), + flow_shift=( + float(_param_any(extra, sp, prompt_data, ("flow_shift", "shift"))) + if _param_any(extra, sp, prompt_data, ("flow_shift", "shift")) is not None + else None + ), + num_video_frames_per_chunk=int( + _param( + extra, + sp, + prompt_data, + "num_video_frames_per_chunk", + TRANSFER_SAMPLE_DEFAULTS["num_video_frames_per_chunk"], + ) + ), + num_conditional_frames=int( + _param(extra, sp, prompt_data, "num_conditional_frames", TRANSFER_SAMPLE_DEFAULTS["num_conditional_frames"]) + ), + max_frames=int(_param(extra, sp, prompt_data, "max_frames", TRANSFER_SAMPLE_DEFAULTS["max_frames"])), + show_control_condition=_as_bool( + _param(extra, sp, prompt_data, "show_control_condition", None), + bool(TRANSFER_SAMPLE_DEFAULTS["show_control_condition"]), + ), + show_input=_as_bool( + _param(extra, sp, prompt_data, "show_input", None), + bool(TRANSFER_SAMPLE_DEFAULTS["show_input"]), + ), + num_first_chunk_conditional_frames=int( + _param( + extra, + sp, + prompt_data, + "num_first_chunk_conditional_frames", + TRANSFER_SAMPLE_DEFAULTS["num_first_chunk_conditional_frames"], + ) + ), + share_vision_temporal_positions=_as_bool( + _param(extra, sp, prompt_data, "share_vision_temporal_positions", None), + bool(TRANSFER_SAMPLE_DEFAULTS["share_vision_temporal_positions"]), + ), + num_frames=( + int(_param(extra, sp, prompt_data, "num_frames")) + if _param(extra, sp, prompt_data, "num_frames", None) is not None + else None + ), + fps=(float(fps_value) if fps_value is not None else None), + ) + + if len(hints) == 1: + hint_key = next(iter(hints)) + for field_name, default_value in TRANSFER_DEFAULTS[hint_key].items(): + if field_name == "guidance_scale": + user_set = _is_user_field(extra, sp, prompt_data, "guidance_scale", "guidance") + elif field_name == "flow_shift": + user_set = _is_user_field(extra, sp, prompt_data, "flow_shift", "shift") + elif field_name == "fps": + user_set = _is_user_field(extra, sp, prompt_data, *_FPS_KEYS) + else: + user_set = _is_user_field(extra, sp, prompt_data, field_name) + if not user_set: + setattr(config, field_name, default_value) + + if config.num_video_frames_per_chunk <= 0: + raise ValueError("Cosmos3 transfer num_video_frames_per_chunk must be positive.") + if config.num_conditional_frames < 0: + raise ValueError("Cosmos3 transfer num_conditional_frames must be non-negative.") + if config.max_frames <= 0: + raise ValueError("Cosmos3 transfer max_frames must be positive.") + if config.num_first_chunk_conditional_frames < 0: + raise ValueError("Cosmos3 transfer num_first_chunk_conditional_frames must be non-negative.") + for hint in hints.values(): + if hint.key == "edge" and hint.preset_edge_threshold not in EDGE_PRESETS: + raise ValueError(f"Unsupported Cosmos3 edge preset: {hint.preset_edge_threshold!r}.") + if hint.key == "blur" and hint.preset_blur_strength not in BLUR_DOWNUP_PRESETS: + raise ValueError(f"Unsupported Cosmos3 blur preset: {hint.preset_blur_strength!r}.") + return config + + +def _import_cv2(hint_key: str): + try: + return importlib.import_module("cv2") + except ImportError as exc: + raise ImportError( + f"Cosmos3 transfer hint '{hint_key}' requires opencv-python for on-the-fly control generation. " + "Install opencv-python in the environment, or provide a precomputed control_path for this hint." + ) from exc + + +def pad_temporal_frames(frames: torch.Tensor, target_frames: int) -> torch.Tensor: + if frames.ndim != 4: + raise ValueError(f"Cosmos3 transfer frames must have shape [C, T, H, W], got {tuple(frames.shape)}.") + target_frames = int(target_frames) + if target_frames <= 0: + raise ValueError("Cosmos3 transfer target frame count must be positive.") + if frames.shape[1] >= target_frames: + return frames + if frames.shape[1] == 0: + raise ValueError("Cannot pad an empty Cosmos3 transfer frame tensor.") + padded = frames + while padded.shape[1] < target_frames: + pad_len = min(padded.shape[1] - 1, target_frames - padded.shape[1]) + if pad_len <= 0: + pad_frame = padded[:, -1:].repeat(1, target_frames - padded.shape[1], 1, 1) + padded = torch.cat([padded, pad_frame], dim=1) + break + padded = torch.cat([padded, padded.flip(dims=[1])[:, :pad_len]], dim=1) + return padded + + +def normalized_video_to_uint8_cthw(video: torch.Tensor) -> torch.Tensor: + tensor = video.detach().cpu() + if tensor.ndim == 5: + if tensor.shape[0] != 1: + raise ValueError("Cosmos3 transfer supports only batch size 1 video inputs.") + tensor = tensor[0] + if tensor.ndim != 4: + raise ValueError( + f"Cosmos3 transfer video must have shape [1, 3, T, H, W] or [3, T, H, W], got {tuple(video.shape)}." + ) + if tensor.shape[0] != 3: + if tensor.shape[-1] in (3, 4): + tensor = tensor[..., :3].permute(3, 0, 1, 2) + else: + raise ValueError(f"Cosmos3 transfer video must have 3 RGB channels, got {tuple(video.shape)}.") + if tensor.is_floating_point(): + if tensor.numel() and (tensor.min().item() < 0.0 or tensor.max().item() > 1.0): + tensor = tensor.clamp(-1.0, 1.0).mul(0.5).add(0.5) + tensor = tensor.clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8) + else: + tensor = tensor.clamp(0, 255).to(torch.uint8) + return tensor.contiguous() + + +def uint8_cthw_to_normalized_5d(frames: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor: + if frames.ndim != 4 or frames.shape[0] != 3: + raise ValueError(f"Cosmos3 transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}.") + return frames.to(dtype=dtype).div(127.5).sub(1.0).unsqueeze(0).contiguous() + + +def _pil_to_uint8_rgb(value: Any) -> np.ndarray: + if isinstance(value, PIL.Image.Image): + return np.asarray(value.convert("RGB"), dtype=np.uint8) + if isinstance(value, str | Path): + return np.asarray(PIL.Image.open(value).convert("RGB"), dtype=np.uint8) + if isinstance(value, torch.Tensor): + tensor = value.detach().cpu() + if tensor.ndim == 3 and tensor.shape[0] in (3, 4): + tensor = tensor[:3].permute(1, 2, 0) + if tensor.is_floating_point(): + if tensor.numel() and (tensor.min().item() < 0.0 or tensor.max().item() > 1.0): + tensor = tensor.clamp(-1.0, 1.0).mul(0.5).add(0.5) + tensor = tensor.clamp(0.0, 1.0).mul(255.0).round().to(torch.uint8) + return tensor[..., :3].numpy().astype(np.uint8) + 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[:3], (1, 2, 0)) + if np.issubdtype(array.dtype, np.floating): + if array.size and (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 array[..., :3].astype(np.uint8) + raise TypeError(f"Cosmos3 transfer expected an RGB frame, got {type(value)!r}.") + + +def media_hw(value: Any) -> tuple[int, int] | None: + if isinstance(value, PIL.Image.Image): + return value.height, value.width + if isinstance(value, str | Path): + media_path = Path(value) + if media_path.suffix.lower() in IMAGE_EXTENSIONS: + with PIL.Image.open(media_path) as image: + return image.height, image.width + try: + import imageio.v3 as iio + except ImportError: + return None + for frame in iio.imiter(media_path): + array = np.asarray(frame) + if array.ndim >= 2: + return int(array.shape[0]), int(array.shape[1]) + return None + if isinstance(value, torch.Tensor): + tensor = value.detach() + if tensor.ndim == 5: + tensor = tensor[0] + if tensor.ndim == 4: + if tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4): + return int(tensor.shape[-2]), int(tensor.shape[-1]) + if tensor.shape[-1] in (3, 4): + return int(tensor.shape[1]), int(tensor.shape[2]) + if tensor.ndim == 3: + if tensor.shape[0] in (3, 4) and tensor.shape[-1] not in (3, 4): + return int(tensor.shape[-2]), int(tensor.shape[-1]) + if tensor.shape[-1] in (3, 4): + return int(tensor.shape[0]), int(tensor.shape[1]) + return None + if isinstance(value, np.ndarray): + array = value + if array.ndim == 5: + array = array[0] + if array.ndim == 4: + if array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + return int(array.shape[-2]), int(array.shape[-1]) + if array.shape[-1] in (3, 4): + return int(array.shape[1]), int(array.shape[2]) + if array.ndim == 3: + if array.shape[0] in (3, 4) and array.shape[-1] not in (3, 4): + return int(array.shape[-2]), int(array.shape[-1]) + if array.shape[-1] in (3, 4): + return int(array.shape[0]), int(array.shape[1]) + return None + if isinstance(value, list | tuple) and value: + return media_hw(value[0]) + return None + + +def resize_center_crop_uint8_cthw(frames: torch.Tensor, height: int, width: int) -> torch.Tensor: + if frames.ndim != 4 or frames.shape[0] != 3: + raise ValueError(f"Cosmos3 transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}.") + orig_h, orig_w = int(frames.shape[2]), int(frames.shape[3]) + scale = max(width / orig_w, height / orig_h) + resize_h = int(np.ceil(scale * orig_h)) + resize_w = int(np.ceil(scale * orig_w)) + frames_tchw = frames.permute(1, 0, 2, 3).to(dtype=torch.float32) + resized = F.interpolate(frames_tchw, size=(resize_h, resize_w), mode="bilinear", align_corners=False) + top = (resize_h - height) // 2 + left = (resize_w - width) // 2 + cropped = resized[:, :, top : top + height, left : left + width] + return cropped.round().clamp(0, 255).to(torch.uint8).permute(1, 0, 2, 3).contiguous() + + +def _path_media_to_uint8_cthw(path: str | Path, max_frames: int | None) -> torch.Tensor: + media_path = Path(path) + if not media_path.exists(): + raise FileNotFoundError(f"Missing Cosmos3 transfer control_path: {media_path}") + if media_path.suffix.lower() in IMAGE_EXTENSIONS: + array = _pil_to_uint8_rgb(media_path) + return torch.from_numpy(array).permute(2, 0, 1).unsqueeze(1).contiguous() + + try: + import imageio.v3 as iio + except ImportError as exc: + raise ImportError( + "Cosmos3 transfer video control_path loading requires imageio. " + "Install imageio[ffmpeg] or provide decoded control frames." + ) from exc + + frames: list[torch.Tensor] = [] + limit = max_frames if max_frames is not None else None + for frame in iio.imiter(media_path): + frames.append(torch.from_numpy(_pil_to_uint8_rgb(frame)).permute(2, 0, 1)) + if limit is not None and len(frames) >= int(limit): + break + if not frames: + raise ValueError(f"Cosmos3 transfer control_path produced no frames: {media_path}") + return torch.stack(frames, dim=1).contiguous() + + +def media_to_uint8_cthw(value: Any, *, height: int, width: int, max_frames: int | None = None) -> torch.Tensor: + if isinstance(value, str | Path): + frames = _path_media_to_uint8_cthw(value, max_frames=max_frames) + elif isinstance(value, torch.Tensor): + frames = normalized_video_to_uint8_cthw(value) + elif isinstance(value, np.ndarray): + frames = normalized_video_to_uint8_cthw(torch.from_numpy(value)) + elif isinstance(value, list | tuple): + if not value: + raise ValueError("Cosmos3 transfer control frames cannot be empty.") + selected = value[:max_frames] if max_frames is not None else value + tensors = [torch.from_numpy(_pil_to_uint8_rgb(frame)).permute(2, 0, 1) for frame in selected] + frames = torch.stack(tensors, dim=1).contiguous() + else: + raise TypeError(f"Unsupported Cosmos3 transfer control payload type: {type(value)!r}.") + if max_frames is not None: + frames = frames[:, : int(max_frames)] + return resize_center_crop_uint8_cthw(frames, int(height), int(width)) + + +def make_edge_control(frames: torch.Tensor, preset: str) -> torch.Tensor: + cv2 = _import_cv2("edge") + try: + lower, upper = EDGE_PRESETS[preset] + except KeyError as exc: + raise ValueError(f"Unsupported Cosmos3 edge preset: {preset!r}.") from exc + frames_np = frames.detach().cpu().numpy().astype(np.uint8) + edge_maps = [] + for idx in range(frames_np.shape[1]): + frame = np.ascontiguousarray(np.transpose(frames_np[:, idx], (1, 2, 0))) + edge_maps.append(cv2.Canny(frame, lower, upper)) + edge = np.stack(edge_maps, axis=0)[None] + return torch.from_numpy(edge).expand(3, -1, -1, -1).contiguous() + + +def _scale_for_bilateral_resolution(value: float, longest_side: int) -> float: + if longest_side <= 0: + return value + return value * (longest_side / BILATERAL_REFERENCE_RESOLUTION) + + +def _scaled_bilateral_params(height: int, width: int) -> tuple[int, float, float]: + longest_side = int(max(height, width)) + diameter = max(1, int(round(_scale_for_bilateral_resolution(float(BILATERAL_D), longest_side)))) + if diameter % 2 == 0: + diameter += 1 + sigma_color = max(1.0, _scale_for_bilateral_resolution(float(BILATERAL_SIGMA_COLOR), longest_side)) + sigma_space = max(1.0, _scale_for_bilateral_resolution(float(BILATERAL_SIGMA_SPACE), longest_side)) + return diameter, sigma_color, sigma_space + + +def make_blur_control(frames: torch.Tensor, preset: str) -> torch.Tensor: + cv2 = _import_cv2("blur") + preset = preset.lower() + if preset not in BLUR_DOWNUP_PRESETS: + raise ValueError(f"Unsupported Cosmos3 blur preset: {preset!r}.") + if preset == "none": + return frames.clone() + + frames_np = frames.detach().cpu().numpy().astype(np.uint8) + _, t, h, w = frames_np.shape + pre_blur_factor = max(1, int(BLUR_PRE_BLUR_DOWNSCALE_PRESETS[preset])) + downup_factor = max(1, int(BLUR_DOWNUP_PRESETS[preset])) + result = np.empty_like(frames_np) + for idx in range(t): + frame = np.ascontiguousarray(np.transpose(frames_np[:, idx], (1, 2, 0))) + if pre_blur_factor > 1: + small_w = max(1, w // pre_blur_factor) + small_h = max(1, h // pre_blur_factor) + frame = cv2.resize(frame, (small_w, small_h), interpolation=cv2.INTER_AREA) + diameter, sigma_color, sigma_space = _scaled_bilateral_params(frame.shape[0], frame.shape[1]) + for _ in range(BILATERAL_ITERATIONS): + frame = cv2.bilateralFilter(frame, diameter, sigma_color, sigma_space) + if pre_blur_factor > 1: + frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_LINEAR) + if downup_factor > 1: + small_w = max(1, w // downup_factor) + small_h = max(1, h // downup_factor) + frame = cv2.resize(frame, (small_w, small_h), interpolation=cv2.INTER_CUBIC) + frame = cv2.resize(frame, (w, h), interpolation=cv2.INTER_CUBIC) + result[:, idx] = np.transpose(frame, (2, 0, 1)) + return torch.from_numpy(result).contiguous() + + +def load_or_compute_control_frames( + hint: Cosmos3TransferHint, + *, + height: int, + width: int, + max_frames: int, + input_frames: torch.Tensor | None, +) -> torch.Tensor: + if hint.control is not None: + return media_to_uint8_cthw(hint.control, height=height, width=width, max_frames=max_frames) + if hint.control_path is not None: + return media_to_uint8_cthw(hint.control_path, height=height, width=width, max_frames=max_frames) + if hint.key == "edge": + if input_frames is None: + raise ValueError( + "Cosmos3 transfer hint 'edge' requires either a video input for on-the-fly control generation " + "or a precomputed control_path." + ) + return make_edge_control(input_frames[:, :max_frames], hint.preset_edge_threshold) + if hint.key == "blur": + if input_frames is None: + raise ValueError( + "Cosmos3 transfer hint 'blur' requires either a video input for on-the-fly control generation " + "or a precomputed control_path." + ) + return make_blur_control(input_frames[:, :max_frames], hint.preset_blur_strength) + raise FileNotFoundError( + f"Cosmos3 transfer hint '{hint.key}' requires a precomputed control_path; " + "on-the-fly generation is supported only for edge and blur." + ) diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index 163158a39dc..de9767d6ff8 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -12,12 +12,15 @@ from __future__ import annotations import math -from typing import Any +from collections.abc import Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import TYPE_CHECKING, Any import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F +from cache_dit import ForwardPattern from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( @@ -30,12 +33,16 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig 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 as _VllmRMSNorm from vllm_omni.platforms import current_omni_platform +if TYPE_CHECKING: + from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook + logger = init_logger(__name__) @@ -389,7 +396,6 @@ def __init__( hidden_size: int, frequency_embedding_size: int = 256, max_period: int = 10000, - target_dtype: torch.dtype = torch.bfloat16, ) -> None: super().__init__() # Following diffusers naming pattern here for checkpoint compatibility. @@ -400,7 +406,7 @@ def __init__( self.hidden_size = hidden_size half = frequency_embedding_size // 2 - freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half) + freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) self.register_buffer("freqs", freqs, persistent=False) def forward(self, t: torch.Tensor) -> torch.Tensor: @@ -580,6 +586,7 @@ def __init__( head_dim: int, rms_norm_eps: float, quant_config: QuantizationConfig | None = None, + qk_norm: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -629,8 +636,10 @@ def __init__( prefix=f"{prefix}.to_out", ) - self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) - self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.qk_norm = qk_norm + if self.qk_norm: + self.norm_q = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.norm_k = RMSNorm(self.head_dim, eps=rms_norm_eps) self.attn = FrameworkAttention( num_heads=self.num_heads, @@ -714,8 +723,9 @@ def forward( v = self.to_v(hidden_states).view(B, S_gen, self.num_kv_heads_local, self.head_dim) # Per-head QK norm - q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon) - k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon) + if self.qk_norm: + q = F.rms_norm(q, (self.head_dim,), self.norm_q.weight, eps=self.norm_q.variance_epsilon) + k = F.rms_norm(k, (self.head_dim,), self.norm_k.weight, eps=self.norm_k.variance_epsilon) # Qwen3-style RoPE q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) @@ -799,6 +809,8 @@ def __init__( head_dim: int, rms_norm_eps: float, quant_config: QuantizationConfig | None = None, + mlp_cls: type[nn.Module] = Cosmos3GatedMLP, + qk_norm: bool = True, prefix: str = "", ) -> None: super().__init__() @@ -810,11 +822,12 @@ def __init__( head_dim=head_dim, rms_norm_eps=rms_norm_eps, quant_config=quant_config, + qk_norm=qk_norm, prefix=f"{prefix}.cross_attention", ) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) - self.mlp = Cosmos3GatedMLP( + self.mlp = mlp_cls( hidden_size=hidden_size, intermediate_size=intermediate_size, quant_config=quant_config, @@ -867,6 +880,8 @@ class Cosmos3LanguageModel(nn.Module): computed once and reused across all sampling steps. """ + _layerwise_offload_blocks_attrs = ["layers"] + def __init__( self, *, @@ -954,9 +969,16 @@ class Cosmos3VFMTransformer(nn.Module): """Cosmos3 VFM Transformer: UND language model + GEN denoising layers. The UND pathway runs once per generation (K/V cached). The GEN pathway - runs at each denoising step. + runs at each denoising step over the target video/image latent stream and + optional transfer-control, action, and sound latent streams. + + Layerwise offloading uses ``gen_layers`` as the GEN block container. The + nested UND language model declares its own ``layers`` container so the two + pathways are offloaded as independent rings. - Layerwise offloading uses ``gen_layers`` as the block container. + Model-level CPU offload is Cosmos3-local: the UND ``reasoner`` and GEN + ``generator`` components are swapped at the phase boundaries marked by + ``with self._offload_context(...)``. Sequence parallelism uses ``_sp_plan`` to shard/gather the GEN pathway at module boundaries. ``Cosmos3CrossAttention`` checks @@ -964,18 +986,37 @@ class Cosmos3VFMTransformer(nn.Module): ``Attention`` layer (with Ulysses all-to-all) or plain SDPA accordingly. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + # Cosmos3 GEN blocks return only hidden_states. Per-layer UND K/V + # conditioning uses the transformer's cache-dit fallback path. + block_forward_patterns={ + "gen_layers": ForwardPattern.Pattern_3, + }, + has_separate_cfg=True, + check_forward_pattern=False, + ) + _repeated_blocks = ["Cosmos3GenDecoderLayer"] _layerwise_offload_blocks_attrs = ["gen_layers"] packed_modules_mapping = {} + _language_model_cls = Cosmos3LanguageModel + _gen_mlp_cls = Cosmos3GatedMLP + @staticmethod def _is_transformer_block(name: str, module) -> bool: return ("gen_layers" in name or "language_model.layers" in name) and name.split(".")[-1].isdigit() _hsdp_shard_conditions = [_is_transformer_block] + # Modules whose parameters must NOT be FSDP-sharded at the root level. + # time_embedder is cast to fp32 by post_load_weights for precision; if it + # were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16), + # the dtype upcast would be silently reverted, causing dtype mismatch in forward. + _hsdp_ignored_modules = ["time_embedder"] + _sp_plan = { "gen_sp_prepare": { 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True), @@ -1000,6 +1041,25 @@ def _validate_supported_config(model_config: Any) -> None: if actual != expected: raise ValueError(f"Unsupported Cosmos3 transformer config: {key}={actual!r}; expected {expected!r}.") + @classmethod + def _resolve_rms_norm_eps(cls, model_config: Any) -> float: + return float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) + + @classmethod + def _resolve_rope_theta(cls, model_config: Any) -> float: + return float(_tf_config_get(model_config, "rope_theta", 5_000_000)) + + @classmethod + def _resolve_mrope_section(cls, model_config: Any) -> list[int]: + rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} + return list(rope_scaling.get("mrope_section", [24, 20, 20])) + + def _language_model_kwargs(self) -> dict[str, Any]: + return {} + + def validate_loaded_weights(self, loaded: set[str]) -> None: + del loaded + def __init__( self, od_config: OmniDiffusionConfig, @@ -1012,7 +1072,6 @@ def __init__( super().__init__() model_config = od_config.tf_model_config self._validate_supported_config(model_config) - rope_scaling = _tf_config_get(model_config, "rope_scaling", {}) or {} self.hidden_size = int(_tf_config_get(model_config, "hidden_size", 4096)) self.num_hidden_layers = int(_tf_config_get(model_config, "num_hidden_layers", 36)) @@ -1021,9 +1080,10 @@ def __init__( self.head_dim = int(_tf_config_get(model_config, "head_dim", 128)) self.intermediate_size = int(_tf_config_get(model_config, "intermediate_size", 12288)) self.vocab_size = int(_tf_config_get(model_config, "vocab_size", 151936)) - self.rms_norm_eps = float(_tf_config_get(model_config, "rms_norm_eps", 1e-6)) - self.rope_theta = float(_tf_config_get(model_config, "rope_theta", 5_000_000)) - self.mrope_section = list(rope_scaling.get("mrope_section", [24, 20, 20])) + self.rms_norm_eps = self._resolve_rms_norm_eps(model_config) + self.rope_theta = self._resolve_rope_theta(model_config) + self.mrope_section = self._resolve_mrope_section(model_config) + self.qk_norm_for_diffusion = bool(_tf_config_get(model_config, "qk_norm_for_diffusion", True)) self.latent_patch_size = int(_tf_config_get(model_config, "latent_patch_size", 2)) self.latent_channel_size = int(_tf_config_get(model_config, "latent_channel", 48)) self.timestep_scale = float(_tf_config_get(model_config, "timestep_scale", 0.001)) @@ -1063,10 +1123,12 @@ def __init__( ) self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size + self.use_k_norm_und_for_gen = _tf_config_get(model_config, "use_k_norm_und_for_gen", None) + dtype = od_config.dtype quant_config = getattr(od_config, "quantization_config", None) if od_config else None - self.language_model = Cosmos3LanguageModel( + self.language_model = self._language_model_cls( hidden_size=self.hidden_size, intermediate_size=self.intermediate_size, num_hidden_layers=self.num_hidden_layers, @@ -1079,12 +1141,13 @@ def __init__( mrope_section=self.mrope_section, quant_config=quant_config, prefix="language_model", + **self._language_model_kwargs(), ) # Video projection layers are small; not worth quantizing. self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size) self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim) - self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype) + self.time_embedder = TimestepEmbedder(self.hidden_size) if self.action_gen: self.action_proj_in = DomainAwareLinear( self.action_dim, @@ -1115,6 +1178,8 @@ def __init__( head_dim=self.head_dim, rms_norm_eps=self.rms_norm_eps, quant_config=quant_config, + mlp_cls=self._gen_mlp_cls, + qk_norm=self.qk_norm_for_diffusion, prefix=f"gen_layers.{i}", ) for i in range(self.num_hidden_layers) @@ -1129,10 +1194,125 @@ def __init__( self.cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None self.cached_freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None + self._model_cpu_offload_enabled = False + self._model_cpu_offload_device: torch.device | None = None + self._model_cpu_offload_mover: SequentialOffloadHook | None = None + self._active_model_cpu_offload_component: str | None = None + @property def device(self) -> torch.device: + offload_device = getattr(self, "_model_cpu_offload_device", None) + if getattr(self, "_model_cpu_offload_enabled", False) and offload_device is not None: + return offload_device return next(self.parameters()).device + def _model_cpu_offload_components(self) -> dict[str, list[nn.Module]]: + """Cosmos3's mutually-exclusive reasoner/generator component sets.""" + return { + "reasoner": [self.language_model.layers], + "generator": [self.gen_layers], + } + + def _model_cpu_offload_component_tensor_ids(self) -> set[int]: + component_tensors: set[int] = set() + for modules in self._model_cpu_offload_components().values(): + for module in modules: + component_tensors.update(id(param) for param in module.parameters()) + component_tensors.update(id(buffer) for buffer in module.buffers()) + return component_tensors + + def _move_model_cpu_offload_residents(self) -> None: + """Keep non-reasoner/non-generator weights resident on the target device.""" + component_tensors = self._model_cpu_offload_component_tensor_ids() + device = self._model_cpu_offload_device + if device is None: + return + with torch.no_grad(): + for param in self.parameters(): + if id(param) not in component_tensors and param.data.device != device: + param.data = param.data.to(device, non_blocking=False) + for buffer in self.buffers(): + if id(buffer) not in component_tensors and buffer.device != device: + buffer.data = buffer.data.to(device, non_blocking=False) + + def _offload_model_cpu_component(self, name: str) -> None: + mover = self._model_cpu_offload_mover + if mover is None: + raise RuntimeError("Cosmos3 model CPU offload is not enabled") + for module in self._model_cpu_offload_components()[name]: + mover._to_cpu(module) + + def _load_model_cpu_component(self, name: str) -> None: + mover = self._model_cpu_offload_mover + if mover is None: + raise RuntimeError("Cosmos3 model CPU offload is not enabled") + for module in self._model_cpu_offload_components()[name]: + mover._to_gpu(module) + + def enable_model_cpu_offload( + self, + *, + device: torch.device, + pin_memory: bool = True, + use_hsdp: bool = False, + ) -> None: + """Enable Cosmos3 reasoner/generator CPU swapping inside ``forward``.""" + if getattr(self, "_model_cpu_offload_enabled", False): + return + + from vllm_omni.diffusion.offloader.sequential_backend import SequentialOffloadHook + + self._model_cpu_offload_device = torch.device(device) + self._model_cpu_offload_mover = SequentialOffloadHook( + offload_targets=[], + device=self._model_cpu_offload_device, + pin_memory=pin_memory, + use_hsdp=use_hsdp, + ) + self._move_model_cpu_offload_residents() + for name in self._model_cpu_offload_components(): + self._offload_model_cpu_component(name) + self._model_cpu_offload_enabled = True + self._active_model_cpu_offload_component = None + logger.info("Cosmos3 component-level CPU offload enabled on %s", self._model_cpu_offload_device) + + def disable_model_cpu_offload(self) -> None: + if not getattr(self, "_model_cpu_offload_enabled", False): + return + for name in self._model_cpu_offload_components(): + self._load_model_cpu_component(name) + if self._model_cpu_offload_device is not None and self._model_cpu_offload_device.type != "cpu": + current_omni_platform.synchronize() + self._model_cpu_offload_enabled = False + self._model_cpu_offload_device = None + self._model_cpu_offload_mover = None + self._active_model_cpu_offload_component = None + + def _activate_model_cpu_offload_component(self, name: str) -> None: + if not getattr(self, "_model_cpu_offload_enabled", False): + return + components = self._model_cpu_offload_components() + if name not in components: + raise ValueError(f"Unknown Cosmos3 offload component: {name!r} (known: {list(components)})") + if self._active_model_cpu_offload_component == name: + return + self._active_model_cpu_offload_component = None + for other in components: + if other != name: + self._offload_model_cpu_component(other) + self._load_model_cpu_component(name) + self._active_model_cpu_offload_component = name + + @contextmanager + def _model_cpu_offload_context(self, name: str) -> Iterator[None]: + self._activate_model_cpu_offload_component(name) + yield + + def _offload_context(self, name: str) -> AbstractContextManager[None]: + if not getattr(self, "_model_cpu_offload_enabled", False): + return nullcontext() + return self._model_cpu_offload_context(name) + # -- Patchify / Unpatchify ----------------------------------------------- def _pad_to_patch_size(self, h: int, w: int) -> tuple[int, int, int, int]: @@ -1216,8 +1396,12 @@ def _compute_rope_freqs( action_start_frame_offset: int = 1, action_fps: float | None = None, t_sound: int | None = None, + num_vision_items: int = 1, + share_vision_temporal_positions: bool = False, ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: """Compute mRoPE cos/sin for UND text and GEN media pathways.""" + if num_vision_items <= 0: + raise ValueError(f"Cosmos3 num_vision_items must be positive, got {num_vision_items}.") B = text_mask.shape[0] S_text = text_mask.shape[1] text_lengths = text_mask.sum(dim=1).long() @@ -1231,17 +1415,33 @@ def _compute_rope_freqs( real_len = int(text_lengths[b].item()) t_pos, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) media_temporal_offset = t_offset + self.temporal_modality_margin - v_pos, _ = compute_mrope_position_ids_vision( - t, - hp, - wp, - temporal_offset=media_temporal_offset, - fps=effective_fps, - base_fps=self.base_fps, - temporal_compression_factor=self.temporal_compression_factor, - enable_fps_modulation=self.enable_fps_modulation, - ) - gen_positions = [v_pos] + gen_positions = [] + if num_vision_items == 1 or share_vision_temporal_positions: + v_pos, _ = compute_mrope_position_ids_vision( + t, + hp, + wp, + temporal_offset=media_temporal_offset, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + gen_positions.extend([v_pos] * num_vision_items) + else: + vision_offset: int | float = media_temporal_offset + for _ in range(num_vision_items): + v_pos, vision_offset = compute_mrope_position_ids_vision( + t, + hp, + wp, + temporal_offset=vision_offset, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + gen_positions.append(v_pos) if action_frames > 0: a_pos, _ = compute_mrope_position_ids_action( action_frames, @@ -1298,16 +1498,21 @@ def _validate_gen_sequence_parallel( *, s_gen: int, s_video: int, + s_control: int, s_action: int, s_sound: int, has_action: bool, has_sound: bool, + has_control: bool, ulysses_size: int, ) -> None: if ulysses_size <= 1 or s_gen % ulysses_size == 0: return - detail_parts = [f"video tokens {s_video}"] + detail_parts = [] + if has_control: + detail_parts.append(f"control tokens {s_control}") + detail_parts.append(f"video tokens {s_video}") if has_action: detail_parts.append(f"action tokens {s_action}") if has_sound: @@ -1317,7 +1522,7 @@ def _validate_gen_sequence_parallel( "Adjust the spatial resolution, frame count, action chunk size, " "sound duration, or sound latent FPS so the combined media sequence is a " "multiple of ulysses_degree." - if has_action or has_sound + if has_control or has_action or has_sound else ( "Adjust the spatial resolution so that " "t * ceil(h/patch) * ceil(w/patch) is a multiple " @@ -1329,6 +1534,30 @@ def _validate_gen_sequence_parallel( f"ulysses_degree ({ulysses_size}). {adjust_detail}" ) + def sound_latent_frames_for_sequence_parallel( + self, + *, + video_shape: tuple[int, int, int], + sound_frames: int, + num_vision_items: int = 1, + ) -> int: + # Sound is the only modality the packed GEN sequence pairs with here: action and + # sound are never generated together (the pipeline rejects action+sound), so the + # base is just the vision tokens. + # + # Note: padded frames go through attention and are only trimmed on decode, so SP + # output is not bit-exact with non-SP (the extra frame perturbs the kept ones). + from vllm_omni.diffusion.distributed.parallel_state import get_ulysses_parallel_world_size + + ulysses_size = get_ulysses_parallel_world_size() + if ulysses_size <= 1 or sound_frames <= 0: + return sound_frames + t, h, w = video_shape + hp, wp, _, _ = self._pad_to_patch_size(h, w) + base = num_vision_items * t * hp * wp + pad = (-(base + sound_frames)) % ulysses_size + return sound_frames + pad + # -- Forward ------------------------------------------------------------- def forward( @@ -1346,6 +1575,8 @@ def forward( action_fps: float | None = None, sound_latents: torch.Tensor | None = None, noisy_frame_mask: torch.Tensor | None = None, + control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """ @@ -1363,13 +1594,19 @@ def forward( sound_latents: Optional [B, C_sound, T_sound] noisy sound latents. noisy_frame_mask: Optional [B, 1, t, 1, 1] mask where 1=noisy (add timestep embedding, predict velocity) and 0=conditioned (clean - context, skip timestep embedding). None means all frames noisy - (T2V mode). + context, skip timestep embedding). None means all target vision + frames are noisy, as in T2I/T2V. + control_latents: Optional transfer-control latents. Controls are + clean vision context and are packed before the noisy target. Returns: [B, C, t, h, w] velocity prediction, or - tuple outputs in video, action, sound order when extra modalities are provided. + tuple outputs in video, action, sound order when action/sound streams + are provided. Transfer-control streams condition the video prediction + and are not returned. """ + if kwargs: + raise TypeError(f"Unexpected Cosmos3 transformer kwargs: {sorted(kwargs)}") t, h, w = video_shape hp, wp, _, _ = self._pad_to_patch_size(h, w) text_lengths = text_mask.sum(dim=1) @@ -1382,6 +1619,15 @@ def forward( ) has_action = action_latents is not None has_sound = sound_latents is not None + if control_latents is None: + control_latent_list: list[torch.Tensor] = [] + elif isinstance(control_latents, torch.Tensor): + control_latent_list = [control_latents] + else: + control_latent_list = list(control_latents) + has_control = len(control_latent_list) > 0 + if has_control and (has_action or has_sound): + raise ValueError("Cosmos3 transfer control latents cannot be combined with action or sound latents.") if has_action and not self.action_gen: raise ValueError( "Cosmos3 action generation was requested, but this transformer " @@ -1398,13 +1644,14 @@ def forward( # Query Ulysses state at runtime ulysses_size, _, _ = _get_ulysses_state() - # Patchify latents and project to hidden space - hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w)) - s_video = hidden_video.shape[1] + # Pack action/sound tokens (no learned weights) up front so the UND + # cache sizing knows their token lengths. The modality projections are + # deferred into the generator offload context below, so model-level + # offload does not stage GEN weights before swapping to the reasoner. + action_tokens = None + sound_tokens = None s_action = 0 - hidden_action = None s_sound = 0 - hidden_sound = None if action_latents is not None: if action_latents.shape[0] != hidden_states.shape[0]: raise ValueError( @@ -1413,60 +1660,16 @@ def forward( ) if action_domain_ids is None: action_domain_ids = torch.zeros(action_latents.shape[0], dtype=torch.long, device=action_latents.device) - hidden_action = self.action_proj_in(self.pack_action(action_latents), action_domain_ids) - hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) - s_action = hidden_action.shape[1] + action_tokens = self.pack_action(action_latents) + s_action = action_tokens.shape[1] if sound_latents is not None: if sound_latents.shape[0] != hidden_states.shape[0]: raise ValueError( "Cosmos3 sound and video batch sizes must match: " f"video={hidden_states.shape[0]}, sound={sound_latents.shape[0]}." ) - hidden_sound = self.audio_proj_in(self.pack_sound(sound_latents)) - hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype) - s_sound = hidden_sound.shape[1] - - # Timestep embedding (fp32 for precision). - # 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(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: - # Build per-token mask from per-frame mask. - # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1] - token_noisy_mask = ( - noisy_frame_mask[:, 0, :, 0, 0] # [B, t] - .unsqueeze(-1) # [B, t, 1] - .expand(-1, -1, hp * wp) # [B, t, hp*wp] - .reshape(hidden_video.shape[0], -1, 1) # [B, t*hp*wp, 1] - ) - hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask - else: - hidden_video = hidden_video + time_embed.unsqueeze(1) - - if hidden_action is not None: - if action_noisy_mask is None: - hidden_action = hidden_action + time_embed.unsqueeze(1) - else: - if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1): - raise ValueError( - "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], " - f"got {tuple(action_noisy_mask.shape)}." - ) - action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device) - hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask - - if hidden_sound is not None: - hidden_sound = hidden_sound + time_embed.unsqueeze(1) - hidden_parts = [hidden_video] - if hidden_action is not None: - hidden_parts.append(hidden_action) - if hidden_sound is not None: - hidden_parts.append(hidden_sound) - hidden_gen = torch.cat(hidden_parts, dim=1) + sound_tokens = self.pack_sound(sound_latents) + s_sound = sound_tokens.shape[1] # Run UND pathway once and cache K/V (replicated across all ranks) if self.cached_kv is None: @@ -1482,81 +1685,165 @@ def forward( action_start_frame_offset=action_start_frame_offset, action_fps=action_fps, t_sound=s_sound, + num_vision_items=len(control_latent_list) + 1, + share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) - cached_kv_full = self.language_model(text_ids, freqs_und) + with self._offload_context("reasoner"): + cached_kv_full = self.language_model(text_ids, freqs_und) self.cached_freqs_gen = freqs_gen # Trim to real text length (remove padding). K/V stay replicated; # the framework Attention layer head-slices them via joint_key/value. self.cached_kv = [(k[:, :max_real_len], v[:, :max_real_len]) for k, v in cached_kv_full] - # Run GEN layers. UND K/V (replicated) is passed to each layer; - # the Cosmos3CrossAttention forwards them as joint_key/value so the - # framework Attention handles the Ulysses head-slicing internally. - if self.cached_kv is None or self.cached_freqs_gen is None: - raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") - self._validate_gen_sequence_parallel( - s_gen=hidden_gen.shape[1], - s_video=s_video, - s_action=s_action, - s_sound=s_sound, - has_action=has_action, - has_sound=has_sound, - ulysses_size=ulysses_size, - ) - freqs_cos, freqs_sin = self.cached_freqs_gen - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) - freqs_gen = (freqs_cos, freqs_sin) - - if len(self.gen_layers) == len(self.cached_kv): - for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): - hidden_gen = layer( - hidden_gen, - k_und=k_und, - v_und=v_und, - freqs_cos=freqs_cos, - freqs_sin=freqs_sin, + with self._offload_context("generator"): + # Patchify latents and project to hidden space after UND cache + # construction, so model-level offload does not stage GEN weights + # before immediately swapping to the reasoner. + hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w)) + s_video = hidden_video.shape[1] + s_control = 0 + hidden_controls: list[torch.Tensor] = [] + for idx, control in enumerate(control_latent_list): + if control.shape != hidden_states.shape: + raise ValueError( + "Cosmos3 transfer control latent shape must match target latent shape: " + f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}." + ) + hidden_control = self.proj_in( + self.patchify(control.to(device=hidden_states.device, dtype=hidden_states.dtype), t, h, w) ) - # Cache-dit's block wrapper may return a tuple; unwrap it. - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] - else: - # Cache-dit patches gen_layers to a grouped wrapper. - for layer in self.gen_layers: - hidden_gen = layer( - hidden_gen, - cached_kv=self.cached_kv, - freqs_gen=freqs_gen, + hidden_controls.append(hidden_control) + s_control += hidden_control.shape[1] + hidden_action = None + hidden_sound = None + if action_tokens is not None: + assert action_domain_ids is not None + hidden_action = self.action_proj_in(action_tokens, action_domain_ids) + hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype) + if sound_tokens is not None: + hidden_sound = self.audio_proj_in(sound_tokens) + hidden_sound = hidden_sound + self.audio_modality_embed.to(hidden_sound.dtype) + + # Timestep embedding (fp32 for precision). + # 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(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: + # Build per-token mask from per-frame mask. + # noisy_frame_mask: [B, 1, t, 1, 1] → token mask: [B, t*hp*wp, 1] + token_noisy_mask = ( + noisy_frame_mask[:, 0, :, 0, 0] # [B, t] + .unsqueeze(-1) # [B, t, 1] + .expand(-1, -1, hp * wp) # [B, t, hp*wp] + .reshape(hidden_video.shape[0], -1, 1) # [B, t*hp*wp, 1] ) - if isinstance(hidden_gen, tuple): - hidden_gen = hidden_gen[0] - - hidden_gen = self.gen_sp_gather(hidden_gen) - - # Final norm and project back to latent space - hidden_gen = self.norm_moe_gen(hidden_gen) - if not has_action and not has_sound: - return self.unpatchify(self.proj_out(hidden_gen), t, h, w) - - split_sizes = [s_video] - if has_action: - split_sizes.append(s_action) - if has_sound: - split_sizes.append(s_sound) - split_hidden = hidden_gen.split(split_sizes, dim=1) - hidden_video = split_hidden[0] - video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) - outputs: list[torch.Tensor] = [video_pred] - split_idx = 1 - if has_action: - hidden_action = split_hidden[split_idx] + hidden_video = hidden_video + time_embed.unsqueeze(1) * token_noisy_mask + else: + hidden_video = hidden_video + time_embed.unsqueeze(1) + + if hidden_action is not None: + if action_noisy_mask is None: + hidden_action = hidden_action + time_embed.unsqueeze(1) + else: + if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1): + raise ValueError( + "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], " + f"got {tuple(action_noisy_mask.shape)}." + ) + action_noisy_mask = action_noisy_mask.to(dtype=hidden_action.dtype, device=hidden_action.device) + hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask + + if hidden_sound is not None: + hidden_sound = hidden_sound + time_embed.unsqueeze(1) + hidden_parts = [*hidden_controls, hidden_video] + if hidden_action is not None: + hidden_parts.append(hidden_action) + if hidden_sound is not None: + hidden_parts.append(hidden_sound) + hidden_gen = torch.cat(hidden_parts, dim=1) + + # Run GEN layers. UND K/V (replicated) is passed to each layer; + # the Cosmos3CrossAttention forwards them as joint_key/value so the + # framework Attention handles the Ulysses head-slicing internally. + if self.cached_kv is None or self.cached_freqs_gen is None: + raise RuntimeError("Cosmos3 GEN cache was not initialized before running GEN layers.") + self._validate_gen_sequence_parallel( + s_gen=hidden_gen.shape[1], + s_video=s_video, + s_control=s_control, + s_action=s_action, + s_sound=s_sound, + has_action=has_action, + has_sound=has_sound, + has_control=has_control, + ulysses_size=ulysses_size, + ) + freqs_cos, freqs_sin = self.cached_freqs_gen + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + freqs_gen = (freqs_cos, freqs_sin) + + if len(self.gen_layers) == len(self.cached_kv): + for layer, (k_und, v_und) in zip(self.gen_layers, self.cached_kv, strict=True): + hidden_gen = layer( + hidden_gen, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + ) + # Cache-dit's block wrapper may return a tuple; unwrap it. + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + else: + # Cache-dit patches gen_layers to a grouped wrapper. + for layer in self.gen_layers: + hidden_gen = layer( + hidden_gen, + cached_kv=self.cached_kv, + freqs_gen=freqs_gen, + ) + if isinstance(hidden_gen, tuple): + hidden_gen = hidden_gen[0] + + hidden_gen = self.gen_sp_gather(hidden_gen) + + # Final norm and project back to latent space + hidden_gen = self.norm_moe_gen(hidden_gen) + if not has_action and not has_sound and not has_control: + return self.unpatchify(self.proj_out(hidden_gen), t, h, w) + + split_sizes = [] + if has_control: + split_sizes.append(s_control) + split_sizes.append(s_video) + if has_action: + split_sizes.append(s_action) + if has_sound: + split_sizes.append(s_sound) + split_hidden = hidden_gen.split(split_sizes, dim=1) + split_idx = 0 + if has_control: + split_idx += 1 + hidden_video = split_hidden[split_idx] split_idx += 1 - assert action_domain_ids is not None - outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids))) - if has_sound: - hidden_sound = split_hidden[split_idx] - outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) - return tuple(outputs) + video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w) + if has_control: + return video_pred + outputs: list[torch.Tensor] = [video_pred] + if has_action: + hidden_action = split_hidden[split_idx] + split_idx += 1 + assert action_domain_ids is not None + outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids))) + if has_sound: + hidden_sound = split_hidden[split_idx] + outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound))) + return tuple(outputs) def post_load_weights(self) -> None: """Post-load processing: ensure correct dtypes.""" diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py new file mode 100644 index 00000000000..6a1b4313bc8 --- /dev/null +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py @@ -0,0 +1,396 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Cosmos3 Edge transformer variant with a Nemotron dense UND backbone.""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention + +from .transformer_cosmos3 import ( + Cosmos3VFMTransformer, + Qwen3VLTextRotaryEmbedding, + RMSNorm, + _apply_rotary_pos_emb, + _nested_get, + _tf_config_get, +) + +COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" + + +def _tf_config_nested_get(config: Any, key: str, default: Any = None) -> Any: + if config is None: + return default + if isinstance(config, dict): + value = _nested_get(config, key) + elif hasattr(config, "to_dict"): + value = _nested_get(config.to_dict(), key) + elif hasattr(config, "params"): + value = _nested_get(config.params, key) + else: + value = getattr(config, key, None) + return default if value is None else value + + +class Cosmos3Relu2MLP(nn.Module): + """Nemotron dense MLP: down_proj(relu(up_proj(x)) ** 2).""" + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.up_proj = ColumnParallelLinear( + hidden_size, + intermediate_size, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + hidden = F.relu(self.up_proj(x)) + return self.down_proj(hidden * hidden) + + +class Cosmos3EdgeCausalAttention(nn.Module): + """Edge UND causal attention. + + Reasoner self-attention uses raw UND Q/K with RoPE because + ``qk_norm_for_text=false``. The GEN-facing UND key cache is a separate + normalized-then-RoPE'd view of the raw K tensor when GEN QK norm is enabled; + no-GEN-QK-norm checkpoints reuse the raw RoPE'd UND K. + """ + + def __init__( + self, + *, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + use_k_norm_und_for_gen: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + + tp_size = get_tensor_model_parallel_world_size() + self.num_heads_local = self.num_heads // tp_size + self.num_kv_heads_local = self.num_kv_heads // tp_size + + self.to_q = ColumnParallelLinear( + hidden_size, + self.num_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_q", + ) + self.to_k = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_k", + ) + self.to_v = ColumnParallelLinear( + hidden_size, + self.num_kv_heads * self.head_dim, + bias=False, + gather_output=False, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_v", + ) + self.to_out = RowParallelLinear( + self.num_heads * self.head_dim, + hidden_size, + bias=False, + input_is_parallel=True, + return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_out", + ) + + self.k_norm_und_for_gen = RMSNorm(self.head_dim, eps=rms_norm_eps) if use_k_norm_und_for_gen else None + self.attn = FrameworkAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=True, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_kv_heads, + skip_sequence_parallel=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, S, _ = hidden_states.shape + + q = self.to_q(hidden_states).view(B, S, self.num_heads_local, self.head_dim) + k = self.to_k(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim) + v = self.to_v(hidden_states).view(B, S, self.num_kv_heads_local, self.head_dim) + + q_reasoner, k_reasoner = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + + if self.k_norm_und_for_gen is not None: + k_gen = F.rms_norm( + k, + (self.head_dim,), + self.k_norm_und_for_gen.weight, + eps=self.k_norm_und_for_gen.variance_epsilon, + ) + _, k_gen = _apply_rotary_pos_emb(q, k_gen, freqs_cos, freqs_sin) + else: + k_gen = k_reasoner + + out = self.attn(q_reasoner, k_reasoner, v).reshape(B, S, -1) + return self.to_out(out), k_gen, v + + +class Cosmos3EdgeUndDecoderLayer(nn.Module): + """Edge UND decoder layer: causal self-attention + ReLU2 MLP.""" + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + rms_norm_eps: float, + use_k_norm_und_for_gen: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.self_attn = Cosmos3EdgeCausalAttention( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + use_k_norm_und_for_gen=use_k_norm_und_for_gen, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + ) + self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) + self.mlp = Cosmos3Relu2MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + cos, sin = freqs + attn_out, k, v = self.self_attn(hidden_states, cos, sin) + hidden_states = residual + attn_out + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = residual + self.mlp(hidden_states) + + return hidden_states, k, v + + +class Cosmos3EdgeLanguageModel(nn.Module): + """Nemotron dense UND tower that returns GEN-facing K plus raw V per layer.""" + + _layerwise_offload_blocks_attrs = ["layers"] + + def __init__( + self, + *, + hidden_size: int, + intermediate_size: int, + num_hidden_layers: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + vocab_size: int, + rms_norm_eps: float, + rope_theta: float, + mrope_section: list[int], + use_k_norm_und_for_gen: bool, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.embed_tokens = nn.Embedding(vocab_size, hidden_size) + self.rotary_emb = Qwen3VLTextRotaryEmbedding( + head_dim=head_dim, + rope_theta=rope_theta, + mrope_section=mrope_section, + ) + self.layers = nn.ModuleList( + [ + Cosmos3EdgeUndDecoderLayer( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + rms_norm_eps=rms_norm_eps, + use_k_norm_und_for_gen=use_k_norm_und_for_gen, + quant_config=quant_config, + prefix=f"{prefix}.layers.{i}", + ) + for i in range(num_hidden_layers) + ] + ) + self.norm = RMSNorm(hidden_size, eps=rms_norm_eps) + + def forward( + self, + text_ids: torch.Tensor, + freqs: tuple[torch.Tensor, torch.Tensor], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + hidden = self.embed_tokens(text_ids) + + cached_kv: list[tuple[torch.Tensor, torch.Tensor]] = [] + for layer in self.layers: + hidden, k, v = layer(hidden, freqs) + cached_kv.append((k, v)) + + return cached_kv + + +class Cosmos3EdgeVFMTransformer(Cosmos3VFMTransformer): + """Cosmos3 Edge variant with Nemotron dense UND and shared GEN diffusion.""" + + _language_model_cls = Cosmos3EdgeLanguageModel + _gen_mlp_cls = Cosmos3Relu2MLP + + @staticmethod + def _validate_supported_config(model_config: Any) -> None: + expected_values = { + "qk_norm_for_text": False, + "position_embedding_type": "unified_3d_mrope", + "unified_3d_mrope_reset_spatial_ids": True, + "joint_attn_implementation": "two_way", + } + for key, expected in expected_values.items(): + actual = _tf_config_get(model_config, key, expected) + if actual != expected: + raise ValueError( + f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}." + ) + + backbone_type = _tf_config_get(model_config, "backbone_type", None) + if backbone_type is None: + raise ValueError( + f"Cosmos3 Edge transformer config must declare backbone_type={COSMOS3_EDGE_BACKBONE_TYPE!r}." + ) + if backbone_type != COSMOS3_EDGE_BACKBONE_TYPE: + raise ValueError( + "Unsupported Cosmos3 Edge transformer config: " + f"backbone_type={backbone_type!r}; expected {COSMOS3_EDGE_BACKBONE_TYPE!r}." + ) + + required_values = { + "latent_channel": 48, + "latent_patch_size": 2, + "temporal_compression_factor": 4, + } + for key, expected in required_values.items(): + actual = _tf_config_get(model_config, key, None) + if actual is None: + raise ValueError(f"Cosmos3 Edge transformer config must declare {key}={expected}.") + if int(actual) != expected: + raise ValueError( + f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}; expected {expected!r}." + ) + + @classmethod + def _resolve_rms_norm_eps(cls, model_config: Any) -> float: + return float( + _tf_config_get( + model_config, + "rms_norm_eps", + _tf_config_get(model_config, "layer_norm_epsilon", 1e-5), + ) + ) + + @classmethod + def _resolve_rope_theta(cls, model_config: Any) -> float: + return float(_tf_config_get(model_config, "rope_theta", 100_000_000)) + + @classmethod + def _resolve_mrope_section(cls, model_config: Any) -> list[int]: + mrope_section = _tf_config_get(model_config, "mrope_section", None) + if mrope_section is None: + mrope_section = _tf_config_nested_get(model_config, "mrope_section", [24, 20, 20]) + return list(mrope_section) + + def _language_model_kwargs(self) -> dict[str, Any]: + return {"use_k_norm_und_for_gen": self.use_k_norm_und_for_gen} + + def validate_loaded_weights(self, loaded: set[str]) -> None: + missing: list[str] = [] + for layer_idx in range(self.num_hidden_layers): + required_markers = ( + f"language_model.layers.{layer_idx}.mlp.up_proj.", + f"language_model.layers.{layer_idx}.mlp.down_proj.", + f"gen_layers.{layer_idx}.mlp.up_proj.", + f"gen_layers.{layer_idx}.mlp.down_proj.", + ) + if self.use_k_norm_und_for_gen: + required_markers = ( + *required_markers, + f"language_model.layers.{layer_idx}.self_attn.k_norm_und_for_gen.", + ) + missing.extend( + marker.rstrip(".") for marker in required_markers if not any(marker in name for name in loaded) + ) + if missing: + preview = ", ".join(missing[:10]) + suffix = " ..." if len(missing) > 10 else "" + raise ValueError( + "Cosmos3 Edge transformer checkpoint is missing required weights for " + f"{preview}{suffix}. Use a converted Edge checkpoint with ReLU2 UND/GEN MLP weights." + ) diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index 8163f785436..94fcb6699c2 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -24,13 +24,32 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.models.diffusers_adapter.pipeline_utils import BasePipelineUtils, get_pipeline_utils +from vllm_omni.diffusion.models.diffusers_adapter.quantization_utils import ( + apply_diffusers_quantization_config, + convert_diffusers_quantization_config, + ensure_supported_diffusers_quantization, +) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniPromptType, OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +_DIFFUSERS_CONFIG_LOAD_KWARGS = { + "cache_dir", + "dduf_entries", + "force_download", + "local_dir", + "local_dir_use_symlinks", + "local_files_only", + "proxies", + "revision", + "subfolder", + "token", + "user_agent", +} + class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): """Black-box adapter that delegates full pipeline execution to a diffusers pipeline. @@ -46,6 +65,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): batching mode. """ + supports_request_batch = False supports_step_execution: bool = False def __init__(self, *, od_config: OmniDiffusionConfig, device: torch.device | None = None): @@ -79,12 +99,19 @@ def load_weights(self) -> None: "torch_dtype": dtype, **self.od_config.diffusers_load_kwargs, } - logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") + convert_diffusers_quantization_config(load_kwargs) pipeline_class = self.od_config.diffusers_pipeline_cls pipeline_class_name = pipeline_class.__name__ if pipeline_class is not None else None self._pipeline_utils = get_pipeline_utils(pipeline_class_name) self._pipeline_utils.update_load_kwargs(self.od_config, load_kwargs) + component_names = ( + self._load_diffusers_component_names(model_id, load_kwargs) + if self.od_config.quantization_config is not None and "quantization_config" not in load_kwargs + else {} + ) + apply_diffusers_quantization_config(self.od_config, load_kwargs, component_names) + logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") self._pipeline = DiffusionPipeline.from_pretrained(model_id, **load_kwargs) self._pipeline_utils.apply_post_load_updates(self._pipeline, self.od_config) @@ -151,9 +178,8 @@ def post_decode(self, **_: Any) -> Any: # Forward pass # ------------------------------------------------------------------ - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Full delegation to diffusers ``pipeline.__call__()``.""" - kwargs = self._build_call_kwargs(req) logger.debug(f"Calling diffusers pipeline with kwargs: {kwargs}") @@ -190,10 +216,38 @@ def _raise_unsupported_features(self) -> None: "Eager execution is not supported with the diffusers backend. " "Use a native pipeline for continuous batching mode." ) - if self.od_config.quantization_config is not None: - raise NotImplementedError( - "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization." + if ( + self.od_config.quantization_config is not None + and "quantization_config" not in self.od_config.diffusers_load_kwargs + ): + ensure_supported_diffusers_quantization(self.od_config.quantization_config) + if self.od_config.diffusers_load_kwargs.get("dduf_file"): + raise NotImplementedError( + "Diffusers backend quantization conversion does not support " + "diffusers_load_kwargs.dduf_file yet. The preflight component " + "discovery would need to mirror Diffusers' DDUF config loading. " + "Use diffusers_load_kwargs.quantization_config for a native " + "Diffusers quantization config, or omit dduf_file for vLLM-Omni " + "quantization conversion." + ) + + def _load_diffusers_component_names( + self, + model_id: str, + load_kwargs: dict[str, Any], + ) -> dict[str, Any]: + config_load_kwargs = {k: load_kwargs[k] for k in _DIFFUSERS_CONFIG_LOAD_KWARGS if k in load_kwargs} + pipeline_config = DiffusionPipeline.load_config(model_id, **config_load_kwargs) + return { + name: value + for name, value in pipeline_config.items() + if ( + isinstance(value, list) + and len(value) > 0 + and value[0] is not None + and (name not in load_kwargs or load_kwargs[name] is not None) ) + } # ------------------------------------------------------------------ # Wrap settings, inputs, and outputs @@ -280,8 +334,8 @@ def _set_attention_backend(self) -> None: f"{dict(zip(attention_backend_attempts, attempt_errors))}" ) - def _build_call_kwargs(self, req: OmniDiffusionRequest) -> dict[str, Any]: - """Translate ``OmniDiffusionRequest`` into diffusers ``__call__`` kwargs.""" + def _build_call_kwargs(self, req: DiffusionRequestBatch) -> dict[str, Any]: + """Translate a ``DiffusionRequestBatch`` into diffusers ``__call__`` kwargs.""" sampling = req.sampling_params input_kwargs = self._extract_input(req.prompts) diff --git a/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py b/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py new file mode 100644 index 00000000000..a7228d21e03 --- /dev/null +++ b/vllm_omni/diffusion/models/diffusers_adapter/quantization_utils.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Quantization helpers for the Diffusers backend.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import Any + +from diffusers import PipelineQuantizationConfig, TorchAoConfig + +logger = logging.getLogger(__name__) + +_DIFFUSERS_TRANSFORMER_COMPONENTS = ("transformer", "transformer_2") +_QuantizationConfigSpec = tuple[Callable[[Any], None], str] + + +def ensure_supported_diffusers_quantization(quant_config: Any) -> None: + """Validate that a vLLM-Omni quantization config has a Diffusers mapping.""" + + method = _get_quant_method_name(quant_config) + validator, _ = _get_quantization_config_spec(method) + validator(quant_config) + + +def convert_diffusers_quantization_config(load_kwargs: dict[str, Any]) -> None: + """Build Diffusers-native quantization config objects from dict kwargs.""" + + quant_config = load_kwargs.get("quantization_config") + if isinstance(quant_config, dict): + load_kwargs["quantization_config"] = _build_diffusers_native_quantization_config(quant_config) + + +def apply_diffusers_quantization_config( + od_config: Any, + load_kwargs: dict[str, Any], + component_names: dict[str, Any], +) -> bool: + """Inject a courtesy-converted quantization_config into load kwargs. + + ``diffusers_load_kwargs`` is the canonical Diffusers backend configuration + path, so an explicit ``quantization_config`` already present in ``load_kwargs`` + is never replaced. Returns whether a config was injected by this helper. + """ + + quant_config = getattr(od_config, "quantization_config", None) + if quant_config is None: + return False + + if "quantization_config" in load_kwargs: + logger.warning( + "Both vLLM-Omni quantization_config and diffusers_load_kwargs.quantization_config " + "were provided for the diffusers backend. Using the Diffusers-native " + "quantization_config from diffusers_load_kwargs." + ) + return False + + transformer_components = _get_diffusers_transformer_components(component_names) + load_kwargs["quantization_config"] = _build_diffusers_quantization_config( + quant_config, + transformer_components, + ) + return True + + +def _normalize_method_name(method: Any) -> str: + method = getattr(method, "value", method) + return str(method).lower().replace("-", "_") + + +def _get_quant_method_name(quant_config: Any) -> str: + get_name = getattr(quant_config, "get_name", None) + if get_name is not None: + return _normalize_method_name(get_name() if callable(get_name) else get_name) + + method = getattr(quant_config, "quant_method", None) + if method is None: + method = getattr(quant_config, "method", None) + if method is None: + raise NotImplementedError( + "Diffusers backend quantization conversion requires a quantization " + "config with get_name(), quant_method, or method." + ) + return _normalize_method_name(method) + + +def _ensure_no_ignored_layers(quant_config: Any) -> None: + ignored_layers = getattr(quant_config, "ignored_layers", None) + if not ignored_layers: + ignored_layers = getattr(quant_config, "modules_to_not_convert", None) + if ignored_layers: + raise NotImplementedError( + "Diffusers backend quantization conversion does not map vLLM " + "ignored_layers/modules_to_not_convert names to Diffusers module " + "names. Use diffusers_load_kwargs for a native Diffusers config." + ) + + +def _get_torchao_quant_type_cls(class_name: str) -> type[Any]: + try: + import torchao.quantization as torchao_quantization + except ImportError as exc: + raise ImportError( + "Diffusers backend quantization conversion for fp8/int8 requires " + "torchao. Install torchao or pass a Diffusers-native " + "quantization_config through diffusers_load_kwargs." + ) from exc + + try: + return getattr(torchao_quantization, class_name) + except AttributeError as exc: + raise ImportError(f"torchao.quantization.{class_name} is required for this quantization mapping.") from exc + + +def _build_diffusers_native_quantization_config(config: dict[str, Any]) -> Any: + config = config.copy() + quant_mapping = config.get("quant_mapping") + if isinstance(quant_mapping, dict): + config["quant_mapping"] = { + component_name: _build_diffusers_component_quant_config(component_config) + for component_name, component_config in quant_mapping.items() + } + return PipelineQuantizationConfig(**config) + + +def _build_diffusers_component_quant_config(config: Any) -> Any: + if not isinstance(config, dict): + return config + + quant_method = config.get("quant_method") + if quant_method is None: + raise ValueError("Diffusers quant_mapping entries must provide quant_method.") + + quant_method = _normalize_method_name(quant_method) + config_cls = _get_diffusers_quantization_config_cls(quant_method) + from_dict = getattr(config_cls, "from_dict", None) + if from_dict is not None: + return from_dict(config) + + init_kwargs = config.copy() + init_kwargs.pop("quant_method", None) + return config_cls(**init_kwargs) + + +def _get_diffusers_quantization_config_cls(quant_method: str) -> type[Any]: + from diffusers.quantizers.auto import AUTO_QUANTIZATION_CONFIG_MAPPING as DIFFUSERS_QUANT_CONFIG_MAPPING + + config_cls = DIFFUSERS_QUANT_CONFIG_MAPPING.get(quant_method) + if config_cls is not None: + return config_cls + + try: + from transformers.quantizers.auto import ( + AUTO_QUANTIZATION_CONFIG_MAPPING as TRANSFORMERS_QUANT_CONFIG_MAPPING, + ) + except ImportError: + TRANSFORMERS_QUANT_CONFIG_MAPPING = {} + + config_cls = TRANSFORMERS_QUANT_CONFIG_MAPPING.get(quant_method) + if config_cls is None: + raise ValueError(f"Diffusers quant_mapping entry uses unknown quant_method={quant_method!r}.") + return config_cls + + +def _build_torchao_pipeline_quant_config( + torchao_quant_type_name: str, + transformer_components: list[str], +) -> Any: + """Build a Diffusers config after the caller has validated the vLLM config.""" + + quant_type_cls = _get_torchao_quant_type_cls(torchao_quant_type_name) + return PipelineQuantizationConfig( + quant_mapping={ + component_name: TorchAoConfig( + quant_type=quant_type_cls(), + ) + for component_name in transformer_components + } + ) + + +def _get_diffusers_transformer_components(component_names: dict[str, Any]) -> list[str]: + transformer_components = [ + component_name for component_name in component_names if component_name in _DIFFUSERS_TRANSFORMER_COMPONENTS + ] + if transformer_components: + return transformer_components + + available_components = ", ".join(sorted(component_names.keys())) or "<none>" + raise NotImplementedError( + "Diffusers backend quantization conversion currently only supports " + "pipelines with transformer or transformer_2 components. " + f"Found pipeline components: {available_components}. Use " + "diffusers_load_kwargs for a native Diffusers quantization config." + ) + + +def _validate_fp8_quant_config(quant_config: Any) -> None: + if getattr(quant_config, "is_checkpoint_fp8_serialized", False): + raise NotImplementedError( + "Diffusers backend fp8 conversion only supports online/dynamic " + "TorchAO quantization; serialized vLLM fp8 checkpoints are not mapped." + ) + + activation_scheme = getattr(quant_config, "activation_scheme", "dynamic") + if activation_scheme != "dynamic": + raise NotImplementedError( + f"Diffusers backend fp8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}." + ) + + weight_block_size = getattr(quant_config, "weight_block_size", None) + if weight_block_size is not None: + raise NotImplementedError( + "Diffusers backend fp8 conversion does not map vLLM weight_block_size " + "to TorchAO. Use diffusers_load_kwargs for a native Diffusers config." + ) + + _ensure_no_ignored_layers(quant_config) + + +def _validate_int8_quant_config(quant_config: Any) -> None: + if getattr(quant_config, "is_checkpoint_int8_serialized", False): + raise NotImplementedError( + "Diffusers backend int8 conversion only supports online/dynamic " + "TorchAO quantization; serialized vLLM int8 checkpoints are not mapped." + ) + + activation_scheme = getattr(quant_config, "activation_scheme", "dynamic") + if activation_scheme != "dynamic": + raise NotImplementedError( + f"Diffusers backend int8 conversion only supports activation_scheme='dynamic'. Got {activation_scheme!r}." + ) + + _ensure_no_ignored_layers(quant_config) + + +_QUANTIZATION_CONFIG_SPECS: dict[str, _QuantizationConfigSpec] = { + "fp8": (_validate_fp8_quant_config, "Float8DynamicActivationFloat8WeightConfig"), + "int8": (_validate_int8_quant_config, "Int8DynamicActivationInt8WeightConfig"), +} + + +def _get_quantization_config_spec(method: str) -> _QuantizationConfigSpec: + spec = _QUANTIZATION_CONFIG_SPECS.get(method) + if spec is None: + raise NotImplementedError( + f"Diffusers backend quantization conversion does not support {method!r}. " + "Use diffusers_load_kwargs for a native Diffusers quantization config, " + "or use a native vLLM-Omni pipeline for this quantization method." + ) + return spec + + +def _build_diffusers_quantization_config( + quant_config: Any, + transformer_components: list[str], +) -> Any: + """Build a Diffusers PipelineQuantizationConfig from a supported config.""" + + method = _get_quant_method_name(quant_config) + _, torchao_quant_type_name = _get_quantization_config_spec(method) + return _build_torchao_pipeline_quant_config( + torchao_quant_type_name, + transformer_components, + ) diff --git a/vllm_omni/diffusion/models/dmd2/mixin.py b/vllm_omni/diffusion/models/dmd2/mixin.py index 48e5719259c..9c98e96b43a 100644 --- a/vllm_omni/diffusion/models/dmd2/mixin.py +++ b/vllm_omni/diffusion/models/dmd2/mixin.py @@ -10,7 +10,7 @@ from vllm_omni.diffusion.models.dmd2.config import DMD2Config from vllm_omni.diffusion.models.schedulers import DMD2EulerScheduler from vllm_omni.diffusion.models.utils import _load_json -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch logger = logging.getLogger(__name__) @@ -35,55 +35,54 @@ def __init_dmd2__(self) -> None: stochastic_sampling=(self.dmd2_config.solver == "sde"), ) - def _sanitize_dmd2_request(self, req: OmniDiffusionRequest) -> None: - """Sanitize CFG-related fields in-place. Mutates req.sampling_params and req.prompts.""" - sp = req.sampling_params - - if sp.num_inference_steps and sp.num_inference_steps != self.dmd2_config.num_inference_steps: - logger.warning( - "DMD2: ignoring num_inference_steps=%d, forcing %d.", - sp.num_inference_steps, - self.dmd2_config.num_inference_steps, - ) - sp.num_inference_steps = self.dmd2_config.num_inference_steps - - if sp.guidance_scale_provided and sp.guidance_scale != self.dmd2_config.guidance_scale: - logger.warning( - "DMD2: ignoring guidance_scale=%.2f, forcing %.2f.", - sp.guidance_scale, - self.dmd2_config.guidance_scale, - ) - sp.guidance_scale = self.dmd2_config.guidance_scale - sp.guidance_scale_provided = False - - if sp.guidance_scale_2 is not None: - logger.warning("DMD2: ignoring guidance_scale_2.") - sp.guidance_scale_2 = None - - if sp.true_cfg_scale is not None: - logger.warning("DMD2: ignoring true_cfg_scale.") - sp.true_cfg_scale = None - - sp.do_classifier_free_guidance = False - sp.is_cfg_negative = False - - # defense: strip scheduler-override extra_args that would let the base pipeline - # (e.g. Wan22Pipeline.forward) rebuild self.scheduler mid-forward and clobber DMD2EulerScheduler. - extra_args = getattr(sp, "extra_args", None) or {} - for key in ("sample_solver", "flow_shift"): - if key in extra_args: - logger.warning("DMD2: ignoring extra_args.%s.", key) - extra_args.pop(key) - - fixed = [] - for p in req.prompts: + def _sanitize_dmd2_request(self, req) -> None: + """Sanitize CFG-related fields in-place. Works with both OmniDiffusionRequest and DiffusionRequestBatch.""" + sampling_params = req.sampling_params_list if isinstance(req, DiffusionRequestBatch) else [req.sampling_params] + for sp in sampling_params: + if sp.num_inference_steps and sp.num_inference_steps != self.dmd2_config.num_inference_steps: + logger.warning( + "DMD2: ignoring num_inference_steps=%d, forcing %d.", + sp.num_inference_steps, + self.dmd2_config.num_inference_steps, + ) + sp.num_inference_steps = self.dmd2_config.num_inference_steps + + if sp.guidance_scale_provided and sp.guidance_scale != self.dmd2_config.guidance_scale: + logger.warning( + "DMD2: ignoring guidance_scale=%.2f, forcing %.2f.", + sp.guidance_scale, + self.dmd2_config.guidance_scale, + ) + sp.guidance_scale = self.dmd2_config.guidance_scale + sp.guidance_scale_provided = False + + if sp.guidance_scale_2 is not None: + logger.warning("DMD2: ignoring guidance_scale_2.") + sp.guidance_scale_2 = None + + if sp.true_cfg_scale is not None: + logger.warning("DMD2: ignoring true_cfg_scale.") + sp.true_cfg_scale = None + + sp.do_classifier_free_guidance = False + sp.is_cfg_negative = False + + # Prevent base pipelines from replacing the DMD2 scheduler. + extra_args = getattr(sp, "extra_args", None) or {} + for key in ("sample_solver", "flow_shift"): + if key in extra_args: + logger.warning("DMD2: ignoring extra_args.%s.", key) + extra_args.pop(key) + + # Strip negative_prompt from each request's prompt in-place. + requests = req.requests if hasattr(req, "requests") else [req] + for request in requests: + p = request.prompt if isinstance(p, dict) and "negative_prompt" in p: logger.warning("DMD2: ignoring negative_prompt.") - p = {k: v for k, v in p.items() if k != "negative_prompt"} - fixed.append(p) - req.prompts = fixed + request.prompt = {k: v for k, v in p.items() if k != "negative_prompt"} - def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch, **kwargs) -> list[DiffusionOutput]: self._sanitize_dmd2_request(req) kwargs.pop("guidance_scale", None) kwargs.pop("num_inference_steps", None) diff --git a/vllm_omni/diffusion/models/dreamid_omni/fusion.py b/vllm_omni/diffusion/models/dreamid_omni/fusion.py index 42393542712..acf24dd2f64 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/fusion.py +++ b/vllm_omni/diffusion/models/dreamid_omni/fusion.py @@ -3,11 +3,13 @@ import torch import torch.nn as nn +from cache_dit import ForwardPattern from vllm.distributed import get_tensor_model_parallel_world_size from vllm.logger import init_logger from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -251,6 +253,13 @@ def forward( class FusionModel(nn.Module): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "fused_blocks": ForwardPattern.Pattern_0, + }, + has_separate_cfg=True, + ) + _layerwise_offload_blocks_attrs = ["fused_blocks"] packed_modules_mapping = { diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py index 47b905de719..5555e370814 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py @@ -5,6 +5,7 @@ import math import os from collections.abc import Iterable +from typing import ClassVar import torch import torch.distributed @@ -23,8 +24,8 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin 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 SupportAudioInput, SupportImageInput -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput, SupportsComponentDiscovery +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch try: from dreamid_omni.utils.divisible_crop import DivisibleCrop @@ -95,9 +96,15 @@ def post_process(output): } -class DreamIDOmniPipeline(nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput): +class DreamIDOmniPipeline( + nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput, SupportsComponentDiscovery +): """DreamID-Omni pipeline for vLLM-Omni.""" + _dit_modules: ClassVar[list[str]] = ["model"] + _encoder_modules: ClassVar[list[str]] = ["text_model"] + _vae_modules: ClassVar[list[str]] = ["vae_model_video", "vae_model_audio"] + def __init__( self, *, @@ -475,7 +482,7 @@ def encode_prompt( def forward( self, - request: OmniDiffusionRequest, + request: DiffusionRequestBatch, **kwargs, ) -> DiffusionOutput: """Main forward pass for DreamID-Omni pipeline for R2AV task.""" diff --git a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py index e187054e2f6..ca92538918f 100644 --- a/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py +++ b/vllm_omni/diffusion/models/dreamzero/causal_wan_model.py @@ -21,11 +21,12 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - get_tp_group, + tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.conv import Conv3dLayer from vllm.model_executor.layers.linear import ( ColumnParallelLinear, + QKVParallelLinear, RowParallelLinear, ) from vllm.model_executor.utils import set_weight_attrs @@ -36,6 +37,20 @@ MultiEmbodimentActionEncoder, ) +# AR-Diffusion paged self-attention (in-tree experimental engine). Import at +# module level so the isinstance check + custom-op call trace cleanly inside +# the fullgraph-compiled DiT block (an import inside the traced region would +# graph-break). The model still works without the engine: the payload type is +# only ever constructed by the AR-Diffusion runner. +try: + from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( + ARDiffusionPagedLayerInputs, + paged_write_attn, + ) +except ImportError: # pragma: no cover - experimental package always ships in-tree + ARDiffusionPagedLayerInputs = None + paged_write_attn = None + # ── RoPE utilities ────────────────────────────────────────────────── @@ -174,22 +189,82 @@ def weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> Non raise ValueError(f"RMSNorm shard shape mismatch: param={tuple(param.shape)}, shard={tuple(shard.shape)}.") param.data.copy_(shard) - def forward(self, x: torch.Tensor) -> torch.Tensor: - tp_size = get_tensor_model_parallel_world_size() + def _local_sum_sq(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]: + """Per-rank float32 activation, local sum-of-squares, and local width.""" x_float = x.float() local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True) - local_count = x.shape[-1] + return x_float, local_sum_sq, x.shape[-1] + + def _scale( + self, + x_float: torch.Tensor, + global_sum_sq: torch.Tensor, + global_count: int, + x: torch.Tensor, + ) -> torch.Tensor: + """Apply the (already reduced) RMS and the per-rank weight shard.""" + mean_sq = global_sum_sq / global_count + return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight + + def forward(self, x: torch.Tensor) -> torch.Tensor: + tp_size = get_tensor_model_parallel_world_size() + x_float, local_sum_sq, local_count = self._local_sum_sq(x) if tp_size > 1: - global_sum_sq = local_sum_sq.clone() - torch.distributed.all_reduce(global_sum_sq, group=get_tp_group().device_group) + # Use vLLM's collective (custom all-reduce / symmetric-mem fast path + # for small tensors) instead of raw torch.distributed.all_reduce, and + # take the return value so the custom-AR path (which may return a new + # buffer) is handled correctly. No .clone() needed. + global_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq) global_count = local_count * tp_size else: global_sum_sq = local_sum_sq global_count = local_count - mean_sq = global_sum_sq / global_count - return (x_float * torch.rsqrt(mean_sq + self.eps)).type_as(x) * self.weight + return self._scale(x_float, global_sum_sq, global_count, x) + + +def fused_qk_rms_norm( + norm_q: nn.Module, + norm_k: nn.Module, + q: torch.Tensor, + k: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Apply q/k :class:`DistributedRMSNorm` with a SINGLE fused TP all-reduce. + + WHY: self-attention norms q and k every step; on TP each norm issues its own + tiny all-reduce of the per-token sum-of-squares. For DreamZero's decode-like + per-step forward these latency-bound micro-collectives dominate, so we pack + both sum-of-squares into one tensor and reduce once (2 collectives → 1). + + NUMERICALLY IDENTICAL to ``norm_q(q), norm_k(k)``: all-reduce is elementwise, + so packing along the last dim reduces each slice independently with the same + fp32 accumulation. Requires q and k to share the same shape (true for + self-attention — both come from the same hidden states). Falls back to + independent application when either norm is not a DistributedRMSNorm + (e.g. nn.Identity when qk_norm=False). + """ + if not (isinstance(norm_q, DistributedRMSNorm) and isinstance(norm_k, DistributedRMSNorm)): + return norm_q(q), norm_k(k) + + # The fused path reduces one packed sum-of-squares and reuses q's token width + # as the RMS count for BOTH q and k, so q and k must share the same shape. + # (Self-attention guarantees this: q and k are projected from the same x.) + assert q.shape == k.shape, "fused_qk_rms_norm requires q and k to have the same shape." + + tp_size = get_tensor_model_parallel_world_size() + q_float, q_sum_sq, count = norm_q._local_sum_sq(q) + k_float, k_sum_sq, _ = norm_k._local_sum_sq(k) + + if tp_size > 1: + packed = torch.cat([q_sum_sq, k_sum_sq], dim=-1) + packed = tensor_model_parallel_all_reduce(packed) + q_sum_sq, k_sum_sq = packed[..., 0:1], packed[..., 1:2] + count = count * tp_size + + q_out = norm_q._scale(q_float, q_sum_sq, count, q) + k_out = norm_k._scale(k_float, k_sum_sq, count, k) + return q_out, k_out # ── Projections ───────────────────────────────────────────────────── @@ -337,22 +412,35 @@ def forward( context = context[:, 257:] n, d = self.tp_num_heads, self.head_dim q = self.norm_q(self.q(x)).unflatten(2, (n, d)) + # context (text) and context_img (clip features) are constant within a + # session, so k/v and k_img/v_img are cached on first call and reused. + # context_img == img_emb(clip feature) is session-invariant (state.clip_feas + # is set once at current_start_frame==0 and only cleared on reset, which also + # rebuilds crossattn_cache). Caching k_img/v_img removes a per-step + # norm_k_img all-reduce plus the k_img/v_img projection GEMMs in steady state. + # Only q depends on the per-step x and is always recomputed. if crossattn_cache is not None: if not crossattn_cache["is_init"]: crossattn_cache["is_init"] = True k = self.norm_k(self.k(context)).unflatten(2, (n, d)) v = self.v(context).unflatten(2, (n, d)) + k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d)) + v_img = self.v_img(context_img).unflatten(2, (n, d)) crossattn_cache["k"] = k crossattn_cache["v"] = v + crossattn_cache["k_img"] = k_img + crossattn_cache["v_img"] = v_img else: k = crossattn_cache["k"] v = crossattn_cache["v"] + k_img = crossattn_cache["k_img"] + v_img = crossattn_cache["v_img"] else: k = self.norm_k(self.k(context)).unflatten(2, (n, d)) v = self.v(context).unflatten(2, (n, d)) + k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d)) + v_img = self.v_img(context_img).unflatten(2, (n, d)) x = self.attn(q, k, v) - k_img = self.norm_k_img(self.k_img(context_img)).unflatten(2, (n, d)) - v_img = self.v_img(context_img).unflatten(2, (n, d)) img_x = self.attn(q, k_img, v_img) x = x.flatten(2) img_x = img_x.flatten(2) @@ -403,9 +491,21 @@ def __init__( self.num_action_per_block = num_action_per_block self.num_state_per_block = num_state_per_block self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen - self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False) - self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False) - self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False, return_bias=False) + # Fused QKV projection: q/k/v all come from x, so a single column-parallel + # GEMM replaces three (fewer launches / better GEMM util in the decode-like + # per-step forward). No GQA here: total_num_kv_heads defaults to num_heads. + self.qkv = QKVParallelLinear( + hidden_size=dim, + head_size=self.head_dim, + total_num_heads=num_heads, + bias=True, + ) + # The forward splits qkv into three equal q/k/v shards, which is only + # correct without GQA. total_num_kv_heads defaults to total_num_heads + # here; fail loud if a future checkpoint introduces GQA so the equal + # split does not silently misalign k/v. + if self.qkv.total_num_kv_heads != self.qkv.total_num_heads: + raise ValueError("Self-attn QKV fusion requires no GQA (total_num_kv_heads == total_num_heads).") self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True, return_bias=False) self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity() @@ -424,16 +524,22 @@ def forward( freqs_action: torch.Tensor, freqs_state: torch.Tensor, action_register_length: int | None, - kv_cache: torch.Tensor | None = None, + kv_cache: torch.Tensor | Any | None = None, current_start_frame: int = 0, is_tf: bool = True, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Inference-only forward (KV cache path).""" n, d = self.tp_num_heads, self.head_dim - q = self.norm_q(self.q(x)).unflatten(2, (n, d)) - k = self.norm_k(self.k(x)).unflatten(2, (n, d)) - v = self.v(x).unflatten(2, (n, d)) + # Single fused QKV GEMM, then split into the per-rank q/k/v shards. + qkv, _ = self.qkv(x) + qk_size = self.tp_num_heads * self.head_dim + q, k, v = qkv.split([qk_size, qk_size, qk_size], dim=-1) + # Fuse q/k qk-norm into a single TP all-reduce (both come from x, same shape). + q, k = fused_qk_rms_norm(self.norm_q, self.norm_k, q, k) + q = q.unflatten(2, (n, d)) + k = k.unflatten(2, (n, d)) + v = v.unflatten(2, (n, d)) updated_kv_cache: torch.Tensor | None = None @@ -475,24 +581,44 @@ def forward( action_v = v[:, -action_register_length:] v = v[:, :-action_register_length] - updated_k = kv_cache[0] - updated_v = kv_cache[1] - new_k = torch.cat([updated_k, roped_key], dim=1) - new_v = torch.cat([updated_v, v], dim=1) - new_k = new_k[:, -self.max_attention_size :] - new_v = new_v[:, -self.max_attention_size :] - - if action_register_length is not None: - q_cat = torch.cat([roped_query, roped_action_query], dim=1) - k_cat = torch.cat([new_k, roped_action_key], dim=1) - v_cat = torch.cat([new_v, action_v], dim=1) + if ARDiffusionPagedLayerInputs is not None and isinstance(kv_cache, ARDiffusionPagedLayerInputs): + # Fused write+attend custom op: one opaque node in the compiled + # graph (slot writes + FlashAttention block-table kernel inside). + # Metadata tensors were prepared once per forward in _forward_blocks. + if action_register_length is not None: + q_cat = torch.cat([roped_query, roped_action_query], dim=1) + k_act, v_act = roped_action_key[0], action_v[0] + else: + q_cat = roped_query + k_act = v_act = None + x = paged_write_attn( + kv_cache, + q_cat[0], + roped_key[0], + v[0], + k_act, + v_act, + self.head_dim**-0.5, + ).unsqueeze(0) else: - q_cat = roped_query - k_cat = new_k - v_cat = new_v + updated_k = kv_cache[0] + updated_v = kv_cache[1] + new_k = torch.cat([updated_k, roped_key], dim=1) + new_v = torch.cat([updated_v, v], dim=1) + new_k = new_k[:, -self.max_attention_size :] + new_v = new_v[:, -self.max_attention_size :] + + if action_register_length is not None: + q_cat = torch.cat([roped_query, roped_action_query], dim=1) + k_cat = torch.cat([new_k, roped_action_key], dim=1) + v_cat = torch.cat([new_v, action_v], dim=1) + else: + q_cat = roped_query + k_cat = new_k + v_cat = new_v - x = self.attn(q_cat, k_cat, v_cat) - updated_kv_cache = torch.stack([new_k, new_v], dim=0) + x = self.attn(q_cat, k_cat, v_cat) + updated_kv_cache = torch.stack([new_k, new_v], dim=0) x = x.flatten(2) x = self.o(x) @@ -554,7 +680,7 @@ def forward( freqs_state: torch.Tensor, context: torch.Tensor, action_register_length: int | None = None, - kv_cache: torch.Tensor | None = None, + kv_cache: torch.Tensor | Any | None = None, crossattn_cache: dict | None = None, current_start_frame: int = 0, is_tf: bool = True, @@ -833,9 +959,10 @@ def _forward_blocks( action: torch.Tensor | None, timestep_action: torch.Tensor | None, state: torch.Tensor | None, - kv_cache: list[torch.Tensor], + kv_cache: list[Any], current_start_frame: int, - ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + crossattn_cache: list[dict] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor | None]]: x = x.flatten(start_dim=2).transpose(1, 2) B = x.shape[0] F_t = timestep.shape[1] @@ -873,7 +1000,28 @@ def _forward_blocks( clip_embedding = self.img_emb(clip_feature) context = torch.cat([clip_embedding, context], dim=1) - updated_kv_caches: list[torch.Tensor] = [] + # AR-Diffusion paged path: host-side prep ONCE per branch forward, outside + # the compiled blocks — allocate current video/action slots, build the + # padded block-table metadata all 40 layers share, and hand the compiled + # region plain tensors (ARDiffusionPagedLayerInputs) instead of a Python + # context object. Lazy-allocation contract preserved: only the branch this + # CFG-parallel rank executes reaches its _forward_blocks. + if kv_cache and getattr(kv_cache[0], "is_ar_diffusion_paged_context", False): + if B != 1: + raise RuntimeError("AR-Diffusion paged self-attention currently expects batch_size=1") + fctx = kv_cache[0].forward_ctx + if seq_len != fctx.seq_len: + raise RuntimeError( + f"AR-Diffusion paged context seq_len={fctx.seq_len} but current video KV has {seq_len} tokens" + ) + fctx.prepare( + device=x.device, + action_len=int(action_register_length or 0), + query_len=int(x.shape[1]), + ) + kv_cache = [c.to_layer_inputs() for c in kv_cache] + + updated_kv_caches: list[torch.Tensor | None] = [] for block_index, block in enumerate(self.blocks): x, updated_kv_cache = block( x=x, @@ -884,6 +1032,7 @@ def _forward_blocks( context=context, action_register_length=action_register_length, kv_cache=kv_cache[block_index] if kv_cache else None, + crossattn_cache=crossattn_cache[block_index] if crossattn_cache else None, current_start_frame=current_start_frame, ) updated_kv_caches.append(updated_kv_cache) @@ -906,9 +1055,9 @@ def _forward_inference( timestep: torch.Tensor, context: torch.Tensor, seq_len: int, - kv_cache: list[torch.Tensor], - crossattn_cache: list[torch.Tensor], + kv_cache: list[Any], current_start_frame: int, + crossattn_cache: list[dict] | None = None, y: torch.Tensor | None = None, clip_feature: torch.Tensor | None = None, action: torch.Tensor | None = None, @@ -941,6 +1090,7 @@ def _forward_inference( timestep_action=timestep_action, state=state, kv_cache=kv_cache, + crossattn_cache=crossattn_cache, current_start_frame=current_start_frame, ) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 30ea2a4cebd..107638f8179 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -3,7 +3,7 @@ """DreamZero pipeline for vllm-omni. -Entry point for DiffusionEngine.step() → pipeline.forward(req) +Entry point for DiffusionEngine.step() -> pipeline.forward(req) """ from __future__ import annotations @@ -23,6 +23,10 @@ from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.cache.stepcache import ( + get_stepcache_state, + is_stepcache_active, +) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( DistributedAutoencoderKLWan, @@ -48,17 +52,12 @@ DEFAULT_SIGMA_SHIFT, ) from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.platforms import current_omni_platform +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 -# --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- - - class VideoActionScheduler: """Wraps video + action schedulers into single .step() interface.""" @@ -94,17 +93,102 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): Multi-output: predict_noise() returns (video_pred, action_pred). CFG: video gets standard CFG, action takes positive branch only. - State: DreamZeroState manages KV cache + frame buffer across forward() calls. + + KV is managed by the AR-Diffusion engine: ``self._ar_diffusion_kv_state`` is set by the + runner before ``forward()`` and the pipeline routes all KV access (get / update / + commit / reset) through the pool-backed state. Purely duck-typed (no engine import). """ + _ar_diffusion_kv_state = None # set by the runner before each forward + + def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False): + return self._ar_diffusion_kv_state.get_kv_caches( + is_negative, + seq_len=seq_len, + commit_current=update_kv_cache, + ) + + def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim): + # The engine owns all KV allocation: self-attn is allocated lazily from + # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross. + # Nothing is created model-side. + return + + def _kv_commit(self, is_negative: bool): + self._ar_diffusion_kv_state.commit_paged_context(is_negative) + + def _kv_get_cross(self, state, is_negative): + """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img).""" + return self._ar_diffusion_kv_state.get_cross_kv_caches(is_negative) + + def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None: + """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool. + + Caches the session-invariant cross-attn projections once, per half: the text + ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets + (prompt unchanged within a session — only session resets clear them), while + the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the + 257 image tokens the forward splits off, cached model-side by #4154) are + re-projected on every window restart from the fresh CLIP features. Must run + after the image is encoded so ``clip_feature`` is available. + """ + s = self._ar_diffusion_kv_state + need_text = not s._cross_text_populated.get(is_negative, False) + need_img = ( + clip_feature is not None + and getattr(self.transformer, "model_type", "t2v") == "i2v" + and not s._cross_img_populated.get(is_negative, False) + ) + if not need_text and not need_img: + return + projected = self.transformer.text_embedding(context) if need_text else None + img_ctx = self.transformer.img_emb(clip_feature) if need_img else None + for i, block in enumerate(self.transformer.blocks): + ca = block.cross_attn + n, d = ca.tp_num_heads, ca.head_dim + k = v = None + if projected is not None: + k = ca.norm_k(ca.k(projected)).unflatten(2, (n, d)) + v = ca.v(projected).unflatten(2, (n, d)) + k_img = v_img = None + if img_ctx is not None: + k_img = ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d)) + v_img = ca.v_img(img_ctx).unflatten(2, (n, d)) + s.kv_cache.write_cross_kv(i, is_negative, k, v, k_img, v_img) + if need_text: + s._cross_text_populated[is_negative] = True + if need_img: + s._cross_img_populated[is_negative] = True + logger.info( + "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s", + "neg" if is_negative else "pos", + len(self.transformer.blocks), + "kept" if projected is None else tuple(context.shape), + None if img_ctx is None else tuple(img_ctx.shape), + ) + + def _kv_reset(self, state, *, clear_video_latents: bool = True): + """Reset the engine's pooled session window plus the model's non-KV state. + + DreamZero resets at the attention-window boundary; the engine pool drops the + same window so the next forward starts fresh. ``clear_video_latents=False`` + keeps the accumulated video latents for export. + + ``clear_video_latents=False`` also marks a window ("inference") reset: the + prompt is unchanged, so the pool keeps the text cross-attn K/V and only the + image half repopulates on the restart forward. + """ + state.reset(clear_video_latents=clear_video_latents) + self._ar_diffusion_kv_state.reset(keep_cross_text=not clear_video_latents) + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: """Initialize pipeline components. DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID): - config.json — root config (action_head_cfg, architectures, etc.) - model-*.safetensors — all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*) - experiment_cfg/metadata.json — per-embodiment action normalization stats - vae/ — symlink to Wan2.1 VAE (diffusers-compatible) + config.json -- root config (action_head_cfg, architectures, etc.) + model-*.safetensors -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*) + experiment_cfg/metadata.json -- per-embodiment action normalization stats + vae/ -- symlink to Wan2.1 VAE (diffusers-compatible) Components are instantiated from config (not from_pretrained), then filled by load_weights() which reads root safetensors and remaps key prefixes. @@ -119,6 +203,18 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: """ super().__init__() + # DreamZero is engine-only: every KV access in forward() routes through + # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale + # or programmatic config that leaves engine_backend="default" would + # otherwise only crash mid-forward on the first KV access. + engine_backend = str(getattr(od_config, "engine_backend", "") or "") + if "ar_diffusion" not in engine_backend.lower().replace("-", "_"): + raise ValueError( + "DreamZeroPipeline requires the AR-Diffusion engine; set " + "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine " + f"in the deploy config (got engine_backend={engine_backend!r})." + ) + model_path = od_config.model model_config = od_config.model_config local_files_only = os.path.exists(model_path) @@ -205,6 +301,9 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: self._max_session_states = MAX_DREAMZERO_SESSIONS self.state = self._get_or_create_state("default") + # DiT step cache is configured by StepCacheBackend + # (cache_backend="step_cache") via pipeline._stepcache_config. + # Keep runtime inference settings separate from the training-time config. self.num_inference_steps: int = model_config.get( "num_inference_steps", @@ -226,8 +325,11 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: self.max_action_dim: int = ah_config["max_action_dim"] self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT) + # The negative prompt is a model constant: encode it once, lazily, and + # reuse across every forward/session (UMT5 encode is deterministic). + self._negative_prompt_embeds_cache: torch.Tensor | None = None - # Embodiment name → numeric ID mapping (model knowledge) + # Embodiment name -> numeric ID mapping (model knowledge) self.embodiment_name_to_id: dict[str, int] = model_config.get( "embodiment_name_to_id", DEFAULT_EMBODIMENT_NAME_TO_ID, @@ -306,6 +408,27 @@ def _load_repo_json(model_path: str, relative_path: str, local_files_only: bool) def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: """Call CausalWanModel, return (video_pred, action_pred).""" + video_pred, action_pred = self._predict_noise_eager(kwargs) + + if is_stepcache_active(self): + video_pred = video_pred.clone() + if action_pred is not None: + action_pred = action_pred.clone() + + if action_pred is None: + batch_size = kwargs["hidden_states"].shape[0] + action_pred = torch.empty( + batch_size, + 0, + self.transformer.action_dim, + device=video_pred.device, + dtype=video_pred.dtype, + ) + return (video_pred, action_pred) + + def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor | None]: + """Eager DiT forward; also handles KV-cache write-back on prefill.""" + self._cudagraph_mark_step_begin() video_pred, action_pred, updated_kv_caches = self.transformer( x=kwargs["hidden_states"], timestep=kwargs["timestep_video"], @@ -321,25 +444,148 @@ def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: state=kwargs.get("state_features"), embodiment_id=kwargs.get("embodiment_id"), ) - if kwargs.get("update_kv_cache", False) and updated_kv_caches: - state = kwargs.get("dreamzero_state", self.state) + if kwargs.get("update_kv_cache", False): is_neg = kwargs.get("is_negative", False) - for i, kv in enumerate(updated_kv_caches): - state.update_kv_cache(i, kv, is_negative=is_neg) + logger.debug( + "AR-Diffusion pipeline predict_noise -> commit paged context: " + "is_neg=%s seq_len=%s current_start_frame=%s layers=%d", + is_neg, + kwargs.get("seq_len"), + kwargs.get("current_start_frame"), + len(updated_kv_caches), + ) + self._kv_commit(is_neg) - video_pred = video_pred.clone() - if action_pred is not None: - action_pred = action_pred.clone() - else: - batch_size = kwargs["hidden_states"].shape[0] - action_pred = torch.empty( - batch_size, - 0, - self.transformer.action_dim, - device=video_pred.device, - dtype=video_pred.dtype, - ) # CFG-parallel-safe dummy action pred - return (video_pred, action_pred) + return video_pred, action_pred + + # ----------------------------------------------------------------------- + # torch.compile setup (paper D.2 extended: encoders + VAE + DiT blocks) + # ----------------------------------------------------------------------- + + def setup_compile(self) -> None: + """Compile DreamZero encoders, VAE, and per-block DiT for inference. + + Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False`` + on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and + compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop). + Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because + Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture. + DiT blocks use per-block ``fullgraph=True``. + """ + if not torch.cuda.is_available(): + logger.info("DreamZero setup_compile skipped: CUDA not available.") + return + + from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import ( + apply_wan_vae_feat_cache_tensor_patch, + ) + + apply_wan_vae_feat_cache_tensor_patch() + + compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False} + # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead. + # The AR-Diffusion paged self-attention is a registered custom op, so the + # block stays fullgraph even on that path. + dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False} + + logger.info( + "DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)." + ) + + try: + self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro) + except Exception as exc: + logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc) + + try: + self.image_encoder.model.visual.forward = torch.compile( + self.image_encoder.model.visual.forward, + **compile_ro, + ) + except Exception as exc: + logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc) + + try: + self.vae._encode = torch.compile(self.vae._encode, **compile_ro) + except Exception as exc: + logger.warning("DreamZero: vae._encode compile failed (%s); skipping.", exc) + + compiled_blocks = 0 + for block in self.transformer.blocks: + try: + block.forward = torch.compile(block.forward, **dit_compile) + compiled_blocks += 1 + except Exception as exc: + logger.warning( + "DreamZero: transformer block %d compile failed (%s); leaving remaining eager.", + compiled_blocks, + exc, + ) + break + if compiled_blocks: + logger.info( + "DreamZero: compiled %d/%d transformer blocks.", + compiled_blocks, + len(self.transformer.blocks), + ) + + self.warmup_compile() + + def warmup_compile(self) -> None: + """Warm up compiled text/image/VAE paths before timed inference.""" + if not torch.cuda.is_available(): + return + + device = next(self.text_encoder.parameters()).device + with torch.inference_mode(): + try: + text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device) + attention_mask = torch.ones_like(text_tokens) + self._encode_text(text_tokens, attention_mask) + except Exception as exc: + logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc) + + try: + image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device) + self._encode_image(image, self.num_frames, 180, 320, state=self.state) + except Exception as exc: + logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc) + + try: + self.state.reset_vae_encoder_stream() + dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device) + self._vae_stream_seed(self.state, dummy_video[:, :, :1]) + for _ in range(4): + self._vae_stream_append_frame(self.state, dummy_video[:, :, :1]) + self._vae_stream_get_observation_latents( + self.state, + self.num_frame_per_block, + dtype=torch.bfloat16, + ) + except Exception as exc: + logger.warning("DreamZero compile warmup (vae encode stream) skipped: %s", exc) + + try: + latent_h, latent_w = 180 // 8, 320 // 8 + dummy_latent = torch.zeros( + 1, + 16, + self.num_frame_per_block, + latent_h * 2, + latent_w * 2, + dtype=self.vae.dtype, + device=device, + ) + mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype) + inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype) + dummy_latent_denorm = dummy_latent / inv_std + mean + self._cudagraph_mark_step_begin() + self.vae.decode(dummy_latent_denorm, return_dict=False) + except Exception as exc: + logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc) + + torch.accelerator.synchronize(device) + logger.info("DreamZero compile warmup finished (text / image / vae decode).") def combine_cfg_noise( self, @@ -364,10 +610,12 @@ def _synchronize_cfg_parallel_step_output( latents: tuple[torch.Tensor, torch.Tensor], do_true_cfg: bool, ) -> tuple[torch.Tensor, torch.Tensor]: - """Post-step sync: .contiguous() + platform synchronize.""" + """Post-step sync: .contiguous() + cuda.synchronize()""" latents = tuple(t.contiguous() for t in latents) if do_true_cfg and get_classifier_free_guidance_world_size() > 1: - current_omni_platform.synchronize() + device = next((t.device for t in latents if t.is_cuda), None) + if device is not None: + torch.cuda.current_stream(device).synchronize() return latents # ----------------------------------------------------------------------- @@ -375,7 +623,7 @@ def _synchronize_cfg_parallel_step_output( # ----------------------------------------------------------------------- def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor: - """uint8 [B,T,H,W,C] → bfloat16 [B,C,T,H,W] normalized to [-1,1].""" + """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1].""" videos = videos.permute(0, 4, 1, 2, 3) if videos.dtype == torch.uint8: videos = videos.float() / 255.0 @@ -388,12 +636,20 @@ def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor: videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) return videos.to(dtype=torch.bfloat16) + @staticmethod + def _cudagraph_mark_step_begin() -> None: + try: + torch.compiler.cudagraph_mark_step_begin() + except Exception: + pass + # ----------------------------------------------------------------------- # Text encoding # ----------------------------------------------------------------------- def _encode_text(self, text_tokens: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: """Encode text prompt via UMT5.""" + self._cudagraph_mark_step_begin() seq_lens = attention_mask.gt(0).sum(dim=1).long() prompt_emb = self.text_encoder( text_tokens, @@ -414,6 +670,8 @@ def _encode_image( num_frames: int, height: int, width: int, + *, + state: DreamZeroState | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Encode first frame via CLIP + VAE. Returns: (clip_feas, ys, image_latent) @@ -422,6 +680,7 @@ def _encode_image( batch_size = image.shape[0] with torch.amp.autocast(dtype=torch.bfloat16, device_type=device.type): + self._cudagraph_mark_step_begin() clip_context = self.image_encoder.encode_image(image) msk = torch.ones(batch_size, num_frames, height // 8, width // 8, device=device) @@ -448,11 +707,147 @@ def _encode_image( new_image = y[:, :, 0:1] y = torch.concat([msk, y], dim=1) + if state is not None: + if not state.vae_stream_initialized: + # Seed the AR streaming encoder after the compiled full encode above; + # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph. + self._vae_stream_seed(state, image_input[:, :, :1]) + else: + # Window ("inference") restart with a live stream: keep the real + # frame history in the Wan feat_cache and append the restart + # observation instead of reseeding from scratch. + self._cudagraph_mark_step_begin() + self._vae_stream_append_frame(state, image_input[:, :, :1]) + return clip_context, y, new_image + def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor: + if self.vae.config.patch_size is not None: + from diffusers.models.autoencoders.autoencoder_kl_wan import patchify + + return patchify(videos, patch_size=self.vae.config.patch_size) + return videos + + @staticmethod + def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]: + return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map] + + def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]: + self.vae.clear_cache() + return self._vae_clone_feat_map(self.vae._enc_feat_map) + + def _vae_encode_encoder_chunk( + self, + chunk: torch.Tensor, + feat_map: list[torch.Tensor | None], + ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: + """Run one Wan encoder chunk while mutating ``feat_map`` in place. + + Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates + conflict with ``reduce-overhead`` CUDAGraph capture. + """ + self.vae._enc_feat_map = feat_map + self.vae._enc_conv_idx = [0] + out = self.vae.encoder( + chunk, + feat_cache=self.vae._enc_feat_map, + feat_idx=self.vae._enc_conv_idx, + ) + return out, self._vae_clone_feat_map(self.vae._enc_feat_map) + + def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor: + enc = self.vae.quant_conv(encoder_out) + mu, _ = enc.chunk(2, dim=1) + mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype) + inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype) + return (mu - mean) * inv_std + + def _vae_stream_seed(self, state: DreamZeroState, first_frame: torch.Tensor) -> None: + """Seed incremental VAE encode with the first observation frame.""" + state.reset_vae_encoder_stream() + feat_map = self._vae_init_enc_feat_map() + chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype)) + self._cudagraph_mark_step_begin() + encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map) + state.vae_enc_feat_map = feat_map + state.vae_encoder_out = encoder_out + state.vae_pending_body_frames = None + state.vae_stream_initialized = True + + def _vae_stream_append_frame(self, state: DreamZeroState, new_frame: torch.Tensor) -> None: + """Append one pixel frame and encode a 4-frame body chunk when ready.""" + if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None: + raise RuntimeError("VAE encoder stream is not initialized.") + + frame = new_frame.to(dtype=self.vae.dtype) + if state.vae_pending_body_frames is None: + state.vae_pending_body_frames = frame + else: + state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2) + + while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4: + body = state.vae_pending_body_frames[:, :, :4] + if state.vae_pending_body_frames.shape[2] > 4: + state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:] + else: + state.vae_pending_body_frames = None + chunk = self._vae_patchify(body) + encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map) + state.vae_encoder_out = torch.cat([state.vae_encoder_out, encoder_chunk], dim=2) + state.vae_enc_feat_map = feat_map + + def _vae_stream_get_observation_latents( + self, + state: DreamZeroState, + num_latent_frames: int, + *, + dtype: torch.dtype, + ) -> torch.Tensor: + if state.vae_encoder_out is None: + raise RuntimeError("VAE encoder stream has no accumulated encoder output.") + latents = self._vae_quantize_encoder_out(state.vae_encoder_out).to(dtype=dtype) + if latents.shape[2] >= num_latent_frames: + return latents[:, :, -num_latent_frames:] + pad_count = num_latent_frames - latents.shape[2] + pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1) + return torch.cat([pad, latents], dim=2) + + def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor: + _, _, num_frames_raw, _, _ = videos.shape + if (num_frames_raw - 1) // 4 == self.num_frame_per_block: + return videos + if num_frames_raw // 4 != self.num_frame_per_block: + repeat_factor = self.num_frame_per_block // (num_frames_raw // 4) + videos = torch.repeat_interleave(videos, repeat_factor, dim=2) + first_frame = videos[:, :, 0:1] + return torch.cat([first_frame, videos], dim=2) + first_frame = videos[:, :, 0:1] + return torch.cat([first_frame, videos], dim=2) + + def _encode_observation_latents( + self, + state: DreamZeroState, + videos: torch.Tensor, + *, + latent_dtype: torch.dtype, + ) -> torch.Tensor: + """Encode current robot observation into normalized VAE latents.""" + if state.vae_stream_initialized: + self._cudagraph_mark_step_begin() + self._vae_stream_append_frame(state, videos[:, :, -1:]) + return self._vae_stream_get_observation_latents( + state, + self.num_frame_per_block, + dtype=latent_dtype, + ) + + videos = self._preprocess_vae_observation_window(videos) + return self._encode_vae_latents(videos).to(dtype=latent_dtype) + def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor: """Encode videos into normalized VAE latents.""" input_dtype = videos.dtype + self._cudagraph_mark_step_begin() hidden = self.vae._encode(videos.to(dtype=self.vae.dtype)) mu, _ = hidden.chunk(2, dim=1) mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype) @@ -469,8 +864,23 @@ def decode_video_latents(self, video_latents: torch.Tensor) -> torch.Tensor: inv_std = self.vae_latents_inv_std.to(device=vae_device, dtype=vae_dtype) latents = latents / inv_std + mean with torch.no_grad(): + self._cudagraph_mark_step_begin() return self.vae.decode(latents, return_dict=False)[0] + def decode_accumulated_video_latents(self, session_id: str | None = None) -> torch.Tensor: + """Decode all AR-chunk latents accumulated for ``session_id``.""" + state = self._get_or_create_state(session_id) + latents = state.get_concatenated_video_latents() + if latents is None: + session_key = str(session_id or "default") + raise RuntimeError(f"No accumulated video latents for session {session_key!r}.") + return self.decode_video_latents(latents) + + def clear_accumulated_video_latents(self, session_id: str | None = None) -> None: + """Clear accumulated video latents for ``session_id`` without resetting KV state.""" + state = self._get_or_create_state(session_id) + state.clear_video_latents() + # ----------------------------------------------------------------------- # KV cache prefill # ----------------------------------------------------------------------- @@ -487,7 +897,7 @@ def _prefill_kv_cache( ) -> None: """Prefill KV cache with first frame and/or current observation. - Uses predict_noise_maybe_with_cfg() for CFG parallel — same path as + Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as the denoise loop. The mixin handles rank dispatch automatically. KV cache update happens as a side effect inside predict_noise(). """ @@ -498,7 +908,8 @@ def _prefill_kv_cache( head_dim = self.transformer.dim // self.transformer.num_heads if state.current_start_frame == 0: - state.create_kv_caches( + self._kv_create( + state, batch_size, dtype, device, @@ -523,16 +934,16 @@ def _prefill_kv_cache( ) positive_kwargs = dict( encoder_hidden_states=prompt_embeds, - kv_cache=state.get_kv_caches(False), - crossattn_cache=state.get_crossattn_caches(False), + kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True), + crossattn_cache=self._kv_get_cross(state, False), is_negative=False, **common, ) negative_kwargs = ( dict( encoder_hidden_states=negative_prompt_embeds, - kv_cache=state.get_kv_caches(True), - crossattn_cache=state.get_crossattn_caches(True), + kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True), + crossattn_cache=self._kv_get_cross(state, True), is_negative=True, **common, ) @@ -573,16 +984,16 @@ def _prefill_kv_cache( ) positive_kwargs = dict( encoder_hidden_states=prompt_embeds, - kv_cache=state.get_kv_caches(False), - crossattn_cache=state.get_crossattn_caches(False), + kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True), + crossattn_cache=self._kv_get_cross(state, False), is_negative=False, **common, ) negative_kwargs = ( dict( encoder_hidden_states=negative_prompt_embeds, - kv_cache=state.get_kv_caches(True), - crossattn_cache=state.get_crossattn_caches(True), + kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True), + crossattn_cache=self._kv_get_cross(state, True), is_negative=True, **common, ) @@ -615,8 +1026,8 @@ def diffuse( For each timestep: 1. Build positive_kwargs / negative_kwargs - 2. predict_noise_maybe_with_cfg() → (video_pred, action_pred) - 3. scheduler_step_maybe_with_cfg() → VideoActionScheduler + 2. predict_noise_maybe_with_cfg() -> (video_pred, action_pred) + 3. scheduler_step_maybe_with_cfg() -> VideoActionScheduler 4. _synchronize_cfg_parallel_step_output() """ seq_len = kwargs["seq_len"] @@ -635,6 +1046,12 @@ def diffuse( noisy_input = video_latents noisy_input_action = action_latents + # ---- step_cache (StepCacheBackend / dreamzero.git) ---- + _cached_flow_pred: torch.Tensor | None = None + _cached_flow_pred_action: torch.Tensor | None = None + _prev_predictions: list[tuple[torch.Tensor]] = [] + _step_cache = get_stepcache_state(self) if is_stepcache_active(self) else None + for index in range(len(timesteps_video)): video_timestep = timesteps_video[index] action_timestep = timesteps_action[index] @@ -663,46 +1080,58 @@ def diffuse( else: y = state.ys[:, :, -self.num_frame_per_block :] - positive_kwargs = dict( - hidden_states=noisy_input.transpose(1, 2), - timestep_video=timestep, - encoder_hidden_states=prompt_embeds, - kv_cache=state.get_kv_caches(False), - crossattn_cache=state.get_crossattn_caches(False), - y=y, - clip_feature=state.clip_feas, - action=noisy_input_action, - timestep_action=timestep_action, - is_negative=False, - **common_kwargs, - ) - - # Negative (uncond) kwargs - if do_true_cfg and negative_prompt_embeds is not None: - negative_kwargs = dict( + run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions) + if run_dit: + positive_kwargs = dict( hidden_states=noisy_input.transpose(1, 2), timestep_video=timestep, - encoder_hidden_states=negative_prompt_embeds, - kv_cache=state.get_kv_caches(True), - crossattn_cache=state.get_crossattn_caches(True), + encoder_hidden_states=prompt_embeds, + kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False), + crossattn_cache=self._kv_get_cross(state, False), y=y, clip_feature=state.clip_feas, action=noisy_input_action, timestep_action=timestep_action, - is_negative=True, + is_negative=False, **common_kwargs, ) - else: - negative_kwargs = None - noise_pred = self.predict_noise_maybe_with_cfg( - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - do_true_cfg=do_true_cfg, - true_cfg_scale=self.cfg_scale, - cfg_normalize=False, - ) - flow_pred, flow_pred_action = noise_pred + if do_true_cfg and negative_prompt_embeds is not None: + negative_kwargs = dict( + hidden_states=noisy_input.transpose(1, 2), + timestep_video=timestep, + encoder_hidden_states=negative_prompt_embeds, + kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False), + crossattn_cache=self._kv_get_cross(state, True), + y=y, + clip_feature=state.clip_feas, + action=noisy_input_action, + timestep_action=timestep_action, + is_negative=True, + **common_kwargs, + ) + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + do_true_cfg=do_true_cfg, + true_cfg_scale=self.cfg_scale, + cfg_normalize=False, + ) + flow_pred, flow_pred_action = noise_pred + _cached_flow_pred = flow_pred + _cached_flow_pred_action = flow_pred_action + + _prev_predictions.append((flow_pred,)) + if _step_cache is not None: + _step_cache.trim_history(_prev_predictions) + else: + # Reuse previous prediction (DiT step-skipping cache). + assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip." + flow_pred = _cached_flow_pred + flow_pred_action = _cached_flow_pred_action latents = (noisy_input, noisy_input_action) t = (video_timestep, action_timestep) @@ -733,7 +1162,7 @@ def _transform_robot_obs(self, robot_obs: dict): return transform, transform.transform_input(robot_obs) @torch.no_grad() - def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput: """Full inference step. Called by DiffusionEngine.step().""" extra_args = req.sampling_params.extra_args or {} robot_obs = extra_args.get("robot_obs") @@ -770,7 +1199,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: device=device, ) - # State: raw from transform → pad to (B, state_horizon=1, max_state_dim) + # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim) raw_state = unified_obs["state"] state_for_postprocess = None if raw_state is not None: @@ -810,35 +1239,66 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: # Explicit reset from OpenPI serving is carried by `extra_args["reset"]` # on the next inference request after websocket reset/session switch. if extra_args.get("reset", False): - state.reset() - # Auto-reset based on model state (before accumulation) - if state.should_reset(text_tokens, 0, self.transformer.local_attn_size): - state.reset() + self._kv_reset(state) + else: + # Auto-reset based on model state (before accumulation). Both a "session" + # reset and an "inference" (window-boundary) reset clear the model-local + # KV window, so route both through _kv_reset to drop the AR-Diffusion pool window + # too and stay parity-exact; "inference" keeps the accumulated video + # latents (reset_inference_state == reset(clear_video_latents=False)). + reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size) + if reset_reason == "session": + self._kv_reset(state) + elif reset_reason == "inference": + self._kv_reset(state, clear_video_latents=False) state.language = text_tokens - # Frame accumulation: stitched single frame → multi-frame video + # Frame accumulation: stitched single frame -> multi-frame video video_frames = state.accumulate_frames(stitched) # (T, H, W, C) videos = torch.from_numpy(video_frames).unsqueeze(0).to(device) # (B=1, T, H, W, C) - videos = self._preprocess_video(videos) # → [B,C,T,H,W] bf16 + videos = self._preprocess_video(videos) # -> [B,C,T,H,W] bf16 _, _, num_frames_raw, height, width = videos.shape - prompt_embeds = self._encode_text(text_tokens, attention_mask) + # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs + # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark + # synchronizes CUDA, so leave it off for timed benchmark runs. + _pt = None + if os.environ.get("DZ_PHASE_TIMING"): + import time as _time + + torch.accelerator.synchronize() + _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []} + + def _pt_mark(name: str) -> None: + if _pt is not None: + torch.accelerator.synchronize() + _pt["marks"].append((name, _pt["time"].perf_counter())) + + # Prompt embeds are constant within a session (a prompt change triggers a + # "session" reset above, which clears this cache alongside state.language). + if state.prompt_embeds is None: + state.prompt_embeds = self._encode_text(text_tokens, attention_mask) + prompt_embeds = state.prompt_embeds # Negative prompt for CFG uncond branch (model constant) negative_prompt_embeds = None if self.cfg_scale > 1.0: - neg_inputs = self.tokenizer( - self.negative_prompt, - max_length=512, - padding="max_length", - truncation=True, - return_tensors="pt", - add_special_tokens=True, - ) - negative_prompt_embeds = self._encode_text( - neg_inputs["input_ids"].to(device), - neg_inputs["attention_mask"].to(device), - ) + if self._negative_prompt_embeds_cache is None: + neg_inputs = self.tokenizer( + self.negative_prompt, + max_length=512, + padding="max_length", + truncation=True, + return_tensors="pt", + add_special_tokens=True, + ) + self._negative_prompt_embeds_cache = self._encode_text( + neg_inputs["input_ids"].to(device), + neg_inputs["attention_mask"].to(device), + ) + negative_prompt_embeds = self._negative_prompt_embeds_cache + + _pt_mark("text_encode") # Extract first/last frame for CLIP + VAE encoding if num_frames_raw == 4 or num_frames_raw == 9: @@ -852,27 +1312,24 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: self.num_frames, height, width, + state=state, ) state.clip_feas = clip_feas.to(dtype=image.dtype) state.ys = ys.to(dtype=image.dtype) - if state.current_start_frame != 0: - # Subsequent calls: encode current observation via VAE - if (num_frames_raw - 1) // 4 == self.num_frame_per_block: - pass - elif num_frames_raw // 4 != self.num_frame_per_block: - repeat_factor = self.num_frame_per_block // (num_frames_raw // 4) - videos = torch.repeat_interleave(videos, repeat_factor, dim=2) - first_frame = videos[:, :, 0:1] - videos = torch.cat([first_frame, videos], dim=2) - else: - first_frame = videos[:, :, 0:1] - videos = torch.cat([first_frame, videos], dim=2) + # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V + # into the pool now that the image is encoded (clip_feas available). + # Runs on the first forward of a session and after each window-boundary + # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry. + self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False) + if negative_prompt_embeds is not None: + self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True) + if state.current_start_frame != 0: latent_dtype = videos.dtype with torch.no_grad(): - image = self._encode_vae_latents(videos) - image = image.to(dtype=latent_dtype) + image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype) + _pt_mark("obs_vae_encode") batch_size = image.shape[0] generator = torch.Generator(device=device).manual_seed(self.seed) @@ -913,6 +1370,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: do_true_cfg, state, ) + _pt_mark("prefill_kv") sample_scheduler = copy.deepcopy(self.scheduler) sample_scheduler_action = copy.deepcopy(self.scheduler) @@ -954,15 +1412,30 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: state_features=state_features, embodiment_id=embodiment_id, ) + if _pt is not None: + _pt_mark("diffuse") + prev_t = _pt["t0"] + parts = [] + for name, t in _pt["marks"]: + parts.append(f"{name}={1000 * (t - prev_t):.1f}ms") + prev_t = t + logger.info( + "DZ_PHASE_TIMING csf=%s total=%.1fms %s", + state.current_start_frame, + 1000 * (prev_t - _pt["t0"]), + " ".join(parts), + ) if state.current_start_frame == 1: video_out = torch.cat([image, video_out], dim=1) state.current_start_frame += self.num_frame_per_block + state.append_video_latents(video_out) + # q99 denorm: [-1,1] → real values action_out = self._denormalize_action(action_out.float(), embodiment_name) - # Relative → absolute: only for relative_action_keys (joint_position only) + # Relative -> absolute: only for relative_action_keys (joint_position only) # gripper_position is NOT relative, so don't add state back to it if self.relative_action and state_for_postprocess is not None: n_relative = self.relative_action_dim # 7 for DROID (joint only) @@ -974,7 +1447,7 @@ def forward(self, req: OmniDiffusionRequest, **kwargs) -> DiffusionOutput: action_out[..., :n_relative] + last_state.unsqueeze(1) # broadcast over horizon ) - # Squeeze batch dim for output: (B, horizon, dim) → (horizon, dim) + # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim) actions_np = action_out.squeeze(0).float().cpu().numpy() # (horizon, max_action_dim) actions_np = transform.transform_action_output(actions_np) @@ -1074,7 +1547,7 @@ def _denormalize_action( q01 = stats["q01"].to(device=action.device, dtype=action.dtype) q99 = stats["q99"].to(device=action.device, dtype=action.dtype) # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim) - # q01/q99 shape: (actual_action_dim,) — only denorm actual dims + # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims actual_dim = q01.shape[0] action_real = action.clone() action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01 @@ -1104,10 +1577,27 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: .replace("img_emb.proj.3.", "img_emb.fc2.") .replace("img_emb.proj.4.", "img_emb.norm2.") ) + + # Self-attn q/k/v are fused into a single QKVParallelLinear; route + # each separate checkpoint weight/bias to the packed `qkv` param + # with its shard id. cross_attn keeps separate q/k/v (q from x, + # k/v from context — not fusible), so it is left untouched here. + qkv_shard_id: str | None = None + for shard_id in ("q", "k", "v"): + needle = f".self_attn.{shard_id}." + if needle in new_name: + new_name = new_name.replace(needle, ".self_attn.qkv.") + qkv_shard_id = shard_id + break + if new_name in params: param = params[new_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, tensor) + if qkv_shard_id is not None: + # QKVParallelLinear.weight_loader needs the shard id. + param.weight_loader(param, tensor, qkv_shard_id) + else: + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, tensor) loaded.add(new_name) elif new_name in buffers: buffers[new_name].data.copy_(tensor) diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py index 0972d1b5a2c..4c9209808a4 100644 --- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py @@ -23,7 +23,8 @@ class DreamZeroState: Lifecycle: - Created once in DreamZeroPipeline.__init__() - Mutated every forward() call (frame append, KV cache grow) - - reset() on new session / language change / local_attn_size exceeded + - reset() on new session / language change + - reset_inference_state() on local_attn_size exceeded (KV only) """ def __init__(self) -> None: @@ -65,97 +66,110 @@ def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray: self.call_count += 1 return np.stack(frames, axis=0) # (T, H, W, C) + def append_video_latents(self, video_out: torch.Tensor) -> None: + """Append one AR chunk of normalized video latents for later decode. + + Args: + video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop + (before the ``transpose(1, 2)`` stored in ``DiffusionOutput``). + """ + if video_out.dim() != 5: + raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}") + # Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W). + chunk = video_out.transpose(1, 2).detach().cpu() + self.video_latents_across_time.append(chunk) + logger.info( + "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d", + tuple(chunk.shape), + len(self.video_latents_across_time), + int(sum(c.shape[2] for c in self.video_latents_across_time)), + ) + + def get_concatenated_video_latents(self) -> torch.Tensor | None: + """Return all accumulated chunks concatenated along the time dimension.""" + if not self.video_latents_across_time: + return None + if len(self.video_latents_across_time) == 1: + return self.video_latents_across_time[0] + return torch.cat(self.video_latents_across_time, dim=2) + + def clear_video_latents(self) -> None: + """Drop accumulated video latents without resetting KV/frame state.""" + self.video_latents_across_time = [] + # ------------------------------------------------------------------ # Reset / should_reset # ------------------------------------------------------------------ - def reset(self) -> None: - """Clear all state.""" - self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK) - self.call_count: int = 0 - - # KV cache once robot-policy diffusion supports that integration. - self.kv_cache: list[torch.Tensor] | None = None - self.kv_cache_neg: list[torch.Tensor] | None = None - self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None - self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None - self.current_start_frame: int = 0 + def reset(self, *, clear_video_latents: bool = True) -> None: + """Clear session state. - self.clip_feas: torch.Tensor | None = None - self.ys: torch.Tensor | None = None - self.language: torch.Tensor | None = None - - def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool: - """Determine if state should be reset before this forward().""" + Args: + clear_video_latents: When ``False``, keep ``video_latents_across_time`` + so offline/online export can decode the full AR rollout. + """ + saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time) + self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK) + self.call_count = 0 + # Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU. + # Concatenated along T before decode (matches upstream video_across_time). + self.video_latents_across_time = saved_video_latents + + self.current_start_frame = 0 + + self.clip_feas = None + self.ys = None + if clear_video_latents: + # Session reset: drop the prompt-embed cache and VAE encoder stream. + self.language = None + self.prompt_embeds = None + self.reset_vae_encoder_stream() + # Window ("inference") resets keep both: the prompt is unchanged, and the + # Wan feat_cache history is independent of the DiT attention window. + + def reset_vae_encoder_stream(self) -> None: + """Clear incremental Wan VAE encoder state used across AR steps.""" + self.vae_stream_initialized = False + self.vae_enc_feat_map: list[torch.Tensor | None] | None = None + self.vae_encoder_out: torch.Tensor | None = None + self.vae_pending_body_frames: torch.Tensor | None = None + + def reset_inference_state(self) -> None: + """Reset KV/frame state after local attention rolls without dropping video latents.""" + self.reset(clear_video_latents=False) + + def reset_reason( + self, + text_tokens: torch.Tensor | None, + num_video_frames: int, + local_attn_size: int, + ) -> str | None: + """Return why state should reset before the next forward(), if any.""" if self.language is None: logger.info("language is None, resetting") - return True + return "session" if text_tokens is not None and not torch.equal(self.language, text_tokens): logger.info("language changed, resetting") - return True + return "session" # NOTE: after accumulate_frames, num_video_frames is the accumulated T # (1 for first call, 4 for subsequent). Only reset on true single-frame # which happens when the stitched_buffer was cleared externally. if num_video_frames == 1 and self.call_count > 1: logger.info("single frame input after first call, resetting") - return True + return "session" if local_attn_size != -1 and self.current_start_frame >= local_attn_size: logger.info( - "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size + "current_start_frame %d >= local_attn_size %d, resetting inference state", + self.current_start_frame, + local_attn_size, ) - return True - - return False + return "inference" - # ------------------------------------------------------------------ - # KV cache management - # ------------------------------------------------------------------ + return None - def create_kv_caches( - self, - batch_size: int, - dtype: torch.dtype, - device: torch.device, - num_layers: int, - num_heads: int, - head_dim: int, - ) -> None: - """Initialize empty KV caches and cross-attention caches.""" - self.kv_cache = [ - torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers) - ] - self.kv_cache_neg = [ - torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers) - ] - - self.crossattn_cache = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)] - self.crossattn_cache_neg = [{"is_init": False, "k": None, "v": None} for _ in range(num_layers)] - - def update_kv_cache( - self, - layer_index: int, - updated_kv: torch.Tensor, - is_negative: bool = False, - ) -> None: - """Update a single layer's KV cache after prefill.""" - cache = self.kv_cache_neg if is_negative else self.kv_cache - if cache is None: - raise RuntimeError("KV caches not initialized, call create_kv_caches first.") - cache[layer_index] = updated_kv.clone() - - def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]: - """Get KV caches for the specified branch.""" - cache = self.kv_cache_neg if is_negative else self.kv_cache - if cache is None: - raise RuntimeError("KV caches not initialized.") - return cache - - def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]: - """Get cross-attention caches for the specified branch.""" - cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache - if cache is None: - raise RuntimeError("Cross-attn caches not initialized.") - return cache + def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool: + """Determine if state should be reset before this forward().""" + return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py index 20c8cd8cdc7..ba980a33e08 100644 --- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py +++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py @@ -9,13 +9,59 @@ class DreamZeroVideoExportWorkerExtension: """DreamZero worker RPCs used by offline example video export.""" + def gpu_mem_stats(self) -> dict: + """Peak GPU memory (GiB) for this worker's device, for profiling. + + ``max_memory_reserved`` is the caching-allocator high-water mark (the best + in-process proxy for "peak VRAM"); ``max_memory_allocated`` is the live + tensor high-water mark. Both are monotonic from process start, so the + end-of-run values capture the whole rollout's peak. + """ + dev = torch.accelerator.current_device_index() + return { + "device": int(dev), + "peak_reserved_gib": torch.accelerator.max_memory_reserved(dev) / (1024**3), + "peak_allocated_gib": torch.accelerator.max_memory_allocated(dev) / (1024**3), + } + + def ar_diffusion_perf_stats(self, reset: bool = True) -> dict: + """Per-request server-side forward E2E times (s) recorded by the AR-Diffusion runner. + + The worker-side compute time per ``execute_model`` — the analog of the + upstream server's per-request E2E — which excludes the engine<->worker IPC + that the client's ``omni.generate`` wall time includes. Returns ``[]`` when + the runner is the base (KV-disabled) runner that does not record timings. + """ + runner = self.model_runner + times = list(getattr(runner, "_perf_e2e_times", []) or []) + if reset and hasattr(runner, "_perf_e2e_times"): + runner._perf_e2e_times = [] + return {"server_e2e_s": times} + + @staticmethod + def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor: + decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous() + decoded = decoded.clamp(-1, 1) * 0.5 + 0.5 + return (decoded * 255.0).round().to(torch.uint8).cpu() + def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor: if self.model_runner is None or self.model_runner.pipeline is None: raise RuntimeError("DreamZero pipeline is not initialized on this worker.") with torch.inference_mode(): decoded = self.model_runner.pipeline.decode_video_latents(video_latents) - decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous() - decoded = decoded.clamp(-1, 1) * 0.5 + 0.5 - decoded = (decoded * 255.0).round().to(torch.uint8).cpu() - return decoded + return self._latents_to_uint8_frames(decoded) + + def decode_accumulated_video_latents_to_uint8(self, session_id: str = "default") -> torch.Tensor: + """Decode all AR-chunk latents accumulated on the server for one session.""" + if self.model_runner is None or self.model_runner.pipeline is None: + raise RuntimeError("DreamZero pipeline is not initialized on this worker.") + + with torch.inference_mode(): + decoded = self.model_runner.pipeline.decode_accumulated_video_latents(session_id) + return self._latents_to_uint8_frames(decoded) + + def clear_accumulated_video_latents(self, session_id: str = "default") -> None: + if self.model_runner is None or self.model_runner.pipeline is None: + raise RuntimeError("DreamZero pipeline is not initialized on this worker.") + self.model_runner.pipeline.clear_accumulated_video_latents(session_id) diff --git a/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py new file mode 100644 index 00000000000..57d9b866c56 --- /dev/null +++ b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Make diffusers Wan VAE feat_cache compile-friendly. + +diffusers ``WanResample`` stores the string sentinel ``"Rep"`` in feat_cache slots. +``tensor == "Rep"`` returns a non-Tensor and breaks ``torch.compile(fullgraph=True)``. +We replace it with a scalar int8 tensor marker (value 0). +""" + +from __future__ import annotations + +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False + + +def _is_feat_rep(entry: Any) -> bool: + if not isinstance(entry, torch.Tensor): + return False + return entry.numel() == 1 and entry.dtype == torch.int8 + + +def _make_feat_rep(x: torch.Tensor) -> torch.Tensor: + return torch.zeros(1, device=x.device, dtype=torch.int8) + + +def _wan_resample_forward_tensorized(self, x, feat_cache=None, feat_idx=[0]): + """``WanResample.forward`` with tensor feat_cache sentinel (no ``"Rep"`` string).""" + from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T + + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = _make_feat_rep(x) + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not _is_feat_rep(feat_cache[idx]): + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and _is_feat_rep(feat_cache[idx]): + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if _is_feat_rep(feat_cache[idx]): + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +def apply_wan_vae_feat_cache_tensor_patch() -> None: + """Monkeypatch diffusers WanResample once per process.""" + global _PATCHED + if _PATCHED: + return + + from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample + + WanResample.forward = _wan_resample_forward_tensorized + _PATCHED = True + logger.info("DreamZero: applied Wan VAE feat_cache tensor sentinel patch.") diff --git a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py index aa9b0d8b79a..8398cb7aba2 100644 --- a/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py +++ b/vllm_omni/diffusion/models/ernie_image/ernie_image_transformer.py @@ -8,6 +8,7 @@ import torch import torch.nn as nn +from cache_dit import ForwardPattern from diffusers.models.embeddings import TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput from vllm.distributed import get_tensor_model_parallel_world_size @@ -21,6 +22,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -459,6 +461,12 @@ def forward(self, x: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor: class ErnieImageTransformer2DModel(nn.Module): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_3, + }, + check_forward_pattern=False, + ) _repeated_blocks = ["ErnieImageSharedAdaLNBlock"] _layerwise_offload_blocks_attrs = ["layers"] diff --git a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py b/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py index cd60cc93b27..82f534c312d 100644 --- a/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py +++ b/vllm_omni/diffusion/models/ernie_image/pipeline_ernie_image.py @@ -27,6 +27,7 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = init_logger(__name__) @@ -392,7 +393,7 @@ def check_inputs( def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = "", height: int = 1024, diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 88d8126c8ad..952ee9eb4f4 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from cache_dit import ForwardPattern from diffusers.models.embeddings import ( CombinedTimestepGuidanceTextProjEmbeddings, CombinedTimestepTextProjEmbeddings, @@ -25,6 +26,8 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig +from vllm_omni.quantization.component_config import safe_quant_config if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -42,23 +45,6 @@ logger = init_logger(__name__) -def _safe_quant_config(quant_config: "QuantizationConfig | None") -> "QuantizationConfig | None": - """Return quant_config only if it is safe to propagate here, else None. - - Dual-stream transformer_blocks, norm modulation layers, and norm_out are - kept at full precision for FP8 (see #2728). Offline quantization (e.g. - INC/AutoRound W4A16) needs the config propagated so packed weights load - correctly. - """ - if quant_config is None: - return None - from vllm.model_executor.layers.quantization.inc import INCConfig - - if isinstance(quant_config, INCConfig): - return quant_config - return None - - class ColumnParallelApproxGELU(nn.Module): def __init__( self, @@ -399,7 +385,7 @@ def __init__( super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) - self.norm = AdaLayerNormZeroSingle(dim, quant_config=_safe_quant_config(quant_config), prefix=f"{prefix}.norm") + self.norm = AdaLayerNormZeroSingle(dim, quant_config=safe_quant_config(quant_config), prefix=f"{prefix}.norm") self.proj_mlp = ReplicatedLinear( dim, self.mlp_hidden_dim, @@ -524,6 +510,13 @@ class FluxTransformer2DModel(nn.Module): The dimensions to use for the rotary positional embeddings. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_1, + "single_transformer_blocks": ForwardPattern.Pattern_1, + } + ) + # the small and frequently-repeated block(s) of a model # -- typically a transformer layer # used for torch compile optimizations @@ -587,7 +580,7 @@ def __init__( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, - quant_config=_safe_quant_config(quant_config), + quant_config=safe_quant_config(quant_config), prefix=f"transformer_blocks.{i}", ) for i in range(num_layers) @@ -612,7 +605,7 @@ def __init__( self.inner_dim, elementwise_affine=False, eps=1e-6, - quant_config=_safe_quant_config(quant_config), + quant_config=safe_quant_config(quant_config), prefix="norm_out", ) self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 03bef7bd605..e7f97597bd2 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -5,7 +5,7 @@ import logging import os from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -28,10 +28,11 @@ from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin from vllm_omni.diffusion.models.flux import FluxTransformer2DModel from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = logging.getLogger(__name__) @@ -66,7 +67,15 @@ def post_process_func(images: torch.Tensor): return post_process_func -class FluxPipeline(nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipelineProfilerMixin): +class FluxPipeline( + nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + supports_request_batch = True + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -493,55 +502,54 @@ def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: boo return False return True - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - prompt_2: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - negative_prompt_2: str | list[str] | None = None, - true_cfg_scale: float = 1.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 28, - sigmas: list[float] | None = None, - guidance_scale: float = 3.5, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.FloatTensor | None = None, - pooled_prompt_embeds: torch.FloatTensor | None = None, - negative_prompt_embeds: torch.FloatTensor | None = None, - negative_pooled_prompt_embeds: torch.FloatTensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - ): + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: """Forward pass for flux.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + sampling_params_list = req.sampling_params_list + common_sampling_params = sampling_params_list[0] + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = common_sampling_params.num_inference_steps or 28 + sigmas = common_sampling_params.sigmas + guidance_scale = common_sampling_params.guidance_scale + true_cfg_scale = common_sampling_params.true_cfg_scale or 1.0 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 + ) + max_sequence_length = common_sampling_params.max_sequence_length or 512 + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": None, + "negative_prompt_embeds": None, + "pooled_prompt_embeds": None, + "negative_pooled_prompt_embeds": None, + }, ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + pooled_prompt_embeds = prompt_fields["pooled_prompt_embeds"] + negative_pooled_prompt_embeds = prompt_fields["negative_pooled_prompt_embeds"] + prompt_2 = None + negative_prompt_2 = None + if prompt_embeds is not None: + prompt = None + if negative_prompt_embeds is not None: + negative_prompt = None + + output_type = common_sampling_params.output_type or "pil" + joint_attention_kwargs: dict[str, Any] | None = None + callback_on_step_end_tensor_inputs = ["latents"] # 1. Check inputs. Raise error if not correct self.check_inputs( @@ -565,9 +573,7 @@ def forward( self._interrupt = False # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): + if prompt is not None: batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] @@ -659,9 +665,17 @@ def forward( latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) + stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None + n = num_images_per_prompt + if req.num_reqs == 1: + return [DiffusionOutput(output=image, stage_durations=stage_durations)] + return [ + DiffusionOutput( + output=image[i * n : (i + 1) * n] if isinstance(image, list) else image[i * n : (i + 1) * n], + stage_durations=stage_durations, + ) + for i in range(req.num_reqs) + ] def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) @@ -674,3 +688,6 @@ class FluxDMD2Pipeline(DMD2PipelineMixin, FluxPipeline): def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) self.__init_dmd2__() + + def forward(self, req: DiffusionRequestBatch, **kwargs) -> list[DiffusionOutput]: + return super().forward(req, **kwargs) diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py index c3bea7dd1c4..15f0e4e0c82 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux_kontext.py @@ -9,7 +9,7 @@ import json import os from collections.abc import Callable, Iterable -from typing import Any, cast +from typing import Any, ClassVar, cast import numpy as np import PIL.Image @@ -24,17 +24,18 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin 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.flux import ( FluxKontextTransformer2DModel, ) from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixin -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.logger import init_logger logger = init_logger(__name__) @@ -70,10 +71,20 @@ def post_process_func(images: torch.Tensor) -> list[PIL.Image.Image]: class FluxKontextPipeline( - nn.Module, FluxPipelineMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin + nn.Module, + FluxPipelineMixin, + CFGParallelMixin, + SupportImageInput, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, ): """FLUX.1-Kontext pipeline for image editing with text guidance.""" + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"] + _vae_modules: ClassVar[list[str]] = ["vae"] + support_image_input = True def __init__( @@ -148,6 +159,7 @@ def __init__( self._attention_kwargs = None self._joint_attention_kwargs = None self._num_timesteps = None + self._current_timestep = None self._interrupt = False self._callback_tensor_inputs = ["latents", "prompt_embeds"] self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16 @@ -456,10 +468,117 @@ def joint_attention_kwargs(self): def num_timesteps(self): return self._num_timesteps + @property + def current_timestep(self) -> torch.Tensor | None: + return self._current_timestep + + def predict_noise( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + guidance: torch.Tensor | None, + pooled_projections: torch.Tensor, + encoder_hidden_states: torch.Tensor, + txt_ids: torch.Tensor, + img_ids: torch.Tensor, + joint_attention_kwargs: dict[str, Any] | None, + return_dict: bool = False, + ) -> torch.Tensor: + return self.transformer( + hidden_states=hidden_states, + timestep=timestep, + guidance=guidance, + pooled_projections=pooled_projections, + encoder_hidden_states=encoder_hidden_states, + txt_ids=txt_ids, + img_ids=img_ids, + joint_attention_kwargs=joint_attention_kwargs, + return_dict=return_dict, + )[0] + + def diffuse( + self, + prompt_embeds: torch.Tensor, + pooled_prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + negative_pooled_prompt_embeds: torch.Tensor | None, + latents: torch.Tensor, + image_latents: torch.Tensor | None, + latent_ids: torch.Tensor, + text_ids: torch.Tensor, + negative_text_ids: torch.Tensor | None, + timesteps: torch.Tensor, + do_true_cfg: bool, + guidance: torch.Tensor | None, + true_cfg_scale: float, + callback_on_step_end: Callable[[int, int, dict], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + ) -> torch.Tensor: + self.scheduler.set_begin_index(0) + self.transformer.do_true_cfg = do_true_cfg + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + latent_model_input = latents + if image_latents is not None: + latent_model_input = torch.cat([latents, image_latents], dim=1) + timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "guidance": guidance, + "pooled_projections": pooled_prompt_embeds, + "encoder_hidden_states": prompt_embeds, + "txt_ids": text_ids, + "img_ids": latent_ids, + "joint_attention_kwargs": self.joint_attention_kwargs, + "return_dict": False, + } + negative_kwargs = None + if do_true_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep / 1000, + "guidance": guidance, + "pooled_projections": negative_pooled_prompt_embeds, + "encoder_hidden_states": negative_prompt_embeds, + "txt_ids": negative_text_ids, + "img_ids": latent_ids, + "joint_attention_kwargs": self.joint_attention_kwargs, + "return_dict": False, + } + + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg, + true_cfg_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize=False, + output_slice=latents.size(1), + ) + + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + pbar.update() + return latents + @torch.no_grad() def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, image: PIL.Image.Image | list[PIL.Image.Image] | None = None, prompt: str | list[str] | None = None, prompt_2: str | list[str] | None = None, @@ -494,8 +613,13 @@ def forward( if isinstance(first_prompt, str): prompt = first_prompt + # Keep negative_prompt from function parameter (passed via API) elif first_prompt: prompt = first_prompt.get("prompt") or "" + # Prefer negative_prompt from prompt dict, fallback to function parameter + np_from_dict = first_prompt.get("negative_prompt") + if np_from_dict is not None: + negative_prompt = np_from_dict # else: keep original prompt unchanged # Handle image from prompt data @@ -515,8 +639,17 @@ def forward( height = req.sampling_params.height or height width = req.sampling_params.width or width num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + sigmas = req.sampling_params.sigmas or sigmas guidance_scale = req.sampling_params.guidance_scale or guidance_scale + true_cfg_scale = ( + req.sampling_params.true_cfg_scale if req.sampling_params.true_cfg_scale is not None else true_cfg_scale + ) generator = req.sampling_params.generator or generator + num_images_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_images_per_prompt + ) latents = ( req.sampling_params.extra_args.get("latents") if req.sampling_params.extra_args.get("latents") is not None @@ -559,6 +692,8 @@ def forward( ) do_true_cfg = true_cfg_scale > 1.0 and has_neg_prompt + do_true_cfg = do_true_cfg and self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) + # 1. Prepare text embeddings prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt( prompt=prompt, @@ -642,57 +777,23 @@ def forward( self._joint_attention_kwargs = {} # 5. Denoising loop - self.scheduler.set_begin_index(0) - with self.progress_bar(total=len(timesteps)) as pbar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - latent_model_input = latents - if image_latents is not None: - latent_model_input = torch.cat([latents, image_latents], dim=1) - timestep = t.expand(latents.shape[0]).to(latents.dtype) - - noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=pooled_prompt_embeds, - encoder_hidden_states=prompt_embeds, - txt_ids=text_ids, - img_ids=latent_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - noise_pred = noise_pred[:, : latents.size(1)] - - if do_true_cfg: - neg_noise_pred = self.transformer( - hidden_states=latent_model_input, - timestep=timestep / 1000, - guidance=guidance, - pooled_projections=negative_pooled_prompt_embeds, - encoder_hidden_states=negative_prompt_embeds, - txt_ids=negative_text_ids, - img_ids=latent_ids, - joint_attention_kwargs=self.joint_attention_kwargs, - return_dict=False, - )[0] - neg_noise_pred = neg_noise_pred[:, : latents.size(1)] - noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred) - - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - if callback_on_step_end is not None: - callback_kwargs = {} - for k in callback_on_step_end_tensor_inputs: - callback_kwargs[k] = locals()[k] - callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) - - latents = callback_outputs.pop("latents", latents) - prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) - - pbar.update() + latents = self.diffuse( + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, + latents=latents, + image_latents=image_latents, + latent_ids=latent_ids, + text_ids=text_ids, + negative_text_ids=negative_text_ids if do_true_cfg else None, + timesteps=timesteps, + do_true_cfg=do_true_cfg, + guidance=guidance, + true_cfg_scale=true_cfg_scale, + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) if output_type == "latent": image = latents else: diff --git a/vllm_omni/diffusion/models/flux2/flux2_transformer.py b/vllm_omni/diffusion/models/flux2/flux2_transformer.py index 0c6896c131d..52e0a44303a 100644 --- a/vllm_omni/diffusion/models/flux2/flux2_transformer.py +++ b/vllm_omni/diffusion/models/flux2/flux2_transformer.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any import torch +from cache_dit import ForwardPattern from diffusers.models.embeddings import ( TimestepEmbedding, Timesteps, @@ -25,6 +26,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -755,6 +757,13 @@ class Flux2Transformer2DModel(nn.Module): Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_1, + "single_transformer_blocks": ForwardPattern.Pattern_2, + } + ) + _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] _sp_plan = { "": { @@ -1025,7 +1034,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name = original_name if name not in params_dict and ".to_out.0." in name: name = name.replace(".to_out.0.", ".to_out.") - # Some GGUF checkpoints include quantized tensors for modules that + # Some quantized checkpoints include tensors for modules that # are intentionally left unquantized in this model. param = params_dict.get(name) if param is None: diff --git a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py index 948f922d468..351d94f7edc 100644 --- a/vllm_omni/diffusion/models/flux2/pipeline_flux2.py +++ b/vllm_omni/diffusion/models/flux2/pipeline_flux2.py @@ -6,7 +6,7 @@ import math import os from collections.abc import Callable, Iterable -from typing import Any, cast +from typing import Any, ClassVar, cast import numpy as np import PIL.Image @@ -30,17 +30,24 @@ 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.flux2 import Flux2Transformer2DModel -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.mistral_encoder import MistralEncoderModel from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +def _resolve_component_quant_config(quant_config, component: str): + if hasattr(quant_config, "resolve"): + return quant_config.resolve(component) + return quant_config + + class Flux2ImageProcessor(VaeImageProcessor): """Image processor to preprocess the reference image for Flux2.""" @@ -336,9 +343,20 @@ def retrieve_latents(encoder_output: torch.Tensor, generator: torch.Generator = raise AttributeError("Could not access latents of provided encoder_output") -class Flux2Pipeline(nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin): +class Flux2Pipeline( + nn.Module, + CFGParallelMixin, + SupportImageInput, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, +): """Flux2 pipeline for text-to-image generation.""" + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _callback_tensor_inputs = ["latents", "prompt_embeds"] support_image_input = True @@ -382,10 +400,14 @@ def __init__( text_encoder_config = AutoConfig.from_pretrained( model, subfolder="text_encoder", local_files_only=local_files_only ) + te_quant_config = _resolve_component_quant_config(od_config.quantization_config, "text_encoder") self.text_encoder = MistralEncoderModel( text_encoder_config, prefix="text_encoder", - ).to(self._execution_device) + quant_config=te_quant_config, + ) + if not any(param.is_meta for param in self.text_encoder.parameters()): + self.text_encoder.to(self._execution_device) self.tokenizer = PixtralProcessor.from_pretrained( model, subfolder="tokenizer", local_files_only=local_files_only ) @@ -398,8 +420,9 @@ def __init__( self._execution_device ) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Flux2Transformer2DModel) + transformer_quant_config = _resolve_component_quant_config(od_config.quantization_config, "transformer") self.transformer = Flux2Transformer2DModel( - quant_config=od_config.quantization_config, od_config=od_config, **transformer_kwargs + quant_config=transformer_quant_config, od_config=od_config, **transformer_kwargs ) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 @@ -848,7 +871,7 @@ def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: boo def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, image: PIL.Image.Image | list[PIL.Image.Image] | None = None, prompt: str | list[str] | None = None, height: int | None = None, @@ -1137,5 +1160,26 @@ def forward( return DiffusionOutput(output=image) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + text_encoder_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "text_encoder") + transformer_quant_config = _resolve_component_quant_config(self.od_config.quantization_config, "transformer") + if ( + self.od_config.enable_cpu_offload + and text_encoder_quant_config is not None + and transformer_quant_config is None + ): + weights = self._offload_transformer_before_text_encoder(weights) + loader = AutoWeightsLoader(self) return loader.load_weights(weights) + + def _offload_transformer_before_text_encoder( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> Iterable[tuple[str, torch.Tensor]]: + transformer_offloaded = False + for name, weight in weights: + if not transformer_offloaded and name.startswith("text_encoder."): + self.transformer.to("cpu") + current_omni_platform.empty_cache() + transformer_offloaded = True + logger.info("Offloaded the unquantized transformer before text encoder quantization") + yield name, weight diff --git a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py index a708a8ece6a..e62fccba0c9 100644 --- a/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py +++ b/vllm_omni/diffusion/models/flux2_klein/flux2_klein_transformer.py @@ -21,6 +21,7 @@ import torch import torch.nn as nn +from cache_dit import ForwardPattern from diffusers.models.embeddings import ( TimestepEmbedding, Timesteps, @@ -40,6 +41,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -768,6 +770,13 @@ class Flux2Transformer2DModel(nn.Module): Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_1, + "single_transformer_blocks": ForwardPattern.Pattern_2, + } + ) + _repeated_blocks = ["Flux2TransformerBlock", "Flux2SingleTransformerBlock"] _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] @@ -1035,7 +1044,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: name = original_name if name not in params_dict and ".to_out.0." in name: name = name.replace(".to_out.0.", ".to_out.") - # Some GGUF checkpoints include quantized tensors for modules that + # Some quantized checkpoints include tensors for modules that # are intentionally left unquantized in this model. param = params_dict.get(name) if param is None: diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index aea6d7e03fd..184b5d40252 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -19,7 +19,7 @@ import math import os from collections.abc import Callable, Iterable -from typing import Any, cast +from typing import Any, ClassVar, cast import numpy as np import PIL.Image @@ -43,10 +43,10 @@ from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import ( Flux2Transformer2DModel, ) -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = init_logger(__name__) @@ -181,9 +181,15 @@ def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float: return float(mu) -class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin): +class Flux2KleinPipeline( + nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): """Flux2 klein pipeline for text-to-image generation.""" + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + support_image_input = True def __init__( @@ -744,108 +750,7 @@ def current_timestep(self): def interrupt(self): return self._interrupt - def forward( - self, - req: OmniDiffusionRequest, - image: PIL.Image.Image | list[PIL.Image.Image] | None = None, - reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None, - mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None, - prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - strength: float = 1.0, - guidance_scale: float | None = 4.0, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end: Callable[[int, int, dict], None] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - text_encoder_out_layers: tuple[int, ...] = (9, 18, 27), - padding_mask_crop: int | None = None, - ) -> DiffusionOutput: - r""" - Function invoked when calling the pipeline for generation. - - Args: - image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, or list of these): - `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both - numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list - or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a - list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image - latents as `image`, but if passing latents directly it is not encoded again. - prompt (`str` or `List[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - guidance_scale (`float`, *optional*, defaults to 4.0): - Guidance scale as defined in [Classifier-Free Diffusion - Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. - of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting - `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to - the text `prompt`, usually at the expense of lower image quality. For step-wise distilled models, - `guidance_scale` is ignored. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for the best results. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for the best results. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - sigmas (`List[float]`, *optional*): - Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in - their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed - will be used. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.Tensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.Tensor`, *optional*): - Pre-generated negative text embeddings. Note that "" is used as the negative prompt in this pipeline. - If not provided, will be generated from "". - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple. - attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`List`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. - text_encoder_out_layers (`Tuple[int]`): - Layer indices to use in the `text_encoder` to derive the final prompt embeddings. - - Examples: - - Returns: - [`~pipelines.flux2.Flux2PipelineOutput`] or `tuple`: [`~pipelines.flux2.Flux2PipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the - generated images. - """ + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: if len(req.prompts) > 1: logger.warning( """This model only supports a single prompt, not a batched request.""", @@ -872,39 +777,32 @@ def forward( else: image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else cast(PIL.Image.Image, raw_image) - height = req.sampling_params.height or height - width = req.sampling_params.width or width + prompt_embeds = None + negative_prompt_embeds = None + + height = req.sampling_params.height + width = req.sampling_params.width num_inference_steps = ( - req.sampling_params.num_inference_steps - if req.sampling_params.num_inference_steps is not None - else num_inference_steps + 50 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps ) - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - generator = req.sampling_params.generator or generator + sigmas = req.sampling_params.sigmas + guidance_scale = req.sampling_params.guidance_scale + generator = req.sampling_params.generator num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - text_encoder_out_layers = req.sampling_params.extra_args.get("text_encoder_out_layers", text_encoder_out_layers) - - req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - # If at list one prompt is provided as an embedding, - # Then assume that the user wants to provide embeddings for all prompts, and enter this if block - # If the user in fact provides mixed input format, req_prompt_embeds will have some None's - # And `torch.stack` automatically raises an exception for us - prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore # intentionally expect TypeError - - req_negative_prompt_embeds = [ - p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts - ] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore # intentionally expect TypeError + max_sequence_length = req.sampling_params.max_sequence_length or 512 + latents = req.sampling_params.latents + strength = req.sampling_params.strength if req.sampling_params.strength is not None else 1.0 + + text_encoder_out_layers: tuple[int, ...] = req.sampling_params.extra_args.get( + "text_encoder_out_layers", (9, 18, 27) + ) + padding_mask_crop: int | None = req.sampling_params.extra_args.get("padding_mask_crop") + attention_kwargs: dict[str, Any] | None = None + callback_on_step_end: Callable[[int, int, dict], None] | None = None + callback_on_step_end_tensor_inputs = ["latents"] + output_type = req.sampling_params.output_type or "pil" # 1. Check inputs. Raise error if not correct self.check_inputs( @@ -924,12 +822,7 @@ def forward( self._interrupt = False # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = 1 device = self._execution_device diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 56a7066e390..c933d4e7f31 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -9,6 +9,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from cache_dit import ForwardPattern from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings from vllm.logger import init_logger @@ -20,6 +21,8 @@ ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -935,6 +938,13 @@ class GlmImageTransformer2DModel(CachedTransformer): `od_config.tf_model_config`. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_0, + }, + has_separate_cfg=True, + ) + _repeated_blocks = ["GlmImageTransformerBlock"] # SP plan using GlmImagePrepare module for sharding hidden_states and RoPE together. # Similar to Qwen-Image's ImageRopePrepare, this creates a module boundary where diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 102456d9f9c..9f3d2940090 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py @@ -17,7 +17,7 @@ import os import re from collections.abc import Iterable -from typing import cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -46,8 +46,10 @@ GlmImageKVCache, GlmImageTransformer2DModel, ) +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -80,58 +82,54 @@ def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig): def pre_process_func(request: OmniDiffusionRequest): """Pre-process condition images for Image Edit mode.""" - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if raw_image is None: - # Text-to-image mode, no preprocessing needed - continue - - if not isinstance(raw_image, list): - raw_image = [raw_image] - images = [ - PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im) - for im in raw_image - ] - - preprocessed = [] - height, width = None, None - - for img in images: - if isinstance(img, PIL.Image.Image): - img_h, img_w = img.size[::-1] # PIL is (width, height) - else: - img_h, img_w = img.shape[:2] - - # Align to multiple of vae_scale_factor * patch_size - multiple_of = vae_scale_factor * patch_size - img_h = (img_h // multiple_of) * multiple_of - img_w = (img_w // multiple_of) * multiple_of - - processed = image_processor.preprocess(img, height=img_h, width=img_w) - preprocessed.append(processed) - - # Use first image dimensions as default - if height is None: - height, width = img_h, img_w - - # Store in request - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt, additional_information={}) - elif "additional_information" not in prompt: - prompt["additional_information"] = {} - prompt["additional_information"]["preprocessed_image"] = processed # type: ignore - prompt["additional_information"]["prompt_image"] = images # type: ignore - request.prompts[i] = prompt - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if raw_image is None: + # Text-to-image mode, no preprocessing needed + return request + + if not isinstance(raw_image, list): + raw_image = [raw_image] + images = [ + PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im) + for im in raw_image + ] + + preprocessed = [] + height, width = None, None + + for img in images: + if isinstance(img, PIL.Image.Image): + img_h, img_w = img.size[::-1] # PIL is (width, height) + else: + img_h, img_w = img.shape[:2] + + # Align to multiple of vae_scale_factor * patch_size + multiple_of = vae_scale_factor * patch_size + img_h = (img_h // multiple_of) * multiple_of + img_w = (img_w // multiple_of) * multiple_of + + processed = image_processor.preprocess(img, height=img_h, width=img_w) + preprocessed.append(processed) + + # Use first image dimensions as default + if height is None: + height, width = img_h, img_w + + # Store in request + prompt["additional_information"]["preprocessed_image"] = processed # type: ignore + prompt["additional_information"]["prompt_image"] = images # type: ignore + request.prompt = prompt + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width return request @@ -238,7 +236,7 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): +class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery): """ GLM-Image Pipeline for text-to-image and image-to-image generation. @@ -255,6 +253,10 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): 4. VAE decodes final latents to image """ + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -555,6 +557,8 @@ def diffuse( noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) # Scheduler step latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + # Re-cast to transformer_dtype before cfg_group.broadcast to avoid dtype mismatches across ranks + latents = latents.to(transformer_dtype) # Broadcast updated latents to all ranks cfg_group.broadcast(latents, src=0) @@ -667,7 +671,7 @@ def _prepare_condition_image_kv_cache( return kv_caches @torch.inference_mode() - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """ Main generation forward pass. diff --git a/vllm_omni/diffusion/models/gr00t/__init__.py b/vllm_omni/diffusion/models/gr00t/__init__.py new file mode 100644 index 00000000000..a10202a40bc --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/__init__.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline + +__all__ = ["Gr00tN1d7Pipeline"] diff --git a/vllm_omni/diffusion/models/gr00t/configs/__init__.py b/vllm_omni/diffusion/models/gr00t/configs/__init__.py new file mode 100644 index 00000000000..9648b859751 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/configs/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +MODEL_CONFIG_TYPES: dict[str, type] = {} + + +def register_model_config(shortname: str, configtype: type) -> None: + MODEL_CONFIG_TYPES[shortname] = configtype diff --git a/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py b/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py new file mode 100644 index 00000000000..7d11a651cf6 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/configs/embodiment_configs.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.diffusion.models.gr00t.dataio.types import ( + ActionConfig, + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) + +__all__ = ["ActionConfig", "ActionFormat", "ActionRepresentation", "ActionType", "ModalityConfig"] diff --git a/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py new file mode 100644 index 00000000000..d7d4527fb93 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from transformers import PretrainedConfig + +from . import register_model_config + + +def _default_diffusion_model_cfg() -> dict: + return { + "positional_embeddings": None, + "num_layers": 16, + "num_attention_heads": 32, + "attention_head_dim": 48, + "norm_type": "ada_norm", + "output_dim": 1024, + "interleave_self_attention": True, + } + + +class Gr00tN1d7Config(PretrainedConfig): + model_type = "Gr00tN1d7" + + def __init__( + self, + model_dtype: str = "bfloat16", + model_name: str = "nvidia/Cosmos-Reason2-2B", + backbone_model_type: str = "qwen", + model_revision: str | None = None, + backbone_embedding_dim: int = 2048, + select_layer: int = 12, + reproject_vision: bool = False, + use_flash_attention: bool = True, + load_bf16: bool = False, + image_crop_size: tuple | None = (230, 230), + image_target_size: tuple | None = (256, 256), + shortest_image_edge: int | None = None, + crop_fraction: float | None = None, + random_rotation_angle: int | None = None, + color_jitter_params: dict | None = None, + formalize_language: bool = True, + apply_sincos_state_encoding: bool = False, + use_percentiles: bool = True, + use_relative_action: bool = False, + max_state_dim: int = 132, + max_action_dim: int = 132, + action_horizon: int = 40, + hidden_size: int = 1024, + input_embedding_dim: int = 1536, + state_history_length: int = 1, + add_pos_embed: bool = True, + attn_dropout: float = 0.2, + use_vlln: bool = True, + max_seq_len: int = 1024, + use_alternate_vl_dit: bool = True, + attend_text_every_n_blocks: int = 2, + diffusion_model_cfg: dict | None = None, + vl_self_attention_cfg: dict | None = None, + num_inference_timesteps: int = 4, + num_timestep_buckets: int = 1000, + exclude_state: bool = False, + use_mean_std: bool = False, + max_num_embodiments: int = 32, + **kwargs, + ): + super().__init__(**kwargs) + self.model_dtype = model_dtype + self.model_name = model_name + self.backbone_model_type = backbone_model_type + self.model_revision = model_revision + self.backbone_embedding_dim = backbone_embedding_dim + self.select_layer = select_layer + self.reproject_vision = reproject_vision + self.use_flash_attention = use_flash_attention + self.load_bf16 = load_bf16 + self.image_crop_size = image_crop_size + self.image_target_size = image_target_size + self.shortest_image_edge = shortest_image_edge + self.crop_fraction = crop_fraction + self.random_rotation_angle = random_rotation_angle + self.color_jitter_params = color_jitter_params + self.formalize_language = formalize_language + self.apply_sincos_state_encoding = apply_sincos_state_encoding + self.use_percentiles = use_percentiles + self.use_relative_action = use_relative_action + self.max_state_dim = max_state_dim + self.max_action_dim = max_action_dim + self.action_horizon = action_horizon + self.hidden_size = hidden_size + self.input_embedding_dim = input_embedding_dim + self.state_history_length = state_history_length + self.add_pos_embed = add_pos_embed + self.attn_dropout = attn_dropout + self.use_vlln = use_vlln + self.max_seq_len = max_seq_len + self.use_alternate_vl_dit = use_alternate_vl_dit + self.attend_text_every_n_blocks = attend_text_every_n_blocks + self.diffusion_model_cfg = ( + diffusion_model_cfg if diffusion_model_cfg is not None else _default_diffusion_model_cfg() + ) + self.vl_self_attention_cfg = vl_self_attention_cfg + self.num_inference_timesteps = num_inference_timesteps + self.num_timestep_buckets = num_timestep_buckets + self.exclude_state = exclude_state + self.use_mean_std = use_mean_std + self.max_num_embodiments = max_num_embodiments + + +register_model_config("Gr00tN1d7", Gr00tN1d7Config) diff --git a/vllm_omni/diffusion/models/gr00t/dataio/__init__.py b/vllm_omni/diffusion/models/gr00t/dataio/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/__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/diffusion/models/gr00t/dataio/embodiment_tags.py b/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py new file mode 100755 index 00000000000..3ed0c9b0360 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/embodiment_tags.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import Enum + +""" +Embodiment tags are used to identify the robot embodiment in the data. + +Naming convention: +<dataset>_<robot_name> + +If using multiple datasets, e.g. sim GR1 and real GR1, we can drop the dataset name and use only the robot name. +""" + + +class EmbodimentTag(Enum): + """Embodiment tags supported by the GR00T N1.7 checkpoint. + + Pretrain tags (baked into the base model nvidia/GR00T-N1.7-3B, inference-ready): + - OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT -> "oxe_droid_relative_eef_relative_joint" + - XDOF -> "xdof_relative_eef_relative_joint" + - XDOF_SUBTASK -> "xdof_relative_eef_relative_joint_subtask" + - REAL_G1 -> "real_g1_relative_eef_relative_joints" + - REAL_R1_PRO_SHARPA -> "real_r1_pro_sharpa_relative_eef" + - REAL_R1_PRO_SHARPA_HUMAN -> "real_r1_pro_sharpa_relative_eef_human" + - REAL_R1_PRO_SHARPA_MAXINSIGHTS -> "real_r1_pro_sharpa_relative_eef_maxinsights" + - REAL_R1_PRO_SHARPA_MECKA -> "real_r1_pro_sharpa_relative_eef_mecka" + + Pre-registered posttrain tags (require finetuned checkpoint): + - UNITREE_G1 -> "unitree_g1_full_body_with_waist_height_nav_cmd" + - UNITREE_G1_SONIC -> "unitree_g1_sonic" + - SIMPLER_ENV_GOOGLE -> "simpler_env_google" + - SIMPLER_ENV_WIDOWX -> "simpler_env_widowx" + - LIBERO_PANDA -> "libero_sim" + + Finetuning tag (for custom robots): + - NEW_EMBODIMENT -> "new_embodiment" + + Use ``EmbodimentTag.resolve(s)`` to look up a tag by name or value, + case-insensitively. + """ + + OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT = "oxe_droid_relative_eef_relative_joint" + """ + The Open-X-Embodiment DROID robot with relative EEF and relative joint position actions. + """ + + XDOF = "xdof_relative_eef_relative_joint" + """ + The generic X-DOF robot with relative EEF and relative joint position actions. + """ + + XDOF_SUBTASK = "xdof_relative_eef_relative_joint_subtask" + """ + The generic X-DOF robot (subtask variant). + """ + + REAL_G1 = "real_g1_relative_eef_relative_joints" + """ + Real-world Unitree G1 with relative EEF and relative joint actions. + """ + + REAL_R1_PRO_SHARPA = "real_r1_pro_sharpa_relative_eef" + """ + Real-world R1 Pro Sharpa with relative EEF actions. + """ + + REAL_R1_PRO_SHARPA_HUMAN = "real_r1_pro_sharpa_relative_eef_human" + """ + Real-world R1 Pro Sharpa with relative EEF actions (human teleop data). + """ + + REAL_R1_PRO_SHARPA_MAXINSIGHTS = "real_r1_pro_sharpa_relative_eef_maxinsights" + """ + Real-world R1 Pro Sharpa with relative EEF actions (MaxInsights data, single-cam). + """ + + REAL_R1_PRO_SHARPA_MECKA = "real_r1_pro_sharpa_relative_eef_mecka" + """ + Real-world R1 Pro Sharpa with relative EEF actions (Mecka data, single-cam). + """ + + UNITREE_G1 = "unitree_g1_full_body_with_waist_height_nav_cmd" + """ + The Unitree G1 robot (sim, full-body with waist height and nav commands). + """ + + UNITREE_G1_SONIC = "unitree_g1_sonic" + """ + The Unitree G1 robot with SONIC whole-body controller. VLA action space is SONIC latents. + """ + + SIMPLER_ENV_GOOGLE = "simpler_env_google" + """ + The SimplerEnv Google robot. + """ + + SIMPLER_ENV_WIDOWX = "simpler_env_widowx" + """ + The SimplerEnv WidowX robot. + """ + + LIBERO_PANDA = "libero_sim" + """ + The LIBERO Panda robot (used for LIBERO-Goal, LIBERO-Object, LIBERO-Spatial, LIBERO-10). + """ + + # New embodiment during post-training + NEW_EMBODIMENT = "new_embodiment" + """ + Any new embodiment. + """ + + @classmethod + def resolve(cls, tag: "str | EmbodimentTag") -> "EmbodimentTag": + """Resolve a string to an EmbodimentTag, case-insensitively. + + Matches by enum **name** first (e.g. ``"xdof"`` -> ``XDOF``), then by + enum **value** (e.g. ``"xdof_relative_eef_relative_joint"`` -> ``XDOF``). + + Raises: + ValueError: If *tag* does not match any known embodiment. + """ + if isinstance(tag, cls): + return tag + key = tag.strip() + key_lower = key.lower() + # Match by enum name (case-insensitive) + for member in cls: + if member.name.lower() == key_lower: + return member + # Match by enum value (case-insensitive) + for member in cls: + if member.value.lower() == key_lower: + return member + + def _fmt(tags): + return "\n".join(f" {m.name:40s} -> {m.value}" for m in tags) + + msg = ( + f"Unknown embodiment tag: {tag!r}\n\n" + f" Base model tags (work with nvidia/GR00T-N1.7-3B):\n" + f"{_fmt(PRETRAIN_TAGS)}\n\n" + f" Posttrain tags (require a finetuned checkpoint):\n" + f"{_fmt(POSTTRAIN_TAGS)}\n\n" + f" Finetuning-only tags (for custom robots):\n" + f"{_fmt(FINETUNE_ONLY_TAGS)}" + ) + raise ValueError(msg) + + @classmethod + def reverse_lookup(cls, value: str) -> "str": + """Map a tag value string back to its enum name, or return the value as-is.""" + for member in cls: + if member.value == value: + return member.name + return value + + +# Module-level tag category sets (cannot be Enum class attributes). +PRETRAIN_TAGS: frozenset[EmbodimentTag] = frozenset( + { + EmbodimentTag.OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT, + EmbodimentTag.XDOF, + EmbodimentTag.XDOF_SUBTASK, + EmbodimentTag.REAL_G1, + EmbodimentTag.REAL_R1_PRO_SHARPA, + EmbodimentTag.REAL_R1_PRO_SHARPA_HUMAN, + EmbodimentTag.REAL_R1_PRO_SHARPA_MAXINSIGHTS, + EmbodimentTag.REAL_R1_PRO_SHARPA_MECKA, + } +) +"""Tags baked into the base model (nvidia/GR00T-N1.7-3B) — usable without finetuning.""" + +POSTTRAIN_TAGS: frozenset[EmbodimentTag] = frozenset( + { + EmbodimentTag.UNITREE_G1, + EmbodimentTag.UNITREE_G1_SONIC, + EmbodimentTag.SIMPLER_ENV_GOOGLE, + EmbodimentTag.SIMPLER_ENV_WIDOWX, + EmbodimentTag.LIBERO_PANDA, + } +) +"""Tags that require a finetuned checkpoint.""" + +FINETUNE_ONLY_TAGS: frozenset[EmbodimentTag] = frozenset( + { + EmbodimentTag.NEW_EMBODIMENT, + } +) +"""Tags for custom robots (finetuning only, not in any shipped checkpoint).""" diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/__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/diffusion/models/gr00t/dataio/state_action/action_chunking.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py new file mode 100644 index 00000000000..4128d6cae6c --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/action_chunking.py @@ -0,0 +1,654 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Sequence +from typing import Generic, TypeVar + +import numpy as np +from numpy.typing import NDArray +from scipy import interpolate +from scipy.spatial.transform import Rotation, Slerp + +from vllm_omni.diffusion.models.gr00t.dataio.state_action.pose import EndEffectorPose, JointPose, Pose +from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat + +PoseType = TypeVar("PoseType", bound=Pose) + + +class ActionChunk(Generic[PoseType]): + """ + Abstract base class for robot action chunking. + + This class provides common functionality for different action chunking types + including relative and delta action chunking computation with optional reference frames, + interpolation, and format conversion. + """ + + def __init__( + self, + poses: Sequence[PoseType], + times: Sequence[float] | NDArray[np.float64] | None = None, + ): + """ + Initialize action chunking from a list of poses. + + Args: + poses: Sequence of Pose objects + times: Optional sequence of timestamps for each pose. If None, assumes + uniform spacing starting from 0 with step 1.0 + + Raises: + ValueError: If action chunking is empty or times length doesn't match poses + """ + if not poses: + raise ValueError("ActionChunk must contain at least one pose") + + self._poses: list[PoseType] = list(poses) + + # Set up times + if times is None: + self._times = np.arange(len(poses), dtype=np.float64) + else: + if len(times) != len(poses): + raise ValueError("Number of times must match number of poses") + self._times = np.array(times, dtype=np.float64) + + @property + def poses(self) -> list[PoseType]: + """Get the list of poses""" + return self._poses.copy() + + @property + def times(self) -> NDArray[np.float64]: + """Get the timestamps""" + return self._times.copy() + + @property + def num_poses(self) -> int: + """Get the number of poses in the action chunking""" + return len(self._poses) + + def relative_chunking(self, reference_frame: PoseType | None = None) -> "ActionChunk[PoseType]": + """ + Compute the relative action chunking with respect to a reference frame. + + If reference_frame is None, uses the first pose in the action chunking as reference. + All poses are transformed to be relative to the reference frame. + + Args: + reference_frame: Optional reference pose. If None, uses first pose. + + Returns: + A new ActionChunk of the same type where all poses are relative to the reference frame. + """ + if not self._poses: + return self.__class__([], times=[]) + + # Use the first pose as the reference if one is not provided. + ref_pose = reference_frame if reference_frame is not None else self._poses[0] + + # Use the polymorphic subtraction defined in the Pose subclasses. + # The subtraction returns the same type as the operands + relative_poses: list[PoseType] = [pose - ref_pose for pose in self._poses] # type: ignore[misc] + + # Return a new instance of the same action chunking class + # (e.g., JointActionChunk or EndEffectorActionChunk) + return self.__class__(relative_poses, times=self.times) + + def delta_chunking(self, reference_frame: PoseType | None = None) -> "ActionChunk[PoseType]": + """ + Compute the delta action chunking where each pose represents the relative + transformation from the previous frame. + + If reference_frame is provided, it is treated as the first frame, and the + first delta will be from reference_frame to the first pose in the action chunking. + Otherwise, the first pose in the delta action chunking will be the identity/zero transformation. + + Args: + reference_frame: Optional reference pose to use as the first frame. + + Returns: + A new ActionChunk of the same type where each pose is relative to the previous pose. + """ + if not self._poses: + return self.__class__([], times=[]) + + delta_poses: list[PoseType] = [] + + # Determine the initial reference for the very first pose. + # If a reference_frame is given, the first delta is pose[0] - reference_frame. + # If not, the first delta is pose[0] - pose[0], resulting in an identity/zero pose. + prev_pose = reference_frame if reference_frame is not None else self._poses[0] + + for current_pose in self._poses: + delta: PoseType = current_pose - prev_pose # type: ignore[assignment] + delta_poses.append(delta) + prev_pose = current_pose # Update the reference for the next step + + return self.__class__(delta_poses, times=self.times.tolist()) + + def to_absolute_chunking(self, reference_frame: PoseType) -> "ActionChunk[PoseType]": + """ + Convert a relative action chunking to an absolute action chunking by applying + the relative poses on top of a reference frame. + + This is the inverse operation of relative_chunking(). Each relative pose + is composed with the reference frame to produce absolute poses. + + Args: + reference_frame: The reference pose to apply the relative action chunking on top of. + + Returns: + A new ActionChunk of the same type with absolute poses. + + Raises: + NotImplementedError: Must be implemented by subclasses + """ + raise NotImplementedError("Subclasses must implement to_absolute_chunking") + + def interpolate( + self, + num_points: int | None = None, + times: NDArray[np.float64] | None = None, + ) -> "ActionChunk": + """ + Interpolate the action chunking to generate intermediate poses. + Must be implemented by subclasses. + + Args: + num_points: Number of evenly-spaced points to generate + times: Specific timestamps at which to interpolate + + Returns: + A new ActionChunk with interpolated poses + """ + raise NotImplementedError("Subclasses must implement interpolate") + + def to(self, action_format: ActionFormat) -> NDArray[np.float64]: + """ + Convert action chunking to the specified action format. + Must be implemented by subclasses. + + Args: + action_format: The desired output format + + Returns: + Array in the requested format + + Raises: + NotImplementedError: If not implemented by subclass + """ + raise NotImplementedError("Subclasses must implement to method") + + def __len__(self) -> int: + """Return the number of poses in the action chunking""" + return len(self._poses) + + def __getitem__(self, index: int) -> PoseType: + """Get a pose by index""" + return self._poses[index] + + def __repr__(self) -> str: + return ( + f"{self.__class__.__name__}(num_poses={len(self._poses)}, " + f"time_range=[{self._times[0]:.2f}, {self._times[-1]:.2f}])" + ) + + +class JointActionChunk(ActionChunk[JointPose]): + """ + Represents action chunking in joint space as a sequence of joint configurations. + + Examples: + # Create a joint action chunking + joint_poses = [ + JointPose([0.0, 0.0, 0.0]), + JointPose([0.5, 0.5, 0.5]), + JointPose([1.0, 1.0, 1.0]), + ] + action_chunking = JointActionChunk(joint_poses) + + # Get relative trajectory (all poses relative to first pose) + relative_traj = action_chunking.relative_chunking() + + # Get relative trajectory with custom reference + reference = JointPose([0.1, 0.1, 0.1]) + relative_traj = action_chunking.relative_chunking(reference_frame=reference) + + # Get delta trajectory (incremental changes) + delta_traj = action_chunking.delta_chunking() + + # Convert relative trajectory back to absolute + reference = JointPose([0.1, 0.1, 0.1]) + absolute_traj = relative_traj.to_absolute_chunking(reference_frame=reference) + + # Interpolate trajectory + interpolated = action_chunking.interpolate(num_points=10) + + # Convert to desired format + from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat + array_data = action_chunking.to(ActionFormat.DEFAULT) # Returns joint array + """ + + def __init__( + self, + poses: Sequence[JointPose], + times: Sequence[float] | NDArray[np.float64] | None = None, + ): + """ + Initialize a joint trajectory from a list of joint poses. + + Args: + poses: Sequence of JointPose objects + times: Optional sequence of timestamps for each pose + + Raises: + TypeError: If poses are not all JointPose objects + """ + # Validate all poses are JointPose + if not all(isinstance(p, JointPose) for p in poses): + raise TypeError("All poses must be JointPose objects for JointActionChunk") + + super().__init__(poses, times) + + def interpolate( + self, + num_points: int | None = None, + times: NDArray[np.float64] | None = None, + ) -> "JointActionChunk": + """ + Interpolate the joint action chunking to generate intermediate configurations. + + Uses linear interpolation for joint values. + + Args: + num_points: Number of evenly-spaced points to generate (including endpoints). + Only used if times is None. + times: Specific timestamps at which to interpolate. If provided, + num_points is ignored. + + Returns: + A new JointActionChunk with interpolated poses + + Raises: + ValueError: If neither num_points nor times is provided, or if + interpolation times are outside the trajectory range + """ + if num_points is None and times is None: + raise ValueError("Must provide either num_points or times") + + if len(self._poses) < 2: + raise ValueError("Need at least 2 poses for interpolation") + + # Prepare data: extract joint values + timestamps = self._times.copy() + joint_values = np.array([pose.joints for pose in self._poses]) # (N, num_joints) + + # Find and remove non-monotonic timestamps + drop_indices = [idx for idx in range(1, len(timestamps)) if timestamps[idx] <= timestamps[idx - 1]] + + if drop_indices: + for idx in drop_indices: + print( + f"Dropping timestamp pair - Previous: {timestamps[idx - 1]}, " + f"Current: {timestamps[idx]} at index {idx}" + ) + timestamps = np.delete(timestamps, drop_indices) + joint_values = np.delete(joint_values, drop_indices, axis=0) + + # Check if we still have enough poses after cleanup + if len(timestamps) < 2: + raise ValueError("Need at least 2 poses with monotonic timestamps for interpolation") + + # Create interpolator + joint_interp = interpolate.interp1d(timestamps, joint_values, kind="linear", axis=0) + + # Generate interpolation times if not provided + if times is None: + assert num_points is not None # Type narrowing for type checker + interp_times = np.linspace(timestamps[0], timestamps[-1], num_points) + else: + interp_times = np.array(times, dtype=np.float64) + + # Check that interpolation times are within bounds + if np.any(interp_times < timestamps[0]) or np.any(interp_times > timestamps[-1]): + raise ValueError(f"Interpolation times must be within [{timestamps[0]}, {timestamps[-1]}]") + + # Interpolate joint values + interp_joint_values = joint_interp(interp_times) + + # Create interpolated poses + joint_names = self._poses[0].joint_names + interpolated_poses = [ + JointPose(joints=interp_joint_values[i], joint_names=joint_names) for i in range(len(interp_times)) + ] + + return JointActionChunk(interpolated_poses, times=interp_times) + + def to_array(self) -> NDArray[np.float64]: + """ + Convert trajectory to array of joint values. + + Returns: + Array with shape (N, num_joints) where N is the number of poses + """ + return np.array([pose.joints for pose in self._poses]) + + def to_absolute_chunking(self, reference_frame: JointPose) -> "JointActionChunk": + """ + Convert a relative joint action chunking to an absolute action chunking by adding + the relative joint positions to the reference frame. + + This is the inverse operation of relative_chunking(). Each relative joint + configuration is added to the reference frame to produce absolute joint positions. + + Args: + reference_frame: The reference joint pose to apply the relative trajectory on top of. + + Returns: + A new JointActionChunk with absolute joint positions. + + Raises: + ValueError: If joint dimensions don't match + """ + if not self._poses: + return JointActionChunk([], times=[]) + + if len(self._poses[0].joints) != len(reference_frame.joints): + raise ValueError( + f"Cannot apply relative trajectory: " + f"joint dimensions don't match ({len(self._poses[0].joints)} vs {len(reference_frame.joints)})" + ) + + # Add each relative pose to the reference frame + absolute_poses: list[JointPose] = [] + for relative_pose in self._poses: + absolute_joints = reference_frame.joints + relative_pose.joints + absolute_pose = JointPose(joints=absolute_joints, joint_names=reference_frame.joint_names) + absolute_poses.append(absolute_pose) + + return JointActionChunk(absolute_poses, times=self.times) + + def to(self, action_format: ActionFormat) -> NDArray[np.float64]: + """ + Convert trajectory to the desired format. + + Args: + action_format: The desired output format + + Returns: + Array in the requested format + + Raises: + ValueError: If the action format is not supported for joint trajectories + """ + if action_format == ActionFormat.DEFAULT: + return self.to_array() + else: + raise ValueError( + f"ActionFormat {action_format} is not supported for JointActionChunk. " + f"Only {ActionFormat.DEFAULT} is supported." + ) + + +class EndEffectorActionChunk(ActionChunk[EndEffectorPose]): + """ + Represents action chunking in Cartesian space as a sequence of end-effector poses. + + Examples: + # Create an end-effector action chunking + ee_poses = [ + EndEffectorPose(translation=[0, 0, 0], rotation=[1, 0, 0, 0], + rotation_type="quat", rotation_order="wxyz"), + EndEffectorPose(translation=[1, 0, 0], rotation=[0.707, 0, 0, 0.707], + rotation_type="quat", rotation_order="wxyz"), + EndEffectorPose(translation=[2, 0, 0], rotation=[0, 0, 0, 1], + rotation_type="quat", rotation_order="wxyz"), + ] + action_chunking = EndEffectorActionChunk(ee_poses) + + # Get relative trajectory (all poses relative to first pose) + relative_traj = action_chunking.relative_chunking() + + # Get relative trajectory with custom reference frame + reference = EndEffectorPose(translation=[0.5, 0, 0], rotation=[1,0,0,0], + rotation_type="quat", rotation_order="wxyz") + relative_traj = action_chunking.relative_chunking(reference_frame=reference) + + # Get delta trajectory + delta_traj = action_chunking.delta_chunking() + + # Convert relative trajectory back to absolute + reference = EndEffectorPose(translation=[0.5, 0, 0], rotation=[1,0,0,0], + rotation_type="quat", rotation_order="wxyz") + absolute_traj = relative_traj.to_absolute_chunking(reference_frame=reference) + + # Interpolate trajectory + interpolated = action_chunking.interpolate(num_points=10) + + # Convert to desired format + from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat + homo_matrices = action_chunking.to(ActionFormat.DEFAULT) # (N, 4, 4) homogeneous matrices + xyz_rot6d = action_chunking.to(ActionFormat.XYZ_ROT6D) # (N, 9) xyz + rot6d + xyz_rotvec = action_chunking.to(ActionFormat.XYZ_ROTVEC) # (N, 6) xyz + rotvec + """ + + def __init__( + self, + poses: Sequence[EndEffectorPose], + times: Sequence[float] | NDArray[np.float64] | None = None, + ): + """ + Initialize an end-effector trajectory from a list of end-effector poses. + + Args: + poses: Sequence of EndEffectorPose objects + times: Optional sequence of timestamps for each pose + + Raises: + TypeError: If poses are not all EndEffectorPose objects + """ + # Validate all poses are EndEffectorPose + if not all(isinstance(p, EndEffectorPose) for p in poses): + raise TypeError("All poses must be EndEffectorPose objects for EndEffectorActionChunk") + + super().__init__(poses, times) + + @classmethod + def from_array(cls, data: np.ndarray, action_format: ActionFormat) -> "EndEffectorActionChunk": + """ + Create an EndEffectorActionChunk from a 2-D array using the specified action format. + + This is the inverse of ``.to(action_format)``. + + Args: + data: Array of shape (N, D) where D depends on the action_format. + action_format: The format that describes the layout of each row. + + Returns: + EndEffectorActionChunk with N poses. + """ + poses = [EndEffectorPose.from_action_format(row, action_format) for row in data] + return cls(poses) + + def interpolate( + self, + num_points: int | None = None, + times: NDArray[np.float64] | None = None, + ) -> "EndEffectorActionChunk": + """ + Interpolate the action chunking to generate intermediate poses. + + Uses linear interpolation for translation and SLERP (Spherical Linear + Interpolation) for rotation. + + Args: + num_points: Number of evenly-spaced points to generate (including endpoints). + Only used if times is None. + times: Specific timestamps at which to interpolate. If provided, + num_points is ignored. + + Returns: + A new EndEffectorActionChunk with interpolated poses + + Raises: + ValueError: If neither num_points nor times is provided, or if + interpolation times are outside the trajectory range + """ + if num_points is None and times is None: + raise ValueError("Must provide either num_points or times") + + if len(self._poses) < 2: + raise ValueError("Need at least 2 poses for interpolation") + + # Prepare data: extract positions and rotations + timestamps = self._times.copy() + homogeneous_matrices = np.array([pose.homogeneous for pose in self._poses]) + positions = homogeneous_matrices[:, :3, 3] + rotations = Rotation.from_matrix(homogeneous_matrices[:, :3, :3]) + + # Find indices where timestamps are not monotonically increasing + drop_indices = [idx for idx in range(1, len(timestamps)) if timestamps[idx] <= timestamps[idx - 1]] + + # Remove the problematic timestamps and corresponding data + if drop_indices: + for idx in drop_indices: + print( + f"Dropping timestamp pair - Previous: {timestamps[idx - 1]}, " + f"Current: {timestamps[idx]} at index {idx}" + ) + timestamps = np.delete(timestamps, drop_indices) + positions = np.delete(positions, drop_indices, axis=0) + rotations = Rotation.from_matrix(np.delete(homogeneous_matrices[:, :3, :3], drop_indices, axis=0)) + + # Check if we still have enough poses after cleanup + if len(timestamps) < 2: + raise ValueError("Need at least 2 poses with monotonic timestamps for interpolation") + + # Create interpolators + pos_interp = interpolate.interp1d(timestamps, positions, kind="linear", axis=0) + rot_interp = Slerp(timestamps, rotations) + + # Generate interpolation times if not provided + if times is None: + assert num_points is not None # Type narrowing for type checker + interp_times = np.linspace(timestamps[0], timestamps[-1], num_points) + else: + interp_times = np.array(times, dtype=np.float64) + + # Check that interpolation times are within bounds + if np.any(interp_times < timestamps[0]) or np.any(interp_times > timestamps[-1]): + raise ValueError(f"Interpolation times must be within [{timestamps[0]}, {timestamps[-1]}]") + + # Interpolate positions and rotations + interp_positions = pos_interp(interp_times) + interp_rotations = rot_interp(interp_times) + + # Create interpolated poses + interpolated_poses = [] + for i in range(len(interp_times)): + pose = EndEffectorPose( + translation=interp_positions[i], + rotation=interp_rotations[i].as_matrix(), + rotation_type="matrix", + ) + interpolated_poses.append(pose) + + return EndEffectorActionChunk(interpolated_poses, times=interp_times) + + def to_homogeneous_matrices(self) -> NDArray[np.float64]: + """ + Convert trajectory to array of homogeneous transformation matrices. + + Returns: + Array of homogeneous matrices with shape (N, 4, 4) where N is the number of poses + """ + return np.array([pose.homogeneous for pose in self._poses]) + + def to_translation_rot6d(self) -> NDArray[np.float64]: + """ + Convert trajectory to array of translations and 6D rotations. + + Returns: + Array with shape (N, 9) - 3 for xyz + 6 for rot6d + """ + translations = np.array([pose.translation for pose in self._poses]) # (N, 3) + rotations = np.array([pose.rot6d for pose in self._poses]) # (N, 6) + + # Concatenate translation and rotation + xyz_rot6d = np.concatenate([translations, rotations], axis=1) # (N, 9) + + return xyz_rot6d + + def to_translation_rotvec(self) -> NDArray[np.float64]: + """ + Convert trajectory to array of translations and rotation vectors. + + Returns: + Array with shape (N, 6) - 3 for xyz + 3 for rotvec + """ + translations = np.array([pose.translation for pose in self._poses]) # (N, 3) + rotations = np.array([pose.rotvec for pose in self._poses]) # (N, 3) + + # Concatenate translation and rotation + xyz_rotvec = np.concatenate([translations, rotations], axis=1) # (N, 6) + + return xyz_rotvec + + def to_absolute_chunking(self, reference_frame: EndEffectorPose) -> "EndEffectorActionChunk": + """ + Convert a relative end-effector action chunking to an absolute action chunking by + composing each relative transformation with the reference frame. + + This is the inverse operation of relative_chunking(). Each relative pose + represents a transformation that is applied on top of the reference frame + to produce absolute poses. + + Args: + reference_frame: The reference end-effector pose to apply the relative trajectory on top of. + + Returns: + A new EndEffectorActionChunk with absolute poses. + """ + if not self._poses: + return EndEffectorActionChunk([], times=[]) + + # Get reference frame as homogeneous matrix + T_ref = reference_frame.homogeneous + + # Compose each relative transformation with the reference frame + absolute_poses: list[EndEffectorPose] = [] + for relative_pose in self._poses: + # Get relative transformation as homogeneous matrix + T_relative = relative_pose.homogeneous + + # Compose transformations: T_absolute = T_ref @ T_relative + T_absolute = T_ref @ T_relative + + # Create absolute pose from composed transformation + absolute_pose = EndEffectorPose(homogeneous=T_absolute) + absolute_poses.append(absolute_pose) + + return EndEffectorActionChunk(absolute_poses, times=self.times) + + def to(self, action_format: ActionFormat) -> NDArray[np.float64]: + """ + Convert trajectory to the desired format. + + Args: + action_format: The desired output format + + Returns: + Array in the requested format + + Raises: + ValueError: If the action format is not supported + """ + if action_format == ActionFormat.DEFAULT: + return self.to_homogeneous_matrices() + elif action_format == ActionFormat.XYZ_ROT6D: + return self.to_translation_rot6d() + elif action_format == ActionFormat.XYZ_ROTVEC: + return self.to_translation_rotvec() + else: + raise ValueError(f"Unsupported action format: {action_format}") diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py new file mode 100644 index 00000000000..0e0849f17b4 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/pose.py @@ -0,0 +1,709 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import Enum +from typing import TypeVar + +import numpy as np +from numpy.typing import NDArray +from scipy.spatial.transform import Rotation + +from vllm_omni.diffusion.models.gr00t.dataio.types import ActionFormat + +# TypeVar for self-type preservation in Pose operations +PoseT = TypeVar("PoseT", bound="Pose") + + +def invert_transformation(transform: NDArray[np.float64]) -> NDArray[np.float64]: + """ + Invert a homogeneous transformation matrix. + + Args: + transform: A 4x4 homogeneous transformation matrix + + Returns: + The inverse of the transformation matrix (4x4) + """ + R = transform[:3, :3] # Extract the rotation matrix + t = transform[:3, 3] # Extract the translation vector + + # Inverse of the rotation matrix is its transpose (since it's orthogonal) + R_inv = R.T + + # Inverse of the translation is -R_inv * t + t_inv = -R_inv @ t + + # Construct the inverse transformation matrix + T_inv = np.eye(4) + T_inv[:3, :3] = R_inv + T_inv[:3, 3] = t_inv + + return T_inv + + +def relative_transformation( + base_transform: NDArray[np.float64], target_transform: NDArray[np.float64] +) -> NDArray[np.float64]: + """ + Compute the relative transformation between two poses. + + Args: + base_transform: Initial 4x4 homogeneous transformation matrix + target_transform: Current 4x4 homogeneous transformation matrix + + Returns: + The relative transformation matrix (4x4) from base_transform to target_transform + """ + # Relative transformation is base_transform^{-1} * target_transform + T_relative = invert_transformation(base_transform) @ target_transform + return T_relative + + +class RotationType(Enum): + """Supported rotation representation types""" + + QUAT = "quat" + EULER = "euler" + ROTVEC = "rotvec" + MATRIX = "matrix" + ROT6D = "rot6d" + + +class EulerOrder(Enum): + """Common Euler angle conventions""" + + XYZ = "xyz" + ZYX = "zyx" + XZY = "xzy" + YXZ = "yxz" + YZX = "yzx" + ZXY = "zxy" + + +class QuatOrder(Enum): + """Quaternion ordering conventions""" + + WXYZ = "wxyz" # scalar-first (w, x, y, z) + XYZW = "xyzw" # scalar-last (x, y, z, w) + + +class Pose: + """ + Abstract base class for robot poses. + + This class provides common functionality for different pose representations + including relative pose computation via the subtraction operator. + """ + + pose_type: str + + def __sub__(self: PoseT, other: PoseT) -> PoseT: + """ + Compute relative transformation between two poses. + + For EndEffectorPose: Computes the relative transformation from other to self. + Result represents the transformation needed to go from other's frame to self's frame. + + For JointPose: Computes the joint-space difference (self - other). + + Args: + other: The reference pose to compute relative transformation from + + Returns: + Relative pose (same type as self) + + Raises: + TypeError: If poses are not of the same type + + Examples: + # End-effector poses + pose1 = EndEffectorPose(translation=[1, 0, 0], rotation=[1,0,0,0], + rotation_type="quat", rotation_order="wxyz") + pose2 = EndEffectorPose(translation=[2, 0, 0], rotation=[1,0,0,0], + rotation_type="quat", rotation_order="wxyz") + relative = pose2 - pose1 # Transformation from pose1 to pose2 + + # Joint poses + joint1 = JointPose([0.0, 0.5, 1.0]) + joint2 = JointPose([0.1, 0.6, 1.2]) + joint_diff = joint2 - joint1 # Joint differences: [0.1, 0.1, 0.2] + """ + if type(self) is not type(other): + raise TypeError( + f"Cannot compute relative transformation between different pose types: " + f"{type(self).__name__} and {type(other).__name__}" + ) + + return self._compute_relative(other) + + def _compute_relative(self: PoseT, other: PoseT) -> PoseT: + """ + Internal method to compute relative transformation. + Must be implemented by subclasses. + + Args: + other: The reference pose + + Returns: + Relative pose + """ + raise NotImplementedError("Subclasses must implement _compute_relative") + + def copy(self: PoseT) -> PoseT: + """ + Create a deep copy of this pose. + Must be implemented by subclasses. + + Returns: + New Pose instance with copied data + """ + raise NotImplementedError("Subclasses must implement copy") + + +class JointPose(Pose): + """ + Represents a robot configuration in joint space. + + This class stores joint angles/positions for a robot manipulator. + Unlike end-effector poses, joint poses represent the configuration + of all joints in the kinematic chain. + + Examples: + # Create a 6-DOF joint configuration + joint_pose = JointPose( + joints=[0.0, -np.pi/4, np.pi/2, 0.0, np.pi/4, 0.0], + joint_names=["shoulder_pan", "shoulder_lift", "elbow", + "wrist_1", "wrist_2", "wrist_3"] + ) + + # Create with default joint names + joint_pose = JointPose(joints=[0.0, 0.5, 1.0]) + + # Get as dictionary + joint_dict = joint_pose.to_dict() # {"joint_0": 0.0, ...} + + # Access individual joints + first_joint = joint_pose.joints[0] + num_joints = joint_pose.num_joints + + # Compute relative joint displacement + joint1 = JointPose([0.0, 0.5, 1.0]) + joint2 = JointPose([0.1, 0.6, 1.2]) + relative = joint2 - joint1 # [0.1, 0.1, 0.2] + """ + + pose_type = "joint" + + def __init__( + self, + joints: list | np.ndarray, + joint_names: list | None = None, + ): + """ + Initialize a joint pose. + + Args: + joints: Joint angles/positions as array-like of shape (n,) + joint_names: Optional list of names for each joint. If None, + defaults to ["joint_0", "joint_1", ...] + """ + super().__init__() + self.joints = np.array(joints, dtype=np.float64) + + # Set defaults and validate joint_names + if joint_names is None: + self.joint_names = [f"joint_{i}" for i in range(len(self.joints))] + else: + if len(joint_names) != len(self.joints): + raise ValueError( + f"Number of joint names ({len(joint_names)}) must match number of joints ({len(self.joints)})" + ) + self.joint_names = joint_names + + @property + def num_joints(self) -> int: + """ + Get the number of joints. + + Returns: + Number of joints in the configuration + """ + return len(self.joints) + + def to_dict(self) -> dict: + """ + Convert joint configuration to dictionary. + + Returns: + Dictionary mapping joint names to joint values + """ + return dict(zip(self.joint_names, self.joints)) + + def _compute_relative(self, other): # type: ignore[override] + """ + Compute relative joint displacement. + + Args: + other: Reference joint pose + + Returns: + JointPose representing the joint-space difference (self - other) + + Raises: + ValueError: If joint dimensions don't match + """ + if len(self.joints) != len(other.joints): + raise ValueError( + f"Cannot compute relative joint pose: " + f"joint dimensions don't match ({len(self.joints)} vs {len(other.joints)})" + ) + + relative_joints = self.joints - other.joints + return JointPose(joints=relative_joints, joint_names=self.joint_names) + + def copy(self) -> "JointPose": + """ + Create a deep copy of this joint pose. + + Returns: + New JointPose instance with copied data + """ + return JointPose( + joints=self.joints.copy(), + joint_names=self.joint_names.copy(), + ) + + def __repr__(self) -> str: + if len(self.joints) <= 6: + joints_str = np.array2string(self.joints, precision=4, suppress_small=True) + else: + joints_str = f"[{self.joints[0]:.4f}, ..., {self.joints[-1]:.4f}] ({len(self.joints)} joints)" + + return f"JointPose(joints={joints_str})" + + def __eq__(self, other) -> bool: + if not isinstance(other, JointPose): + return False + return np.allclose(self.joints, other.joints) and self.joint_names == other.joint_names + + def __getitem__(self, index) -> float | NDArray[np.float64]: + """Allow indexing: joint_pose[0] returns first joint value""" + return self.joints[index] + + def __len__(self) -> int: + """Allow len(): len(joint_pose) returns number of joints""" + return len(self.joints) + + +class EndEffectorPose(Pose): + """ + Represents a single end-effector pose with translation and rotation components. + + This class handles Cartesian space representations of robot end-effector poses, + supporting multiple rotation representations (quaternions, Euler angles, rotation + vectors, rotation matrices, etc.). + + Examples: + # Create with quaternion (wxyz order) + pose = EndEffectorPose( + translation=[1.0, 2.0, 3.0], + rotation=[1.0, 0.0, 0.0, 0.0], + rotation_type="quat", + rotation_order="wxyz" + ) + + # Create with Euler angles (degrees by default) + pose = EndEffectorPose( + translation=[1, 2, 3], + rotation=[0, 0, 90], + rotation_type="euler", + rotation_order="xyz" + ) + + # Create with Euler angles in radians + pose = EndEffectorPose( + translation=[1, 2, 3], + rotation=[0, 0, np.pi/2], + rotation_type="euler", + rotation_order="xyz", + degrees=False + ) + + # Create from homogeneous matrix + H = np.eye(4) + H[:3, 3] = [1, 2, 3] + pose = EndEffectorPose(homogeneous=H) + + # Convert between representations + quat_wxyz = pose.to_rotation("quat", "wxyz") + euler_zyx = pose.to_rotation("euler", "zyx") + rot6d = pose.to_rotation("rot6d") + + # Compute relative transformation + pose1 = EndEffectorPose(translation=[1, 0, 0], rotation=[1,0,0,0], + rotation_type="quat", rotation_order="wxyz") + pose2 = EndEffectorPose(translation=[2, 0, 0], rotation=[1,0,0,0], + rotation_type="quat", rotation_order="wxyz") + relative = pose2 - pose1 # Transformation from pose1's frame to pose2's frame + """ + + pose_type = "end_effector" + + def __init__( + self, + translation: list | np.ndarray | None = None, + rotation: list | np.ndarray | None = None, + rotation_type: str | None = None, + rotation_order: str | None = None, + homogeneous: np.ndarray | None = None, + degrees: bool = True, + ): + """ + Initialize an end-effector pose. + + Args: + translation: Translation vector [x, y, z] + rotation: Rotation in specified format + rotation_type: Type of rotation ("quat", "euler", "rotvec", "matrix", "rot6d") + rotation_order: Order/convention for the rotation type + homogeneous: Homogeneous transformation matrix (4, 4) + If provided, overrides translation and rotation + degrees: For Euler angles, whether the input is in degrees (default True) + """ + super().__init__() + + # Cache for homogeneous matrix + self._homogeneous_cache: NDArray[np.float64] | None = None + self._cache_valid = False + + # Handle homogeneous matrix input + if homogeneous is not None: + self._from_homogeneous(homogeneous) + return + + # Store translation + self._translation = np.array(translation) if translation is not None else np.zeros(3) + + # Store rotation as scipy Rotation object internally + if rotation is not None: + if rotation_type is None: + raise ValueError("rotation_type must be specified when rotation is provided") + self._set_rotation(rotation, rotation_type, rotation_order, degrees) + else: + self._rotation = Rotation.identity() + + def _from_homogeneous(self, homogeneous: np.ndarray): + """Initialize from homogeneous transformation matrix""" + homogeneous = np.array(homogeneous) + + # Extract translation (last column, first 3 rows) + self._translation = homogeneous[:3, 3] + + # Extract rotation matrix (top-left 3x3) + rotation_matrix = homogeneous[:3, :3] + + # Create Rotation object from matrix + self._rotation = Rotation.from_matrix(rotation_matrix) + + @staticmethod + def _rot6d_to_matrix(rot6d: np.ndarray) -> np.ndarray: + """ + Convert 6D rotation representation to rotation matrix. + + Args: + rot6d: 6D rotation as (6,) array - first two rows of rotation matrix flattened + + Returns: + Rotation matrix (3, 3) + """ + rot6d = rot6d.reshape(2, 3) + + # First two rows + row1 = rot6d[0] + row2 = rot6d[1] + + # Normalize first row + row1 = row1 / np.linalg.norm(row1) + + # Gram-Schmidt orthogonalization for second row + row2 = row2 - np.dot(row1, row2) * row1 + row2 = row2 / np.linalg.norm(row2) + + # Third row is cross product + row3 = np.cross(row1, row2) + + # Construct rotation matrix + rotation_matrix = np.vstack([row1, row2, row3]) + + return rotation_matrix + + @staticmethod + def _matrix_to_rot6d(rotation_matrix: np.ndarray) -> np.ndarray: + """ + Convert rotation matrix to 6D rotation representation. + + Args: + rotation_matrix: Rotation matrix (3, 3) + + Returns: + 6D rotation - (6,) array (first two rows flattened) + """ + return rotation_matrix[:2, :].flatten() + + def _set_rotation( + self, + rotation: list | np.ndarray, + rotation_type: str, + rotation_order: str | None = None, + degrees: bool = True, + ): + """Internal method to set rotation from various representations""" + rotation = np.array(rotation) + rot_type = RotationType(rotation_type.lower()) + + if rot_type == RotationType.QUAT: + quat_order = QuatOrder(rotation_order.lower()) if rotation_order else QuatOrder.WXYZ + if quat_order == QuatOrder.WXYZ: + # scipy uses xyzw order, so convert + quat_xyzw = np.array([rotation[1], rotation[2], rotation[3], rotation[0]]) + else: + quat_xyzw = rotation + self._rotation = Rotation.from_quat(quat_xyzw) + + elif rot_type == RotationType.EULER: + euler_order = EulerOrder(rotation_order.lower()) if rotation_order else EulerOrder.XYZ + self._rotation = Rotation.from_euler(euler_order.value, rotation, degrees=degrees) + + elif rot_type == RotationType.ROTVEC: + self._rotation = Rotation.from_rotvec(rotation) + + elif rot_type == RotationType.MATRIX: + self._rotation = Rotation.from_matrix(rotation) + + elif rot_type == RotationType.ROT6D: + rotation_matrix = self._rot6d_to_matrix(rotation) + self._rotation = Rotation.from_matrix(rotation_matrix) + + else: + raise ValueError(f"Unknown rotation type: {rotation_type}") + + # Invalidate cache + self._cache_valid = False + + @property + def translation(self) -> np.ndarray: + """ + Get translation vector. + + Returns: + Translation array - shape (3,) + """ + return self._translation.copy() + + @property + def quat_wxyz(self) -> np.ndarray: + """Get rotation as quaternion in wxyz order (w, x, y, z)""" + return self.to_rotation("quat", "wxyz") + + @property + def quat_xyzw(self) -> np.ndarray: + """Get rotation as quaternion in xyzw order (x, y, z, w)""" + return self.to_rotation("quat", "xyzw") + + @property + def euler_xyz(self) -> np.ndarray: + """Get rotation as Euler angles in xyz order (degrees)""" + return self.to_rotation("euler", "xyz") + + @property + def rotvec(self) -> np.ndarray: + """Get rotation as rotation vector (axis-angle)""" + return self.to_rotation("rotvec") + + @property + def rotation_matrix(self) -> np.ndarray: + """Get rotation as 3x3 rotation matrix""" + return self.to_rotation("matrix") + + @property + def rot6d(self) -> np.ndarray: + """Get rotation as 6D representation (first two rows of rotation matrix)""" + return self.to_rotation("rot6d") + + @property + def xyz_rot6d(self) -> np.ndarray: + """Get pose as concatenated translation and 6D rotation (9,)""" + return np.concatenate([self._translation, self.rot6d]) + + @property + def xyz_rotvec(self) -> np.ndarray: + """Get pose as concatenated translation and rotation vector (6,)""" + return np.concatenate([self._translation, self.rotvec]) + + @property + def homogeneous(self) -> np.ndarray: + """ + Get homogeneous transformation matrix. + + Returns: + Homogeneous matrix - shape (4, 4) + """ + if not self._cache_valid: + self._homogeneous_cache = self._compute_homogeneous() + self._cache_valid = True + assert self._homogeneous_cache is not None + return self._homogeneous_cache.copy() + + def _compute_homogeneous(self) -> np.ndarray: + """Compute homogeneous transformation matrix""" + H = np.eye(4) + H[:3, :3] = self._rotation.as_matrix() + H[:3, 3] = self._translation + return H + + def to_rotation( + self, + rotation_type: str, + rotation_order: str | None = None, + degrees: bool = True, + ) -> np.ndarray: + """ + Get rotation in specified representation. + + Args: + rotation_type: Desired type ("quat", "euler", "rotvec", "matrix", "rot6d") + rotation_order: Order/convention for the rotation type + degrees: For Euler angles, return in degrees (default True) + + Returns: + Rotation in requested format + - Shape (4,) for quat + - Shape (3,) for euler/rotvec + - Shape (6,) for rot6d + - Shape (3, 3) for matrix + """ + rot_type = RotationType(rotation_type.lower()) + + if rot_type == RotationType.ROT6D: + rotation_matrix = self._rotation.as_matrix() + return self._matrix_to_rot6d(rotation_matrix) + + if rot_type == RotationType.QUAT: + quat_order = QuatOrder(rotation_order.lower()) if rotation_order else QuatOrder.WXYZ + quat_xyzw = self._rotation.as_quat() + if quat_order == QuatOrder.WXYZ: + return np.array([quat_xyzw[3], quat_xyzw[0], quat_xyzw[1], quat_xyzw[2]]) + else: + return quat_xyzw + + elif rot_type == RotationType.EULER: + euler_order = EulerOrder(rotation_order.lower()) if rotation_order else EulerOrder.XYZ + return self._rotation.as_euler(euler_order.value, degrees=degrees) + + elif rot_type == RotationType.ROTVEC: + return self._rotation.as_rotvec() + + elif rot_type == RotationType.MATRIX: + return self._rotation.as_matrix() + + else: + raise ValueError(f"Unknown rotation type: {rotation_type}") + + def to_homogeneous(self) -> np.ndarray: + """ + Convert pose to homogeneous transformation matrix. + (Alias for the homogeneous property) + + Returns: + Homogeneous matrix - shape (4, 4) + """ + return self.homogeneous + + def set_rotation( + self, + rotation: list | np.ndarray, + rotation_type: str, + rotation_order: str | None = None, + degrees: bool = True, + ): + """ + Set rotation from specified representation. + + Args: + rotation: Rotation data + rotation_type: Type of rotation ("quat", "euler", "rotvec", "matrix", "rot6d") + rotation_order: Order/convention for the rotation type + degrees: For Euler angles, whether the input is in degrees (default True) + """ + self._set_rotation(rotation, rotation_type, rotation_order, degrees) + + def _compute_relative(self, other): # type: ignore[override] + """ + Compute relative transformation from other to self. + + The result represents the transformation needed to go from other's frame to self's frame. + Mathematically: T_relative = T_other^{-1} * T_self + + Args: + other: Reference end-effector pose + + Returns: + EndEffectorPose representing the relative transformation + """ + # Get homogeneous matrices + T_self = self.homogeneous + T_other = other.homogeneous + + # Compute relative transformation: T_other^{-1} * T_self + T_relative = relative_transformation(T_other, T_self) + + # Create new EndEffectorPose from relative transformation + return EndEffectorPose(homogeneous=T_relative) + + @classmethod + def from_action_format(cls, data: np.ndarray, action_format: ActionFormat) -> "EndEffectorPose": + """ + Create an EndEffectorPose from a flat array using the specified action format. + + This is the inverse of the xyz_rot6d / xyz_rotvec / homogeneous properties. + + Args: + data: Flat array whose layout depends on action_format. + action_format: One of ActionFormat.XYZ_ROT6D, XYZ_ROTVEC, or DEFAULT. + + Returns: + EndEffectorPose instance. + """ + if action_format == ActionFormat.XYZ_ROT6D: + return cls(translation=data[:3], rotation=data[3:], rotation_type="rot6d") + elif action_format == ActionFormat.XYZ_ROTVEC: + return cls(translation=data[:3], rotation=data[3:], rotation_type="rotvec") + elif action_format == ActionFormat.DEFAULT: + return cls(homogeneous=data.reshape(4, 4)) + else: + raise ValueError(f"Unsupported ActionFormat: {action_format}") + + def copy(self) -> "EndEffectorPose": + """ + Create a deep copy of this end-effector pose. + + Returns: + New EndEffectorPose instance with copied data + """ + return EndEffectorPose( + translation=self._translation.copy(), + rotation=self._rotation.as_quat(), + rotation_type="quat", + rotation_order="xyzw", + ) + + def __repr__(self) -> str: + quat = self.to_rotation("quat", "wxyz") + return f"EndEffectorPose(translation={self.translation}, rotation_quat_wxyz={quat})" + + def __eq__(self, other) -> bool: + if not isinstance(other, EndEffectorPose): + return False + return np.allclose(self._translation, other._translation) and np.allclose( + self._rotation.as_quat(), other._rotation.as_quat() + ) diff --git a/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py b/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py new file mode 100644 index 00000000000..14e2b0b3a37 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/state_action/state_action_processor.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unified state and action processor for robotics.""" + +from copy import deepcopy + +import numpy as np + +from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ( + ActionFormat, + ActionRepresentation, + ActionType, + ModalityConfig, +) +from vllm_omni.diffusion.models.gr00t.dataio.state_action.action_chunking import ( + EndEffectorActionChunk, + JointActionChunk, +) +from vllm_omni.diffusion.models.gr00t.dataio.state_action.pose import EndEffectorPose, JointPose +from vllm_omni.diffusion.models.gr00t.dataio.utils import ( + apply_sin_cos_encoding, + nested_dict_to_numpy, + normalize_values_meanstd, + normalize_values_minmax, + parse_modality_configs, + unnormalize_values_meanstd, + unnormalize_values_minmax, +) + + +class StateActionProcessor: + """ + Unified processor for robot state and action data. + + Handles: + - State normalization (min/max, mean/std, sin/cos encoding) + - Action normalization + - Absolute <-> Relative action representation conversion + - Action processing with state dependency + """ + + def __init__( + self, + modality_configs: dict[str, dict[str, ModalityConfig]], + statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None, + use_percentiles: bool = False, + clip_outliers: bool = True, + apply_sincos_state_encoding: bool = False, + use_relative_action: bool = False, + ): + """ + Initialize unified state and action processor. + + Args: + modality_configs: Nested dict with structure: + {embodiment_tag: {modality: ModalityConfig}} + where modality in ["state", "action"] + Example: {"gr1": {"state": ModalityConfig(...), "action": ModalityConfig(...)}} + statistics: Optional nested dict with structure: + {embodiment_tag: {modality: {joint_group: {stat_type: values}}}} + where modality in ["state", "action", "relative_action"] + and stat_type in ["min", "max", "mean", "std", "q01", "q99"] + Example: {"gr1": {"state": {"left_arm": {"min": [...], "max": [...], ...}}}} + use_percentiles: Whether to use percentiles (q01/q99) instead of min/max + clip_outliers: Whether to clip normalized values to [-1, 1] + apply_sincos_state_encoding: Global flag to enable sin/cos encoding for states + """ + self.modality_configs = parse_modality_configs(modality_configs) + self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {} + self.use_percentiles = use_percentiles + self.clip_outliers = clip_outliers + self.apply_sincos_state_encoding = apply_sincos_state_encoding + self.use_relative_action = use_relative_action + + # Normalization parameters computed from statistics + self.norm_params: dict[str, dict[str, dict[str, dict[str, np.ndarray]]]] = {} + # Format: norm_params[embodiment_tag][modality][joint_group][stat_type] + # where stat_type in ["min", "max", "mean", "std", "dim"] + + if statistics is not None: + self.set_statistics(statistics) + + def set_statistics( + self, + statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]], + override: bool = False, + ) -> None: + """ + Set dataset statistics for normalization. + + Args: + statistics: Nested dict with structure: + {embodiment_tag: {modality: {joint_group: {stat_type: values}}}} + """ + for key in statistics: + if key not in self.statistics or override: + self.statistics[key] = deepcopy(statistics[key]) + else: + print(f"Embodiment tag {key} already in statistics, skipping updating") + self._compute_normalization_parameters() + + def _compute_normalization_parameters(self) -> None: + """Compute and cache normalization parameters from statistics for all embodiments and modalities.""" + for embodiment_tag in self.statistics: + self.norm_params[embodiment_tag] = {} + + for modality in ["state", "action"]: + if modality not in self.statistics[embodiment_tag]: + continue + + self.norm_params[embodiment_tag][modality] = {} + + for joint_group, stats in self.statistics[embodiment_tag][modality].items(): + if self.use_percentiles: + min_vals = np.array(stats["q01"]) + max_vals = np.array(stats["q99"]) + else: + min_vals = np.array(stats["min"]) + max_vals = np.array(stats["max"]) + + mean_vals = np.array(stats["mean"]) + std_vals = np.array(stats["std"]) + + # Compute range, ensuring it's not zero + range_vals = max_vals - min_vals + range_vals = np.maximum(range_vals, 1e-8) + + self.norm_params[embodiment_tag][modality][joint_group] = { + "min": min_vals, + "max": max_vals, + "dim": np.array(range_vals.shape[0]), + "mean": mean_vals, + "std": std_vals, + } + + # Override absolute action stats with relative stats where specified + if "action" in self.modality_configs[embodiment_tag]: + modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys + action_configs = self.modality_configs[embodiment_tag]["action"].action_configs + + if action_configs is not None: + for key, action_config in zip(modality_keys, action_configs): + if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action: + if "relative_action" not in self.statistics[embodiment_tag]: + raise ValueError( + f"Relative action statistics required for embodiment '{embodiment_tag}' " + f"but 'relative_action' not found in statistics" + ) + if key not in self.statistics[embodiment_tag]["relative_action"]: + raise ValueError( + f"Relative action statistics required for key '{key}' " + f"in embodiment '{embodiment_tag}' but not found" + ) + action_dim = self.norm_params[embodiment_tag]["action"][key]["dim"] + self.norm_params[embodiment_tag]["action"][key] = nested_dict_to_numpy( + self.statistics[embodiment_tag]["relative_action"][key] + ) + self.norm_params[embodiment_tag]["action"][key]["dim"] = action_dim + + def apply_state( + self, + state: dict[str, np.ndarray], + embodiment_tag: str, + ) -> dict[str, np.ndarray]: + """ + Apply state processing (normalization, encoding). + + Args: + state: Dict mapping joint_group -> raw state values + Shape per group: (..., D) where D is state dimension + embodiment_tag: Embodiment identifier (e.g., "gr1") + + Returns: + Dict mapping joint_group -> processed state values + - Sin/cos encoded groups: (..., 2*D) + - Other groups: (..., D) + """ + normalized_values = {} + state = deepcopy(state) # Avoid modifying input + + # Get sin/cos embedding keys if enabled + sin_cos_keys = None + if self.apply_sincos_state_encoding: + state_config = self.modality_configs[embodiment_tag].get("state") + if state_config and hasattr(state_config, "sin_cos_embedding_keys"): + sin_cos_keys = state_config.sin_cos_embedding_keys + + for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys: + if joint_group not in state: + raise KeyError(f"Joint group '{joint_group}' not found in state dict for embodiment '{embodiment_tag}'") + + # Strategy 1: Sin/cos encoding (doubles dimension) + if sin_cos_keys and joint_group in sin_cos_keys: + normalized_values[joint_group] = apply_sin_cos_encoding(state[joint_group]) + + # Strategy 2: Mean/std normalization + elif ( + hasattr( + self.modality_configs[embodiment_tag]["state"], + "mean_std_embedding_keys", + ) + and self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys + and joint_group in self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys + ): + params = self.norm_params[embodiment_tag]["state"][joint_group] + normalized = normalize_values_meanstd(state[joint_group], params) + normalized_values[joint_group] = normalized + + # Strategy 3: Min/max normalization to [-1, 1] + else: + params = self.norm_params[embodiment_tag]["state"][joint_group] + normalized = normalize_values_minmax(state[joint_group], params) + + if self.clip_outliers: + normalized = np.clip(normalized, -1.0, 1.0) + + normalized_values[joint_group] = normalized + + return normalized_values + + def unapply_state( + self, + state: dict[str, np.ndarray], + embodiment_tag: str, + ) -> dict[str, np.ndarray]: + """ + Reverse state processing (denormalization). + + Args: + state: Dict mapping joint_group -> processed state values + embodiment_tag: Embodiment identifier + + Returns: + Dict mapping joint_group -> raw state values + + Raises: + ValueError: If attempting to reverse sin/cos encoding (not reversible) + """ + unnormalized_values = {} + + # Get sin/cos embedding keys if enabled + sin_cos_keys = None + if self.apply_sincos_state_encoding: + state_config = self.modality_configs[embodiment_tag].get("state") + if state_config and hasattr(state_config, "sin_cos_embedding_keys"): + sin_cos_keys = state_config.sin_cos_embedding_keys + + for joint_group in self.modality_configs[embodiment_tag]["state"].modality_keys: + if joint_group not in state: + raise KeyError(f"Joint group '{joint_group}' not found in state dict for embodiment '{embodiment_tag}'") + + # Sin/cos encoding is not reversible + if sin_cos_keys and joint_group in sin_cos_keys: + raise ValueError( + f"Cannot unapply sin/cos encoding for joint group '{joint_group}' " + f"in embodiment '{embodiment_tag}'. This transformation is not reversible." + ) + + # Reverse mean/std normalization + elif ( + hasattr( + self.modality_configs[embodiment_tag]["state"], + "mean_std_embedding_keys", + ) + and self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys + and joint_group in self.modality_configs[embodiment_tag]["state"].mean_std_embedding_keys + ): + params = self.norm_params[embodiment_tag]["state"][joint_group] + unnormalized = unnormalize_values_meanstd(state[joint_group], params) + unnormalized_values[joint_group] = unnormalized + + # Reverse min/max normalization + else: + params = self.norm_params[embodiment_tag]["state"][joint_group] + unnormalized_values[joint_group] = unnormalize_values_minmax(state[joint_group], params) + + return unnormalized_values + + def apply_action( + self, + action: dict[str, np.ndarray], + embodiment_tag: str, + state: dict[str, np.ndarray] | None = None, + ) -> dict[str, np.ndarray]: + """ + Apply action processing (absolute->relative conversion, normalization). + + Processing order: + 1. Convert absolute actions to relative (if configured) + 2. Normalize actions + + Args: + action: Dict mapping joint_group -> raw action values + Shape per group: (T, D) where T is action horizon, D is action dimension + embodiment_tag: Embodiment identifier + state: Optional dict mapping joint_group -> raw state values + Required if any action group uses ActionRepresentation.RELATIVE + Shape per group: (T_state, D) where last timestep is used as reference + + Returns: + Dict mapping joint_group -> processed action values + Shape per group: (T, D) + + Raises: + ValueError: If state is None but required for relative action conversion + """ + action = deepcopy(action) # Avoid modifying input + + # Step 1: Convert absolute actions to relative (if needed) + modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys + action_configs = self.modality_configs[embodiment_tag]["action"].action_configs + + if action_configs is not None: + for key, action_config in zip(modality_keys, action_configs): + if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action: + if state is None: + raise ValueError( + f"State dict required for relative action processing of key '{key}' " + f"in embodiment '{embodiment_tag}'" + ) + + # Determine which state key to use as reference + state_key = action_config.state_key if action_config.state_key else key + + if state_key not in state: + raise KeyError( + f"Reference state key '{state_key}' not found in state dict " + f"for embodiment '{embodiment_tag}'" + ) + + # Use last state as reference frame + reference_state = state[state_key][-1] + + # Convert absolute to relative + action[key] = self._convert_to_relative_action( + action=action[key], + reference_state=reference_state, + action_type=action_config.type, + action_format=action_config.format, + ) + + # Step 2: Normalize actions + normalized_values = {} + for joint_group in modality_keys: + if joint_group not in action: + raise KeyError( + f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'" + ) + + params = self.norm_params[embodiment_tag]["action"][joint_group] + if ( + self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None + and joint_group in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys + ): + normalized = normalize_values_meanstd(action[joint_group], params) + else: + normalized = normalize_values_minmax(action[joint_group], params) + + if self.clip_outliers: + normalized = np.clip(normalized, -1.0, 1.0) + + normalized_values[joint_group] = normalized + + return normalized_values + + def unapply_action( + self, + action: dict[str, np.ndarray], + embodiment_tag: str, + state: dict[str, np.ndarray] | None = None, + ) -> dict[str, np.ndarray]: + """ + Reverse action processing (denormalization, relative->absolute conversion). + + Processing order: + 1. Denormalize actions + 2. Convert relative actions to absolute (if configured) + + Args: + action: Dict mapping joint_group -> processed action values + Shape per group: (T, D) or (B, T, D) for batched + embodiment_tag: Embodiment identifier + state: Optional dict mapping joint_group -> raw state values + Required if any action group uses ActionRepresentation.RELATIVE + Shape per group: (T_state, D) or (B, T_state, D) for batched + + Returns: + Dict mapping joint_group -> raw absolute action values + Shape per group: (T, D) or (B, T, D) for batched + + Raises: + ValueError: If state is None but required for relative->absolute conversion + """ + # Step 1: Unnormalize actions + unnormalized_values = {} + modality_keys = self.modality_configs[embodiment_tag]["action"].modality_keys + + for joint_group in modality_keys: + if joint_group not in action: + raise KeyError( + f"Joint group '{joint_group}' not found in action dict for embodiment '{embodiment_tag}'" + ) + + params = self.norm_params[embodiment_tag]["action"][joint_group] + group_values = action[joint_group] + + if ( + self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys is not None + and joint_group in self.modality_configs[embodiment_tag]["action"].mean_std_embedding_keys + ): + unnormalized = unnormalize_values_meanstd(group_values, params) + else: + unnormalized = unnormalize_values_minmax(group_values, params) + + unnormalized_values[joint_group] = unnormalized + + # Step 2: Convert relative actions to absolute (if needed) + action_configs = self.modality_configs[embodiment_tag]["action"].action_configs + + if action_configs is not None: + for key, action_config in zip(modality_keys, action_configs): + if action_config.rep == ActionRepresentation.RELATIVE and self.use_relative_action: + if state is None: + raise ValueError( + f"State dict required for relative->absolute conversion of key '{key}' " + f"in embodiment '{embodiment_tag}'" + ) + + # Determine which state key to use as reference + state_key = action_config.state_key if action_config.state_key else key + + if state_key not in state: + raise KeyError( + f"Reference state key '{state_key}' not found in state dict " + f"for embodiment '{embodiment_tag}'" + ) + + relative_action = unnormalized_values[key] + + # Handle batched and unbatched cases + is_batched = relative_action.ndim == 3 + if not is_batched: + assert relative_action.ndim == 2 + reference_state = state[state_key] + if reference_state.ndim == 2: + reference_state = reference_state[None, :] + relative_action = relative_action[None, :] + else: + reference_state = state[state_key] + if reference_state.ndim == 2: + reference_state = reference_state[None, :] + + # Convert batched relative actions to absolute + absolute_actions = [] + for s, a in zip(reference_state, relative_action): + # Use last timestep of state as reference + absolute_action = self._convert_to_absolute_action( + action=a, + reference_state=s[-1], + action_type=action_config.type, + action_format=action_config.format, + ) + absolute_actions.append(absolute_action) + + if is_batched: + unnormalized_values[key] = np.stack(absolute_actions, axis=0) + else: + unnormalized_values[key] = absolute_actions[0] + + return unnormalized_values + + def apply( + self, + state: dict[str, np.ndarray], + action: dict[str, np.ndarray], + embodiment_tag: str, + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + """ + Apply both state and action processing together. + + Convenience method that processes state and action in one call, + automatically passing raw state to action processor for relative conversion. + + Args: + state: Dict mapping joint_group -> raw state values + action: Dict mapping joint_group -> raw action values + embodiment_tag: Embodiment identifier + + Returns: + Tuple of (processed_state, processed_action) + """ + processed_state = self.apply_state(state, embodiment_tag) + if action: + processed_action = self.apply_action(action, embodiment_tag, state=state) + else: + processed_action = {} + return processed_state, processed_action + + def unapply( + self, + state: dict[str, np.ndarray], + action: dict[str, np.ndarray], + embodiment_tag: str, + raw_state: dict[str, np.ndarray] | None = None, + ) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]: + """ + Reverse both state and action processing together. + + Args: + state: Dict mapping joint_group -> processed state values + action: Dict mapping joint_group -> processed action values + embodiment_tag: Embodiment identifier + raw_state: Optional dict of raw states for relative->absolute conversion + If None, will use unapplied state (but won't work for sin/cos encoded states) + + Returns: + Tuple of (raw_state, raw_action) + """ + # Unapply state first + try: + unapplied_state = self.unapply_state(state, embodiment_tag) + except ValueError as e: + if "sin/cos encoding" in str(e) and raw_state is None: + raise ValueError("Cannot unapply sin/cos encoded state. Please provide raw_state parameter.") from e + raise + + # Use provided raw_state if available, otherwise use unapplied state + state_for_action = raw_state if raw_state is not None else unapplied_state + + # Unapply action + unapplied_action = self.unapply_action(action, embodiment_tag, state=state_for_action) + + return unapplied_state, unapplied_action + + def get_state_dim(self, embodiment_tag: str, include_sincos_expansion: bool = False) -> int: + """ + Get total state dimension after processing. + + Args: + embodiment_tag: Embodiment identifier + include_sincos_expansion: If True, accounts for sin/cos encoding doubling dimensions + + Returns: + Total state dimension across all joint groups + """ + total_dim = 0 + state_config = self.modality_configs[embodiment_tag]["state"] + + # Get sin/cos embedding keys if enabled + sin_cos_keys = set() + if self.apply_sincos_state_encoding and hasattr(state_config, "sin_cos_embedding_keys"): + sin_cos_keys = set(state_config.sin_cos_embedding_keys) + + for joint_group in state_config.modality_keys: + base_dim = self.norm_params[embodiment_tag]["state"][joint_group]["dim"].item() + + # Sin/cos encoding doubles the dimension + if include_sincos_expansion and joint_group in sin_cos_keys: + total_dim += base_dim * 2 + else: + total_dim += base_dim + + return total_dim + + def get_action_dim(self, embodiment_tag: str) -> int: + """ + Get total action dimension. + + Args: + embodiment_tag: Embodiment identifier + + Returns: + Total action dimension across all joint groups + """ + total_dim = 0 + for joint_group in self.modality_configs[embodiment_tag]["action"].modality_keys: + total_dim += self.norm_params[embodiment_tag]["action"][joint_group]["dim"].item() + return total_dim + + def _convert_to_relative_action( + self, + action: np.ndarray, + reference_state: np.ndarray, + action_type: ActionType, + action_format: ActionFormat, + ) -> np.ndarray: + """Convert absolute action to relative action using reference state.""" + assert action.ndim == 2, f"Expected action shape (T, D), got {action.shape}" + assert reference_state.ndim == 1, f"Expected state shape (D,), got {reference_state.shape}" + + if action_type == ActionType.EEF: + action_chunking = EndEffectorActionChunk.from_array(action, action_format) + reference_frame = EndEffectorPose.from_action_format(reference_state, action_format) + + elif action_type == ActionType.NON_EEF: + action_chunking = JointActionChunk([JointPose(m) for m in action]) + reference_frame = JointPose(reference_state) + + else: + raise ValueError(f"Unknown ActionType: {action_type}") + + relative_action_chunking = action_chunking.relative_chunking(reference_frame=reference_frame) + return relative_action_chunking.to(action_format) + + def _convert_to_absolute_action( + self, + action: np.ndarray, + reference_state: np.ndarray, + action_type: ActionType, + action_format: ActionFormat, + ) -> np.ndarray: + """Convert relative action to absolute action using reference state.""" + assert action.ndim == 2, f"Expected action shape (T, D), got {action.shape}" + assert reference_state.ndim == 1, f"Expected state shape (D,), got {reference_state.shape}" + assert reference_state.shape[0] == action.shape[1], ( + f"State dim {reference_state.shape[0]} != action dim {action.shape[1]}" + ) + + if action_type == ActionType.EEF: + rel_action = EndEffectorActionChunk.from_array(action, action_format) + reference_frame = EndEffectorPose.from_action_format(reference_state, action_format) + + elif action_type == ActionType.NON_EEF: + rel_action = JointActionChunk([JointPose(pose) for pose in action]) + reference_frame = JointPose(reference_state) + + else: + raise ValueError(f"Unknown ActionType: {action_type}") + + abs_action = rel_action.to_absolute_chunking(reference_frame=reference_frame) + return abs_action.to(action_format) + + def __str__(self) -> str: + return ( + "StateActionProcessor(" + f"modality_configs={self.modality_configs}, " + f"statistics={self.statistics}, " + f"use_percentiles={self.use_percentiles}, " + f"clip_outliers={self.clip_outliers}, " + f"apply_sincos_state_encoding={self.apply_sincos_state_encoding}, " + f"use_relative_action={self.use_relative_action})" + ) diff --git a/vllm_omni/diffusion/models/gr00t/dataio/types.py b/vllm_omni/diffusion/models/gr00t/dataio/types.py new file mode 100644 index 00000000000..d4948a1c234 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/types.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any + +import numpy as np + +from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag + + +class MessageType(Enum): + START_OF_EPISODE = "start_of_episode" + END_OF_EPISODE = "end_of_episode" + EPISODE_STEP = "episode_step" + IMAGE = "image" + TEXT = "text" + + +class ActionRepresentation(Enum): + RELATIVE = "relative" + DELTA = "delta" + ABSOLUTE = "absolute" + + +class ActionType(Enum): + EEF = "eef" + NON_EEF = "non_eef" + + +class ActionFormat(Enum): + DEFAULT = "default" + XYZ_ROT6D = "xyz+rot6d" + XYZ_ROTVEC = "xyz+rotvec" + + +@dataclass +class VLAStepData: + """ + Represents a single step of VLA (Vision-Language-Action) data. + + This is the core data structure returned by datasets, containing raw observation + and action data that will be processed by the SequenceVLAProcessor. + """ + + # Core data + images: dict[str, list[np.ndarray]] # view_name -> list[np.ndarray] (for temporal stacking) + states: dict[str, np.ndarray] # state_name -> np.ndarray (dim,) for single step or (horizon, dim) for trajectory + actions: dict[str, np.ndarray] # action_name -> np.ndarray (horizon, dim) for action chunk + masks: dict[str, list[np.ndarray]] | None = None # view_name -> list[np.ndarray] (H, W) + text: str | None = None # Optional task description or instruction + embodiment: EmbodimentTag = EmbodimentTag.NEW_EMBODIMENT # Optional embodiment tag for cross-embodiment training + is_demonstration: bool = ( + False # Whether the step is a demonstration. If True, no loss should be computed for this step. + ) + + # Flexible metadata that can be extended by users + metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ActionConfig: + rep: ActionRepresentation + type: ActionType + format: ActionFormat + state_key: str | None = None + + +@dataclass +class ModalityConfig: + """Configuration for a modality defining how data should be sampled and loaded. + + This class specifies which indices to sample relative to a base index and which + keys to load for a particular modality (e.g., video, state, action). + """ + + delta_indices: list[int] + """Delta indices to sample relative to the current index. The returned data will + correspond to the original data at a sampled base index + delta indices.""" + modality_keys: list[str] + """The keys to load for the modality in the dataset.""" + sin_cos_embedding_keys: list[str] | None = None + """Optional list of keys to apply sin/cos encoding. If None or empty, use + min/max normalization for all keys.""" + mean_std_embedding_keys: list[str] | None = None + """Optional list of keys to apply mean/std normalization. If None or empty, + use min/max normalization for all keys.""" + action_configs: list[ActionConfig] | None = None + + def __post_init__(self): + """Validate fields and set default values.""" + if self.delta_indices is None or not isinstance(self.delta_indices, list): + raise ValueError(f"delta_indices must be a non-None list, got {self.delta_indices!r}") + if self.modality_keys is None or not isinstance(self.modality_keys, list) or len(self.modality_keys) == 0: + raise ValueError(f"modality_keys must be a non-empty list, got {self.modality_keys!r}") + if self.action_configs is not None: + assert len(self.action_configs) == len(self.modality_keys), ( + f"Number of action configs ({len(self.action_configs)}) must match " + f"number of modality keys ({len(self.modality_keys)})" + ) + parsed_action_configs = [] + for action_config in self.action_configs: + if isinstance(action_config, dict): + action_config = ActionConfig( + rep=ActionRepresentation[action_config["rep"]], + type=ActionType[action_config["type"]], + format=ActionFormat[action_config["format"]], + state_key=action_config.get("state_key", None), + ) + parsed_action_configs.append(action_config) + self.action_configs = parsed_action_configs diff --git a/vllm_omni/diffusion/models/gr00t/dataio/utils.py b/vllm_omni/diffusion/models/gr00t/dataio/utils.py new file mode 100644 index 00000000000..d3846d44228 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/dataio/utils.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from dataclasses import asdict, is_dataclass +from enum import Enum +from typing import Any + +import numpy as np + +from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig + + +def apply_sin_cos_encoding(values: np.ndarray) -> np.ndarray: + """Apply sin/cos encoding to values. + + Args: + values: Array of shape (..., D) containing values to encode + + Returns: + Array of shape (..., 2*D) with [sin, cos] concatenated + + Note: This DOUBLES the dimension. For example: + Input: [v₁, v₂, v₃] with shape (..., 3) + Output: [sin(v₁), sin(v₂), sin(v₃), cos(v₁), cos(v₂), cos(v₃)] with shape (..., 6) + """ + sin_values = np.sin(values) + cos_values = np.cos(values) + # Concatenate sin and cos: [sin(v1), sin(v2), ..., cos(v1), cos(v2), ...] + return np.concatenate([sin_values, cos_values], axis=-1) + + +def nested_dict_to_numpy(data: Any) -> Any: + """Recursively convert bottom-level list-of-lists to NumPy arrays. + + Args: + data: A nested dictionary whose leaves are lists (or lists of lists). + + Returns: + The same dictionary structure with leaf lists converted to ``np.ndarray``. + """ + if isinstance(data, dict): + return {key: nested_dict_to_numpy(value) for key, value in data.items()} + elif isinstance(data, list): + # Convert lists to numpy arrays + # NumPy will handle both 1D and 2D cases appropriately + return np.array(data) + else: + return data + + +def normalize_values_minmax(values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray: + """Min-max normalize ``values`` to ``[-1, 1]`` using ``params['min']`` and ``params['max']``. + + Accepts 2D ``(T, D)`` or 3D ``(B, T, D)`` arrays. Features with ``min == max`` + are emitted as 0. + """ + min_vals = params["min"] + max_vals = params["max"] + normalized = np.zeros_like(values) + + mask = ~np.isclose(max_vals, min_vals) + + normalized[..., mask] = (values[..., mask] - min_vals[..., mask]) / (max_vals[..., mask] - min_vals[..., mask]) + normalized[..., mask] = 2 * normalized[..., mask] - 1 + + return normalized + + +def unnormalize_values_minmax(normalized_values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray: + """Inverse of :func:`normalize_values_minmax`. + + Input values are clipped to ``[-1, 1]`` before being mapped back to + ``[params['min'], params['max']]``. + """ + + min_vals = params["min"] + max_vals = params["max"] + range_vals = max_vals - min_vals + + # Unnormalize from [-1, 1] + unnormalized = (np.clip(normalized_values, -1.0, 1.0) + 1.0) / 2.0 * range_vals + min_vals + return unnormalized + + +def normalize_values_meanstd(values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray: + """Z-score normalize ``values`` using ``params['mean']`` and ``params['std']``. + + Features whose ``std == 0`` are passed through unchanged. + """ + mean_vals = params["mean"] + std_vals = params["std"] + + # Create mask for non-zero standard deviations + mask = std_vals != 0 + + # Initialize normalized array + normalized = np.zeros_like(values) + + # Normalize only features with non-zero std + normalized[..., mask] = (values[..., mask] - mean_vals[..., mask]) / std_vals[..., mask] + + # Keep original values for zero-std features + normalized[..., ~mask] = values[..., ~mask] + + return normalized + + +def unnormalize_values_meanstd(normalized_values: np.ndarray, params: dict[str, np.ndarray]) -> np.ndarray: + """Inverse of :func:`normalize_values_meanstd` (``x * std + mean``). + + Features whose ``std == 0`` are passed through unchanged. + """ + mean_vals = params["mean"] + std_vals = params["std"] + + # Create mask for non-zero standard deviations + mask = std_vals != 0 + + # Initialize unnormalized array + unnormalized = np.zeros_like(normalized_values) + + # Unnormalize only features with non-zero std + unnormalized[..., mask] = normalized_values[..., mask] * std_vals[..., mask] + mean_vals[..., mask] + + # Keep normalized values for zero-std features + unnormalized[..., ~mask] = normalized_values[..., ~mask] + + return unnormalized + + +def to_json_serializable(obj: Any) -> Any: + """ + Recursively convert dataclasses and numpy arrays to JSON-serializable format. + + Args: + obj: Object to convert (can be dataclass, numpy array, dict, list, etc.) + + Returns: + JSON-serializable representation of the object + """ + if is_dataclass(obj) and not isinstance(obj, type): + # Convert dataclass to dict, then recursively process the dict + return to_json_serializable(asdict(obj)) + elif isinstance(obj, np.ndarray): + # Convert numpy array to list + return obj.tolist() + elif isinstance(obj, np.integer): + # Convert numpy integers to Python int + return int(obj) + elif isinstance(obj, np.floating): + # Convert numpy floats to Python float + return float(obj) + elif isinstance(obj, np.bool_): + # Convert numpy bool to Python bool + return bool(obj) + elif isinstance(obj, dict): + # Recursively process dictionary values + return {key: to_json_serializable(value) for key, value in obj.items()} + elif isinstance(obj, (list, tuple)): + # Recursively process list/tuple elements + return [to_json_serializable(item) for item in obj] + elif isinstance(obj, set): + # Convert set to list + return [to_json_serializable(item) for item in obj] + elif isinstance(obj, (str, int, float, bool, type(None))): + # Already JSON-serializable + return obj + elif isinstance(obj, Enum): + return obj.name + else: + # For other types, try to convert to string as fallback + # You might want to handle specific types differently + return str(obj) + + +def parse_modality_configs( + modality_configs: dict[str, dict[str, ModalityConfig]], +) -> dict[str, dict[str, ModalityConfig]]: + parsed_modality_configs = {} + for embodiment_tag, modality_config in modality_configs.items(): + parsed_modality_configs[embodiment_tag] = {} + for modality, config in modality_config.items(): + if isinstance(config, dict): + parsed_modality_configs[embodiment_tag][modality] = ModalityConfig(**config) + else: + parsed_modality_configs[embodiment_tag][modality] = config + return parsed_modality_configs diff --git a/vllm_omni/diffusion/models/gr00t/modeling/__init__.py b/vllm_omni/diffusion/models/gr00t/modeling/__init__.py new file mode 100644 index 00000000000..7aecb405ff9 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/modeling/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7 +from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7Processor + +__all__ = ["Gr00tN1d7", "Gr00tN1d7Processor"] diff --git a/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py new file mode 100644 index 00000000000..3fcbce203f5 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/modeling/gr00t_n1d7.py @@ -0,0 +1,491 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +import os +from typing import Any + +import torch +from torch import nn +from transformers import AutoConfig, AutoModel, PreTrainedModel +from transformers.feature_extraction_utils import BatchFeature + +from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config +from vllm_omni.diffusion.models.gr00t.modeling.modules.dit import AlternateVLDiT, DiT, SelfAttentionTransformer +from vllm_omni.diffusion.models.gr00t.modeling.modules.embodiment_conditioned_mlp import ( + CategorySpecificMLP, + MultiEmbodimentActionEncoder, +) +from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7DataCollator + +logger = logging.getLogger(__name__) + + +class Gr00tN1d7ActionHead(nn.Module): + """Action head component for flow matching diffusion policy.""" + + supports_gradient_checkpointing = True + + def __init__(self, config: Gr00tN1d7Config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.input_embedding_dim = config.input_embedding_dim + + if config.use_alternate_vl_dit: + self.model = AlternateVLDiT( + **config.diffusion_model_cfg, + cross_attention_dim=config.backbone_embedding_dim, + attend_text_every_n_blocks=config.attend_text_every_n_blocks, + ) + logger.info("Using AlternateVLDiT for diffusion model") + else: + self.model = DiT( + **config.diffusion_model_cfg, + cross_attention_dim=config.backbone_embedding_dim, + ) + logger.info("Using DiT for diffusion model") + self.action_dim = config.max_action_dim + self.action_horizon = config.action_horizon + self.num_inference_timesteps = config.num_inference_timesteps + + self.state_encoder = CategorySpecificMLP( + num_categories=config.max_num_embodiments, + input_dim=config.max_state_dim * config.state_history_length, + hidden_dim=self.hidden_size, + output_dim=self.input_embedding_dim, + ) + self.action_encoder = MultiEmbodimentActionEncoder( + action_dim=self.action_dim, + hidden_size=self.input_embedding_dim, + num_embodiments=config.max_num_embodiments, + ) + self.action_decoder = CategorySpecificMLP( + num_categories=config.max_num_embodiments, + input_dim=self.hidden_size, + hidden_dim=self.hidden_size, + output_dim=self.action_dim, + ) + + self.vlln = nn.LayerNorm(config.backbone_embedding_dim) if config.use_vlln else nn.Identity() + + vl_self_attention_cfg = config.vl_self_attention_cfg + if vl_self_attention_cfg and vl_self_attention_cfg.get("num_layers", 0) > 0: + self.vl_self_attention = SelfAttentionTransformer(**vl_self_attention_cfg) + else: + self.vl_self_attention = nn.Identity() + + if config.add_pos_embed: + self.position_embedding = nn.Embedding(config.max_seq_len, self.input_embedding_dim) + nn.init.normal_(self.position_embedding.weight, mean=0.0, std=0.02) + + self.num_timestep_buckets = config.num_timestep_buckets + + def process_backbone_output(self, backbone_output: BatchFeature) -> BatchFeature: + backbone_features = backbone_output["backbone_features"] + backbone_features = self.vlln(backbone_features) + backbone_features = self.vl_self_attention(backbone_features) + backbone_output["backbone_features"] = backbone_features + return backbone_output + + def _encode_features(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + """ + Encode features for the action head. + + Args: + backbone_output: Output from the backbone model containing: + - backbone_features: [B, seq_len, backbone_embedding_dim] + - backbone_attention_mask: [B, seq_len] + action_input: Input containing: + - state: [B, state_history_length, max_state_dim] + - embodiment_id: [B] (embodiment IDs) + + Returns: + BatchFeature containing: + - backbone_features: [B, seq_len, backbone_embedding_dim] + - state_features: [B, 1, input_embedding_dim] + """ + backbone_output = self.process_backbone_output(backbone_output) + + vl_embeds = backbone_output.backbone_features + embodiment_id = action_input.embodiment_id + + state = action_input.state + current_T = state.shape[1] + assert current_T == self.config.state_history_length, "current_T != state_history_length" + # [B, state_history_length, max_state_dim] -> [B, 1, state_history_length * max_state_dim] + state = state.view(state.shape[0], 1, -1) + + state_features = self.state_encoder(state, embodiment_id) + + return BatchFeature(data={"backbone_features": vl_embeds, "state_features": state_features}) + + @torch.no_grad() + def get_action_with_features( + self, + backbone_features: torch.Tensor, + state_features: torch.Tensor, + embodiment_id: torch.Tensor, + backbone_output: BatchFeature, + action_input: BatchFeature, + options: dict[str, Any] | None = None, + ) -> BatchFeature: + """ + Generate actions using the flow matching diffusion process. + + Args: + backbone_features: [B, seq_len, backbone_embedding_dim] + state_features: [B, state_horizon, input_embedding_dim] + embodiment_id: [B] (embodiment IDs) + backbone_output: Output from the backbone model + """ + vl_embeds = backbone_features + + batch_size = vl_embeds.shape[0] + device = vl_embeds.device + _seed_env = os.environ.get("GR00T_NOISE_SEED") + if _seed_env is not None: + _gen = torch.Generator(device=device).manual_seed(int(_seed_env)) + actions = torch.randn( + size=(batch_size, self.config.action_horizon, self.action_dim), + dtype=vl_embeds.dtype, + device=device, + generator=_gen, + ) + else: + actions = torch.randn( + size=(batch_size, self.config.action_horizon, self.action_dim), + dtype=vl_embeds.dtype, + device=device, + ) + + dt = 1.0 / self.num_inference_timesteps + vel_strength = torch.ones_like(actions) + + if "action" in action_input: + # If action in input when doing get action, it means we want to use RTC. + # action_horizon is the action horizon of the input action. + # rtc_overlap_steps is the number of steps to overlap with the previous action chunks. + # rtc_frozen_steps is the policy inference latency expressed as frozen action steps. + # rtc_ramp_rate is the rate of the ramp of denoising the actions. + assert options is not None, "options is not None" + assert "action_horizon" in options, "action_horizon is not in options" + assert "rtc_overlap_steps" in options, "rtc_overlap_steps is not in options" + assert "rtc_frozen_steps" in options, "rtc_frozen_steps is not in options" + assert "rtc_ramp_rate" in options, "rtc_ramp_rate is not in options" + + action_horizon_before_padding = options["action_horizon"] + + # Use previous action instead of pure noise to do inpainting + actions[:, : options["rtc_overlap_steps"], :] = action_input["action"][ + :, + action_horizon_before_padding - options["rtc_overlap_steps"] : action_horizon_before_padding, + :, + ] + vel_strength[:, : options["rtc_frozen_steps"], :] = 0.0 + # NOTE: use an exponential ramp strength to set the remaining unfrozen rtc_steps + intermediate_steps = options["rtc_overlap_steps"] - options["rtc_frozen_steps"] + # Create exponential ramp from 0 to 1 over intermediate steps + t = torch.linspace(0.0, 1.0, intermediate_steps + 2, device=device) + ramp = 1 - torch.exp(-options["rtc_ramp_rate"] * t) + ramp = ramp / ramp[-1].clamp_min(1e-8) # normalize to [0,1] + ramp = ramp[1:-1] # we will only take the middle part of the ramp, ignore the 0.0 and 1.0 + # Apply ramp to the intermediate steps [batch, intermediate_steps, action_dim] + vel_strength[ + :, + options["rtc_frozen_steps"] : options["rtc_overlap_steps"], + :, + ] = ramp[None, :, None].to(device) + + # Run denoising steps. + for t in range(self.num_inference_timesteps): + t_discretized = t * self.num_timestep_buckets // self.num_inference_timesteps + + # Embed noised action trajectory. + timesteps_tensor = torch.full(size=(batch_size,), fill_value=t_discretized, device=device) + action_features = self.action_encoder(actions, timesteps_tensor, embodiment_id) + # Add position embedding. + if self.config.add_pos_embed: + pos_ids = torch.arange(action_features.shape[1], dtype=torch.long, device=device) + pos_embs = self.position_embedding(pos_ids).unsqueeze(0) + action_features = action_features + pos_embs + + # Join vision, language, state and action embedding along sequence dimension. + sa_embs = torch.cat((state_features, action_features), dim=1) + + # Run model forward. + if self.config.use_alternate_vl_dit: + model_output = self.model( + hidden_states=sa_embs, + encoder_hidden_states=vl_embeds, + timestep=timesteps_tensor, + image_mask=backbone_output.image_mask, + backbone_attention_mask=backbone_output.backbone_attention_mask, + ) + else: + model_output = self.model( + hidden_states=sa_embs, + encoder_hidden_states=vl_embeds, + timestep=timesteps_tensor, + ) + pred = self.action_decoder(model_output, embodiment_id) + + pred_velocity = pred[:, -self.action_horizon :] + + # Update actions using euler integration. + actions = actions + dt * pred_velocity * vel_strength + + return BatchFeature( + data={ + "action_pred": actions, + "backbone_features": vl_embeds, + "state_features": state_features, + } + ) + + @torch.no_grad() + def get_action( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + options: dict[str, Any] | None = None, + ) -> BatchFeature: + """ + Generate actions using the flow matching diffusion process. + + Args: + backbone_output: Output from the backbone model containing: + - backbone_features: [B, seq_len, backbone_embedding_dim] + - backbone_attention_mask: [B, seq_len] + action_input: Input containing: + - state: [B, state_dim] + - embodiment_id: [B] (embodiment IDs) + + Returns: + BatchFeature containing: + - action_pred: [B, action_horizon, action_dim] predicted actions + """ + features = self._encode_features(backbone_output, action_input) + return self.get_action_with_features( + backbone_features=features.backbone_features, + state_features=features.state_features, + embodiment_id=action_input.embodiment_id, + backbone_output=backbone_output, + action_input=action_input, + options=options, + ) + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + def prepare_input(self, batch: dict) -> BatchFeature: + """Prepare input batch for the action head.""" + return BatchFeature(data=batch) + + +def _load_qwen3_vl_cls() -> type: + """Import Qwen3-VL lazily so the missing-dep failure is clear and deferred.""" + try: + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLForConditionalGeneration + except ImportError as exc: + raise ImportError( + "GR00T-N1.7 requires transformers>=4.57.1 for transformers.models.qwen3_vl " + "(the repo's declared floor is 4.56.0). Please upgrade transformers." + ) from exc + return Qwen3VLForConditionalGeneration + + +class _Qwen3VLBackbone(nn.Module): + """GR00T adapter around the shared Qwen3-VL implementation.""" + + def __init__( + self, + model_name: str, + select_layer: int, + backbone_embedding_dim: int, + load_bf16: bool, + transformers_loading_kwargs: dict[str, Any] | None = None, + ): + super().__init__() + backbone_config = AutoConfig.from_pretrained( + model_name, **(transformers_loading_kwargs or {"trust_remote_code": True}) + ) + backbone_config.text_config._attn_implementation = "sdpa" + backbone_config.vision_config._attn_implementation = "sdpa" + self.model = _load_qwen3_vl_cls()(backbone_config).eval() + if load_bf16: + self.model.to(dtype=torch.bfloat16) + + target_layers = select_layer if select_layer >= 0 else len(self.model.model.language_model.layers) + while len(self.model.model.language_model.layers) > target_layers: + self.model.model.language_model.layers.pop(-1) + + # Drop the unused lm_head: feature-extractor only, and .to(dtype) above unties it (~0.58 GiB). + self.model.lm_head = None + + def set_frozen_modules_to_eval_mode(self) -> None: + self.eval() + + def prepare_input(self, batch: dict) -> BatchFeature: + return BatchFeature(data=batch) + + def forward(self, vl_input: BatchFeature) -> BatchFeature: + self.set_frozen_modules_to_eval_mode() + keys_to_use = [ + "input_ids", + "attention_mask", + "pixel_values", + "image_grid_thw", + "mm_token_type_ids", + ] + vl_input = {key: vl_input[key] for key in keys_to_use if key in vl_input} + # GR00T was trained against the **pre-final-RMSNorm** output of the last + # decoder layer (this was `outer.hidden_states[-1]` under HF 4.57.x semantics, + # where the @check_model_inputs tie-logic skipped overwriting because + # Qwen3VLCausalLMOutputWithPast has no `last_hidden_state` field). + # HF >=5.x changed the output_hidden_states tuple plumbing so that + # `hidden_states[-1]` is post-norm, which is a different tensor and breaks + # the trained action head. Capture the pre-norm tensor via a hook on the + # final RMSNorm regardless of transformers version. + _captured: list[torch.Tensor] = [] + + def _pre_norm_hook(_module, args, _out): + _captured.append(args[0]) + + norm = self.model.model.language_model.norm + handle = norm.register_forward_hook(_pre_norm_hook) + try: + self.model.model(**vl_input, return_dict=True) + finally: + handle.remove() + if not _captured: + raise RuntimeError("Failed to capture pre-norm hidden states from Qwen3-VL backbone") + backbone_features = _captured[-1] + image_mask = vl_input["input_ids"] == self.model.config.image_token_id + attention_mask = vl_input["attention_mask"] == 1 + return BatchFeature( + data={ + "backbone_features": backbone_features, + "backbone_attention_mask": attention_mask, + "image_mask": image_mask, + } + ) + + +def get_backbone_cls(config: Gr00tN1d7Config): + if "nvidia/Cosmos-Reason2" in config.model_name: + return _Qwen3VLBackbone + raise ValueError(f"Unsupported model name: {config.model_name}") + + +class Gr00tN1d7(PreTrainedModel): + """Gr00tN1d7: VLA model with Cosmos-Reason2-2B (Qwen3-VL) backbone.""" + + config_class = Gr00tN1d7Config + supports_gradient_checkpointing = True + _tp_plan = {} + + @property + def all_tied_weights_keys(self) -> dict[str, Any]: + return {} + + def __init__( + self, + config: Gr00tN1d7Config, + transformers_loading_kwargs: dict | None = None, + ): + super().__init__(config) + if transformers_loading_kwargs is None: + transformers_loading_kwargs = {"trust_remote_code": True} + self.config = config + + backbone_cls = get_backbone_cls(config) + self.backbone = backbone_cls( + model_name=config.model_name, + select_layer=config.select_layer, + backbone_embedding_dim=config.backbone_embedding_dim, + load_bf16=config.load_bf16, + transformers_loading_kwargs=transformers_loading_kwargs, + ) + + # Initialize action head + self.action_head = Gr00tN1d7ActionHead(config) + + self.collator = Gr00tN1d7DataCollator( + model_name=config.model_name, + model_type=config.backbone_model_type, + transformers_loading_kwargs=transformers_loading_kwargs, + ) + + def prepare_input(self, inputs: dict) -> tuple[BatchFeature, BatchFeature]: + """Prepare inputs for backbone and action head.""" + + # NOTE -- currently the eval code doesn't use collator, so we need to add it here + if "vlm_content" in inputs: + # Fix for n_envs > 1: Process all environments' VLM content, not just the first + vlm_content_list = inputs["vlm_content"] + # Ensure vlm_content_list is always a list for consistent processing + if not isinstance(vlm_content_list, list): + vlm_content_list = [vlm_content_list] + + # Process all VLM contents through the collator + prep = self.collator([{"vlm_content": vlm} for vlm in vlm_content_list])["inputs"] + inputs.pop("vlm_content") + inputs.update(prep) + + backbone_inputs = self.backbone.prepare_input(inputs) + action_inputs = self.action_head.prepare_input(inputs) + + backbone_inputs = backbone_inputs.to(device=self.device, dtype=self.dtype) + action_inputs = action_inputs.to(device=self.device, dtype=self.dtype) + + return backbone_inputs, action_inputs + + def forward(self, inputs: dict) -> BatchFeature: + """ + Forward pass through the complete model. + + Args: + inputs: Dictionary containing: + - Action inputs (state, action, embodiment_id, etc.) + + Returns: + BatchFeature containing loss and other outputs + """ + # Prepare inputs for backbone and action head + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_outputs = self.action_head(backbone_outputs, action_inputs) + + return action_outputs + + def get_action(self, inputs: dict, options: dict[str, Any] | None = None) -> BatchFeature: + """ + Generate actions using the complete model. + """ + # Prepare inputs for backbone and action head + backbone_inputs, action_inputs = self.prepare_input(inputs) + + # Forward through backbone + backbone_outputs = self.backbone(backbone_inputs) + action_outputs = self.action_head.get_action(backbone_outputs, action_inputs, options) + + return action_outputs + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype + + +# Register the model with HuggingFace +AutoConfig.register("Gr00tN1d7", Gr00tN1d7Config) +AutoModel.register(Gr00tN1d7Config, Gr00tN1d7) diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/modeling/modules/__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/diffusion/models/gr00t/modeling/modules/dit.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py new file mode 100644 index 00000000000..54fbe770d0d --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/modeling/modules/dit.py @@ -0,0 +1,344 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn.functional as F +from diffusers import ConfigMixin, ModelMixin +from diffusers.configuration_utils import register_to_config +from diffusers.models.attention import Attention, FeedForward +from diffusers.models.embeddings import SinusoidalPositionalEmbedding, TimestepEmbedding, Timesteps +from torch import nn + + +class AdaLayerNorm(nn.Module): + def __init__( + self, + embedding_dim: int, + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + chunk_dim: int = 0, + ): + super().__init__() + self.chunk_dim = chunk_dim + output_dim = embedding_dim * 2 + self.silu = nn.SiLU() + self.linear = nn.Linear(embedding_dim, output_dim) + self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine) + + def forward( + self, + x: torch.Tensor, + temb: torch.Tensor | None = None, + ) -> torch.Tensor: + temb = self.linear(self.silu(temb)) + scale, shift = temb.chunk(2, dim=1) + x = self.norm(x) * (1 + scale[:, None]) + shift[:, None] + return x + + +class BasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + dropout=0.0, + cross_attention_dim: int | None = None, + activation_fn: str = "geglu", + attention_bias: bool = False, + upcast_attention: bool = False, + norm_elementwise_affine: bool = True, + norm_type: str = "layer_norm", + norm_eps: float = 1e-5, + final_dropout: bool = False, + attention_type: str = "default", + positional_embeddings: str | None = None, + num_positional_embeddings: int | None = None, + ff_inner_dim: int | None = None, + ff_bias: bool = True, + attention_out_bias: bool = True, + ): + super().__init__() + self.norm_type = norm_type + + if positional_embeddings and (num_positional_embeddings is None): + raise ValueError( + "If `positional_embedding` type is defined, `num_position_embeddings` must also be defined." + ) + + if positional_embeddings == "sinusoidal": + self.pos_embed = SinusoidalPositionalEmbedding(dim, max_seq_length=num_positional_embeddings) + else: + self.pos_embed = None + + if norm_type == "ada_norm": + self.norm1 = AdaLayerNorm(dim) + else: + self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps) + + self.attn1 = Attention( + query_dim=dim, + heads=num_attention_heads, + dim_head=attention_head_dim, + dropout=dropout, + bias=attention_bias, + cross_attention_dim=cross_attention_dim, + upcast_attention=upcast_attention, + out_bias=attention_out_bias, + ) + + self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine) + self.ff = FeedForward( + dim, + dropout=dropout, + activation_fn=activation_fn, + final_dropout=final_dropout, + inner_dim=ff_inner_dim, + bias=ff_bias, + ) + self.final_dropout = nn.Dropout(dropout) if final_dropout else None + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + temb: torch.LongTensor | None = None, + ) -> torch.Tensor: + if self.norm_type == "ada_norm": + norm_hidden_states = self.norm1(hidden_states, temb) + else: + norm_hidden_states = self.norm1(hidden_states) + + if self.pos_embed is not None: + norm_hidden_states = self.pos_embed(norm_hidden_states) + + attn_output = self.attn1( + norm_hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=(encoder_attention_mask if encoder_hidden_states is not None else attention_mask), + ) + if self.final_dropout: + attn_output = self.final_dropout(attn_output) + + hidden_states = attn_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + + norm_hidden_states = self.norm3(hidden_states) + ff_output = self.ff(norm_hidden_states) + + hidden_states = ff_output + hidden_states + if hidden_states.ndim == 4: + hidden_states = hidden_states.squeeze(1) + return hidden_states + + +class TimestepEncoder(nn.Module): + def __init__(self, embedding_dim: int): + super().__init__() + self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1) + self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + dtype = next(self.parameters()).dtype + return self.timestep_embedder(self.time_proj(timesteps).to(dtype)) + + +class DiT(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 8, + attention_head_dim: int = 64, + output_dim: int = 26, + num_layers: int = 12, + dropout: float = 0.1, + attention_bias: bool = True, + activation_fn: str = "gelu-approximate", + num_embeds_ada_norm: int | None = 1000, + upcast_attention: bool = False, + norm_type: str = "ada_norm", + norm_elementwise_affine: bool = False, + norm_eps: float = 1e-5, + max_num_positional_embeddings: int = 512, + compute_dtype=torch.float32, + final_dropout: bool = True, + positional_embeddings: str | None = "sinusoidal", + interleave_self_attention=False, + cross_attention_dim: int | None = None, + ): + super().__init__() + + self.inner_dim = num_attention_heads * attention_head_dim + self.gradient_checkpointing = False + self.timestep_encoder = TimestepEncoder(embedding_dim=self.inner_dim) + + all_blocks = [] + for idx in range(num_layers): + use_self_attn = idx % 2 == 1 and interleave_self_attention + curr_cross_attention_dim = cross_attention_dim if not use_self_attn else None + all_blocks.append( + BasicTransformerBlock( + self.inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + activation_fn=activation_fn, + attention_bias=attention_bias, + upcast_attention=upcast_attention, + norm_type=norm_type, + norm_elementwise_affine=norm_elementwise_affine, + norm_eps=norm_eps, + positional_embeddings=positional_embeddings, + num_positional_embeddings=max_num_positional_embeddings, + final_dropout=final_dropout, + cross_attention_dim=curr_cross_attention_dim, + ) + ) + self.transformer_blocks = nn.ModuleList(all_blocks) + self.norm_out = nn.LayerNorm(self.inner_dim, elementwise_affine=False, eps=1e-6) + self.proj_out_1 = nn.Linear(self.inner_dim, 2 * self.inner_dim) + self.proj_out_2 = nn.Linear(self.inner_dim, output_dim) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + return_all_hidden_states: bool = False, + ): + temb = self.timestep_encoder(timestep) + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + all_hidden_states = [hidden_states] + + for idx, block in enumerate(self.transformer_blocks): + if idx % 2 == 1 and self.config.interleave_self_attention: + hidden_states = block(hidden_states, temb=temb) + else: + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + temb=temb, + ) + all_hidden_states.append(hidden_states) + + shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + if return_all_hidden_states: + return self.proj_out_2(hidden_states), all_hidden_states + return self.proj_out_2(hidden_states) + + +class AlternateVLDiT(DiT): + def __init__(self, *args, attend_text_every_n_blocks: int = 2, **kwargs): + super().__init__(*args, **kwargs) + self.attend_text_every_n_blocks = attend_text_every_n_blocks + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.LongTensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + return_all_hidden_states: bool = False, + image_mask: torch.Tensor | None = None, + backbone_attention_mask: torch.Tensor | None = None, + ): + assert image_mask is not None, "Image mask is required" + assert self.config.interleave_self_attention, "Interleave self attention must be enabled" + + temb = self.timestep_encoder(timestep) + hidden_states = hidden_states.contiguous() + encoder_hidden_states = encoder_hidden_states.contiguous() + + image_attention_mask = image_mask & backbone_attention_mask + non_image_attention_mask = (~image_mask) & backbone_attention_mask + + all_hidden_states = [hidden_states] + for idx, block in enumerate(self.transformer_blocks): + if idx % 2 == 1: + hidden_states = block(hidden_states, temb=temb) + else: + curr_encoder_attention_mask = ( + non_image_attention_mask + if idx % (2 * self.attend_text_every_n_blocks) == 0 + else image_attention_mask + ) + hidden_states = block( + hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_attention_mask=curr_encoder_attention_mask, + temb=temb, + ) + all_hidden_states.append(hidden_states) + + shift, scale = self.proj_out_1(F.silu(temb)).chunk(2, dim=1) + hidden_states = self.norm_out(hidden_states) * (1 + scale[:, None]) + shift[:, None] + if return_all_hidden_states: + return self.proj_out_2(hidden_states), all_hidden_states + return self.proj_out_2(hidden_states) + + +class SelfAttentionTransformer(ModelMixin, ConfigMixin): + _supports_gradient_checkpointing = True + + @register_to_config + def __init__( + self, + num_attention_heads: int = 8, + attention_head_dim: int = 64, + output_dim: int = 26, + num_layers: int = 12, + dropout: float = 0.1, + attention_bias: bool = True, + activation_fn: str = "gelu-approximate", + num_embeds_ada_norm: int | None = 1000, + upcast_attention: bool = False, + max_num_positional_embeddings: int = 512, + compute_dtype=torch.float32, + final_dropout: bool = True, + positional_embeddings: str | None = "sinusoidal", + interleave_self_attention=False, + ): + super().__init__() + + self.inner_dim = num_attention_heads * attention_head_dim + self.gradient_checkpointing = False + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + self.inner_dim, + num_attention_heads, + attention_head_dim, + dropout=dropout, + activation_fn=activation_fn, + attention_bias=attention_bias, + upcast_attention=upcast_attention, + positional_embeddings=positional_embeddings, + num_positional_embeddings=max_num_positional_embeddings, + final_dropout=final_dropout, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + return_all_hidden_states: bool = False, + ): + hidden_states = hidden_states.contiguous() + all_hidden_states = [hidden_states] + for block in self.transformer_blocks: + hidden_states = block(hidden_states) + all_hidden_states.append(hidden_states) + if return_all_hidden_states: + return hidden_states, all_hidden_states + return hidden_states diff --git a/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py b/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py new file mode 100644 index 00000000000..81fcb4726f6 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/modeling/modules/embodiment_conditioned_mlp.py @@ -0,0 +1,178 @@ +# 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 vllm_omni.diffusion.utils.flow_matching import SinusoidalPositionalEncoding, swish + + +class CategorySpecificLinear(nn.Module): + """Linear layer with category-specific weights and biases for multi-embodiment support.""" + + def __init__(self, num_categories: int, input_dim: int, hidden_dim: int) -> None: + super().__init__() + self.num_categories = num_categories + # For each category, we have separate weights and biases. + self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) + self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) + + def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor: + """Apply per-category linear projection. + + Args: + x: ``[B, T, input_dim]`` input tensor. + cat_ids: ``[B]`` category / embodiment IDs. + + Returns: + ``[B, T, hidden_dim]`` output tensor. + """ + selected_W = self.W[cat_ids] + selected_b = self.b[cat_ids] + return torch.bmm(x, selected_W) + selected_b.unsqueeze(1) + + def expand_action_dimension( + self, + old_action_dim: int, + new_action_dim: int, + expand_input: bool = False, + expand_output: bool = False, + ) -> None: + """Safely expand action dimension with explicit targeting. + + Args: + old_action_dim: Original action dimension. + new_action_dim: New (larger) action dimension. + expand_input: Whether to expand input dimension (dim=1). + expand_output: Whether to expand output dimension (dim=2). + """ + if new_action_dim <= old_action_dim: + raise ValueError(f"New action dim {new_action_dim} must be larger than old action dim {old_action_dim}") + + # Expand input dimension (dim=1) only if explicitly requested AND dimensions match + if expand_input and self.W.shape[1] == old_action_dim: + repeat_times = new_action_dim // old_action_dim + remainder = new_action_dim % old_action_dim + + new_W_parts = [self.W] * repeat_times + if remainder > 0: + new_W_parts.append(self.W[:, :remainder, :]) + + new_W = torch.cat(new_W_parts, dim=1) + self.W = nn.Parameter(new_W) + + # Expand output dimension (dim=2) only if explicitly requested AND dimensions match + if expand_output and self.W.shape[2] == old_action_dim: + repeat_times = new_action_dim // old_action_dim + remainder = new_action_dim % old_action_dim + + new_W_parts = [self.W] * repeat_times + if remainder > 0: + new_W_parts.append(self.W[:, :, :remainder]) + + new_W = torch.cat(new_W_parts, dim=2) + self.W = nn.Parameter(new_W) + + # Expand bias for output dimension + if self.b.shape[1] == old_action_dim: + new_b_parts = [self.b] * repeat_times + if remainder > 0: + new_b_parts.append(self.b[:, :remainder]) + + new_b = torch.cat(new_b_parts, dim=1) + self.b = nn.Parameter(new_b) + + +class CategorySpecificMLP(nn.Module): + """Two-layer MLP with category-specific weights for multi-embodiment support.""" + + def __init__(self, num_categories: int, input_dim: int, hidden_dim: int, output_dim: int) -> None: + super().__init__() + self.num_categories = num_categories + self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) + self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) + + def forward(self, x: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor: + """Apply two-layer per-category MLP. + + Args: + x: ``[B, T, input_dim]`` input tensor. + cat_ids: ``[B]`` category / embodiment IDs. + + Returns: + ``[B, T, output_dim]`` output tensor. + """ + hidden = F.relu(self.layer1(x, cat_ids)) + return self.layer2(hidden, cat_ids) + + def expand_action_dimension(self, old_action_dim: int, new_action_dim: int) -> None: + """Expand action dimension by copying weights from existing dimensions. + + Args: + old_action_dim: Original action dimension. + new_action_dim: New (larger) action dimension. + """ + # self.layer1 does not take action_dim as input, so no expansion needed + self.layer2.expand_action_dimension(old_action_dim, new_action_dim, expand_input=False, expand_output=True) + + +class MultiEmbodimentActionEncoder(nn.Module): + """Action encoder with multi-embodiment support and sinusoidal positional encoding.""" + + def __init__(self, action_dim: int, hidden_size: int, num_embodiments: int) -> None: + super().__init__() + self.hidden_size = hidden_size + self.num_embodiments = num_embodiments + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) + self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions: torch.Tensor, timesteps: torch.Tensor, cat_ids: torch.Tensor) -> torch.Tensor: + """Encode actions with embodiment-conditioned MLP and time embedding. + + Args: + actions: ``[B, T, action_dim]`` action tensor. + timesteps: ``[B]`` timesteps (one scalar per batch item). + cat_ids: ``[B]`` category / embodiment IDs. + + Returns: + ``[B, T, hidden_size]`` encoded action features. + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError("Expected `timesteps` to have shape (B,) so we can replicate across T.") + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions, cat_ids) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x, cat_ids)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x, cat_ids) + return x + + def expand_action_dimension(self, old_action_dim: int, new_action_dim: int) -> None: + """Expand action dimension by copying weights from existing dimensions. + + Args: + old_action_dim: Original action dimension. + new_action_dim: New (larger) action dimension. + """ + # Only W1 takes action_dim as input, so only expand its input dimension + self.W1.expand_action_dimension(old_action_dim, new_action_dim, expand_input=True, expand_output=False) diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py new file mode 100644 index 00000000000..ade5eac83b1 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py @@ -0,0 +1,561 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +import os +import re +import warnings +from copy import deepcopy +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torchvision.transforms.v2 as transforms +from PIL import Image +from transformers import AutoProcessor, ProcessorMixin +from transformers.feature_extraction_utils import BatchFeature +from transformers.utils import cached_file +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig +from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config +from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag +from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor +from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable + + +class LetterBoxTransform: + """Pad image to square dimensions by adding black bars to the smaller side.""" + + def __call__(self, img: torch.Tensor) -> torch.Tensor: + import math + + *leading_dims, c, h, w = img.shape + if h == w: + return img + max_dim = max(h, w) + pad_h, pad_w = max_dim - h, max_dim - w + pad_top, pad_left = pad_h // 2, pad_w // 2 + pad_bottom, pad_right = pad_h - pad_top, pad_w - pad_left + if leading_dims: + batch_size = math.prod(leading_dims) + img_r = img.reshape(batch_size, c, h, w) + padded = transforms.functional.pad(img_r, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) + return padded.reshape(leading_dims + [c, max_dim, max_dim]) + return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) + + +def _build_eval_image_transform( + image_target_size: list[int], + image_crop_size: list[int], +) -> transforms.Compose: + """Deterministic eval/inference image transform (letterbox → resize → centercrop → resize).""" + return transforms.Compose( + [ + transforms.ToImage(), + LetterBoxTransform(), + transforms.Resize(size=image_target_size), + transforms.CenterCrop(size=image_crop_size), + transforms.Resize(size=image_target_size), + ] + ) + + +logger = init_logger(__name__) + +# Suppress protobuf deprecation warnings +warnings.filterwarnings("ignore", category=DeprecationWarning, module="google.protobuf") + +EMBODIMENT_TAG_TO_PROJECTOR_INDEX = { + # Pretrain embodiment ids + "oxe_droid_relative_eef_relative_joint": 24, + "xdof_relative_eef_relative_joint": 27, + "xdof_relative_eef_relative_joint_subtask": 27, + "real_g1_relative_eef_relative_joints": 25, + "real_r1_pro_sharpa_relative_eef": 26, + "real_r1_pro_sharpa_relative_eef_human": 26, + "real_r1_pro_sharpa_relative_eef_maxinsights": 26, + "real_r1_pro_sharpa_relative_eef_mecka": 26, + # Posttrain embodiment ids + "unitree_g1_full_body_with_waist_height_nav_cmd": 25, + "unitree_g1_sonic": 11, + "simpler_env_google": 0, + "simpler_env_widowx": 1, + "libero_sim": 2, + "new_embodiment": 10, +} + +QWEN3_VL_2B_PROCESSOR = "Qwen/Qwen3-VL-2B-Instruct" + + +def build_processor(model_name: str, transformers_loading_kwargs: dict) -> ProcessorMixin: + try: + from transformers import Qwen3VLProcessor + except ImportError as exc: + raise ImportError( + "GR00T-N1.7 requires transformers>=4.57.1 for Qwen3VLProcessor " + "(the repo's declared floor is 4.56.0). Please upgrade transformers." + ) from exc + if model_name == "nvidia/Cosmos-Reason2-2B": + # Cosmos-Reason2-2B lacks a Qwen3VLProcessor; fall back to upstream artifacts. + logger.warning_once( + "Substituting processor from %s because %s does not ship one. " + "If you fine-tune Cosmos-Reason2-2B's tokenizer/image processor, " + "load the processor explicitly instead of relying on this fallback.", + QWEN3_VL_2B_PROCESSOR, + model_name, + ) + model_name = QWEN3_VL_2B_PROCESSOR + return Qwen3VLProcessor.from_pretrained(model_name, **transformers_loading_kwargs) + + +class Gr00tN1d7DataCollator: + def __init__( + self, + model_name: str, + model_type: str = "qwen", + transformers_loading_kwargs: dict | None = None, + ): + if transformers_loading_kwargs is None: + transformers_loading_kwargs = {} + self.processor = build_processor(model_name, transformers_loading_kwargs) + self.processor.tokenizer.padding_side = "left" + self.model_type = model_type + self.model_name = model_name + + def __call__(self, features: list[dict[str, Any]]) -> BatchFeature: + batch = {} + keys = list(set().union(*(elem.keys() for elem in features))) + + for key in keys: + values = [elem[key] for elem in features if key in elem] + if key == "vlm_content": + text_list = [] + image_inputs = [] + for v in values: + curr_text_list = [v["text"]] + + text_list += curr_text_list + curr_image_inputs = v["images"] + image_inputs += curr_image_inputs + + vlm_inputs = self.processor( + text=text_list, + images=image_inputs, + return_tensors="pt", + padding=True, + ) + for k, v in vlm_inputs.items(): + batch[k] = v + elif key in ( + "pixel_values", + "image_grid_thw", + "attention_mask", + "input_ids", + ): + raise Exception("Not implemented") + else: + batch[key] = torch.from_numpy(np.stack(values)) + return BatchFeature(data={"inputs": batch}) + + def __str__(self): + return f"Gr00tN1d7DataCollator(model_name={self.model_name}, model_type={self.model_type})" + + +class Gr00tN1d7Processor(ProcessorMixin): + data_collator_class = Gr00tN1d7DataCollator + + def __init__( + self, + modality_configs: dict[str, dict[str, ModalityConfig]], + statistics: (dict[str, dict[str, dict[str, dict[str, list[float]]]]] | None) = None, + use_percentiles: bool = False, + clip_outliers: bool = True, + image_crop_size: list[int] = None, + image_target_size: list[int] = None, + shortest_image_edge: int = 256, + crop_fraction: float = 0.95, + random_rotation_angle: int | None = None, + color_jitter_params: dict[str, float] | None = None, + formalize_language: bool = True, + model_name: str = "nvidia/Cosmos-Reason2-2B", + model_type: str = "qwen", + max_state_dim: int = 29, + max_action_dim: int = 29, + max_action_horizon: int = 50, + apply_sincos_state_encoding: bool = False, + use_relative_action: bool = False, + embodiment_id_mapping: dict[str, int] | None = None, + transformers_loading_kwargs: dict | None = None, + exclude_state: bool = False, + # Normalization + use_mean_std: bool = False, + **kwargs, # absorb deprecated training-only keys from saved processor_config.json + ): + if transformers_loading_kwargs is None: + transformers_loading_kwargs = {"trust_remote_code": True} + if kwargs: + logger.debug("Gr00tN1d7Processor: ignoring unknown keys: %s", list(kwargs)) + self.modality_configs = parse_modality_configs(modality_configs) + + # Initialize StateActionProcessor for state/action normalization + self.state_action_processor = StateActionProcessor( + modality_configs=modality_configs, + statistics=statistics, + use_percentiles=use_percentiles, + clip_outliers=clip_outliers, + apply_sincos_state_encoding=apply_sincos_state_encoding, + use_relative_action=use_relative_action, + ) + + self.use_percentiles = use_percentiles + self.use_mean_std = use_mean_std + self.clip_outliers = clip_outliers + self.apply_sincos_state_encoding = apply_sincos_state_encoding + self.use_relative_action = use_relative_action + + self.exclude_state = exclude_state + + self.formalize_language = formalize_language + self.model_name = model_name + self.model_type = model_type + + self.max_state_dim = max_state_dim + self.max_action_dim = max_action_dim + self.max_action_horizon = max_action_horizon + + self.image_crop_size = image_crop_size + self.image_target_size = image_target_size + self.random_rotation_angle = random_rotation_angle + self.color_jitter_params = color_jitter_params + self.processor = build_processor(model_name, transformers_loading_kwargs) + self.processor.tokenizer.padding_side = "left" + self.embodiment_id_mapping = embodiment_id_mapping or EMBODIMENT_TAG_TO_PROJECTOR_INDEX + for k, v in EMBODIMENT_TAG_TO_PROJECTOR_INDEX.items(): + if k not in self.embodiment_id_mapping: + self.embodiment_id_mapping[k] = v + self.shortest_image_edge = shortest_image_edge + self.crop_fraction = crop_fraction + + self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {} + + # Eval/inference image transform + self.eval_image_transform = _build_eval_image_transform( + image_target_size, + image_crop_size, + ) + self._collator = self.data_collator_class( + model_name=model_name, + model_type=model_type, + transformers_loading_kwargs=transformers_loading_kwargs, + ) + + @property + def collator(self): + return self._collator + + def set_statistics( + self, + statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]], + override: bool = False, + ) -> None: + """Set dataset statistics for normalization.""" + for key in statistics: + if key not in self.statistics or override: + if override: + logger.info("Overriding statistics for %s", key) + self.statistics[key] = deepcopy(statistics[key]) + else: + logger.debug("Embodiment tag %s already in statistics, skipping update", key) + + self.state_action_processor.set_statistics(statistics, override=override) + + # Compute action dimensions for convenience + self.action_dim = {} + for embodiment_tag in self.state_action_processor.statistics: + self.action_dim[embodiment_tag] = self.state_action_processor.get_action_dim(embodiment_tag) + + def decode_action( + self, + action: np.ndarray, + embodiment_tag: EmbodimentTag, + state: dict[str, np.ndarray] | None = None, + ): + """Undo action normalization and convert relative actions to absolute.""" + # Split concatenated action into joint groups + out_dict = {} + start_idx = 0 + joint_groups = self.modality_configs[embodiment_tag.value]["action"].modality_keys + action_horizon = len(self.modality_configs[embodiment_tag.value]["action"].delta_indices) + for key in joint_groups: + joint_dim = self.state_action_processor.norm_params[embodiment_tag.value]["action"][key]["dim"].item() + out_dict[key] = action[..., :action_horizon, start_idx : start_idx + joint_dim] + start_idx += joint_dim + + # Use StateActionProcessor to unnormalize and convert to absolute + return self.state_action_processor.unapply_action(out_dict, embodiment_tag.value, state=state) + + def _apply_vlm_processing(self, images: np.ndarray, language: str) -> BatchFeature: + pil_images = [Image.fromarray(np.transpose(v, (1, 2, 0))) for v in images] + conversation = [ + { + "role": "user", + "content": [ + *[{"type": "image", "image": img} for img in pil_images], + {"type": "text", "text": language}, + ], + } + ] + + text = self.processor.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False) + return { + "vlm_content": { + "text": text, + "images": pil_images, + "conversation": conversation, + } + } + + def __call__( + self, + messages: list[dict[str, Any]], + ): + assert len(messages) == 1 + content = messages[0]["content"] + embodiment_tag = content.embodiment + action_data = content.actions + state_data = content.states + + norm_state_dict, normalized_actions = self.state_action_processor.apply( + state=state_data, + action=action_data, + embodiment_tag=embodiment_tag.value, + ) + + if normalized_actions: + action_keys = self.modality_configs[embodiment_tag.value]["action"].modality_keys + normalized_actions = torch.cat( + [torch.from_numpy(normalized_actions[key]) for key in action_keys], + dim=-1, + ) # (t, d) + action_dim = normalized_actions.shape[1] + normalized_actions = torch.cat( + [ + normalized_actions, + torch.zeros( + normalized_actions.shape[0], + self.max_action_dim - normalized_actions.shape[1], + ), + ], + dim=-1, + ) # (t, max_action_dim) + action_horizon = normalized_actions.shape[0] + assert action_horizon <= self.max_action_horizon, ( + f"Action sequence length {action_horizon} exceeds max_action_horizon" + f" {self.max_action_horizon}. Increase model config action_horizon to" + f" >= {action_horizon}." + ) + normalized_actions = torch.cat( + [ + normalized_actions, + torch.zeros( + self.max_action_horizon - normalized_actions.shape[0], + self.max_action_dim, + ), + ], + dim=0, + ) # (max_action_horizon, max_action_dim) + action_mask = torch.ones_like(normalized_actions) + action_mask[action_horizon:] = 0 + action_mask[:, action_dim:] = 0 + else: + normalized_actions = None + action_mask = None + + state_keys = self.modality_configs[embodiment_tag.value]["state"].modality_keys + exclude_state = self.exclude_state or getattr( + self.modality_configs[embodiment_tag.value]["state"], "exclude_state", False + ) + if exclude_state: + normalized_states = torch.cat( + [torch.from_numpy(np.zeros_like(state_data[key])) for key in state_keys], dim=-1 + ) + else: + normalized_states = torch.cat([torch.from_numpy(norm_state_dict[key]) for key in state_keys], dim=-1) + normalized_states = torch.cat( + [ + normalized_states, + torch.zeros( + normalized_states.shape[0], + self.max_state_dim - normalized_states.shape[1], + ), + ], + dim=-1, + ) + + image_transform = self.eval_image_transform + image_keys = self.modality_configs[embodiment_tag.value]["video"].modality_keys + + if self.formalize_language: + language = content.text.lower() + language = re.sub(r"[^\w\s]", "", language) + else: + language = content.text + + vlm_inputs = self._get_vlm_inputs( + image_keys=image_keys, + images=content.images, + masks=content.masks, + image_transform=image_transform, + language=language, + ) + + transformed_inputs = { + "state": normalized_states.to(torch.get_default_dtype()), + } + if normalized_actions is not None: + transformed_inputs["action"] = normalized_actions.to(torch.get_default_dtype()) + # Add VLM inputs + transformed_inputs.update(vlm_inputs) + if action_mask is not None: + transformed_inputs["action_mask"] = action_mask + transformed_inputs["embodiment_id"] = self.embodiment_id_mapping[embodiment_tag.value] + return transformed_inputs + + def _get_vlm_inputs( + self, + image_keys: list[str], + images: list[Image.Image], + masks: dict[str, list[np.ndarray]] | None, + image_transform: transforms.Compose, + language: str, + ): + temporal_stacked_images = {} + + if masks is not None: + raise ValueError("Mask-based transforms are not supported at inference.") + for view in image_keys: + assert view in images, f"{view} not in {images}" + temporal_stacked_images[view] = torch.stack([image_transform(img) for img in images[view]]) # (T, C, H, W) + + for k, v in temporal_stacked_images.items(): + assert isinstance(k, str), f"{k} is not a string" + assert isinstance(v, torch.Tensor), f"{v} is not a torch tensor" + assert v.ndim == 4, f"{v} is not a 4D tensor" + assert v.dtype == torch.uint8, f"{v} is not a uint8 tensor" + assert v.shape[1] == 3, f"{v} is not a 3 channel tensor" + + stacked_images = ( + torch.stack([temporal_stacked_images[view] for view in image_keys], dim=1).flatten(0, 1).numpy() + ) # (T*V, C, H, W), processor expects numpy array + + vlm_inputs = self._apply_vlm_processing(stacked_images, language) + return vlm_inputs + + def save_pretrained(self, save_directory: str | Path) -> list[Path]: + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + main_config_file = save_directory / "processor_config.json" + statistics_file = save_directory / "statistics.json" + embodiment_id_file = save_directory / "embodiment_id.json" + + config = { + "processor_class": self.__class__.__name__, + "processor_kwargs": { + "modality_configs": to_json_serializable(self.modality_configs), + "image_crop_size": self.image_crop_size, + "image_target_size": self.image_target_size, + "random_rotation_angle": self.random_rotation_angle, + "color_jitter_params": self.color_jitter_params, + "shortest_image_edge": self.shortest_image_edge, + "crop_fraction": self.crop_fraction, + "model_name": self.model_name, + "model_type": self.model_type, + "formalize_language": self.formalize_language, + "max_state_dim": self.max_state_dim, + "max_action_dim": self.max_action_dim, + "max_action_horizon": self.max_action_horizon, + "use_percentiles": self.use_percentiles, + "use_mean_std": self.use_mean_std, + "clip_outliers": self.clip_outliers, + "apply_sincos_state_encoding": self.apply_sincos_state_encoding, + "use_relative_action": self.use_relative_action, + "exclude_state": self.exclude_state, + }, + } + with open(main_config_file, "w") as f: + json.dump(config, f, indent=2) + with open(statistics_file, "w") as f: + json.dump( + to_json_serializable(self.state_action_processor.statistics), + f, + indent=2, + ) + with open(embodiment_id_file, "w") as f: + json.dump(self.embodiment_id_mapping, f, indent=2) + return [main_config_file, statistics_file, embodiment_id_file] + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str | Path, **kwargs): + transformers_loading_kwargs = kwargs.pop("transformers_loading_kwargs", {"trust_remote_code": True}) + pretrained_model_name_or_path = Path(pretrained_model_name_or_path) + config_file = pretrained_model_name_or_path / "processor_config.json" + statistics_file = pretrained_model_name_or_path / "statistics.json" + embodiment_id_file = pretrained_model_name_or_path / "embodiment_id.json" + is_local = os.path.isdir(pretrained_model_name_or_path) + if not is_local: + config_file = Path(cached_file(pretrained_model_name_or_path, "processor_config.json")) + statistics_file = Path(cached_file(pretrained_model_name_or_path, "statistics.json")) + embodiment_id_file = Path(cached_file(pretrained_model_name_or_path, "embodiment_id.json")) + + with open(config_file) as f: + config = json.load(f) + with open(statistics_file) as f: + statistics = json.load(f) + if embodiment_id_file.exists(): + with open(embodiment_id_file) as f: + embodiment_id_mapping = json.load(f) + else: + embodiment_id_mapping = None + processor_kwargs = config["processor_kwargs"] + processor_kwargs["statistics"] = statistics + processor_kwargs["embodiment_id_mapping"] = embodiment_id_mapping + + # Backfill missing fields from older checkpoints. + processor_kwargs.setdefault("model_name", "nvidia/Cosmos-Reason2-2B") + processor_kwargs.setdefault("model_type", "qwen") + processor_kwargs.setdefault("clip_outliers", True) + + # Directly override other processor kwargs + if kwargs: + # Override modality configs while keeping pretrained embodiment configs + modality_configs = kwargs.pop("modality_configs", {}) + for embodiment_tag, modality_config in modality_configs.items(): + processor_kwargs["modality_configs"][embodiment_tag] = modality_config + override_keys = [ + "random_rotation_angle", + "color_jitter_params", + "use_relative_action", + "exclude_state", + "use_mean_std", + "model_name", + "model_type", + "max_action_horizon", + "max_state_dim", + "max_action_dim", + ] + for key in override_keys: + if key in kwargs: + override = kwargs.pop(key) + if override is not None: + processor_kwargs[key] = override + return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs) + + +# transformers >= 5.5 requires the config *class* (not a model_type string) as the +# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing +# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute +# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in +# modeling/gr00t_n1d7.py. +AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor) diff --git a/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py new file mode 100644 index 00000000000..1154dcf7145 --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/pipeline_gr00t.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable, Mapping +from typing import Any + +import numpy as np +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.gr00t.policy import Gr00tPolicy +from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +logger = init_logger(__name__) + + +def _to_float32_action_dict(actions: Mapping[str, Any]) -> dict[str, np.ndarray]: + converted = {str(key): np.asarray(value, dtype=np.float32) for key, value in actions.items()} + if not converted: + raise RuntimeError("GR00T policy returned an empty action dict.") + return converted + + +class Gr00tN1d7Pipeline(nn.Module): + """GR00T N1.7 policy pipeline backed by vLLM-Omni's local GR00T port. + + vLLM-Omni owns the serving integration: OpenPI observations arrive through + `sampling_params.extra_args["robot_obs"]`, this pipeline runs GR00T policy + inference, and actions are returned through `DiffusionOutput.output["actions"]`. + """ + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: + super().__init__() + model_config = od_config.model_config + self.model_path = od_config.model + self.embodiment_tag = str(model_config.get("embodiment_tag") or "OXE_DROID_RELATIVE_EEF_RELATIVE_JOINT") + self.strict = bool(model_config.get("strict", True)) + self.device = "cuda" if torch.cuda.is_available() else "cpu" + + logger.info("Loading GR00T N1.7 policy from %s with embodiment_tag=%s", self.model_path, self.embodiment_tag) + self.policy = Gr00tPolicy( + model_path=self.model_path, + embodiment_tag=self.embodiment_tag, + device=self.device, + strict=self.strict, + ) + self._validate_policy_server_config(model_config.get("policy_server_config")) + + def _validate_policy_server_config(self, psc: Mapping[str, Any] | None) -> None: + """Fail fast if the deploy handshake drifts from the loaded checkpoint. + + ``policy_server_config`` is sent verbatim to the OpenPI client, so its + model/embodiment-specific values must match what the loaded policy actually + produces; otherwise the client is handed the wrong action contract. + """ + if not isinstance(psc, Mapping): + return + action_config = self.policy.modality_configs["action"] + expected_horizon = len(action_config.delta_indices) + expected_keys = set(action_config.modality_keys) + + if "action_horizon" in psc and psc["action_horizon"] != expected_horizon: + raise ValueError( + f"policy_server_config.action_horizon={psc['action_horizon']} != loaded model's " + f"action horizon {expected_horizon}." + ) + if "action_keys" in psc and set(psc["action_keys"]) != expected_keys: + raise ValueError( + f"policy_server_config.action_keys={list(psc['action_keys'])} != loaded model's " + f"action keys {sorted(expected_keys)}." + ) + psc_embodiment = psc.get("embodiment_tag") + if psc_embodiment is not None and psc_embodiment != self.embodiment_tag: + raise ValueError( + f"policy_server_config.embodiment_tag={psc_embodiment!r} != " + f"model_config.embodiment_tag={self.embodiment_tag!r}." + ) + + def reset(self) -> dict[str, Any]: + return self.policy.reset() or {} + + @property + def weights_sources(self) -> tuple[Any, ...]: + return () + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + consumed = list(weights) + if consumed: + raise RuntimeError( + f"Gr00tN1d7Pipeline.load_weights received {len(consumed)} weight tensors; " + "weights_sources=() should prevent this. GR00T weights are loaded directly by Gr00tPolicy." + ) + return set() + + def _dummy_actions(self) -> dict[str, np.ndarray]: + embodiment_value = self.policy.embodiment_tag.value + action_config = self.policy.modality_configs["action"] + horizon = len(action_config.delta_indices) + norm_params = self.policy.processor.state_action_processor.norm_params[embodiment_value]["action"] + actions = {} + for key in action_config.modality_keys: + dim = norm_params[key]["dim"] + dim = int(dim.item() if hasattr(dim, "item") else dim) + actions[key] = np.zeros((1, horizon, dim), dtype=np.float32) + return actions + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput: + del kwargs + extra_args = req.sampling_params.extra_args or {} + robot_obs = extra_args.get("robot_obs") + if robot_obs is None: + if req.request_id == DUMMY_DIFFUSION_REQUEST_ID: + return DiffusionOutput(output={"actions": self._dummy_actions()}) + return DiffusionOutput(error="Gr00tN1d7Pipeline.forward expects sampling_params.extra_args['robot_obs'].") + if not isinstance(robot_obs, Mapping): + return DiffusionOutput(error=f"robot_obs must be a dict, got {type(robot_obs).__name__}.") + + if extra_args.get("reset"): + self.reset() + + policy_obs = _normalize_observation(robot_obs, language_key=self.policy.language_key) + result = self.policy.get_action(policy_obs) + actions = result[0] if isinstance(result, tuple) else result + if not isinstance(actions, Mapping): + return DiffusionOutput(error=f"GR00T policy returned {type(actions).__name__}; expected dict actions.") + # Return actions via output.output (like the DreamZero OpenPI policy) so the engine's + # empty-output guard passes. + return DiffusionOutput(output={"actions": _to_float32_action_dict(actions)}) + + +def _normalize_observation(robot_obs: Mapping[str, Any], *, language_key: str) -> dict[str, Any]: + obs: dict[str, Any] = {} + if "video" in robot_obs: + obs["video"] = robot_obs["video"] + elif "images" in robot_obs: + obs["video"] = robot_obs["images"] + if "state" in robot_obs: + obs["state"] = robot_obs["state"] + if "language" in robot_obs: + obs["language"] = robot_obs["language"] + else: + prompt = robot_obs.get("prompt") + if prompt is not None: + obs["language"] = {language_key: [[str(prompt)]]} + return obs diff --git a/vllm_omni/diffusion/models/gr00t/policy.py b/vllm_omni/diffusion/models/gr00t/policy.py new file mode 100644 index 00000000000..c16b344632f --- /dev/null +++ b/vllm_omni/diffusion/models/gr00t/policy.py @@ -0,0 +1,434 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from transformers import AutoModel, AutoProcessor + +from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import FINETUNE_ONLY_TAGS, POSTTRAIN_TAGS, EmbodimentTag +from vllm_omni.diffusion.models.gr00t.dataio.types import MessageType, VLAStepData +from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7 +from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import Gr00tN1d7Processor + + +def _rec_to_dtype(value: Any, dtype: torch.dtype) -> Any: + """Recursively convert floating-point tensors in nested collator output.""" + if isinstance(value, torch.Tensor): + return value.to(dtype=dtype) if torch.is_floating_point(value) else value + if isinstance(value, dict) or hasattr(value, "items"): + return {key: _rec_to_dtype(item, dtype) for key, item in value.items()} + if isinstance(value, list): + return [_rec_to_dtype(item, dtype) for item in value] + return value + + +class Gr00tPolicy: + """Core policy class for Gr00t model inference. + + This policy handles the end-to-end inference pipeline: + 1. Validates input observations + 2. Processes observations with pretrained VLA processor + 3. Runs model inference + 4. Decodes and returns actions + + The policy expects observations with specific modalities (video, state, language) + and returns actions in the format defined by the model's modality configuration. + """ + + def __init__( + self, + embodiment_tag: EmbodimentTag | str, + model_path: str, + *, + device: int | str, + strict: bool = True, + ): + """Initialize the Gr00t Policy. + + Args: + embodiment_tag: The embodiment tag defining the robot/environment type. + Accepts an EmbodimentTag enum or a string (resolved case-insensitively). + model_path: Path to the pretrained model checkpoint directory + device: Device to run the model on (e.g., 'cuda:0', 0, 'cpu') + strict: Whether to enforce strict input validation (default: True) + """ + self.strict = strict + if isinstance(embodiment_tag, str): + embodiment_tag = EmbodimentTag.resolve(embodiment_tag) + model_dir = Path(model_path) + + # Load directly into bfloat16 to avoid a transient fp32 copy at load time. + self.model: Gr00tN1d7 = AutoModel.from_pretrained(model_dir, torch_dtype=torch.bfloat16) + self.model.eval() + self.model.to(device=device) + + # Training saves processor files under a "processor/" subdirectory, but + # AutoProcessor expects them at the model root. Fall back to the + # subdirectory when the root lacks a processor_config.json. + processor_dir = ( + model_dir / "processor" + if (model_dir / "processor").is_dir() and not (model_dir / "processor_config.json").exists() + else model_dir + ) + self.processor: Gr00tN1d7Processor = AutoProcessor.from_pretrained(processor_dir) + + # Store embodiment-specific configurations + self.embodiment_tag = embodiment_tag + all_modality_configs = self.processor.modality_configs + if self.embodiment_tag.value not in all_modality_configs: + # Map raw checkpoint tag values to user-friendly enum names where possible. + supported_lines = [] + for tag_value in sorted(all_modality_configs.keys()): + enum_name = EmbodimentTag.reverse_lookup(tag_value) + if enum_name != tag_value: + supported_lines.append(f" {enum_name:30s} (--embodiment-tag {enum_name})") + else: + supported_lines.append(f" {tag_value:30s} (internal, no public enum)") + supported_str = "\n".join(supported_lines) + + hint = "" + if self.embodiment_tag in POSTTRAIN_TAGS: + hint = ( + f"\n\nHint: '{self.embodiment_tag.name}' is a posttrain tag that requires " + f"a finetuned checkpoint, not the base model. " + f"See the example READMEs for how to finetune and download checkpoints." + ) + elif self.embodiment_tag in FINETUNE_ONLY_TAGS: + hint = ( + f"\n\nHint: '{self.embodiment_tag.name}' is for finetuning custom robots. " + f"Use it with launch_finetune.py, not with the base model directly." + ) + + raise ValueError( + f"Embodiment tag '{self.embodiment_tag.name}' " + f"(value='{self.embodiment_tag.value}') is not supported " + f"by this checkpoint.\n\n" + f"Supported tags in this checkpoint:\n{supported_str}" + f"{hint}" + ) + self.modality_configs = { + k: v for k, v in all_modality_configs[self.embodiment_tag.value].items() if k != "rl_info" + } + self.collate_fn = self.processor.collator + + # Extract and validate language configuration + # Embodiments may define multiple language keys (e.g. paraphrases); use only the first at inference. + language_keys = self.modality_configs["language"].modality_keys + language_delta_indices = self.modality_configs["language"].delta_indices + assert len(language_keys) >= 1, "At least one language key is required" + assert len(language_delta_indices) == 1, "Only one language delta index is supported" + self.language_key = language_keys[0] + + def _unbatch_observation(self, value: dict[str, Any]) -> list[dict[str, Any]]: + """Unbatch a batched observation into a list of single observations. + + Args: + value: Batched observation with shape (B, ...) for each modality + + Returns: + List of B observations, each with the batch dimension removed + """ + unbatched_obs = [] + # Infer batch size from the first video key + batch_size = value["video"][list(value["video"].keys())[0]].shape[0] + + # Split each modality along the batch dimension + for i in range(batch_size): + unbatched_value = { + "video": {k: v[i] for k, v in value["video"].items()}, + "state": {k: v[i] for k, v in value["state"].items()}, + "language": {k: v[i] for k, v in value["language"].items()}, + } + unbatched_obs.append(unbatched_value) + return unbatched_obs + + def _to_vla_step_data(self, observation: dict[str, Any]) -> VLAStepData: + """Convert a single observation into a VLAStepData object for processing. + + Args: + observation: Single observation dict with video, state, and language + + Returns: + VLAStepData object ready for processor input + """ + return VLAStepData( + images=observation["video"], + states=observation["state"], + actions={}, # No ground truth actions during inference + text=observation["language"][self.language_key][0], + embodiment=self.embodiment_tag, + ) + + def check_observation(self, observation: dict[str, Any]) -> None: + """Validate that the observation has the correct structure and types. + + This method ensures that all required modalities are present and that their + data types, shapes, and dimensions match the model's expectations. + + Expected observation structure: + - video: dict[str, np.ndarray[np.uint8, (B, T, H, W, C)]] + - B: batch size + - T: temporal horizon (number of frames) + - H, W: image height and width + - C: number of channels (must be 3 for RGB) + - state: dict[str, np.ndarray[np.float32, (B, T, D)]] + - B: batch size + - T: temporal horizon (number of state observations) + - D: state dimension + - language: dict[str, list[list[str]]] + - Shape: (B, T) where each element is a string + - T: temporal horizon (typically 1 for language) + + Args: + observation: Dictionary containing video, state, and language modalities + + Raises: + AssertionError: If any validation check fails + """ + # Check that observation contains all required top-level modality keys + for modality in ["video", "state", "language"]: + assert modality in observation, f"Observation must contain a '{modality}' key" + assert isinstance(observation[modality], dict), ( + f"Observation '{modality}' must be a dictionary. " + f"Got {type(observation[modality])}: {observation[modality]}" + ) + + # Track batch size across modalities to ensure consistency + bs = -1 + + for video_key in self.modality_configs["video"].modality_keys: + assert video_key in observation["video"], f"Video key '{video_key}' must be in observation" + + # Set or verify batch size consistency across all video keys + if bs == -1: + bs = len(observation["video"][video_key]) + else: + assert len(observation["video"][video_key]) == bs, ( + f"Video key '{video_key}' must have batch size {bs}. Got {len(observation['video'][video_key])}" + ) + + batched_video = observation["video"][video_key] + + # Verify data type is numpy array + assert isinstance(batched_video, np.ndarray), ( + f"Video key '{video_key}' must be a numpy array. Got {type(batched_video)}" + ) + + # Verify dtype is uint8 (standard for image data, range 0-255) + assert batched_video.dtype == np.uint8, ( + f"Video key '{video_key}' must be a numpy array of type np.uint8. Got {batched_video.dtype}" + ) + + # Verify shape has 5 dimensions: (B, T, H, W, C) + assert batched_video.ndim == 5, ( + f"Video key '{video_key}' must be a numpy array of shape (B, T, H, W, C), got {batched_video.shape}" + ) + + # Verify temporal dimension matches the expected horizon from config + assert batched_video.shape[1] == len(self.modality_configs["video"].delta_indices), ( + f"Video key '{video_key}'s horizon must be " + f"{len(self.modality_configs['video'].delta_indices)}. Got {batched_video.shape[1]}" + ) + + # Verify channel dimension is 3 (RGB images) + assert batched_video.shape[-1] == 3, ( + f"Video key '{video_key}'s channel 'C' must be 3. Got {batched_video.shape[-1]}" + ) + + for state_key in self.modality_configs["state"].modality_keys: + # Check that the expected state key exists in the observation + # Must happen before indexing; see video validation above. + assert state_key in observation["state"], f"State key '{state_key}' must be in observation" + + # Set or verify batch size consistency across all state keys + if bs == -1: + bs = len(observation["state"][state_key]) + else: + assert len(observation["state"][state_key]) == bs, ( + f"State key '{state_key}' must have batch size {bs}. Got {len(observation['state'][state_key])}" + ) + + batched_state = observation["state"][state_key] + + # Verify data type is numpy array + assert isinstance(batched_state, np.ndarray), ( + f"State key '{state_key}' must be a numpy array. Got {type(batched_state)}" + ) + + # Verify dtype is float32 (standard for continuous state values) + assert batched_state.dtype == np.float32, ( + f"State key '{state_key}' must be a numpy array of type np.float32. Got {batched_state.dtype}" + ) + + # Verify shape has 3 dimensions: (B, T, D) + assert batched_state.ndim == 3, ( + f"State key '{state_key}' must be a numpy array of shape (B, T, D), got {batched_state.shape}" + ) + + # Verify temporal dimension matches the expected horizon from config + assert batched_state.shape[1] == len(self.modality_configs["state"].delta_indices), ( + f"State key '{state_key}'s horizon must be " + f"{len(self.modality_configs['state'].delta_indices)}. Got {batched_state.shape[1]}" + ) + + for language_key in self.modality_configs["language"].modality_keys: + # Check that the expected language key exists in the observation + # Must happen before indexing; see video validation above. + assert language_key in observation["language"], f"Language key '{language_key}' must be in observation" + + # Set or verify batch size consistency (language uses len instead of .shape) + if bs == -1: + bs = len(observation["language"][language_key]) + else: + assert len(observation["language"][language_key]) == bs, ( + f"Language key '{language_key}' must have batch size {bs}. " + f"Got {len(observation['language'][language_key])}" + ) + + batched_language: list[list[str]] = observation["language"][language_key] + + # Verify outer structure is a list (batch dimension) + assert isinstance(batched_language, list), ( + f"Language key '{language_key}' must be a list. Got {type(batched_language)}" + ) + + # Validate each batch item + for batch_item in batched_language: + # Verify temporal dimension matches expected horizon + assert len(batch_item) == len(self.modality_configs["language"].delta_indices), ( + f"Language key '{language_key}'s horizon must be " + f"{len(self.modality_configs['language'].delta_indices)}. Got {len(batched_language)}" + ) + + # Verify inner structure is also a list (temporal dimension) + assert isinstance(batch_item, list), f"Language batch item must be a list. Got {type(batch_item)}" + + # Current implementation expects exactly one language instruction per timestep + assert len(batch_item) == 1, f"Language batch item must have exactly one item. Got {len(batch_item)}" + + # Verify the instruction itself is a string + assert isinstance(batch_item[0], str), ( + f"Language batch item must be a string. Got {type(batch_item[0])}" + ) + + def _get_action( + self, observation: dict[str, Any], options: dict[str, Any] | None = None + ) -> tuple[dict[str, Any], dict[str, Any]]: + """Internal method to compute actions from observations. + + Pipeline: + 1. Unbatch observations into individual samples + 2. Convert each to VLAStepData and process + 3. Collate into model input batch + 4. Run model inference + 5. Decode and unnormalize actions + + Args: + observation: Batched observation dictionary + options: Optional parameters (currently unused) + + Returns: + Tuple of (actions_dict, info_dict) + """ + # Step 1: Split batched observation into individual observations + unbatched_observations = self._unbatch_observation(observation) + processed_inputs = [] + + # Step 2: Process each observation through the VLA processor + states = [] + for obs in unbatched_observations: + vla_step_data = self._to_vla_step_data(obs) + states.append(vla_step_data.states) # dict[str, np.ndarray[np.float32, (T, D)]] + messages = [{"type": MessageType.EPISODE_STEP.value, "content": vla_step_data}] + processed_inputs.append(self.processor(messages)) + + # Step 3: Collate processed inputs into a single batch for model + collated_inputs = self.collate_fn(processed_inputs) + collated_inputs = _rec_to_dtype(collated_inputs, dtype=torch.bfloat16) + + # Step 4: Run model inference to predict actions + with torch.inference_mode(): + model_pred = self.model.get_action(**collated_inputs) + normalized_action = model_pred["action_pred"].float() + + # Step 5: Decode actions from normalized space back to physical units + batched_states = {} + for k in self.modality_configs["state"].modality_keys: + batched_states[k] = np.stack([s[k] for s in states], axis=0) # (B, T, D) + unnormalized_action = self.processor.decode_action( + normalized_action.cpu().numpy(), self.embodiment_tag, batched_states + ) + + # Cast all actions to float32 for consistency + casted_action = {key: value.astype(np.float32) for key, value in unnormalized_action.items()} + return casted_action, {} + + def check_action(self, action: dict[str, Any]) -> None: + """Validate that the action has the correct structure and types. + + This method ensures that all required action keys are present and that their + data types, shapes, and dimensions match the model's action space. + + Expected action structure: + - action: dict[str, np.ndarray[np.float32, (B, T, D)]] + - B: batch size + - T: action horizon (number of future action steps) + - D: action dimension (e.g., joint positions, velocities, gripper state) + + Args: + action: Dictionary containing action arrays for each action key + + Raises: + AssertionError: If any validation check fails + """ + # Validate each action key defined in the modality config + for action_key in self.modality_configs["action"].modality_keys: + # Check that the expected action key exists + assert action_key in action, f"Action key '{action_key}' must be in action" + + action_arr = action[action_key] + + # Verify data type is numpy array + assert isinstance(action_arr, np.ndarray), ( + f"Action key '{action_key}' must be a numpy array. Got {type(action_arr)}" + ) + + # Verify dtype is float32 (standard for continuous actions) + assert action_arr.dtype == np.float32, ( + f"Action key '{action_key}' must be a numpy array of type np.float32. Got {action_arr.dtype}" + ) + + # Verify shape has 3 dimensions: (B, T, D) + assert action_arr.ndim == 3, ( + f"Action key '{action_key}' must be a numpy array of shape (B, T, D), got {action_arr.shape}" + ) + + # Verify action horizon matches the expected temporal dimension from config + assert action_arr.shape[1] == len(self.modality_configs["action"].delta_indices), ( + f"Action key '{action_key}'s horizon must be " + f"{len(self.modality_configs['action'].delta_indices)}. Got {action_arr.shape[1]}" + ) + + def get_action( + self, observation: dict[str, Any], options: dict[str, Any] | None = None + ) -> tuple[dict[str, Any], dict[str, Any]]: + if self.strict: + self.check_observation(observation) + action, info = self._get_action(observation, options) + if self.strict: + self.check_action(action) + return action, info + + def reset(self, options: dict[str, Any] | None = None) -> dict[str, Any]: + """Reset the policy to its initial state. + + Args: + options: Dictionary containing the options for the reset + + Returns: + Dictionary containing the info after resetting the policy + """ + return {} diff --git a/vllm_omni/diffusion/models/helios/helios_transformer.py b/vllm_omni/diffusion/models/helios/helios_transformer.py index ed0a50fee38..730c10f1c78 100644 --- a/vllm_omni/diffusion/models/helios/helios_transformer.py +++ b/vllm_omni/diffusion/models/helios/helios_transformer.py @@ -11,6 +11,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from cache_dit import ForwardPattern from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import FP32LayerNorm @@ -25,6 +26,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, SequenceParallelOutput, @@ -602,6 +604,13 @@ class HeliosTransformer3DModel(nn.Module): guidance cross-attention, and chunked video generation support. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "blocks": ForwardPattern.Pattern_2, + }, + has_separate_cfg=True, + ) + _repeated_blocks = ["HeliosTransformerBlock"] _layerwise_offload_blocks_attrs = ["blocks"] packed_modules_mapping = { diff --git a/vllm_omni/diffusion/models/helios/pipeline_helios.py b/vllm_omni/diffusion/models/helios/pipeline_helios.py index c3a47ebaab5..12ecda3683a 100644 --- a/vllm_omni/diffusion/models/helios/pipeline_helios.py +++ b/vllm_omni/diffusion/models/helios/pipeline_helios.py @@ -8,8 +8,8 @@ import logging import math import os -from collections.abc import Iterable -from typing import TYPE_CHECKING, Any +from collections.abc import Iterable, Sequence +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch @@ -26,15 +26,18 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.helios.helios_transformer import HeliosTransformer3DModel from vllm_omni.diffusion.models.helios.scheduling_helios import HeliosScheduler +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.base_config import ( - QuantizationConfig, - ) + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import DiffusionRequestState logger = logging.getLogger(__name__) @@ -151,13 +154,20 @@ def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: return pre_process_func -class HeliosPipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin): +class HeliosPipeline( + nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): """Helios text-to-video / image-to-video / video-to-video pipeline for vllm-omni. Supports T2V, I2V (with image input), and V2V (with video input). Implements chunked video generation with multi-term memory history context. """ + supports_step_execution: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -260,9 +270,650 @@ def num_timesteps(self): def current_timestep(self): return self._current_timestep + def _stage1_sigmas(self, num_steps: int) -> np.ndarray: + # DMD drops the last timestep in set_timesteps(). Only compensate for + # the one-step dummy-run edge case; otherwise preserve the historical + # stage-1 schedule used by request-level forward(). + sigma_count = num_steps + 2 if self.is_distilled and num_steps == 1 else num_steps + 1 + return np.linspace(0.999, 0.0, sigma_count)[:-1] + + def prepare_encode( + self, + state: DiffusionRequestState, + **kwargs: Any, + ) -> DiffusionRequestState: + """Initialize Helios request state for chunk-wise step execution.""" + del kwargs + # Wrap the single request in a DiffusionRequestBatch so the batch + # compatibility properties (`prompts`, etc.) used below are available; + # OmniDiffusionRequest itself only exposes a singular `prompt`. + req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=state.prompt, + sampling_params=state.sampling, + request_id=state.request_id, + ) + ] + ) + extra = getattr(state.sampling, "extra_args", {}) or {} + + history_sizes = sorted(extra.get("history_sizes", [16, 2, 1]), reverse=True) + num_latent_frames_per_chunk = int(extra.get("num_latent_frames_per_chunk", 9)) + keep_first_frame = bool(extra.get("keep_first_frame", True)) + frame_num = int(extra.get("frame_num", 132)) + height = (int(state.sampling.height or extra.get("height", 384)) // 16) * 16 + width = (int(state.sampling.width or extra.get("width", 640)) // 16) * 16 + num_frames = max(int(state.sampling.num_frames or frame_num), 1) + num_steps = int(state.sampling.num_inference_steps or extra.get("num_inference_steps", 50)) + output_type = extra.get("output_type", "np") + + image = extra.get("image") + video = extra.get("video") + if image is not None and video is not None: + raise ValueError("image and video cannot be provided simultaneously") + if len(req.prompts) > 1: + raise ValueError("Helios step execution supports a single prompt, not a batched request.") + + prompt = None + negative_prompt = None + if len(req.prompts) == 1: + prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") + negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") + if prompt is None: + raise ValueError("Prompt is required for Helios generation.") + + guidance_scale = float(extra.get("guidance_scale", 5.0)) + if state.sampling.guidance_scale_provided: + guidance_scale = state.sampling.guidance_scale + self._guidance_scale = guidance_scale + + device = self.device + dtype = self.transformer.dtype + generator = state.sampling.generator + if generator is None and state.sampling.seed is not None: + generator = torch.Generator(device=device).manual_seed(state.sampling.seed) + state.sampling.generator = generator + + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=state.sampling.num_outputs_per_prompt or 1, + max_sequence_length=state.sampling.max_sequence_length or 226, + device=device, + dtype=dtype, + ) + batch_size = prompt_embeds.shape[0] + + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(self.vae.device, self.vae.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + self.vae.device, self.vae.dtype + ) + + add_noise_to_image_latents = bool(extra.get("add_noise_to_image_latents", True)) + image_noise_sigma_min = float(extra.get("image_noise_sigma_min", 0.111)) + image_noise_sigma_max = float(extra.get("image_noise_sigma_max", 0.135)) + add_noise_to_video_latents = bool(extra.get("add_noise_to_video_latents", True)) + video_noise_sigma_min = float(extra.get("video_noise_sigma_min", 0.111)) + video_noise_sigma_max = float(extra.get("video_noise_sigma_max", 0.135)) + + fake_image_latents = None + image_latents = None + if image is not None: + image_latents, fake_image_latents = self.prepare_image_latents( + image, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + ) + + if image_latents is not None and add_noise_to_image_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + fake_image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + fake_image_latents = ( + fake_image_noise_sigma * randn_tensor(fake_image_latents.shape, generator=generator, device=device) + + (1 - fake_image_noise_sigma) * fake_image_latents + ) + + video_latents = None + if video is not None: + image_latents, video_latents = self.prepare_video_latents( + video, + latents_mean=latents_mean, + latents_std=latents_std, + num_latent_frames_per_chunk=num_latent_frames_per_chunk, + dtype=torch.float32, + device=device, + generator=generator, + ) + + if video_latents is not None and add_noise_to_video_latents: + image_noise_sigma = ( + torch.rand(1, device=device, generator=generator) * (image_noise_sigma_max - image_noise_sigma_min) + + image_noise_sigma_min + ) + image_latents = ( + image_noise_sigma * randn_tensor(image_latents.shape, generator=generator, device=device) + + (1 - image_noise_sigma) * image_latents + ) + noisy_latents_chunks = [] + num_latent_chunks = video_latents.shape[2] // num_latent_frames_per_chunk + for i in range(num_latent_chunks): + chunk_start = i * num_latent_frames_per_chunk + chunk_end = chunk_start + num_latent_frames_per_chunk + latent_chunk = video_latents[:, :, chunk_start:chunk_end, :, :] + chunk_frames = latent_chunk.shape[2] + frame_sigmas = ( + torch.rand(chunk_frames, device=device, generator=generator) + * (video_noise_sigma_max - video_noise_sigma_min) + + video_noise_sigma_min + ) + frame_sigmas = frame_sigmas.view(1, 1, chunk_frames, 1, 1) + noisy_chunk = ( + frame_sigmas * randn_tensor(latent_chunk.shape, generator=generator, device=device) + + (1 - frame_sigmas) * latent_chunk + ) + noisy_latents_chunks.append(noisy_chunk) + video_latents = torch.cat(noisy_latents_chunks, dim=2) + + num_channels_latents = self.transformer.config.in_channels + window_num_frames = (num_latent_frames_per_chunk - 1) * self.vae_scale_factor_temporal + 1 + num_latent_chunk = max(1, (num_frames + window_num_frames - 1) // window_num_frames) + num_history_latent_frames = sum(history_sizes) + total_generated_latent_frames = 0 + if not keep_first_frame: + history_sizes[-1] = history_sizes[-1] + 1 + + history_latents = torch.zeros( + batch_size, + num_channels_latents, + num_history_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + device=device, + dtype=torch.float32, + ) + if fake_image_latents is not None: + history_latents = torch.cat([history_latents[:, :, :-1, :, :], fake_image_latents], dim=2) + total_generated_latent_frames += 1 + if video_latents is not None: + history_frames = history_latents.shape[2] + video_frames = video_latents.shape[2] + if video_frames < history_frames: + keep_frames = history_frames - video_frames + history_latents = torch.cat([history_latents[:, :, :keep_frames, :, :], video_latents], dim=2) + else: + history_latents = video_latents + total_generated_latent_frames += video_latents.shape[2] + + if keep_first_frame: + indices = torch.arange(0, sum([1, *history_sizes, num_latent_frames_per_chunk])) + ( + indices_prefix, + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_1x, + indices_hidden_states, + ) = indices.split([1, *history_sizes, num_latent_frames_per_chunk], dim=0) + indices_latents_history_short = torch.cat([indices_prefix, indices_latents_history_1x], dim=0) + else: + indices = torch.arange(0, sum([*history_sizes, num_latent_frames_per_chunk])) + ( + indices_latents_history_long, + indices_latents_history_mid, + indices_latents_history_short, + indices_hidden_states, + ) = indices.split([*history_sizes, num_latent_frames_per_chunk], dim=0) + + state.prompt_embeds = prompt_embeds + state.negative_prompt_embeds = negative_prompt_embeds + state.do_true_cfg = self.do_classifier_free_guidance + state.chunk_index = 0 + state.step_in_chunk = 0 + state.total_chunks = num_latent_chunk + state.extra.update( + { + "attention_kwargs": extra.get("attention_kwargs", {}) or {}, + "batch_size": batch_size, + "dtype": dtype, + "generator": generator, + "guidance_scale": guidance_scale, + "height": height, + "history_latents": history_latents, + "history_sizes": history_sizes, + "history_video": None, + "image_latents": image_latents, + "indices_hidden_states": indices_hidden_states.unsqueeze(0), + "indices_latents_history_long": indices_latents_history_long.unsqueeze(0), + "indices_latents_history_mid": indices_latents_history_mid.unsqueeze(0), + "indices_latents_history_short": indices_latents_history_short.unsqueeze(0), + "is_amplify_first_chunk": bool(extra.get("is_amplify_first_chunk", False)), + "is_enable_stage2": bool(extra.get("is_enable_stage2", False)), + "is_skip_first_chunk": bool(extra.get("is_skip_first_chunk", False)), + "keep_first_frame": keep_first_frame, + "latents_mean": latents_mean, + "latents_std": latents_std, + "num_channels_latents": num_channels_latents, + "num_history_latent_frames": num_history_latent_frames, + "num_latent_frames_per_chunk": num_latent_frames_per_chunk, + "num_steps": num_steps, + "output_type": output_type, + "pyramid_num_inference_steps_list": extra.get("pyramid_num_inference_steps_list", [10, 10, 10]), + "pyramid_num_stages": int(extra.get("pyramid_num_stages", 3)), + "total_generated_latent_frames": total_generated_latent_frames, + "use_cfg_zero_star": bool(extra.get("use_cfg_zero_star", False)), + "use_zero_init": bool(extra.get("use_zero_init", True)), + "vae_dtype": self.vae.dtype, + "width": width, + "window_num_frames": window_num_frames, + "zero_steps": int(extra.get("zero_steps", 1)), + } + ) + self._prepare_next_chunk(state) + return state + + def _prepare_next_chunk(self, state: DiffusionRequestState) -> None: + extra = state.extra + k = state.chunk_index + is_first_chunk = k == 0 + is_second_chunk = k == 1 + history_sizes = extra["history_sizes"] + history_latents = extra["history_latents"] + keep_first_frame = extra["keep_first_frame"] + image_latents = extra["image_latents"] + num_history_latent_frames = extra["num_history_latent_frames"] + + if keep_first_frame: + latents_history_long, latents_history_mid, latents_history_1x = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + if image_latents is None and is_first_chunk: + latents_prefix = torch.zeros( + ( + extra["batch_size"], + extra["num_channels_latents"], + 1, + latents_history_1x.shape[-2], + latents_history_1x.shape[-1], + ), + device=latents_history_1x.device, + dtype=latents_history_1x.dtype, + ) + else: + latents_prefix = image_latents + latents_history_short = torch.cat([latents_prefix, latents_history_1x], dim=2) + else: + latents_history_long, latents_history_mid, latents_history_short = history_latents[ + :, :, -num_history_latent_frames: + ].split(history_sizes, dim=2) + + state.latents = self.prepare_latents( + extra["batch_size"], + extra["num_channels_latents"], + extra["height"], + extra["width"], + extra["window_num_frames"], + dtype=torch.float32, + device=self.device, + generator=extra["generator"], + latents=None, + ) + extra["stage1_start_latents"] = state.latents + extra.update( + { + "is_first_chunk": is_first_chunk, + "is_second_chunk": is_second_chunk, + "latents_history_long": latents_history_long, + "latents_history_mid": latents_history_mid, + "latents_history_short": latents_history_short, + } + ) + + if not extra["is_enable_stage2"]: + patch_size = self.transformer.config.patch_size + image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + sigmas = self._stage1_sigmas(extra["num_steps"]) + mu = calculate_shift(image_seq_len) + self.scheduler.set_timesteps(extra["num_steps"], device=self.device, sigmas=sigmas, mu=mu) + state.timesteps = self.scheduler.timesteps + state.chunk_num_steps = len(state.timesteps) + else: + self._prepare_stage2_chunk(state) + + state.step_in_chunk = 0 + state.step_index = 0 + self._num_timesteps = state.chunk_num_steps + + def _prepare_stage2_chunk(self, state: DiffusionRequestState) -> None: + extra = state.extra + batch_size, num_channel, num_frames_lat, height, width = state.latents.shape + latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames_lat, num_channel, height, width + ) + for _ in range(extra["pyramid_num_stages"] - 1): + height //= 2 + width //= 2 + latents_flat = F.interpolate(latents_flat, size=(height, width), mode="bilinear") * 2 + state.latents = latents_flat.reshape(batch_size, num_frames_lat, num_channel, height, width).permute( + 0, 2, 1, 3, 4 + ) + extra["stage2_height"] = height + extra["stage2_width"] = width + extra["stage2_start_point_list"] = [state.latents] if self.is_distilled else None + extra["stage_index"] = 0 + extra["stage_step_index"] = 0 + amplify_first_chunk = extra["is_amplify_first_chunk"] and extra["is_first_chunk"] + state.chunk_num_steps = sum( + self._stage2_effective_num_steps(int(num_steps), amplify_first_chunk) + for num_steps in extra["pyramid_num_inference_steps_list"] + ) + self._set_stage2_timesteps(state) + + def _stage2_effective_num_steps(self, num_steps: int, is_amplify_first_chunk: bool) -> int: + if self.is_distilled and is_amplify_first_chunk: + return num_steps * 2 + return num_steps + + def _set_stage2_timesteps(self, state: DiffusionRequestState) -> None: + extra = state.extra + patch_size = self.transformer.config.patch_size + image_seq_len = (state.latents.shape[-1] * state.latents.shape[-2] * state.latents.shape[-3]) // ( + patch_size[0] * patch_size[1] * patch_size[2] + ) + mu = calculate_shift(image_seq_len) + stage_index = extra["stage_index"] + self.scheduler.set_timesteps( + extra["pyramid_num_inference_steps_list"][stage_index], + stage_index=stage_index, + device=self.device, + mu=mu, + is_amplify_first_chunk=extra["is_amplify_first_chunk"] and extra["is_first_chunk"], + ) + state.timesteps = self.scheduler.timesteps + state.step_index = extra["stage_step_index"] + + def denoise_step( + self, + input_batch: InputBatch, + states: Sequence[DiffusionRequestState], + **kwargs: Any, + ) -> torch.Tensor | None: + del kwargs + if len(states) != 1: + raise ValueError("Helios step execution supports a single request, not a batched request.") + state = states[0] + if state.extra["is_enable_stage2"]: + return self._denoise_stage2_step(state) + return self._denoise_stage1_step(state, input_batch.latents, input_batch.timesteps) + + def _denoise_stage1_step( + self, + state: DiffusionRequestState, + latents: torch.Tensor, + timesteps: torch.Tensor, + ) -> torch.Tensor: + extra = state.extra + batch_size = latents.shape[0] + t = timesteps + self._current_timestep = t + timestep = t.expand(batch_size) + transformer_kwargs = { + "hidden_states": latents.to(extra["dtype"]), + "timestep": timestep, + "indices_hidden_states": extra["indices_hidden_states"], + "indices_latents_history_short": extra["indices_latents_history_short"], + "indices_latents_history_mid": extra["indices_latents_history_mid"], + "indices_latents_history_long": extra["indices_latents_history_long"], + "latents_history_short": extra["latents_history_short"].to(extra["dtype"]), + "latents_history_mid": extra["latents_history_mid"].to(extra["dtype"]), + "latents_history_long": extra["latents_history_long"].to(extra["dtype"]), + "attention_kwargs": extra["attention_kwargs"], + "return_dict": False, + } + if extra["use_cfg_zero_star"] and state.do_true_cfg: + noise_pred = self.transformer( + encoder_hidden_states=state.prompt_embeds, + **transformer_kwargs, + )[0] + noise_uncond = self.transformer( + encoder_hidden_states=state.negative_prompt_embeds, + **transformer_kwargs, + )[0] + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))).to(noise_pred.dtype) + if (state.step_in_chunk <= extra["zero_steps"]) and extra["use_zero_init"]: + return noise_pred * 0.0 + return noise_uncond * alpha_cfg + extra["guidance_scale"] * (noise_pred - noise_uncond * alpha_cfg) + + positive_kwargs = { + "encoder_hidden_states": state.prompt_embeds, + **transformer_kwargs, + } + negative_kwargs = ( + { + "encoder_hidden_states": state.negative_prompt_embeds, + **transformer_kwargs, + } + if state.do_true_cfg + else None + ) + return self.predict_noise_maybe_with_cfg( + do_true_cfg=state.do_true_cfg, + true_cfg_scale=extra["guidance_scale"], + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + def _denoise_stage2_step(self, state: DiffusionRequestState) -> torch.Tensor: + extra = state.extra + latents = state.latents + assert latents is not None + batch_size = latents.shape[0] + t = state.current_timestep + assert t is not None + self._current_timestep = t + timestep = t.expand(batch_size).to(torch.int64) + transformer_kwargs = { + "hidden_states": latents.to(extra["dtype"]), + "timestep": timestep, + "indices_hidden_states": extra["indices_hidden_states"], + "indices_latents_history_short": extra["indices_latents_history_short"], + "indices_latents_history_mid": extra["indices_latents_history_mid"], + "indices_latents_history_long": extra["indices_latents_history_long"], + "latents_history_short": extra["latents_history_short"].to(extra["dtype"]), + "latents_history_mid": extra["latents_history_mid"].to(extra["dtype"]), + "latents_history_long": extra["latents_history_long"].to(extra["dtype"]), + "attention_kwargs": extra["attention_kwargs"], + "return_dict": False, + } + noise_pred = self.transformer(encoder_hidden_states=state.prompt_embeds, **transformer_kwargs)[0] + if state.do_true_cfg: + noise_uncond = self.transformer(encoder_hidden_states=state.negative_prompt_embeds, **transformer_kwargs)[0] + if extra["use_cfg_zero_star"]: + positive_flat = noise_pred.view(batch_size, -1) + negative_flat = noise_uncond.view(batch_size, -1) + alpha_cfg = optimized_scale(positive_flat, negative_flat) + alpha_cfg = alpha_cfg.view(batch_size, *([1] * (len(noise_pred.shape) - 1))).to(noise_pred.dtype) + if ( + extra["stage_index"] == 0 + and extra["stage_step_index"] <= extra["zero_steps"] + and extra["use_zero_init"] + ): + noise_pred = noise_pred * 0.0 + else: + noise_pred = noise_uncond * alpha_cfg + extra["guidance_scale"] * ( + noise_pred - noise_uncond * alpha_cfg + ) + else: + noise_pred = noise_uncond + extra["guidance_scale"] * (noise_pred - noise_uncond) + return noise_pred + + def step_scheduler( + self, + state: DiffusionRequestState, + noise_pred: torch.Tensor, + **kwargs: Any, + ) -> None: + del kwargs + if state.extra["is_enable_stage2"]: + self._step_scheduler_stage2(state, noise_pred) + else: + t = state.current_timestep + assert t is not None + if self.is_distilled: + state.latents = self.scheduler.step( + noise_pred, + t, + state.latents, + return_dict=False, + cur_sampling_step=state.step_in_chunk, + dmd_noisy_tensor=state.extra["stage1_start_latents"], + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=state.timesteps, + )[0] + else: + state.latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, state.latents, state.do_true_cfg) + state.step_in_chunk += 1 + state.step_index = state.step_in_chunk + + def _step_scheduler_stage2(self, state: DiffusionRequestState, noise_pred: torch.Tensor) -> None: + extra = state.extra + t = state.current_timestep + assert t is not None and state.latents is not None + start_points = extra["stage2_start_point_list"] + state.latents = self.scheduler.step( + noise_pred, + t, + state.latents, + return_dict=False, + cur_sampling_step=extra["stage_step_index"], + dmd_noisy_tensor=start_points[extra["stage_index"]] if start_points is not None else None, + dmd_sigmas=self.scheduler.sigmas, + dmd_timesteps=self.scheduler.timesteps, + all_timesteps=state.timesteps, + )[0] + extra["stage_step_index"] += 1 + state.step_in_chunk += 1 + state.step_index = extra["stage_step_index"] + + if extra["stage_step_index"] < len(state.timesteps): + return + if extra["stage_index"] >= extra["pyramid_num_stages"] - 1: + return + + extra["stage_index"] += 1 + extra["stage_step_index"] = 0 + extra["stage2_height"] *= 2 + extra["stage2_width"] *= 2 + batch_size, num_channel, num_frames_cur = state.latents.shape[:3] + latents_flat = state.latents.permute(0, 2, 1, 3, 4).reshape( + batch_size * num_frames_cur, + num_channel, + extra["stage2_height"] // 2, + extra["stage2_width"] // 2, + ) + latents_flat = F.interpolate(latents_flat, size=(extra["stage2_height"], extra["stage2_width"]), mode="nearest") + state.latents = latents_flat.reshape( + batch_size, + num_frames_cur, + num_channel, + extra["stage2_height"], + extra["stage2_width"], + ).permute(0, 2, 1, 3, 4) + + ori_sigma = 1 - self.scheduler.ori_start_sigmas[extra["stage_index"]] + gamma = self.scheduler.config.gamma + alpha = 1 / (math.sqrt(1 + (1 / gamma)) * (1 - ori_sigma) + ori_sigma) + beta = alpha * (1 - ori_sigma) / math.sqrt(gamma) + patch_size = self.transformer.config.patch_size + noise = self.sample_block_noise( + batch_size, + num_channel, + state.latents.shape[2], + extra["stage2_height"], + extra["stage2_width"], + patch_size, + generator=extra["generator"], + ).to(device=self.device, dtype=extra["dtype"]) + state.latents = alpha * state.latents + beta * noise + if self.is_distilled and extra["stage2_start_point_list"] is not None: + extra["stage2_start_point_list"].append(state.latents) + self._set_stage2_timesteps(state) + + def post_decode( + self, + state: DiffusionRequestState, + **kwargs: Any, + ) -> DiffusionOutput: + del kwargs + extra = state.extra + is_first_chunk = extra["is_first_chunk"] + is_second_chunk = extra["is_second_chunk"] + if extra["keep_first_frame"] and ( + (is_first_chunk and extra["image_latents"] is None) or (extra["is_skip_first_chunk"] and is_second_chunk) + ): + extra["image_latents"] = state.latents[:, :, 0:1, :, :] + + extra["total_generated_latent_frames"] += state.latents.shape[2] + extra["history_latents"] = torch.cat([extra["history_latents"], state.latents], dim=2) + real_history_latents = extra["history_latents"][:, :, -extra["total_generated_latent_frames"] :] + current_latents = ( + real_history_latents[:, :, -extra["num_latent_frames_per_chunk"] :].to(extra["vae_dtype"]) + / extra["latents_std"] + + extra["latents_mean"] + ) + current_video = self.vae.decode(current_latents, return_dict=False)[0] + if extra["history_video"] is None: + extra["history_video"] = current_video + else: + extra["history_video"] = torch.cat([extra["history_video"], current_video], dim=2) + + output = current_latents if extra["output_type"] == "latent" else current_video + completed_chunk_index = state.chunk_index + state.chunk_index += 1 + finished = state.request_denoise_completed + if not finished: + self._prepare_next_chunk(state) + else: + self._current_timestep = None + if current_omni_platform.is_available(): + current_omni_platform.empty_cache() + + return DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else {}, + chunk_index=completed_chunk_index, + total_chunks=state.total_chunks, + finished=finished, + ) + def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | None = None, negative_prompt: str | None = None, height: int = 384, diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py deleted file mode 100644 index abc50eb8a22..00000000000 --- a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py +++ /dev/null @@ -1,1059 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import inspect -import json -import math -import os -from collections.abc import Callable, Iterable -from typing import Any - -import numpy as np -import torch -from diffusers.image_processor import VaeImageProcessor -from diffusers.models import AutoencoderKL -from diffusers.schedulers import FlowMatchEulerDiscreteScheduler -from diffusers.utils import deprecate, logging -from diffusers.utils.torch_utils import randn_tensor -from torch import nn -from transformers import ( - CLIPTextModelWithProjection, - CLIPTokenizer, - LlamaForCausalLM, - PreTrainedTokenizerFast, - T5EncoderModel, - T5Tokenizer, -) -from vllm.model_executor.models.utils import AutoWeightsLoader - -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig -from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size -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.hidream_image import HiDreamImageTransformer2DModel -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs -from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific - -logger = logging.get_logger(__name__) - - -def get_hidream_image_post_process_func( - od_config: OmniDiffusionConfig, -): - model_name = od_config.model - if os.path.exists(model_name): - model_path = model_name - else: - model_path = download_weights_from_hf_specific(model_name, None, ["*"]) - vae_config_path = os.path.join(model_path, "vae/config.json") - if not os.path.exists(vae_config_path): - raise FileNotFoundError( - f"VAE config not found at {vae_config_path}. " - "Please ensure the model path contains a valid VAE configuration." - ) - with open(vae_config_path) as f: - vae_config = json.load(f) - vae_scale_factor = 2 ** (len(vae_config["block_out_channels"]) - 1) if "block_out_channels" in vae_config else 8 - - image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) - - def post_process_func( - images: torch.Tensor, - ): - return image_processor.postprocess(images) - - return post_process_func - - -# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift -def calculate_shift( - image_seq_len, - base_seq_len: int = 256, - max_seq_len: int = 4096, - base_shift: float = 0.5, - max_shift: float = 1.15, -): - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - b = base_shift - m * base_seq_len - mu = image_seq_len * m + b - return mu - - -# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps -def retrieve_timesteps( - scheduler, - num_inference_steps: int | None = None, - device: str | torch.device | None = None, - timesteps: list[int] | None = None, - sigmas: list[float] | None = None, - **kwargs, -): - r""" - Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles - custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. - - Args: - scheduler (`SchedulerMixin`): - The scheduler to get timesteps from. - num_inference_steps (`int`): - The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` - must be `None`. - device (`str` or `torch.device`, *optional*): - The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. - timesteps (`list[int]`, *optional*): - Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, - `num_inference_steps` and `sigmas` must be `None`. - sigmas (`list[float]`, *optional*): - Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, - `num_inference_steps` and `timesteps` must be `None`. - - Returns: - timesteps (`torch.Tensor`): The timestep schedule from the scheduler. - num_inference_steps (`int`): The number of inference steps. - """ - if timesteps is not None and sigmas is not None: - raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") - if timesteps is not None: - accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accepts_timesteps: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" timestep schedules. Please check whether you are using the correct scheduler." - ) - scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - elif sigmas is not None: - accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) - if not accept_sigmas: - raise ValueError( - f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" - f" sigmas schedules. Please check whether you are using the correct scheduler." - ) - scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) - timesteps = scheduler.timesteps - num_inference_steps = len(timesteps) - else: - scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) - timesteps = scheduler.timesteps - return timesteps, num_inference_steps - - -class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, ProgressBarMixin): - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - dtype = getattr(od_config, "dtype", torch.bfloat16) - llama_path = self.od_config.extras["auxiliary_text_encoder"] - if llama_path is None: - logger.warning( - f"auxiliary_text_encoder is not provided. " - f"Attempting to load default LLaMA model: {self.od_config.extras['default_llama_model_id']}" - ) - llama_path = self.od_config.extras["default_llama_model_id"] - logger.info(f"Loading LLAMA model from: {llama_path}") - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - self.device = get_local_device() - model = od_config.model - # Check if model is a local path - local_files_only = os.path.exists(model) - - # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder - # race; prefetch the in-repo component set before any from_pretrained - # (``text_encoder_4`` lives in a separate Llama repo and is unaffected). - hidream_subfolders = [ - "scheduler", - "vae", - "text_encoder", - "tokenizer", - "text_encoder_2", - "tokenizer_2", - "text_encoder_3", - "tokenizer_3", - ] - prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only) - - self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - model, subfolder="scheduler", local_files_only=local_files_only - ) - self.vae = from_pretrained_with_prefetch( - AutoencoderKL.from_pretrained, - model, - subfolder="vae", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ).to(self.device) - self.text_encoder = from_pretrained_with_prefetch( - CLIPTextModelWithProjection.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ) - self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - self.text_encoder_2 = from_pretrained_with_prefetch( - CLIPTextModelWithProjection.from_pretrained, - model, - subfolder="text_encoder_2", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ) - self.tokenizer_2 = CLIPTokenizer.from_pretrained( - model, subfolder="tokenizer_2", local_files_only=local_files_only - ) - self.text_encoder_3 = from_pretrained_with_prefetch( - T5EncoderModel.from_pretrained, - model, - subfolder="text_encoder_3", - prefetch_list=hidream_subfolders, - local_files_only=local_files_only, - ) - self.tokenizer_3 = T5Tokenizer.from_pretrained( - model, subfolder="tokenizer_3", local_files_only=local_files_only - ) - self.text_encoder_4 = LlamaForCausalLM.from_pretrained(llama_path, output_hidden_states=True, dtype=dtype).to( - self.device - ) - self.tokenizer_4 = PreTrainedTokenizerFast.from_pretrained(llama_path, use_fast=False) - transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, HiDreamImageTransformer2DModel) - self.transformer = HiDreamImageTransformer2DModel( - od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs - ) - - self.stage = None - - self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 - # HiDreamImage latents are turned into 2x2 patches and packed. - # This means the latent width and height has to be divisible - # by the patch size. So the vae scale factor is multiplied by the patch size to account for this - self.default_sample_size = 128 - if getattr(self, "tokenizer_4", None) is not None: - self.tokenizer_4.pad_token = self.tokenizer_4.eos_token - - self.setup_diffusion_pipeline_profiler( - enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler - ) - - def check_inputs( - self, - prompt, - prompt_2, - prompt_3, - prompt_4, - negative_prompt=None, - negative_prompt_2=None, - negative_prompt_3=None, - negative_prompt_4=None, - prompt_embeds_t5=None, - prompt_embeds_llama3=None, - negative_prompt_embeds_t5=None, - negative_prompt_embeds_llama3=None, - pooled_prompt_embeds=None, - negative_pooled_prompt_embeds=None, - callback_on_step_end_tensor_inputs=None, - ): - if prompt is not None and pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `pooled_prompt_embeds`: {pooled_prompt_embeds}. " - "Please make sure to only forward one of the two." - ) - elif prompt_2 is not None and pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt_2`: {prompt_2} and `pooled_prompt_embeds`: {pooled_prompt_embeds}." - "Please make sure to only forward one of the two." - ) - elif prompt_3 is not None and prompt_embeds_t5 is not None: - raise ValueError( - f"Cannot forward both `prompt_3`: {prompt_3} and `prompt_embeds_t5`: {prompt_embeds_t5}." - "Please make sure to only forward one of the two." - ) - elif prompt_4 is not None and prompt_embeds_llama3 is not None: - raise ValueError( - f"Cannot forward both `prompt_4`: {prompt_4} and `prompt_embeds_llama3`: {prompt_embeds_llama3}." - "Please make sure to only forward one of the two." - ) - elif prompt is None and pooled_prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `pooled_prompt_embeds`." - "Cannot leave both `prompt` and `pooled_prompt_embeds` undefined." - ) - elif prompt is None and prompt_embeds_t5 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_t5`." - "Cannot leave both `prompt` and `prompt_embeds_t5` undefined." - ) - elif prompt is None and prompt_embeds_llama3 is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds_llama3`." - "Cannot leave both `prompt` and `prompt_embeds_llama3` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - elif prompt_2 is not None and (not isinstance(prompt_2, str) and not isinstance(prompt_2, list)): - raise ValueError(f"`prompt_2` has to be of type `str` or `list` but is {type(prompt_2)}") - elif prompt_3 is not None and (not isinstance(prompt_3, str) and not isinstance(prompt_3, list)): - raise ValueError(f"`prompt_3` has to be of type `str` or `list` but is {type(prompt_3)}") - elif prompt_4 is not None and (not isinstance(prompt_4, str) and not isinstance(prompt_4, list)): - raise ValueError(f"`prompt_4` has to be of type `str` or `list` but is {type(prompt_4)}") - - if negative_prompt is not None and negative_pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_pooled_prompt_embeds`:" - f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_2 is not None and negative_pooled_prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_2`: {negative_prompt_2} and `negative_pooled_prompt_embeds`:" - f" {negative_pooled_prompt_embeds}. Please make sure to only forward one of the two." - ) - elif negative_prompt_3 is not None and negative_prompt_embeds_t5 is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_3`: {negative_prompt_3} and `negative_prompt_embeds_t5`:" - f" {negative_prompt_embeds_t5}. Please make sure to only forward one of the two." - ) - elif negative_prompt_4 is not None and negative_prompt_embeds_llama3 is not None: - raise ValueError( - f"Cannot forward both `negative_prompt_4`: {negative_prompt_4} and `negative_prompt_embeds_llama3`:" - f" {negative_prompt_embeds_llama3}. Please make sure to only forward one of the two." - ) - - if pooled_prompt_embeds is not None and negative_pooled_prompt_embeds is not None: - if pooled_prompt_embeds.shape != negative_pooled_prompt_embeds.shape: - raise ValueError( - "`pooled_prompt_embeds` and `negative_pooled_prompt_embeds`" - " must have the same shape when passed directly, but" - f" got: `pooled_prompt_embeds` {pooled_prompt_embeds.shape} != `negative_pooled_prompt_embeds`" - f" {negative_pooled_prompt_embeds.shape}." - ) - if prompt_embeds_t5 is not None and negative_prompt_embeds_t5 is not None: - if prompt_embeds_t5.shape != negative_prompt_embeds_t5.shape: - raise ValueError( - "`prompt_embeds_t5` and `negative_prompt_embeds_t5`" - " must have the same shape when passed directly, but" - f" got: `prompt_embeds_t5` {prompt_embeds_t5.shape} != `negative_prompt_embeds_t5`" - f" {negative_prompt_embeds_t5.shape}." - ) - if prompt_embeds_llama3 is not None and negative_prompt_embeds_llama3 is not None: - if prompt_embeds_llama3.shape != negative_prompt_embeds_llama3.shape: - raise ValueError( - "`prompt_embeds_llama3` and `negative_prompt_embeds_llama3`" - " must have the same shape when passed directly, but" - f" got: `prompt_embeds_llama3` {prompt_embeds_llama3.shape} != `negative_prompt_embeds_llama3`" - f" {negative_prompt_embeds_llama3.shape}." - ) - - def _get_t5_prompt_embeds( - self, - prompt: str | list[str] = None, - max_sequence_length: int = 128, - dtype: torch.dtype | None = None, - ): - dtype = dtype or self.text_encoder_3.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_3( - prompt, - padding="max_length", - max_length=min(max_sequence_length, self.tokenizer_3.model_max_length), - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - untruncated_ids = self.tokenizer_3(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer_3.batch_decode( - untruncated_ids[:, min(max_sequence_length, self.tokenizer_3.model_max_length) - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because `max_sequence_length` is set to " - f" {min(max_sequence_length, self.tokenizer_3.model_max_length)} tokens: {removed_text}" - ) - - prompt_embeds = self.text_encoder_3( - text_input_ids.to(self.device), attention_mask=attention_mask.to(self.device) - )[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) - return prompt_embeds - - def _get_clip_prompt_embeds( - self, - tokenizer, - text_encoder, - prompt: str | list[str], - max_sequence_length: int = 128, - dtype: torch.dtype | None = None, - ): - dtype = dtype or text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - model_max_length = self.tokenizer.model_max_length - text_inputs = tokenizer( - prompt, - padding="max_length", - max_length=min(max_sequence_length, model_max_length), - truncation=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - untruncated_ids = tokenizer(prompt, padding="longest", return_tensors="pt").input_ids - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = tokenizer.batch_decode(untruncated_ids[:, model_max_length - 1 : -1]) - logger.warning( - "The following part of your input was truncated because CLIP can only handle sequences up to" - f" {model_max_length} tokens: {removed_text}" - ) - prompt_embeds = text_encoder(text_input_ids.to(self.device), output_hidden_states=True) - - # Use pooled output of CLIPTextModel - prompt_embeds = prompt_embeds[0] - prompt_embeds = prompt_embeds.to(dtype=dtype, device=self.device) - return prompt_embeds - - def _get_llama3_prompt_embeds( - self, - prompt: str | list[str] = None, - max_sequence_length: int = 128, - dtype: torch.dtype | None = None, - ): - dtype = dtype or self.text_encoder_4.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - - text_inputs = self.tokenizer_4( - prompt, - padding="max_length", - max_length=min(max_sequence_length, self.tokenizer_4.model_max_length), - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - attention_mask = text_inputs.attention_mask - untruncated_ids = self.tokenizer_4(prompt, padding="longest", return_tensors="pt").input_ids - - if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids): - removed_text = self.tokenizer_4.batch_decode( - untruncated_ids[:, min(max_sequence_length, self.tokenizer_4.model_max_length) - 1 : -1] - ) - logger.warning( - "The following part of your input was truncated because `max_sequence_length` is set to " - f" {min(max_sequence_length, self.tokenizer_4.model_max_length)} tokens: {removed_text}" - ) - - outputs = self.text_encoder_4( - text_input_ids.to(self.device), - attention_mask=attention_mask.to(self.device), - output_hidden_states=True, - ) - - prompt_embeds = outputs.hidden_states[1:] - prompt_embeds = torch.stack(prompt_embeds, dim=0) - return prompt_embeds - - def encode_prompt( - self, - prompt: str | list[str] | None = None, - prompt_2: str | list[str] | None = None, - prompt_3: str | list[str] | None = None, - prompt_4: str | list[str] | None = None, - dtype: torch.dtype | None = None, - num_images_per_prompt: int = 1, - do_classifier_free_guidance: bool = True, - negative_prompt: str | list[str] | None = None, - negative_prompt_2: str | list[str] | None = None, - negative_prompt_3: str | list[str] | None = None, - negative_prompt_4: str | list[str] | None = None, - prompt_embeds_t5: list[torch.FloatTensor] | None = None, - prompt_embeds_llama3: list[torch.FloatTensor] | None = None, - negative_prompt_embeds_t5: list[torch.FloatTensor] | None = None, - negative_prompt_embeds_llama3: list[torch.FloatTensor] | None = None, - pooled_prompt_embeds: torch.FloatTensor | None = None, - negative_pooled_prompt_embeds: torch.FloatTensor | None = None, - max_sequence_length: int = 128, - lora_scale: float | None = None, - ): - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = pooled_prompt_embeds.shape[0] - - if pooled_prompt_embeds is None: - pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( - self.tokenizer, self.text_encoder, prompt, max_sequence_length, dtype - ) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt = [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - - if len(negative_prompt) > 1 and len(negative_prompt) != batch_size: - raise ValueError(f"negative_prompt must be of length 1 or {batch_size}") - - negative_pooled_prompt_embeds_1 = self._get_clip_prompt_embeds( - self.tokenizer, self.text_encoder, negative_prompt, max_sequence_length, dtype - ) - - if negative_pooled_prompt_embeds_1.shape[0] == 1 and batch_size > 1: - negative_pooled_prompt_embeds_1 = negative_pooled_prompt_embeds_1.repeat(batch_size, 1) - - if pooled_prompt_embeds is None: - prompt_2 = prompt_2 or prompt - prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 - - if len(prompt_2) > 1 and len(prompt_2) != batch_size: - raise ValueError(f"prompt_2 must be of length 1 or {batch_size}") - - pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( - self.tokenizer_2, self.text_encoder_2, prompt_2, max_sequence_length, dtype - ) - - if pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: - pooled_prompt_embeds_2 = pooled_prompt_embeds_2.repeat(batch_size, 1) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_prompt_2 = negative_prompt_2 or negative_prompt - negative_prompt_2 = [negative_prompt_2] if isinstance(negative_prompt_2, str) else negative_prompt_2 - - if len(negative_prompt_2) > 1 and len(negative_prompt_2) != batch_size: - raise ValueError(f"negative_prompt_2 must be of length 1 or {batch_size}") - - negative_pooled_prompt_embeds_2 = self._get_clip_prompt_embeds( - self.tokenizer_2, self.text_encoder_2, negative_prompt_2, max_sequence_length, dtype - ) - - if negative_pooled_prompt_embeds_2.shape[0] == 1 and batch_size > 1: - negative_pooled_prompt_embeds_2 = negative_pooled_prompt_embeds_2.repeat(batch_size, 1) - - if pooled_prompt_embeds is None: - pooled_prompt_embeds = torch.cat([pooled_prompt_embeds_1, pooled_prompt_embeds_2], dim=-1) - - if do_classifier_free_guidance and negative_pooled_prompt_embeds is None: - negative_pooled_prompt_embeds = torch.cat( - [negative_pooled_prompt_embeds_1, negative_pooled_prompt_embeds_2], dim=-1 - ) - - if prompt_embeds_t5 is None: - prompt_3 = prompt_3 or prompt - prompt_3 = [prompt_3] if isinstance(prompt_3, str) else prompt_3 - - if len(prompt_3) > 1 and len(prompt_3) != batch_size: - raise ValueError(f"prompt_3 must be of length 1 or {batch_size}") - - prompt_embeds_t5 = self._get_t5_prompt_embeds(prompt_3, max_sequence_length, dtype) - - if prompt_embeds_t5.shape[0] == 1 and batch_size > 1: - prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) - - if do_classifier_free_guidance and negative_prompt_embeds_t5 is None: - negative_prompt_3 = negative_prompt_3 or negative_prompt - negative_prompt_3 = [negative_prompt_3] if isinstance(negative_prompt_3, str) else negative_prompt_3 - - if len(negative_prompt_3) > 1 and len(negative_prompt_3) != batch_size: - raise ValueError(f"negative_prompt_3 must be of length 1 or {batch_size}") - - negative_prompt_embeds_t5 = self._get_t5_prompt_embeds(negative_prompt_3, max_sequence_length, dtype) - - if negative_prompt_embeds_t5.shape[0] == 1 and batch_size > 1: - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) - - if prompt_embeds_llama3 is None: - prompt_4 = prompt_4 or prompt - prompt_4 = [prompt_4] if isinstance(prompt_4, str) else prompt_4 - - if len(prompt_4) > 1 and len(prompt_4) != batch_size: - raise ValueError(f"prompt_4 must be of length 1 or {batch_size}") - - prompt_embeds_llama3 = self._get_llama3_prompt_embeds(prompt_4, max_sequence_length, dtype) - - if prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - - if do_classifier_free_guidance and negative_prompt_embeds_llama3 is None: - negative_prompt_4 = negative_prompt_4 or negative_prompt - negative_prompt_4 = [negative_prompt_4] if isinstance(negative_prompt_4, str) else negative_prompt_4 - - if len(negative_prompt_4) > 1 and len(negative_prompt_4) != batch_size: - raise ValueError(f"negative_prompt_4 must be of length 1 or {batch_size}") - - negative_prompt_embeds_llama3 = self._get_llama3_prompt_embeds( - negative_prompt_4, max_sequence_length, dtype - ) - - if negative_prompt_embeds_llama3.shape[0] == 1 and batch_size > 1: - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - - # duplicate pooled_prompt_embeds for each generation per prompt - pooled_prompt_embeds = pooled_prompt_embeds.repeat(1, num_images_per_prompt) - pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - # duplicate t5_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len, _ = prompt_embeds_t5.shape - if bs_embed == 1 and batch_size > 1: - prompt_embeds_t5 = prompt_embeds_t5.repeat(batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate prompt_embeds_t5 of batch size {bs_embed}") - prompt_embeds_t5 = prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) - prompt_embeds_t5 = prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) - - # duplicate llama3_prompt_embeds for batch_size and num_images_per_prompt - _, bs_embed, seq_len, dim = prompt_embeds_llama3.shape - if bs_embed == 1 and batch_size > 1: - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate prompt_embeds_llama3 of batch size {bs_embed}") - prompt_embeds_llama3 = prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) - prompt_embeds_llama3 = prompt_embeds_llama3.view(-1, batch_size * num_images_per_prompt, seq_len, dim) - - if do_classifier_free_guidance: - # duplicate negative_pooled_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len = negative_pooled_prompt_embeds.shape - if bs_embed == 1 and batch_size > 1: - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(batch_size, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_pooled_prompt_embeds of batch size {bs_embed}") - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.repeat(1, num_images_per_prompt) - negative_pooled_prompt_embeds = negative_pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) - - # duplicate negative_t5_prompt_embeds for batch_size and num_images_per_prompt - bs_embed, seq_len, _ = negative_prompt_embeds_t5.shape - if bs_embed == 1 and batch_size > 1: - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_prompt_embeds_t5 of batch size {bs_embed}") - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.repeat(1, num_images_per_prompt, 1) - negative_prompt_embeds_t5 = negative_prompt_embeds_t5.view(batch_size * num_images_per_prompt, seq_len, -1) - - # duplicate negative_prompt_embeds_llama3 for batch_size and num_images_per_prompt - _, bs_embed, seq_len, dim = negative_prompt_embeds_llama3.shape - if bs_embed == 1 and batch_size > 1: - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, batch_size, 1, 1) - elif bs_embed > 1 and bs_embed != batch_size: - raise ValueError(f"cannot duplicate negative_prompt_embeds_llama3 of batch size {bs_embed}") - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.repeat(1, 1, num_images_per_prompt, 1) - negative_prompt_embeds_llama3 = negative_prompt_embeds_llama3.view( - -1, batch_size * num_images_per_prompt, seq_len, dim - ) - - return ( - prompt_embeds_t5, - negative_prompt_embeds_t5, - prompt_embeds_llama3, - negative_prompt_embeds_llama3, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - def enable_vae_slicing(self): - r""" - Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to - compute decoding in several steps. This is useful to save some memory and allow larger batch sizes. - """ - depr_message = ( - f"Calling `enable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" - " method will be removed in a future version. Please use `pipe.vae.enable_slicing()`." - ) - deprecate( - "enable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.enable_slicing() - - def disable_vae_slicing(self): - r""" - Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = ( - f"Calling `disable_vae_slicing()` on a `{self.__class__.__name__}` is deprecated and this" - " method will be removed in a future version. Please use `pipe.vae.disable_slicing()`." - ) - deprecate( - "disable_vae_slicing", - "0.40.0", - depr_message, - ) - self.vae.disable_slicing() - - def enable_vae_tiling(self): - r""" - Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to - compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow - processing larger images. - """ - depr_message = ( - f"Calling `enable_vae_tiling()` on a `{self.__class__.__name__}`" - " is deprecated and this method will be removed in a future version. Please use `pipe.vae.enable_tiling()`." - ) - deprecate( - "enable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.enable_tiling() - - def disable_vae_tiling(self): - r""" - Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to - computing decoding in one step. - """ - depr_message = ( - f"Calling `disable_vae_tiling()` on a `{self.__class__.__name__}` is" - " deprecated and this method will be removed in a future version. Please use `pipe.vae.disable_tiling()`." - ) - deprecate( - "disable_vae_tiling", - "0.40.0", - depr_message, - ) - self.vae.disable_tiling() - - def prepare_latents( - self, - batch_size, - num_channels_latents, - height, - width, - dtype, - generator, - latents=None, - ): - # VAE applies 8x compression on images but we must also account for packing which requires - # latent height and width to be divisible by 2. - height = 2 * (int(height) // (self.vae_scale_factor * 2)) - width = 2 * (int(width) // (self.vae_scale_factor * 2)) - - shape = (batch_size, num_channels_latents, height, width) - - if latents is None: - latents = randn_tensor(shape, generator=generator, device=self.device, dtype=dtype) - else: - if latents.shape != shape: - raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}") - latents = latents.to(self.device) - return latents - - def prepare_timesteps(self, num_inference_steps, sigmas, image_seq_len): - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - mu = calculate_shift( - image_seq_len, - self.scheduler.config.get("base_image_seq_len", 256), - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_shift", 0.5), - self.scheduler.config.get("max_shift", 1.15), - ) - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - self.device, - sigmas=sigmas, - mu=mu, - ) - return timesteps, num_inference_steps - - def _extract_prompts(self, prompts): - """Extract prompt and negative_prompt from OmniPromptType list.""" - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None - if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): - negative_prompt = None - elif prompts: - negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] - else: - negative_prompt = None - return prompt, negative_prompt - - def check_cfg_parallel_validity(self, true_cfg_scale: float, has_neg_prompt: bool): - if get_classifier_free_guidance_world_size() == 1: - return True - - if true_cfg_scale <= 1: - logger.warning("CFG parallel is NOT working correctly when true_cfg_scale <= 1.") - return False - - if not has_neg_prompt: - logger.warning( - "CFG parallel is NOT working correctly when there is no negative prompt or negative prompt embeddings." - ) - return False - return True - - def diffuse( - self, - prompt_embeds_t5: torch.Tensor, - prompt_embeds_llama3: torch.Tensor, - pooled_prompt_embeds: torch.Tensor, - latents: torch.Tensor, - timesteps: torch.Tensor, - do_true_cfg: bool, - ) -> torch.Tensor: - with self.progress_bar(total=len(timesteps)) as pbar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - # expand the latents if we are doing classifier free guidance - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - # broadcast to batch dimension in a way that's compatible with ONNX/Core ML - timestep = t.expand(latent_model_input.shape[0]) - noise_pred = self.transformer( - hidden_states=latent_model_input, - timesteps=timestep, - encoder_hidden_states_t5=prompt_embeds_t5, - encoder_hidden_states_llama3=prompt_embeds_llama3, - pooled_embeds=pooled_prompt_embeds, - return_dict=False, - )[0] - noise_pred = -noise_pred - - # TODO: Modify CFG guidance - # perform guidance - if self.do_classifier_free_guidance: - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - - # compute the previous noisy sample x_t -> x_t-1 - latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg) - - pbar.update() - - return latents - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1 - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def interrupt(self): - return self._interrupt - - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] = None, - prompt_2: str | list[str] | None = None, - prompt_3: str | list[str] | None = None, - prompt_4: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 5.0, - negative_prompt: str | list[str] | None = None, - negative_prompt_2: str | list[str] | None = None, - negative_prompt_3: str | list[str] | None = None, - negative_prompt_4: str | list[str] | None = None, - num_images_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds_t5: torch.FloatTensor | None = None, - prompt_embeds_llama3: torch.FloatTensor | None = None, - negative_prompt_embeds_t5: torch.FloatTensor | None = None, - negative_prompt_embeds_llama3: torch.FloatTensor | None = None, - pooled_prompt_embeds: torch.FloatTensor | None = None, - negative_pooled_prompt_embeds: torch.FloatTensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end: Callable[[int, int], None] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 128, - **kwargs, - ): - extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) - prompt = extracted_prompt or prompt - - height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor - - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or guidance_scale - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt - ) - - prompt_embeds = kwargs.get("prompt_embeds", None) - negative_prompt_embeds = kwargs.get("negative_prompt_embeds", None) - - if prompt_embeds is not None: - deprecation_message = ( - "The `prompt_embeds` argument is deprecated." - " Please use `prompt_embeds_t5` and `prompt_embeds_llama3` instead." - ) - deprecate("prompt_embeds", "0.35.0", deprecation_message) - prompt_embeds_t5 = prompt_embeds[0] - prompt_embeds_llama3 = prompt_embeds[1] - - if negative_prompt_embeds is not None: - deprecation_message = ( - "The `negative_prompt_embeds` argument is deprecated." - "Please use `negative_prompt_embeds_t5` and `negative_prompt_embeds_llama3` instead." - ) - deprecate("negative_prompt_embeds", "0.35.0", deprecation_message) - negative_prompt_embeds_t5 = negative_prompt_embeds[0] - negative_prompt_embeds_llama3 = negative_prompt_embeds[1] - - division = self.vae_scale_factor * 2 - S_max = (self.default_sample_size * self.vae_scale_factor) ** 2 - scale = S_max / (width * height) - scale = math.sqrt(scale) - width, height = int(width * scale // division * division), int(height * scale // division * division) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - prompt_3, - prompt_4, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - negative_prompt_4=negative_prompt_4, - prompt_embeds_t5=prompt_embeds_t5, - prompt_embeds_llama3=prompt_embeds_llama3, - negative_prompt_embeds_t5=negative_prompt_embeds_t5, - negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - - # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - elif pooled_prompt_embeds is not None: - batch_size = pooled_prompt_embeds.shape[0] - - # TODO: CFG guidance configuration - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None - ) - do_true_cfg = true_cfg_scale > 1 and has_neg_prompt - self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) - - # 3. Encode prompt - lora_scale = self.attention_kwargs.get("scale", None) if self.attention_kwargs is not None else None - ( - prompt_embeds_t5, - negative_prompt_embeds_t5, - prompt_embeds_llama3, - negative_prompt_embeds_llama3, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - prompt_3=prompt_3, - prompt_4=prompt_4, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - negative_prompt_3=negative_prompt_3, - negative_prompt_4=negative_prompt_4, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds_t5=prompt_embeds_t5, - prompt_embeds_llama3=prompt_embeds_llama3, - negative_prompt_embeds_t5=negative_prompt_embeds_t5, - negative_prompt_embeds_llama3=negative_prompt_embeds_llama3, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - num_images_per_prompt=num_images_per_prompt, - max_sequence_length=max_sequence_length, - lora_scale=lora_scale, - ) - - if self.do_classifier_free_guidance: - prompt_embeds_t5 = torch.cat([negative_prompt_embeds_t5, prompt_embeds_t5], dim=0) - prompt_embeds_llama3 = torch.cat([negative_prompt_embeds_llama3, prompt_embeds_llama3], dim=1) - pooled_prompt_embeds = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds], dim=0) - - # 4. Prepare latent variables - num_channels_latents = self.transformer.in_channels - latents = self.prepare_latents( - batch_size * num_images_per_prompt, - num_channels_latents, - height, - width, - pooled_prompt_embeds.dtype, - generator, - latents, - ) - - # 5. Prepare timesteps - timesteps, num_inference_steps = self.prepare_timesteps(num_inference_steps, sigmas, latents.shape[1]) - self._num_timesteps = len(timesteps) - - # 6. Denoising loop - latents = self.diffuse( - prompt_embeds_t5, - prompt_embeds_llama3, - pooled_prompt_embeds, - latents, - timesteps, - do_true_cfg, - ) - - if output_type == "latent": - image = latents - - else: - latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor - - image = self.vae.decode(latents, return_dict=False)[0] - - return DiffusionOutput( - output=image, 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) - return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py index 135e21e1a73..af156b34bf8 100644 --- a/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py +++ b/vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py @@ -35,8 +35,8 @@ from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = logging.get_logger(__name__) @@ -866,7 +866,7 @@ def interrupt(self): def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] = None, prompt_2: str | list[str] | None = None, prompt_3: str | list[str] | None = None, @@ -896,7 +896,7 @@ def forward( callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 128, **kwargs, - ): + ) -> DiffusionOutput: extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) prompt = extracted_prompt or prompt diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py index f9f114513b5..fc1317a48e6 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py @@ -3,6 +3,8 @@ from typing import Any +import torch +import vllm.distributed.parallel_state as _vllm_ps import vllm.forward_context as _vllm_fc from vllm.utils.import_utils import resolve_obj_by_qualname @@ -24,6 +26,30 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None: forward_context.in_profile_run = False +def _set_forward_context_dp_metadata(num_tokens: int) -> None: + """Populate vLLM MoE DP token metadata for diffusion-only DP mappings.""" + if not _vllm_fc.is_forward_context_available(): + return + forward_context = _vllm_fc.get_forward_context() + if getattr(forward_context, "dp_metadata", None) is not None: + return + + dp_group = getattr(_vllm_ps, "_DP", None) + if dp_group is None or getattr(dp_group, "world_size", 1) <= 1: + return + + gathered_num_tokens: list[int | None] = [None] * dp_group.world_size + torch.distributed.all_gather_object( + gathered_num_tokens, + int(num_tokens), + group=dp_group.cpu_group, + ) + if any(count is None for count in gathered_num_tokens): + raise RuntimeError(f"Failed to gather MoE DP token counts: {gathered_num_tokens}") + + forward_context.dp_metadata = _vllm_fc.DPMetadata(torch.tensor(gathered_num_tokens, dtype=torch.int64)) + + class HunyuanFusedMoEDefault: """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3. @@ -62,6 +88,7 @@ def _num_tokens_pre_hook(module: Any, args: Any, kwargs: Any) -> None: hidden_states = args[0] if hidden_states is not None: _set_forward_context_num_tokens(hidden_states.shape[0]) + _set_forward_context_dp_metadata(hidden_states.shape[0]) moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True) 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 20ed54b4f18..f38d1ca7cd8 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -11,6 +11,7 @@ import numpy as np import regex as re import torch +from cache_dit import ForwardPattern from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback from diffusers.image_processor import VaeImageProcessor from diffusers.pipelines.pipeline_utils import DiffusionPipeline @@ -28,7 +29,7 @@ CausalLMOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel -from vllm.config import CacheConfig +from vllm.config import CacheConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_world_size, ) @@ -59,6 +60,7 @@ AttentionMetadata, ) from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.parallel_state import ( get_cfg_group, get_classifier_free_guidance_rank, @@ -1065,9 +1067,10 @@ def _reuse_prompt_kv( assert self.image_kv_cache_lens is not None if position_ids is not None: assert position_ids.shape == (bs, q_len) - assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens), ( - "The first current position must immediately follow each sample's cached prompt KV." - ) + if logger.isEnabledFor(logging.DEBUG): + assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens.to(position_ids.device)), ( + "The first current position must immediately follow each sample's cached prompt KV." + ) new_key = torch.cat([cached_key, key], dim=1) new_value = torch.cat([cached_value, value], dim=1) return new_key.contiguous(), new_value.contiguous() @@ -1096,6 +1099,12 @@ def _build_neg_ar_kv( neg_kv = (key.reshape(-1, num_kv_heads, head_dim), value.reshape(-1, num_kv_heads, head_dim)) self._injected_ar_kv = [self._injected_ar_kv[0], neg_kv] + logger.debug( + "[AR KV Reuse] neg branch: reused shared_prefix_len=%d, neg_kv_total_len=%d (q_len_actual=%d)", + shared_prefix_len, + key.shape[1], + q_len_actual, + ) return key, value def __call__( @@ -1621,6 +1630,7 @@ def __init__( else: self.shared_mlp = None + enable_expert_parallel = get_current_vllm_config().parallel_config.enable_expert_parallel self.experts = HunyuanFusedMoE( shared_experts=self.shared_mlp, num_experts=self.n_routed_experts, @@ -1632,7 +1642,7 @@ def __init__( prefix=f"{prefix}.experts", enable_eplb=self.enable_eplb, num_redundant_experts=self.n_redundant_experts, - pcp_size=1, + pcp_size=None if enable_expert_parallel else 1, ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: @@ -2017,6 +2027,13 @@ def forward(self, x): class HunyuanImage3Model(nn.Module): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_4, + }, + check_forward_pattern=False, + ) + _sp_plan = { # Split custom_pos_emb tuple elements (cos, sin) at model forward input "pre_processor": { @@ -2224,13 +2241,8 @@ def get_loaded_weight_shard( # processed with quantization, LoRA, fine-tuning, etc. if self.config.tie_word_embeddings and "lm_head.weight" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - # Loading kv cache scales for compressed-tensors quantization - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight[0] - weight_loader(param, loaded_weight) - continue + # KV-cache scales are renamed via maybe_remap_kv_scale_name below; + # quant_config.get_cache_scale was removed in vLLM v0.23.0 (see #4810). is_found = False for param_name, weight_name, shard_id in stacked_params_mapping: @@ -2913,6 +2925,10 @@ def _maybe_handle_ar_kv_reuse( ): ar_kv_data = model_kwargs.pop("ar_kv_data", None) if ar_kv_data is None: + logger.debug( + "[AR KV Reuse] cfg_rank=%s: no AR KV received, fallback to full recompute (reuse_len=0)", + cfg_rank, + ) return input_ids, 0 # 1. positive prefix len 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 6635168b983..c0476b87972 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -24,11 +24,12 @@ 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.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer @@ -300,32 +301,32 @@ def _build_cond_joint_image(raw_image: Any) -> dict[str, Any]: ) def pre_process_func(request: OmniDiffusionRequest): - for i, prompt in enumerate(request.prompts): - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - multi_modal_data = prompt.get("multi_modal_data") or {} - raw_images = multi_modal_data.get("image") - if raw_images is None: - raw_images = prompt.get("pil_image") - has_images = raw_images is not None and (not isinstance(raw_images, list) or len(raw_images) > 0) - if has_images: - image_list = raw_images if isinstance(raw_images, list) else [raw_images] - cond_image_infos = [_build_cond_joint_image(image) for image in image_list] - prompt["additional_information"]["batch_cond_image_info"] = cond_image_infos - - bridge_h = prompt.get("height") if isinstance(prompt, dict) else None - bridge_w = prompt.get("width") if isinstance(prompt, dict) else None - first_image_w, first_image_h = _to_pil_image(image_list[0]).size - if request.sampling_params.width is None: - request.sampling_params.width = int(bridge_w or first_image_w) - if request.sampling_params.height is None: - request.sampling_params.height = int(bridge_h or first_image_h) - - request.prompts[i] = prompt + prompt = request.prompt + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + multi_modal_data = prompt.get("multi_modal_data") or {} + raw_images = multi_modal_data.get("image") + if raw_images is None: + raw_images = prompt.get("pil_image") + has_images = raw_images is not None and (not isinstance(raw_images, list) or len(raw_images) > 0) + if has_images: + image_list = raw_images if isinstance(raw_images, list) else [raw_images] + cond_image_infos = [_build_cond_joint_image(image) for image in image_list] + prompt["additional_information"]["batch_cond_image_info"] = cond_image_infos + + bridge_h = prompt.get("height") if isinstance(prompt, dict) else None + bridge_w = prompt.get("width") if isinstance(prompt, dict) else None + first_image_w, first_image_h = _to_pil_image(image_list[0]).size + if request.sampling_params.width is None: + request.sampling_params.width = int(bridge_w or first_image_w) + if request.sampling_params.height is None: + request.sampling_params.height = int(bridge_h or first_image_h) + + request.prompt = prompt return request @@ -336,10 +337,16 @@ class HunyuanImage3Pipeline( HunyuanImage3PreTrainedModel, GenerationMixin, SupportImageInput, + SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, ): supports_step_execution: ClassVar[bool] = True + supports_request_batch = False support_image_input = True + _dit_modules: ClassVar[list[str]] = ["model"] + _encoder_modules: ClassVar[list[str]] = ["vision_model"] + _vae_modules: ClassVar[list[str]] = ["vae"] + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "model.": "", @@ -480,9 +487,9 @@ def pipeline(self): return self._pipeline def _validate_step_request(self, state: "DiffusionRequestState") -> None: - prompts = state.prompts or [] + prompt = state.prompt sampling = state.sampling - if len(prompts) != 1: + if prompt is None: raise ValueError("HunyuanImage3 step execution currently requires exactly one prompt per request.") if sampling.timesteps is not None or sampling.sigmas is not None: raise ValueError("HunyuanImage3 step execution does not support custom timesteps or sigmas yet.") @@ -581,7 +588,7 @@ def _extract_step_prompt_inputs( ) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]: sampling = state.sampling return self._extract_prompt_inputs( - state.prompts or [], + [state.prompt] if state.prompt is not None else [], getattr(sampling, "extra_args", {}) or {}, request_id=state.request_id, allow_cond_image=False, @@ -987,7 +994,7 @@ def instantiate_timestep_tokens( timestep_scatter_src = self.timestep_emb(t.reshape(-1)).reshape(batch_size, -1, n_embd) x.scatter_( dim=1, - index=timestep_scatter_index.unsqueeze(-1).repeat(1, 1, n_embd), + index=timestep_scatter_index.to(x.device).unsqueeze(-1).repeat(1, 1, n_embd), src=timestep_scatter_src, ) @@ -1370,7 +1377,7 @@ def prepare_model_inputs( num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, image_mask=to_device(output.gen_image_mask, device), - gen_timestep_scatter_index=to_device(output.gen_timestep_scatter_index, device), + gen_timestep_scatter_index=output.gen_timestep_scatter_index, cond_vae_images=to_device(cond_vae_images, device), cond_timestep=to_device(cond_timestep, device), cond_vae_image_mask=to_device(output.cond_vae_image_mask, device), @@ -2202,7 +2209,6 @@ def post_decode( if output_type == "latent": return DiffusionOutput( output=latents, - custom_output={}, stage_durations=getattr(self, "stage_durations", None), ) @@ -2225,18 +2231,20 @@ def post_decode( ) cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or [] - custom_output = {} + metadata = {} if any(text is not None for text in cot_text_list): - custom_output["ar_generated_text"] = cot_text_list[0] + metadata["text"] = {"ar_generated_text": cot_text_list[0]} return DiffusionOutput( - output=image[0], - custom_output=custom_output, + output={ + "payload": {"image": image[0]}, + "metadata": metadata, + }, stage_durations=getattr(self, "stage_durations", None), ) def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] = "", image_size="auto", height: int = 1024, @@ -2294,11 +2302,15 @@ def forward( model_inputs.update(ar_kv_kwargs) outputs = self._generate(**model_inputs, **kwargs) - custom_output = {} + image = outputs[0] + metadata = {} if any(t is not None for t in cot_text_list): - custom_output["ar_generated_text"] = cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list + metadata["text"] = {"ar_generated_text": cot_text_list[0] if len(cot_text_list) == 1 else cot_text_list} + stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None return DiffusionOutput( - output=outputs[0], - custom_output=custom_output, - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + output={ + "payload": {"image": image}, + "metadata": metadata, + }, + stage_durations=stage_durations, ) 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 9af9e848875..9224973dda8 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 @@ -6,6 +6,7 @@ import torch import torch.nn.functional as F +from cache_dit import ForwardPattern from diffusers.models.embeddings import ( CombinedTimestepTextProjEmbeddings, TimestepEmbedding, @@ -22,6 +23,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.distributed.parallel_state import get_sequence_parallel_world_size @@ -591,6 +593,13 @@ class HunyuanVideo15Transformer3DModel(nn.Module): tensor-parallel layers for the 54 main transformer blocks. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_0, + }, + has_separate_cfg=True, + ) + _repeated_blocks = ["HunyuanVideo15TransformerBlock"] _layerwise_offload_blocks_attrs = ["transformer_blocks"] packed_modules_mapping = { 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 0684f67492c..942e0b28435 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 @@ -7,7 +7,7 @@ import os import re from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -27,11 +27,12 @@ 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.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -86,7 +87,13 @@ def post_process_func(video: torch.Tensor, output_type: str = "pil"): return post_process_func -class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin): +class HunyuanVideo15Pipeline( + nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -387,7 +394,7 @@ def predict_noise(self, **kwargs: Any) -> torch.Tensor: def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, num_inference_steps: int = 50, guidance_scale: float = 6.0, height: int = 480, 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 583ef87ede6..1281a9e39d6 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 @@ -6,7 +6,7 @@ import logging import os from collections.abc import Iterable -from typing import Any, cast +from typing import Any, ClassVar, cast import numpy as np import PIL.Image @@ -40,12 +40,13 @@ get_hunyuan_video_15_post_process_func, retrieve_latents, ) -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -62,10 +63,10 @@ def get_hunyuan_video_15_i2v_pre_process_func(od_config: OmniDiffusionConfig): divisor = 16 # Must be divisible by VAE spatial compression def pre_process_func(req: OmniDiffusionRequest) -> OmniDiffusionRequest: - if not req.prompts: + if not req.prompt: return req - prompt_data = req.prompts[0] + prompt_data = req.prompt if isinstance(prompt_data, str): return req @@ -104,8 +105,17 @@ def pre_process_func(req: OmniDiffusionRequest) -> OmniDiffusionRequest: class HunyuanVideo15I2VPipeline( - nn.Module, CFGParallelMixin, SupportImageInput, ProgressBarMixin, DiffusionPipelineProfilerMixin + nn.Module, + CFGParallelMixin, + SupportImageInput, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, ): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2", "image_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + support_image_input = True color_format = "RGB" @@ -465,7 +475,7 @@ def predict_noise(self, **kwargs: Any) -> torch.Tensor: def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, num_inference_steps: int = 50, guidance_scale: float = 6.0, height: int = 480, diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index c9d013ff390..657a45bf882 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -56,15 +56,19 @@ class SupportsStepExecution(Protocol): def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState: """Prepare request-level inputs and return initialized state.""" + ... def denoise_step(self, input_batch: InputBatch, **kwargs: Any) -> torch.Tensor | None: """Run one denoise forward on the runner-assembled batch.""" + ... def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: """Run one scheduler step.""" + ... def post_decode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionOutput: - """Decode output after denoise loop.""" + """Decode output after denoise loop or at a partial chunk boundary.""" + ... @runtime_checkable diff --git a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py index b761ae8c668..10498620c5e 100644 --- a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py +++ b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py @@ -16,7 +16,7 @@ DiffusionPipelineProfilerMixin, wrap_methods_by_paths, ) -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from .config import ( DEFAULT_QWEN3_VL_MODEL, @@ -228,7 +228,7 @@ def _predict_actions( ) @torch.inference_mode() - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: if len(req.prompts) > 1: logger.warning("InternVLAA1Pipeline only supports a single prompt/request; taking the first sample.") extra_args = getattr(req.sampling_params, "extra_args", {}) or {} @@ -247,11 +247,13 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: noise=extra_args.get("noise"), decode_image=bool(extra_args.get("decode_image", False)), ) - custom_output: dict[str, Any] = {} + metadata: dict[str, Any] = {} if decoded is not None: - custom_output["decoded"] = decoded + metadata["actions"] = {"decoded": decoded} return DiffusionOutput( - output=output, - custom_output=custom_output, + output={ + "payload": {"actions": output}, + "metadata": metadata, + }, post_process_func=get_internvla_a1_post_process_func(self.od_config), ) diff --git a/vllm_omni/diffusion/models/krea2/__init__.py b/vllm_omni/diffusion/models/krea2/__init__.py new file mode 100644 index 00000000000..e802b09848a --- /dev/null +++ b/vllm_omni/diffusion/models/krea2/__init__.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Krea 2 diffusion model components.""" + +from vllm_omni.diffusion.models.krea2.krea2_transformer import Krea2Transformer2DModel +from vllm_omni.diffusion.models.krea2.pipeline_krea2 import ( + Krea2Pipeline, + get_krea2_post_process_func, +) + +__all__ = [ + "Krea2Pipeline", + "Krea2Transformer2DModel", + "get_krea2_post_process_func", +] diff --git a/vllm_omni/diffusion/models/krea2/krea2_transformer.py b/vllm_omni/diffusion/models/krea2/krea2_transformer.py new file mode 100644 index 00000000000..e8c0247b5b7 --- /dev/null +++ b/vllm_omni/diffusion/models/krea2/krea2_transformer.py @@ -0,0 +1,571 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from collections.abc import Iterable +from typing import TYPE_CHECKING + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import apply_rotary_emb, get_1d_rotary_pos_embed +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + + from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + +# Devices that do not support float64; RoPE frequencies fall back to float32 there. +_FP64_UNSUPPORTED_DEVICES = frozenset({"mps", "npu", "neuron"}) + + +def _rope_freqs_dtype(device: torch.device) -> torch.dtype: + return torch.float32 if device.type in _FP64_UNSUPPORTED_DEVICES else torch.float64 + + +def _join_prefix(prefix: str, name: str) -> str: + return f"{prefix}.{name}" if prefix else name + + +def _linear( + in_features: int, + out_features: int, + bias: bool, + quant_config: "QuantizationConfig | None", + prefix: str, +) -> ReplicatedLinear: + # ReplicatedLinear (not nn.Linear) so the diffusion LoRA manager can wrap these projections. + return ReplicatedLinear( + in_features, + out_features, + bias=bias, + quant_config=quant_config, + prefix=prefix, + return_bias=False, + ) + + +class Krea2RMSNorm(nn.Module): + """RMSNorm with a zero-centered scale: the effective multiplier is ``1 + weight``, matching the Krea 2 checkpoint + format. Activations are upcast so the normalization runs in float32.""" + + def __init__(self, dim: int, eps: float = 1e-5) -> None: + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.zeros(dim)) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + dtype = hidden_states.dtype + hidden_states = F.rms_norm(hidden_states.float(), (self.dim,), weight=self.weight + 1.0, eps=self.eps) + return hidden_states.to(dtype) + + +class Krea2Attention(nn.Module): + """Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings and a sigmoid output gate. + + Q/K/V layout is ``[B, seq, heads, head_dim]``; ``attention_mask`` is a 2D boolean key-padding mask + ``(batch, key_seq_len)`` that the backend broadcasts. + """ + + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int | None = None, + eps: float = 1e-5, + role: str = "self", + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: + super().__init__() + if hidden_size % num_heads != 0: + raise ValueError(f"hidden_size={hidden_size} must be divisible by num_heads={num_heads}") + self.hidden_size = hidden_size + self.num_heads = num_heads + self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads + self.head_dim = hidden_size // num_heads + + q_dim = self.head_dim * self.num_heads + kv_dim = self.head_dim * self.num_kv_heads + self.to_q = _linear(hidden_size, q_dim, False, quant_config, _join_prefix(prefix, "to_q")) + self.to_k = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_k")) + self.to_v = _linear(hidden_size, kv_dim, False, quant_config, _join_prefix(prefix, "to_v")) + self.to_gate = _linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_gate")) + self.norm_q = Krea2RMSNorm(self.head_dim, eps=eps) + self.norm_k = Krea2RMSNorm(self.head_dim, eps=eps) + self.to_out = nn.ModuleList( + [_linear(hidden_size, hidden_size, False, quant_config, _join_prefix(prefix, "to_out.0")), nn.Dropout(0.0)] + ) + + self.attn = Attention( + num_heads=self.num_heads, + head_size=self.head_dim, + softmax_scale=1.0 / (self.head_dim**0.5), + causal=False, + num_kv_heads=self.num_kv_heads, + role=role, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + query = self.to_q(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)) + key = self.to_k(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim)) + value = self.to_v(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim)) + gate = self.to_gate(hidden_states) + + query = self.norm_q(query) + key = self.norm_k(key) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + query = query.to(value.dtype) + key = key.to(value.dtype) + + attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None + hidden_states = self.attn(query, key, value, attn_metadata) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states * torch.sigmoid(gate) + return self.to_out[0](hidden_states) + + +class Krea2SwiGLU(nn.Module): + """SwiGLU feed-forward network.""" + + def __init__( + self, dim: int, hidden_dim: int, quant_config: "QuantizationConfig | None" = None, prefix: str = "" + ) -> None: + super().__init__() + self.gate = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "gate")) + self.up = _linear(dim, hidden_dim, False, quant_config, _join_prefix(prefix, "up")) + self.down = _linear(hidden_dim, dim, False, quant_config, _join_prefix(prefix, "down")) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.down(F.silu(self.gate(hidden_states)) * self.up(hidden_states)) + + +class Krea2TextFusionBlock(nn.Module): + """Pre-norm transformer block (no rotary embeddings, no time modulation) used by the text fusion stage.""" + + def __init__( + self, + dim: int, + num_heads: int, + num_kv_heads: int, + intermediate_size: int, + eps: float, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: + super().__init__() + self.norm1 = Krea2RMSNorm(dim, eps=eps) + self.norm2 = Krea2RMSNorm(dim, eps=eps) + self.attn = Krea2Attention( + dim, + num_heads, + num_kv_heads, + eps=eps, + role="self", + quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), + ) + self.ff = Krea2SwiGLU(dim, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff")) + + def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask=attention_mask) + hidden_states = hidden_states + self.ff(self.norm2(hidden_states)) + return hidden_states + + +class Krea2TextFusion(nn.Module): + """Fuses the stack of tapped text-encoder hidden states into a single sequence of text features. + + Two ``layerwise_blocks`` attend across the ``num_text_layers`` axis independently for every token, a linear + ``projector`` collapses that axis, and two ``refiner_blocks`` attend across the token sequence. + """ + + def __init__( + self, + num_text_layers: int, + dim: int, + num_heads: int, + num_kv_heads: int, + intermediate_size: int, + num_layerwise_blocks: int, + num_refiner_blocks: int, + eps: float, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layerwise_blocks = nn.ModuleList( + [ + Krea2TextFusionBlock( + dim, + num_heads, + num_kv_heads, + intermediate_size, + eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"layerwise_blocks.{i}"), + ) + for i in range(num_layerwise_blocks) + ] + ) + self.projector = _linear(num_text_layers, 1, False, quant_config, _join_prefix(prefix, "projector")) + self.refiner_blocks = nn.ModuleList( + [ + Krea2TextFusionBlock( + dim, + num_heads, + num_kv_heads, + intermediate_size, + eps, + quant_config=quant_config, + prefix=_join_prefix(prefix, f"refiner_blocks.{i}"), + ) + for i in range(num_refiner_blocks) + ] + ) + + def forward(self, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor: + batch_size, seq_len, num_text_layers, dim = encoder_hidden_states.shape + + hidden_states = encoder_hidden_states.reshape(batch_size * seq_len, num_text_layers, dim) + for block in self.layerwise_blocks: + hidden_states = block(hidden_states.contiguous()) + + hidden_states = hidden_states.reshape(batch_size, seq_len, num_text_layers, dim).permute(0, 1, 3, 2) + hidden_states = self.projector(hidden_states).squeeze(-1) + + for block in self.refiner_blocks: + hidden_states = block(hidden_states, attention_mask=attention_mask) + + return hidden_states + + +class Krea2TransformerBlock(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + norm_eps: float, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: + super().__init__() + self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size)) + self.norm1 = Krea2RMSNorm(hidden_size, eps=norm_eps) + self.norm2 = Krea2RMSNorm(hidden_size, eps=norm_eps) + self.attn = Krea2Attention( + hidden_size, + num_heads, + num_kv_heads, + eps=norm_eps, + role="self", + quant_config=quant_config, + prefix=_join_prefix(prefix, "attn"), + ) + self.ff = Krea2SwiGLU( + hidden_size, intermediate_size, quant_config=quant_config, prefix=_join_prefix(prefix, "ff") + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an additive table. + modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table + prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2) + + attn_out = self.attn( + (1.0 + prescale) * self.norm1(hidden_states) + preshift, + attention_mask=attention_mask, + image_rotary_emb=image_rotary_emb, + ) + hidden_states = hidden_states + pregate * attn_out + ff_out = self.ff((1.0 + postscale) * self.norm2(hidden_states) + postshift) + hidden_states = hidden_states + postgate * ff_out + return hidden_states + + +class Krea2TimestepEmbedding(nn.Module): + """Sinusoidal flow-time embedding (cos-first, input scaled by 1000) followed by a two-layer MLP. + + Keeps the sequence dimension at size 1 so the per-block modulations broadcast over tokens. + """ + + def __init__( + self, embed_dim: int, hidden_size: int, quant_config: "QuantizationConfig | None" = None, prefix: str = "" + ) -> None: + super().__init__() + self.embed_dim = embed_dim + self.linear_1 = _linear(embed_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1")) + self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2")) + + def forward(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + half = self.embed_dim // 2 + freqs = torch.exp(-math.log(1e4) * torch.arange(half, dtype=torch.float32, device=timestep.device) / half) + args = (timestep.float() * 1e3)[:, None, None] * freqs + emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype) + return self.linear_2(F.gelu(self.linear_1(emb), approximate="tanh")) + + +class Krea2TextProjection(nn.Module): + """Projects the fused text features into the transformer width.""" + + def __init__( + self, + text_dim: int, + hidden_size: int, + eps: float, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: + super().__init__() + self.norm = Krea2RMSNorm(text_dim, eps=eps) + self.linear_1 = _linear(text_dim, hidden_size, True, quant_config, _join_prefix(prefix, "linear_1")) + self.linear_2 = _linear(hidden_size, hidden_size, True, quant_config, _join_prefix(prefix, "linear_2")) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.linear_1(self.norm(hidden_states)) + return self.linear_2(F.gelu(hidden_states, approximate="tanh")) + + +class Krea2FinalLayer(nn.Module): + """Final adaptive RMSNorm and output projection.""" + + def __init__( + self, + hidden_size: int, + out_channels: int, + eps: float, + quant_config: "QuantizationConfig | None" = None, + prefix: str = "", + ) -> None: + super().__init__() + self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size)) + self.norm = Krea2RMSNorm(hidden_size, eps=eps) + self.linear = _linear(hidden_size, out_channels, True, quant_config, _join_prefix(prefix, "linear")) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + modulation = temb + self.scale_shift_table + scale, shift = modulation.chunk(2, dim=1) + hidden_states = (1.0 + scale) * self.norm(hidden_states) + shift + return self.linear(hidden_states) + + +class Krea2RotaryPosEmbed(nn.Module): + """Multi-axis (t, h, w) rotary position embedding, following the Flux/Krea 2 convention.""" + + def __init__(self, theta: int, axes_dim: list[int]): + super().__init__() + self.theta = theta + self.axes_dim = axes_dim + + def forward(self, ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + n_axes = ids.shape[-1] + cos_out = [] + sin_out = [] + pos = ids.float() + freqs_dtype = _rope_freqs_dtype(ids.device) + for i in range(n_axes): + cos, sin = get_1d_rotary_pos_embed( + self.axes_dim[i], + pos[:, i], + theta=self.theta, + repeat_interleave_real=True, + use_real=True, + freqs_dtype=freqs_dtype, + ) + cos_out.append(cos) + sin_out.append(sin) + freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device) + freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device) + return freqs_cos, freqs_sin + + +class Krea2Transformer2DModel(nn.Module): + r"""The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline (vLLM-Omni port). + + Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A + small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated + with the patchified image latents into a single ``[text, image]`` sequence processed by the transformer blocks. + The timestep conditions every block through one shared modulation vector plus per-block learned tables. + """ + + _repeated_blocks = ["Krea2TransformerBlock"] + _layerwise_offload_blocks_attrs = ["transformer_blocks"] + + @staticmethod + def _is_transformer_block(name: str, module) -> bool: + return "transformer_blocks" in name and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + + def __init__( + self, + in_channels: int = 64, + num_layers: int = 28, + attention_head_dim: int = 128, + num_attention_heads: int = 48, + num_key_value_heads: int = 12, + intermediate_size: int = 16384, + timestep_embed_dim: int = 256, + text_hidden_dim: int = 2560, + num_text_layers: int = 12, + text_num_attention_heads: int = 20, + text_num_key_value_heads: int = 20, + text_intermediate_size: int = 6912, + num_layerwise_text_blocks: int = 2, + num_refiner_text_blocks: int = 2, + axes_dims_rope: tuple[int, int, int] = (32, 48, 48), + rope_theta: float = 1000.0, + norm_eps: float = 1e-5, + od_config: "OmniDiffusionConfig | None" = None, + quant_config: "QuantizationConfig | None" = None, + ) -> None: + super().__init__() + # The pipeline reads ``self.dtype`` to cast inputs before the transformer forward. + self.dtype = od_config.dtype if od_config is not None else torch.get_default_dtype() + self.od_config = od_config + + hidden_size = attention_head_dim * num_attention_heads + if sum(axes_dims_rope) != attention_head_dim: + raise ValueError( + f"sum(axes_dims_rope)={sum(axes_dims_rope)} must equal attention_head_dim={attention_head_dim}" + ) + + self.in_channels = in_channels + self.out_channels = in_channels + self.hidden_size = hidden_size + self.num_layers = num_layers + + self.img_in = _linear(in_channels, hidden_size, True, quant_config, "img_in") + self.time_embed = Krea2TimestepEmbedding(timestep_embed_dim, hidden_size, quant_config, "time_embed") + self.time_mod_proj = _linear(hidden_size, 6 * hidden_size, True, quant_config, "time_mod_proj") + self.text_fusion = Krea2TextFusion( + num_text_layers=num_text_layers, + dim=text_hidden_dim, + num_heads=text_num_attention_heads, + num_kv_heads=text_num_key_value_heads, + intermediate_size=text_intermediate_size, + num_layerwise_blocks=num_layerwise_text_blocks, + num_refiner_blocks=num_refiner_text_blocks, + eps=norm_eps, + quant_config=quant_config, + prefix="text_fusion", + ) + self.txt_in = Krea2TextProjection(text_hidden_dim, hidden_size, norm_eps, quant_config, "txt_in") + self.rotary_emb = Krea2RotaryPosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope)) + + self.transformer_blocks = nn.ModuleList( + [ + Krea2TransformerBlock( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_heads=num_attention_heads, + num_kv_heads=num_key_value_heads, + norm_eps=norm_eps, + quant_config=quant_config, + prefix=f"transformer_blocks.{i}", + ) + for i in range(num_layers) + ] + ) + + self.final_layer = Krea2FinalLayer(hidden_size, in_channels, norm_eps, quant_config, "final_layer") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + position_ids: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + r"""Predict the flow-matching velocity for the image tokens. + + Args: + hidden_states: Packed (patchified) noisy image latents ``(batch, image_seq_len, in_channels)``. + encoder_hidden_states: Tapped text-encoder hidden states + ``(batch, text_seq_len, num_text_layers, text_hidden_dim)``. + timestep: Flow-matching time in ``[0, 1]`` of shape ``(batch,)``. + position_ids: ``(t, h, w)`` rotary coordinates ``(text_seq_len + image_seq_len, 3)``. + encoder_attention_mask: Boolean mask marking valid text tokens ``(batch, text_seq_len)`` or ``None``. + """ + if position_ids.ndim != 2 or position_ids.shape[-1] != 3: + raise ValueError(f"`position_ids` must have shape (sequence_length, 3), got {tuple(position_ids.shape)}.") + + batch_size, image_seq_len, _ = hidden_states.shape + + temb = self.time_embed(timestep, dtype=hidden_states.dtype) + temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh")) + + # 2D key-padding masks: padded text tokens are excluded as keys; their lanes are dropped at the output slice. + text_attention_mask = None + attention_mask = None + if encoder_attention_mask is not None: + text_attention_mask = encoder_attention_mask.bool() + image_mask = text_attention_mask.new_ones((batch_size, image_seq_len)) + attention_mask = torch.cat([text_attention_mask, image_mask], dim=1) + + encoder_hidden_states = self.text_fusion(encoder_hidden_states, attention_mask=text_attention_mask) + encoder_hidden_states = self.txt_in(encoder_hidden_states) + + text_seq_len = encoder_hidden_states.shape[1] + hidden_states = self.img_in(hidden_states) + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + image_rotary_emb = self.rotary_emb(position_ids) + + for block in self.transformer_blocks: + hidden_states = block(hidden_states, temb_mod, image_rotary_emb, attention_mask) + + hidden_states = hidden_states[:, text_seq_len:] + return self.final_layer(hidden_states, temb) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # Module/parameter names mirror the diffusers ``Krea2Transformer2DModel`` exactly, so weights map 1:1. + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params diff --git a/vllm_omni/diffusion/models/krea2/pipeline_krea2.py b/vllm_omni/diffusion/models/krea2/pipeline_krea2.py new file mode 100644 index 00000000000..4aa0617c7b0 --- /dev/null +++ b/vllm_omni/diffusion/models/krea2/pipeline_krea2.py @@ -0,0 +1,551 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import inspect +import logging +import os +from collections.abc import Iterable +from typing import ClassVar + +import numpy as np +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import AutoTokenizer, Qwen3VLModel +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.transformers_utils.config import get_hf_file_to_dict + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage +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.krea2.krea2_transformer import Krea2Transformer2DModel +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs + +logger = logging.getLogger(__name__) + +# Text-encoder layer taps used when model_index.json does not specify ``text_encoder_select_layers``. +DEFAULT_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35) + + +def get_krea2_post_process_func(od_config: OmniDiffusionConfig): + model_name = od_config.model + # Read only the VAE config (``get_hf_file_to_dict`` resolves a local dir via + # ``Path(model)/file`` and a hub repo via a single-file download) instead of + # pulling all weights just to compute ``vae_scale_factor``. + vae_config = get_hf_file_to_dict("vae/config.json", model_name) or {} + vae_scale_factor = 2 ** len(vae_config["temperal_downsample"]) if "temperal_downsample" in vae_config else 8 + + patch_size = 2 + try: + model_index = get_hf_file_to_dict("model_index.json", model_name) or {} + patch_size = int(model_index.get("patch_size", patch_size)) + except Exception: # noqa: BLE001 + pass + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * patch_size) + + def post_process_func(images: torch.Tensor): + return image_processor.postprocess(images) + + return post_process_func + + +# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift +def calculate_shift( + image_seq_len, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +): + m = (max_shift - base_shift) / (max_seq_len - base_seq_len) + b = base_shift - m * base_seq_len + mu = image_seq_len * m + b + return mu + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +) -> tuple[torch.Tensor, int]: + if timesteps is not None and sigmas is not None: + raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") + if timesteps is not None: + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + if not accept_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigmas schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + return timesteps, num_inference_steps + + +class Krea2Pipeline(nn.Module, DiffusionPipelineProfilerMixin, ProgressBarMixin, SupportsComponentDiscovery): + """The Krea 2 text-to-image pipeline (vLLM-Omni port of ``diffusers.Krea2Pipeline``). + + Components: + - ``scheduler``: ``FlowMatchEulerDiscreteScheduler`` (resolution-aware exponential time shift). + - ``vae``: the Qwen-Image VAE (f8, 16 latent channels). + - ``text_encoder``: a Qwen3-VL model; the pipeline consumes a stack of hidden states tapped from several + decoder layers rather than the last hidden state. + - ``tokenizer``: the tokenizer paired with the text encoder. + - ``transformer``: the Krea 2 single-stream MMDiT that predicts the flow-matching velocity. + """ + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + _PROFILER_TARGETS: ClassVar[list[str]] = ["text_encoder.forward", "diffuse", "vae.decode"] + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + # Only the transformer goes through the standard weights loader; text encoder and VAE load eagerly below. + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=od_config.revision, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + self.device = get_local_device() + model = od_config.model + local_files_only = os.path.isdir(model) + + # Pipeline-level config lives in model_index.json, not the transformer config. + model_index = get_hf_file_to_dict("model_index.json", model) or {} + select_layers = model_index.get("text_encoder_select_layers") or DEFAULT_TEXT_ENCODER_SELECT_LAYERS + self.text_encoder_select_layers = tuple(int(i) for i in select_layers) + self.is_distilled = bool(model_index.get("is_distilled", False)) + self.patch_size = int(model_index.get("patch_size", 2)) + + subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"] + # See ``hub_prefetch.py`` for the transformers v5 subfolder race. + prefetch_subfolders(model, subfolders, local_files_only=local_files_only) + + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + + self.text_encoder = from_pretrained_with_prefetch( + Qwen3VLModel.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=subfolders, + local_files_only=local_files_only, + torch_dtype=od_config.dtype, + ) + # Drop the unused Qwen3-VL vision tower before moving to GPU so it never consumes GPU memory. + if hasattr(self.text_encoder, "visual"): + del self.text_encoder.visual + else: + logger.warning("Krea2: vision tower not found on text encoder; skipping drop") + self.text_encoder = self.text_encoder.to(self.device) + + self.vae = from_pretrained_with_prefetch( + DistributedAutoencoderKLQwenImage.from_pretrained, + model, + subfolder="vae", + prefetch_list=subfolders, + local_files_only=local_files_only, + torch_dtype=od_config.dtype, + ).to(self.device) + + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, Krea2Transformer2DModel) + self.transformer = Krea2Transformer2DModel( + od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + ) + + self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + + self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size) + self.default_sample_size = 1024 + + # Chat template tokenized as a fixed-length block: prompt padded to a fixed length first, assistant suffix + # appended after the padding; the first ``prompt_template_encode_start_idx`` system tokens are later dropped. + self.prompt_template_encode_prefix = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, " + "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n" + ) + self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n" + self.prompt_template_encode_start_idx = 34 + self.prompt_template_encode_num_suffix_tokens = 5 + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def get_text_hidden_states( + self, + prompt: str | list[str], + max_sequence_length: int = 512, + device: torch.device | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Tokenize ``prompt`` into the fixed-length Krea 2 layout and tap the selected encoder hidden states. + + Returns ``(hidden_states, attention_mask)`` of shapes ``(batch, text_seq_len, num_text_layers, + text_hidden_dim)`` and ``(batch, text_seq_len)`` (bool). + """ + device = device or self.device + prompt = [prompt] if isinstance(prompt, str) else prompt + prefix_idx = self.prompt_template_encode_start_idx + text = [self.prompt_template_encode_prefix + e for e in prompt] + text_tokens = self.tokenizer( + text, + truncation=True, + padding="max_length", + max_length=max_sequence_length + prefix_idx - self.prompt_template_encode_num_suffix_tokens, + return_tensors="pt", + ).to(device) + suffix_tokens = self.tokenizer([self.prompt_template_encode_suffix] * len(text), return_tensors="pt").to(device) + + input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1) + attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool() + + # Padding sits in the middle (``[prefix | prompt | PAD | suffix]``), so positions must count only real tokens + # (padding consumes no position). Broadcast across the 3 mRoPE axes (T/H/W equal for text). + position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0) + position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) + + outputs = self.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + output_hidden_states=True, + ) + hidden_states = torch.stack([outputs.hidden_states[i] for i in self.text_encoder_select_layers], dim=2) + + hidden_states = hidden_states[:, prefix_idx:] + attention_mask = attention_mask[:, prefix_idx:] + return hidden_states, attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + device: torch.device | None = None, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + prompt_embeds_mask: torch.Tensor | None = None, + max_sequence_length: int = 512, + ) -> tuple[torch.Tensor, torch.Tensor]: + device = device or self.device + + if prompt_embeds is None: + prompt_embeds, prompt_embeds_mask = self.get_text_hidden_states(prompt, max_sequence_length, device) + + batch_size, seq_len, num_text_layers, dim = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, num_text_layers, dim) + prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt) + prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len) + + return prompt_embeds, prompt_embeds_mask + + def _pack_latents(self, latents, batch_size, num_channels_latents, height, width): + p = self.patch_size + latents = latents.view(batch_size, num_channels_latents, height // p, p, width // p, p) + latents = latents.permute(0, 2, 4, 1, 3, 5) + latents = latents.reshape(batch_size, (height // p) * (width // p), num_channels_latents * p * p) + return latents + + def _unpack_latents(self, latents, height, width): + batch_size, _, channels = latents.shape + p = self.patch_size + + height = p * (int(height) // (self.vae_scale_factor * p)) + width = p * (int(width) // (self.vae_scale_factor * p)) + + latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p) + latents = latents.permute(0, 3, 1, 4, 2, 5) + latents = latents.reshape(batch_size, channels // (p * p), 1, height, width) + return latents + + @staticmethod + def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device): + """Build the ``(text_seq_len + grid_height * grid_width, 3)`` rotary coordinates: text tokens sit at the + origin, image tokens carry their ``(0, h, w)`` latent-grid coordinates.""" + text_ids = torch.zeros(text_seq_len, 3, device=device) + image_ids = torch.zeros(grid_height, grid_width, 3, device=device) + image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None] + image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :] + image_ids = image_ids.reshape(grid_height * grid_width, 3) + return torch.cat([text_ids, image_ids], dim=0) + + def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None): + if latents is not None: + return latents.to(device=device, dtype=dtype) + + latent_height = height // self.vae_scale_factor + latent_width = width // self.vae_scale_factor + shape = (batch_size, num_channels_latents, latent_height, latent_width) + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width) + return latents + + def _extract_prompts(self, prompts): + """Extract prompt and negative_prompt from the request's OmniPromptType list.""" + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] or None + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in prompts): + negative_prompt = None + elif prompts: + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in prompts] + else: + negative_prompt = None + return prompt, negative_prompt + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale > 0 + + def diffuse( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + position_ids: torch.Tensor, + prompt_embeds: torch.Tensor, + prompt_embeds_mask: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + negative_prompt_embeds_mask: torch.Tensor | None, + guidance_scale: float, + ) -> torch.Tensor: + """Run the flow-matching denoising loop and return the final latents.""" + self.scheduler.set_begin_index(0) + transformer_dtype = self.transformer.dtype + with self.progress_bar(total=len(timesteps)) as progress_bar: + for _, t in enumerate(timesteps): + timestep = (t / self.scheduler.config.num_train_timesteps).expand(latents.shape[0]) + timestep = timestep.to(transformer_dtype) + latent_input = latents.to(transformer_dtype) + + noise_pred = self.transformer( + hidden_states=latent_input, + encoder_hidden_states=prompt_embeds.to(transformer_dtype), + timestep=timestep, + position_ids=position_ids, + encoder_attention_mask=prompt_embeds_mask, + ) + + if self.do_classifier_free_guidance: + neg_noise_pred = self.transformer( + hidden_states=latent_input, + encoder_hidden_states=negative_prompt_embeds.to(transformer_dtype), + timestep=timestep, + position_ids=position_ids, + encoder_attention_mask=negative_prompt_embeds_mask, + ) + noise_pred = noise_pred + guidance_scale * (noise_pred - neg_noise_pred) + + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + progress_bar.update() + + return latents + + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int = 1024, + width: int = 1024, + num_inference_steps: int = 28, + sigmas: list[float] | None = None, + guidance_scale: float = 4.5, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + output_type: str | None = "pil", + max_sequence_length: int = 512, + ) -> DiffusionOutput: + extracted_prompt, extracted_negative = self._extract_prompts(req.prompts) + prompt = extracted_prompt or prompt + if extracted_negative is not None: + negative_prompt = extracted_negative + + sp = req.sampling_params + height = sp.height or height + width = sp.width or width + num_inference_steps = sp.num_inference_steps or num_inference_steps + sigmas = sp.sigmas or sigmas + max_sequence_length = sp.max_sequence_length or max_sequence_length + generator = sp.generator or generator + if sp.guidance_scale_provided: + guidance_scale = sp.guidance_scale + else: + # Distilled checkpoint is trained guidance-free (and the request layer coerces an explicit + # guidance_scale=0 to "not provided"), so default it to no CFG. + guidance_scale = 0.0 if self.is_distilled else guidance_scale + num_images_per_prompt = sp.num_outputs_per_prompt if sp.num_outputs_per_prompt > 0 else num_images_per_prompt + + multiple = self.vae_scale_factor * self.patch_size + if height % multiple != 0 or width % multiple != 0: + rounded_height = ((height + multiple - 1) // multiple) * multiple + rounded_width = ((width + multiple - 1) // multiple) * multiple + logger.warning( + "`height` and `width` must be multiples of %d; rounding up from %dx%d to %dx%d.", + multiple, + height, + width, + rounded_height, + rounded_width, + ) + height, width = rounded_height, rounded_width + + device = self.device + self._guidance_scale = guidance_scale + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = 1 + + prompt_embeds, prompt_embeds_mask = self.encode_prompt( + prompt=prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + negative_prompt_embeds = None + negative_prompt_embeds_mask = None + if self.do_classifier_free_guidance: + if negative_prompt is None: + negative_prompt = "" + if isinstance(negative_prompt, str): + negative_prompt = [negative_prompt] * batch_size + negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt( + prompt=negative_prompt, + device=device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + ) + + num_channels_latents = self.transformer.in_channels // (self.patch_size**2) + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + device, + generator, + latents, + ) + grid_height = height // (self.vae_scale_factor * self.patch_size) + grid_width = width // (self.vae_scale_factor * self.patch_size) + position_ids = self.prepare_position_ids(prompt_embeds.shape[1], grid_height, grid_width, device) + + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + image_seq_len = latents.shape[1] + if self.is_distilled: + mu = 1.15 + else: + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("max_image_seq_len", 6400), + self.scheduler.config.get("base_shift", 0.5), + self.scheduler.config.get("max_shift", 1.15), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, device, sigmas=sigmas, mu=mu + ) + + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + position_ids=position_ids, + prompt_embeds=prompt_embeds, + prompt_embeds_mask=prompt_embeds_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_embeds_mask=negative_prompt_embeds_mask, + guidance_scale=guidance_scale, + ) + + if output_type == "latent": + image = latents + else: + latents = self._unpack_latents(latents, height, width) + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.z_dim, 1, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to( + latents.device, latents.dtype + ) + latents = latents / latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False)[0][:, :, 0] + + return DiffusionOutput( + output=image, + 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) + loaded_weights = loader.load_weights(weights) + # VAE and text encoder load eagerly in __init__; record their params so strict load checks pass. + loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()} + if self.text_encoder is not None: + loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()} + return loaded_weights diff --git a/vllm_omni/diffusion/models/lance/lance_transformer.py b/vllm_omni/diffusion/models/lance/lance_transformer.py index 012ba3f4248..12959eb6472 100644 --- a/vllm_omni/diffusion/models/lance/lance_transformer.py +++ b/vllm_omni/diffusion/models/lance/lance_transformer.py @@ -278,6 +278,11 @@ class LanceBagel(Bagel): ``forward_cache_update_vae`` — is reused unchanged. """ + # Upstream Lance samples ``num_timesteps + 1`` schedule points (one extra + # over official BAGEL) so the denoise loop runs exactly ``num_timesteps`` + # Euler steps. See Bagel.generate_image and issue #4470. + _denoise_schedule_extra_step: bool = True + @staticmethod def _lance_compute_vit_buckets(): """Replicate upstream Lance's ``BucketResize.init_buckets`` for the ViT path. diff --git a/vllm_omni/diffusion/models/lance/pipeline_lance.py b/vllm_omni/diffusion/models/lance/pipeline_lance.py index 6f7e9b95e7e..dc885ac44f9 100644 --- a/vllm_omni/diffusion/models/lance/pipeline_lance.py +++ b/vllm_omni/diffusion/models/lance/pipeline_lance.py @@ -836,8 +836,10 @@ def _forward_t2v(self, req): frames = [Image.fromarray(f) for f in frames_np] logger.info("Lance t2v: decoded %d frames at %dx%d", len(frames), frames[0].width, frames[0].height) return DiffusionOutput( - output=frames[0], - custom_output={"video_frames": frames, "video_shape": video_shape}, + output={ + "payload": {"video": frames}, + "metadata": {"video": {"shape": video_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1069,13 +1071,11 @@ def vit_transforms(img): ) img = self._decode_image_from_latent(self.bagel, self.vae, latents[0], image_shape) - # Put the PIL image into ``custom_output`` too — the orchestrator - # serializes ``custom_output`` across the IPC boundary but strips - # the bare ``output`` field, so consumers that only see - # ``outputs[0].output`` (e.g. gradio_demo) get ``None`` otherwise. return DiffusionOutput( - output=img, - custom_output={"image": img, "image_shape": image_shape}, + output={ + "payload": {"image": img}, + "metadata": {"image": {"shape": image_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1361,8 +1361,10 @@ def _forward_i2v(self, req): frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], out_shape) frames = [Image.fromarray(f) for f in frames_np] return DiffusionOutput( - output=frames[0], - custom_output={"video_frames": frames, "video_shape": out_shape}, + output={ + "payload": {"video": frames}, + "metadata": {"video": {"shape": out_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1622,8 +1624,10 @@ def _vae_transforms(_): frames_np = self._decode_video_from_latent(self.bagel, self.vae, latents[0], video_shape) frames = [Image.fromarray(f) for f in frames_np] return DiffusionOutput( - output=frames[0], - custom_output={"video_frames": frames, "video_shape": video_shape}, + output={ + "payload": {"video": frames}, + "metadata": {"video": {"shape": video_shape}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1705,8 +1709,10 @@ def _forward_x2t_video(self, req): text_output = text_output.lstrip("\n") logger.info("Lance x2t_video: generated %d tokens", token_ids.shape[0]) return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) @@ -1801,8 +1807,10 @@ def vit_transforms(img): text_output = text_output.lstrip("\n") logger.info("Lance x2t_image: generated %d tokens", token_ids.shape[0]) return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) diff --git a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py index 8f0ff446afd..744c1806a5a 100644 --- a/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py +++ b/vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py @@ -6,6 +6,7 @@ import torch import torch.nn as nn +from cache_dit import ForwardPattern from diffusers.models.embeddings import TimestepEmbedding, Timesteps, apply_rotary_emb, get_1d_rotary_pos_embed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle @@ -18,6 +19,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelInput, @@ -581,6 +583,14 @@ class LongCatImageTransformer2DModel(nn.Module): Supports Sequence Parallelism (Ulysses and Ring) when configured via OmniDiffusionConfig. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_1, + "single_transformer_blocks": ForwardPattern.Pattern_1, + }, + has_separate_cfg=True, + ) + _repeated_blocks = ["LongCatImageTransformerBlock", "LongCatImageSingleTransformerBlock"] _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"] diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py index ec942c25ff1..b95f5922839 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py @@ -9,7 +9,7 @@ import re from collections.abc import Iterable from functools import partial -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -28,9 +28,10 @@ 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.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -201,7 +202,13 @@ def get_prompt_language(prompt): return "en" -class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): +class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery): + supports_request_batch = False + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -488,66 +495,35 @@ def cfg_normalize_function(self, noise_pred, comb_pred, cfg_renorm_min=0.0): noise_pred = comb_pred * scale return noise_pred - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 4.5, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - enable_cfg_renorm: bool | None = True, - cfg_renorm_min: float | None = 0.0, - enable_prompt_rewrite: bool | None = True, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - height = req.sampling_params.height or height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or width or self.default_sample_size * self.vae_scale_factor - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - generator = req.sampling_params.generator or generator - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) + prompt_embeds = None + negative_prompt_embeds = None + + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + generator = req.sampling_params.generator + guidance_scale = req.sampling_params.guidance_scale num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt is not None - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) - enable_prompt_rewrite = req.sampling_params.extra_args.get("enable_prompt_rewrite", enable_prompt_rewrite) - enable_cfg_renorm = req.sampling_params.extra_args.get("enable_cfg_renorm", enable_cfg_renorm) - cfg_renorm_min = req.sampling_params.extra_args.get("cfg_renorm_min", cfg_renorm_min) - - req_prompt_embeds = [p.get("prompt_embeds") if not isinstance(p, str) else None for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - # If at list one prompt is provided as an embedding, - # Then assume that the user wants to provide embeddings for all prompts, and enter this if block - # If the user in fact provides mixed input format, req_prompt_embeds will have some None's - # And `torch.stack` automatically raises an exception for us - prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore # intentionally expect TypeError - - req_negative_prompt_embeds = [ - p.get("negative_prompt_embeds") if not isinstance(p, str) else None for p in req.prompts - ] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore # intentionally expect TypeError + latents = req.sampling_params.latents + + extra_args = getattr(req.sampling_params, "extra_args", {}) or {} + enable_prompt_rewrite: bool = extra_args.get("enable_prompt_rewrite", True) + enable_cfg_renorm: bool = extra_args.get("enable_cfg_renorm", True) + cfg_renorm_min: float = extra_args.get("cfg_renorm_min", 0.0) + joint_attention_kwargs: dict[str, Any] | None = None + output_type = req.sampling_params.output_type or "pil" self.check_inputs( prompt, @@ -564,12 +540,7 @@ def forward( self._interrupt = False # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = len(prompt) device = self.device if enable_prompt_rewrite and prompt is not None: @@ -688,9 +659,8 @@ def forward( image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) + stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None + return DiffusionOutput(output=image, stage_durations=stage_durations) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights using AutoWeightsLoader for vLLM integration.""" diff --git a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py index 4bfa00a976d..4821425b0cd 100644 --- a/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py +++ b/vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py @@ -6,7 +6,7 @@ import os import re from collections.abc import Iterable -from typing import Any, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -30,13 +30,14 @@ 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 SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import ( LongCatImageTransformer2DModel, ) from vllm_omni.diffusion.models.longcat_image.pipeline_longcat_image import calculate_shift from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -66,50 +67,48 @@ def pre_process_func( request: OmniDiffusionRequest, ): """Pre-process requests for LongCatImageEditPipeline.""" - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if not raw_image: # None or empty list - raise ValueError("""Received no input image. This model requires one input image to run.""") - elif isinstance(raw_image, list): - if len(raw_image) > 1: - raise ValueError( - """Received multiple input images. Only a single image is supported by this model.""" - ) - else: - raw_image = raw_image[0] - - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if not raw_image: # None or empty list + raise ValueError("""Received no input image. This model requires one input image to run.""") + elif isinstance(raw_image, list): + if len(raw_image) > 1: + raise ValueError("""Received multiple input images. Only a single image is supported by this model.""") else: - image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) - - image_size = image.size - calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1]) - height = request.sampling_params.height or calculated_height - width = request.sampling_params.width or calculated_width - - # Store calculated dimensions in request - prompt["additional_information"]["calculated_height"] = calculated_height - prompt["additional_information"]["calculated_width"] = calculated_width - request.sampling_params.height = height - request.sampling_params.width = width - - # Preprocess image - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels): - image = image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image_processor.resize(image, calculated_height // 2, calculated_width // 2) - image = image_processor.preprocess(image, calculated_height, calculated_width) - - # Store preprocessed image and prompt image in request - prompt["additional_information"]["preprocessed_image"] = image - prompt["additional_information"]["prompt_image"] = prompt_image - request.prompts[i] = prompt + raw_image = raw_image[0] + + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image) + else: + image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) + + image_size = image.size + calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1]) + height = request.sampling_params.height or calculated_height + width = request.sampling_params.width or calculated_width + + # Store calculated dimensions in request + prompt["additional_information"]["calculated_height"] = calculated_height + prompt["additional_information"]["calculated_width"] = calculated_width + request.sampling_params.height = height + request.sampling_params.width = width + + # Preprocess image + if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels): + image = image_processor.resize(image, calculated_height, calculated_width) + prompt_image = image_processor.resize(image, calculated_height // 2, calculated_width // 2) + image = image_processor.preprocess(image, calculated_height, calculated_width) + + # Store preprocessed image and prompt image in request + prompt["additional_information"]["preprocessed_image"] = image + prompt["additional_information"]["prompt_image"] = prompt_image + request.prompt = prompt return request return pre_process_func @@ -222,7 +221,13 @@ def split_quotation(prompt, quote_pairs=None): return result -class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin): +class LongCatImageEditPipeline( + nn.Module, CFGParallelMixin, SupportImageInput, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -537,24 +542,7 @@ def check_inputs( f" {negative_prompt_embeds}. Please make sure to only forward one of the two." ) - def forward( - self, - req: OmniDiffusionRequest, - image: PIL.Image.Image | torch.Tensor | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 3.5, - num_images_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - ): + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. if len(req.prompts) > 1: @@ -565,27 +553,22 @@ def forward( first_prompt = req.prompts[0] prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") - prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("prompt_embeds") - negative_prompt_embeds = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt_embeds") # type: ignore # Why it is list[torch.Tensor] in OmniTokenInputs or OmniEmbedsPrompt? Doesn't make sense - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + sigmas = req.sampling_params.sigmas + guidance_scale = req.sampling_params.guidance_scale + num_inference_steps = req.sampling_params.num_inference_steps or 50 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt is not None - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) - generator = req.sampling_params.generator or generator + generator = req.sampling_params.generator height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + output_type = req.sampling_params.output_type or "pil" - if prompt is not None: - batch_size = 1 if isinstance(prompt, str) else len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = 1 if isinstance(prompt, str) else len(prompt) if not isinstance(first_prompt, str) and "preprocessed_image" in ( additional_information := first_prompt.get("additional_information", {}) @@ -595,13 +578,7 @@ def forward( calculated_height = additional_information.get("calculated_height", height) calculated_width = additional_information.get("calculated_width", width) else: - image_size = image[0].size if isinstance(image, list) else image.size - calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] * 1.0 / image_size[1]) - - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = self.image_processor.resize(image, calculated_height, calculated_width) - prompt_image = self.image_processor.resize(image, calculated_height // 2, calculated_width // 2) - image = self.image_processor.preprocess(image, calculated_height, calculated_width) + raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") self.check_inputs( prompt, @@ -616,7 +593,8 @@ def forward( prompt=prompt, image=prompt_image, prompt_embeds=prompt_embeds, num_images_per_prompt=num_images_per_prompt ) - if guidance_scale > 1: + do_true_cfg = guidance_scale > 1 and negative_prompt is not None + if do_true_cfg: (negative_prompt_embeds, negative_text_ids) = self.encode_prompt( prompt=negative_prompt, image=prompt_image, @@ -675,7 +653,6 @@ def forward( latent_model_input = torch.cat([latents, image_latents], dim=1) timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) - do_true_cfg = guidance_scale > 1 positive_kwargs = { "hidden_states": latent_model_input, "timestep": timestep / 1000, diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index bf57d3f9e84..ca34e0390c2 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -1,25 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( - LTX2Pipeline, - LTX2T2VDMD2Pipeline, - LTX2TwoStagesPipeline, +from vllm_omni.diffusion.models.ltx2.ltx2_components import ( create_transformer_from_config, get_ltx2_post_process_func, load_transformer_config, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import ( +from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2I2VDMD2Pipeline, + LTX2ImageToVideoPipeline, + LTX2Pipeline, + LTX2T2VDMD2Pipeline, LTX23ImageToVideoPipeline, LTX23Pipeline, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import ( - LTX2I2VDMD2Pipeline, - LTX2ImageToVideoPipeline, +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( LTX2ImageToVideoTwoStagesPipeline, + LTX2TwoStagesPipeline, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline __all__ = [ "LTX2Pipeline", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py new file mode 100644 index 00000000000..cfe990a4e09 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared component construction helpers for the LTX model family.""" + +from __future__ import annotations + +import inspect +import json +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler +from diffusers.pipelines.ltx2 import LTX2TextConnectors +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder +from diffusers.video_processor import VideoProcessor +from huggingface_hub import hf_hub_download +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video +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.offloader.module_collector import ModuleDiscovery + +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + +from .ltx2_transformer import LTX2VideoTransformer3DModel + +try: + from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE +except ImportError: + LTX2VocoderWithBWE = None + + +_LTX_COMPONENT_SUBFOLDERS = ( + "tokenizer", + "text_encoder", + "connectors", + "vae", + "audio_vae", + "vocoder", + "scheduler", +) + + +@dataclass(frozen=True) +class LTXComponentProfile: + """Component construction and discovery contract for an LTX variant.""" + + name: str + dit_modules: tuple[str, ...] + encoder_modules: tuple[str, ...] + vae_modules: tuple[str, ...] + resident_modules: tuple[str, ...] = () + video_vae_cls: type = AutoencoderKLLTX2Video + vocoder_cls: type = LTX2Vocoder + vocoder_fallback_cls: type | None = None + + +LTX2_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder",), + video_vae_cls=DistributedAutoencoderKLLTX2Video, +) + +LTX23_COMPONENT_PROFILE = LTXComponentProfile( + name="ltx2_3", + dit_modules=("transformer",), + encoder_modules=("text_encoder", "connectors"), + vae_modules=("vae", "audio_vae"), + resident_modules=("vocoder",), + video_vae_cls=DistributedAutoencoderKLLTX2Video, + vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder, + vocoder_fallback_cls=LTX2Vocoder, +) + + +def _detect_vocoder_output_sample_rate(model: str) -> int | None: + """Read the generated waveform sample rate from the vocoder config.""" + vocoder_config_path = os.path.join(model, "vocoder", "config.json") + if not os.path.exists(vocoder_config_path): + try: + vocoder_config_path = hf_hub_download(model, "vocoder/config.json") + except Exception: + return None + try: + with open(vocoder_config_path) as config_file: + return json.load(config_file).get("output_sampling_rate") + except Exception: + return None + + +def get_ltx2_post_process_func(od_config: Any): + """Build the common LTX engine-output adapter.""" + output_sample_rate = _detect_vocoder_output_sample_rate(od_config.model) + + def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): + if not (isinstance(output, tuple) and len(output) == 2): + return output + video, audio = output + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu() + result: dict[str, Any] = {"video": video, "audio": audio} + if output_sample_rate is not None: + result["audio_sample_rate"] = output_sample_rate + return result + + return post_process_func + + +def _load_component( + component_cls: type, + model: str, + subfolder: str, + *, + local_files_only: bool, + dtype: torch.dtype, +) -> Any: + return from_pretrained_with_prefetch( + component_cls.from_pretrained, + model, + subfolder=subfolder, + prefetch_list=_LTX_COMPONENT_SUBFOLDERS, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + + +def _place_aux_components(pipeline: Any) -> None: + parallel_config = getattr(pipeline.od_config, "parallel_config", None) + use_managed_placement = bool( + getattr(pipeline.od_config, "enable_cpu_offload", False) + or getattr(pipeline.od_config, "enable_layerwise_offload", False) + or getattr(parallel_config, "use_hsdp", False) + ) + if use_managed_placement: + return + + modules = ModuleDiscovery.discover(pipeline) + for module in (*modules.encoders, *modules.vaes, *modules.resident_modules): + module.to(pipeline.device) + + +def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None: + """Build the common LTX component graph selected by ``component_profile``.""" + profile: LTXComponentProfile = pipeline.component_profile + pipeline.od_config = od_config + pipeline.device = get_local_device() + dtype = getattr(od_config, "dtype", torch.bfloat16) + model = od_config.model + local_files_only = os.path.exists(model) + + pipeline.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only) + + pipeline.tokenizer = AutoTokenizer.from_pretrained( + model, + subfolder="tokenizer", + local_files_only=local_files_only, + ) + with torch.device("cpu"): + pipeline.text_encoder = _load_component( + Gemma3ForConditionalGeneration, + model, + "text_encoder", + local_files_only=local_files_only, + dtype=dtype, + ) + pipeline.connectors = _load_component( + LTX2TextConnectors, + model, + "connectors", + local_files_only=local_files_only, + dtype=dtype, + ) + pipeline.vae = _load_component( + profile.video_vae_cls, + model, + "vae", + local_files_only=local_files_only, + dtype=dtype, + ) + pipeline.audio_vae = _load_component( + AutoencoderKLLTX2Audio, + model, + "audio_vae", + local_files_only=local_files_only, + dtype=dtype, + ) + try: + pipeline.vocoder = _load_component( + profile.vocoder_cls, + model, + "vocoder", + local_files_only=local_files_only, + dtype=dtype, + ) + except (TypeError, OSError, ValueError): + if profile.vocoder_fallback_cls is None or profile.vocoder_fallback_cls is profile.vocoder_cls: + raise + pipeline.vocoder = _load_component( + profile.vocoder_fallback_cls, + model, + "vocoder", + local_files_only=local_files_only, + dtype=dtype, + ) + + transformer_config = load_transformer_config(model, "transformer", local_files_only) + quant_config = getattr(od_config, "quantization_config", None) + pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config) + _place_aux_components(pipeline) + pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, + subfolder="scheduler", + local_files_only=local_files_only, + ) + + pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio + pipeline.vae_temporal_compression_ratio = pipeline.vae.temporal_compression_ratio + pipeline.audio_vae_mel_compression_ratio = pipeline.audio_vae.mel_compression_ratio + pipeline.audio_vae_temporal_compression_ratio = pipeline.audio_vae.temporal_compression_ratio + pipeline.transformer_spatial_patch_size = pipeline.transformer.config.patch_size + pipeline.transformer_temporal_patch_size = pipeline.transformer.config.patch_size_t + pipeline.audio_sampling_rate = pipeline.audio_vae.config.sample_rate + pipeline.audio_hop_length = pipeline.audio_vae.config.mel_hop_length + pipeline.video_processor = VideoProcessor(vae_scale_factor=pipeline.vae_spatial_compression_ratio) + + tokenizer_max_length = pipeline.tokenizer.model_max_length + if tokenizer_max_length is None or tokenizer_max_length > 100000: + encoder_config = getattr(pipeline.text_encoder, "config", None) + tokenizer_max_length = getattr(encoder_config, "max_position_embeddings", None) + if tokenizer_max_length is None: + tokenizer_max_length = getattr(encoder_config, "max_seq_len", None) + pipeline.tokenizer_max_length = int(tokenizer_max_length or 1024) + + pipeline._guidance_scale = None + pipeline._guidance_rescale = None + pipeline._interrupt = False + + +def load_transformer_config( + model_path: str, + subfolder: str = "transformer", + local_files_only: bool = True, +) -> dict: + """Load an LTX transformer config from a local model or the HF Hub.""" + if local_files_only: + config_path = os.path.join(model_path, subfolder, "config.json") + if os.path.exists(config_path): + with open(config_path) as config_file: + return json.load(config_file) + else: + try: + from huggingface_hub import hf_hub_download + + config_path = hf_hub_download( + repo_id=model_path, + filename=f"{subfolder}/config.json", + ) + with open(config_path) as config_file: + return json.load(config_file) + except Exception: + pass + return {} + + +def create_transformer_from_config( + config: dict, + quant_config: QuantizationConfig | None = None, +) -> LTX2VideoTransformer3DModel: + """Construct the shared LTX transformer from a Diffusers config.""" + if not config and quant_config is None: + return LTX2VideoTransformer3DModel() + + signature = inspect.signature(LTX2VideoTransformer3DModel.__init__) + allowed_keys = set(signature.parameters) + kwargs = {key: value for key, value in config.items() if key in allowed_keys} + if quant_config is not None: + kwargs["quant_config"] = quant_config + + return LTX2VideoTransformer3DModel(**kwargs) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py new file mode 100644 index 00000000000..48d319f6dcf --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -0,0 +1,578 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Task-specific conditioning shared by LTX model versions.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import PIL.Image +import torch +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor + +from . import ltx2_latents as latent_ops +from .ltx2_denoise import I2VVideoAudioScheduler + +if TYPE_CHECKING: + from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext + from .ltx2_request import LTXRequestInputs + + +@dataclass +class LTXPromptContext: + """Connector outputs consumed by an LTX denoise phase.""" + + batch_size: int + connector_prompt_embeds: torch.Tensor + connector_audio_prompt_embeds: torch.Tensor + connector_attention_mask: torch.Tensor + positive_connector_prompt_embeds: torch.Tensor + positive_connector_audio_prompt_embeds: torch.Tensor + positive_connector_attention_mask: torch.Tensor + negative_connector_prompt_embeds: torch.Tensor | None + negative_connector_audio_prompt_embeds: torch.Tensor | None + negative_connector_attention_mask: torch.Tensor | None + + +def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_outputs: int) -> torch.Tensor: + if num_outputs == 1: + return tensor + return tensor.repeat_interleave(num_outputs, dim=0) + + +class LTXTextConditioningMixin: + """Shared Gemma encoding and text connector orchestration.""" + + connector_batches_cfg = False + + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + del scale_factor + device = device or self.device + dtype = dtype or self.text_encoder.dtype + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + text_inputs = self.tokenizer( + [text.strip() for text in prompt], + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(device) + prompt_attention_mask = text_inputs.attention_mask.to(device) + hidden_states = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, + ).hidden_states + + prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) + prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + prompt_attention_mask, + num_videos_per_prompt, + ) + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + device = device or self.device + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) if prompt is not None else prompt_embeds.shape[0] + negative_prompt_embeds_provided = negative_prompt_embeds is not None + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + elif num_videos_per_prompt > 1: + prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) + prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + prompt_attention_mask, + num_videos_per_prompt, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + if isinstance(negative_prompt, list) and batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size" + f" {batch_size}." + ) + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + scale_factor=scale_factor, + device=device, + dtype=dtype, + ) + elif do_classifier_free_guidance and negative_prompt_embeds_provided and num_videos_per_prompt > 1: + negative_prompt_embeds = _repeat_prompt_tensor_for_outputs( + negative_prompt_embeds, + num_videos_per_prompt, + ) + negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + negative_prompt_attention_mask, + num_videos_per_prompt, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + def _prepare_prompt_context( + self, + *, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None, + prompt_embeds: torch.Tensor | None, + negative_prompt_embeds: torch.Tensor | None, + prompt_attention_mask: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None, + num_videos_per_prompt: int, + max_sequence_length: int, + ) -> LTXPromptContext: + if isinstance(prompt, str): + batch_size = 1 + elif isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = ( + self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=self.device, + ) + ) + padding_side = getattr(self.tokenizer, "padding_side", "left") + + if self.do_classifier_free_guidance and self.connector_batches_cfg: + 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) + + video_context, audio_context, attention_mask = self.connectors( + prompt_embeds, + prompt_attention_mask, + padding_side=padding_side, + ) + positive_video_context = video_context + positive_audio_context = audio_context + positive_attention_mask = attention_mask + negative_video_context = None + negative_audio_context = None + negative_attention_mask = None + + if self.do_classifier_free_guidance and self.connector_batches_cfg: + split_batch = batch_size * num_videos_per_prompt + negative_video_context = video_context[:split_batch] + positive_video_context = video_context[split_batch:] + negative_audio_context = audio_context[:split_batch] + positive_audio_context = audio_context[split_batch:] + negative_attention_mask = attention_mask[:split_batch] + positive_attention_mask = attention_mask[split_batch:] + elif self.do_classifier_free_guidance: + negative_video_context, negative_audio_context, negative_attention_mask = self.connectors( + negative_prompt_embeds, + negative_prompt_attention_mask, + padding_side=padding_side, + ) + + return LTXPromptContext( + batch_size=batch_size, + connector_prompt_embeds=video_context, + connector_audio_prompt_embeds=audio_context, + connector_attention_mask=attention_mask, + positive_connector_prompt_embeds=positive_video_context, + positive_connector_audio_prompt_embeds=positive_audio_context, + positive_connector_attention_mask=positive_attention_mask, + negative_connector_prompt_embeds=negative_video_context, + negative_connector_audio_prompt_embeds=negative_audio_context, + negative_connector_attention_mask=negative_attention_mask, + ) + + +class LTXI2VConditioningMixin: + """First-frame conditioning behavior common to LTX2 and LTX2.3.""" + + support_image_input = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.video_processor = VideoProcessor( + vae_scale_factor=self.vae_spatial_compression_ratio, + resample="bilinear", + ) + + def forward(self, req: Any, image: Any | None = None, *args: Any, **kwargs: Any) -> Any: + """Preserve the public I2V positional API while sharing the T2V runner.""" + return super().forward(req, *args, image=image, **kwargs) + + @staticmethod + def _resolve_single_prompt_image(raw_image: Any) -> Any: + if isinstance(raw_image, list): + if len(raw_image) != 1: + raise ValueError( + "LTX I2V prompt dictionaries support exactly one image per prompt. " + "Pass one image per prompt for batched I2V requests." + ) + return raw_image[0] + return raw_image + + @staticmethod + def _resolve_additional_image(additional: dict[str, Any]) -> Any: + for field_name in ("preprocessed_image", "pixel_values", "image"): + raw_image = additional.get(field_name) + if raw_image is not None: + return raw_image + return None + + def _resolve_request_image( + self, + req: Any, + image: Any | None, + request_inputs: LTXRequestInputs, + ) -> Any | None: + if image is not None or not req.prompts: + return image + + raw_images = [] + for prompt_item in req.prompts: + if isinstance(prompt_item, str): + raw_image = None + else: + multi_modal_data = prompt_item.get("multi_modal_data") or {} + raw_image = multi_modal_data.get("image") + if raw_image is None: + raw_image = self._resolve_additional_image(prompt_item.get("additional_information") or {}) + raw_image = self._resolve_single_prompt_image(raw_image) + if isinstance(raw_image, str): + raw_image = PIL.Image.open(raw_image).convert("RGB") + raw_images.append(raw_image) + + if any(raw_image is None for raw_image in raw_images) and request_inputs.latents is None: + raise ValueError("Image is required for LTX I2V generation.") + if len(raw_images) == 1: + return raw_images[0] + return raw_images or image + + def _check_forward_inputs( + self, + request_inputs: LTXRequestInputs, + image: Any | None = None, + ) -> None: + if image is None and request_inputs.latents is None: + raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.") + super()._check_forward_inputs(request_inputs, image=image) + + def prepare_latents( + self, + image: torch.Tensor | None = None, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + num_frames, height, width = latent_ops.resolve_video_latent_shape( + height, + width, + num_frames, + vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, + vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, + ) + shape = (batch_size, num_channels_latents, num_frames, height, width) + mask_shape = (batch_size, 1, num_frames, height, width) + + if latents is not None: + if latents.ndim == 5: + batch_size, _, num_frames, height, width = latents.shape + mask_shape = (batch_size, 1, num_frames, height, width) + conditioning_mask = latents.new_zeros(mask_shape) + conditioning_mask[:, :, 0] = 1.0 + latents = latent_ops.normalize_latents( + latents, + self.vae.latents_mean, + self.vae.latents_std, + self.vae.config.scaling_factor, + ) + latents = latent_ops.create_noised_state( + latents, + noise_scale * (1 - conditioning_mask), + generator, + ) + latents = latent_ops.pack_latents( + latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + else: + conditioning_mask = latents.new_zeros(mask_shape) + conditioning_mask[:, :, 0] = 1.0 + + conditioning_mask = latent_ops.pack_latents( + conditioning_mask, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ).squeeze(-1) + if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape: + raise ValueError( + "Provided `latents` tensor has shape" + f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." + ) + return latents.to(device=device, dtype=dtype), conditioning_mask + + if image is None: + raise ValueError("`image` must be provided when `latents` is None.") + + image_batch_size = image.shape[0] + if image_batch_size == 0: + raise ValueError("`image` batch is empty.") + if batch_size % image_batch_size != 0: + raise ValueError( + f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) " + "for image-to-video outputs." + ) + num_videos_per_prompt = batch_size // image_batch_size + + if isinstance(generator, list): + if len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective" + f" batch size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)] + init_latents = [ + retrieve_latents( + self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), + image_generators[i], + "argmax", + ) + for i in range(image_batch_size) + ] + else: + init_latents = [ + retrieve_latents( + self.vae.encode(img.unsqueeze(0).unsqueeze(2)), + generator, + "argmax", + ) + for img in image + ] + + init_latents = torch.cat(init_latents, dim=0).to(dtype) + if num_videos_per_prompt > 1: + init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0) + init_latents = latent_ops.normalize_latents( + init_latents, + self.vae.latents_mean, + self.vae.latents_std, + ) + init_latents = init_latents.repeat(1, 1, num_frames, 1, 1) + + conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype) + conditioning_mask[:, :, 0] = 1.0 + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask) + + conditioning_mask = latent_ops.pack_latents( + conditioning_mask, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ).squeeze(-1) + latents = latent_ops.pack_latents( + latents, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + return latents, conditioning_mask + + def _step_video_latents_i2v( + self, + noise_pred_video: torch.Tensor, + latents: torch.Tensor, + timestep: torch.Tensor, + latent_num_frames: int, + latent_height: int, + latent_width: int, + ) -> torch.Tensor: + noise_pred_video = latent_ops.unpack_latents( + noise_pred_video, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents_unpacked = latent_ops.unpack_latents( + latents, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + noise_pred_video = noise_pred_video[:, :, 1:] + noise_latents = latents_unpacked[:, :, 1:] + pred_latents = self.scheduler.step( + noise_pred_video, + timestep, + noise_latents, + return_dict=False, + )[0] + latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2) + return latent_ops.pack_latents( + latents_unpacked, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + + def _prepare_video_latents_stage( + self, + request_inputs: LTXRequestInputs, + prompt_context: LTXPromptContext, + *, + device: torch.device, + noise_scale: float, + image: Any | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if request_inputs.latents is None: + if isinstance(image, torch.Tensor): + if image.ndim == 3: + image = image.unsqueeze(0) + elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor): + image = torch.stack(image, dim=0) + else: + image = self.video_processor.preprocess( + image, + height=request_inputs.height, + width=request_inputs.width, + ) + image = image.to( + device=device, + dtype=prompt_context.positive_connector_prompt_embeds.dtype, + ) + + return self.prepare_latents( + image, + prompt_context.batch_size * request_inputs.num_videos_per_prompt, + self.transformer.config.in_channels, + request_inputs.height, + request_inputs.width, + request_inputs.num_frames, + noise_scale, + torch.float32, + device, + request_inputs.generator, + request_inputs.latents, + ) + + def _make_video_audio_scheduler( + self, + audio_scheduler: Any, + latent_num_frames: int, + latent_height: int, + latent_width: int, + ) -> Any: + return I2VVideoAudioScheduler( + self, + audio_scheduler, + latent_num_frames, + latent_height, + latent_width, + ) + + def _prepare_denoise_context_for_cfg( + self, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> LTXDenoiseContext: + denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + if denoise_ctx.conditioning_mask is None: + raise ValueError("LTX I2V denoising requires a conditioning mask.") + + mask_batch = denoise_ctx.conditioning_mask.shape[0] + model_batch = denoise_ctx.video_coords.shape[0] + if model_batch % mask_batch != 0: + raise ValueError( + "I2V conditioning-mask batch must divide the Transformer input batch, " + f"but got {mask_batch} and {model_batch}." + ) + repeats = model_batch // mask_batch + denoise_ctx.conditioning_mask_for_model = ( + denoise_ctx.conditioning_mask if repeats == 1 else torch.cat([denoise_ctx.conditioning_mask] * repeats) + ) + return denoise_ctx + + def _denoise_timestep_kwargs( + self, + ts: torch.Tensor, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> dict[str, torch.Tensor]: + kwargs = super()._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx) + conditioning_mask = ( + denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model + ) + if conditioning_mask is None: + raise ValueError("LTX I2V denoising requires a conditioning mask.") + kwargs.update( + timestep=ts.unsqueeze(-1) * (1 - conditioning_mask), + audio_timestep=ts, + ) + return kwargs diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py new file mode 100644 index 00000000000..af8df4b4182 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -0,0 +1,407 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared denoise execution primitives for LTX pipelines.""" + +from __future__ import annotations + +import copy +from collections.abc import Callable, Iterable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +import numpy as np +import torch +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps + +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +from .ltx2_latents import LTXAVState + +if TYPE_CHECKING: + from .ltx2_conditioning import LTXPromptContext + from .ltx2_request import LTXRequestInputs + + +@dataclass +class LTXForwardContext: + """Immutable metadata and schedulers for one LTX denoise phase.""" + + req: DiffusionRequestBatch + request_inputs: LTXRequestInputs + prompt_context: LTXPromptContext + device: torch.device + cfg_parallel_ready: bool + attention_kwargs: dict[str, Any] | None + latent_num_frames: int + latent_height: int + latent_width: int + latent_mel_bins: int + original_audio_num_frames: int + padded_audio_num_frames: int + timesteps: torch.Tensor + audio_scheduler: Any + video_audio_scheduler: Any + + @property + def batch_size(self) -> int: + return self.prompt_context.batch_size + + @property + def num_videos_per_prompt(self) -> int: + return self.request_inputs.num_videos_per_prompt + + +@dataclass +class LTXDenoiseContext: + """Mutable AV state and positional metadata for a denoise phase.""" + + latents: torch.Tensor + audio_latents: torch.Tensor + video_coords: torch.Tensor + audio_coords: torch.Tensor + conditioning_mask: torch.Tensor | None = None + conditioning_mask_for_model: torch.Tensor | None = None + + +@dataclass +class LTXPhaseResult: + """Denoised, unpacked AV latents and the context used to produce them.""" + + forward_context: LTXForwardContext + video: torch.Tensor + audio: torch.Tensor + + +class LTXDenoisePipeline(Protocol): + """Pipeline state required by :class:`LTXDenoiseExecutor`.""" + + @property + def interrupt(self) -> bool: ... + + def progress_bar(self, iterable=None, total=None): ... + + +LTXDenoiseStep = Callable[[int, torch.Tensor, LTXAVState], LTXAVState] + + +class LTXDenoiseExecutor: + """Run the one shared LTX denoise loop. + + Prediction and scheduler math remain injectable so structural refactors do + not change the existing LTX2/LTX2.3 numerical paths. Guidance will replace + that step policy independently. + """ + + @staticmethod + def run( + pipeline: LTXDenoisePipeline, + state: LTXAVState, + timesteps: Iterable[torch.Tensor], + step: LTXDenoiseStep, + ) -> LTXAVState: + timesteps = tuple(timesteps) + with pipeline.progress_bar(total=len(timesteps)) as progress_bar: + for index, timestep in enumerate(timesteps): + if pipeline.interrupt: + continue + state = step(index, timestep, state) + progress_bar.update() + return state + + +def calculate_shift( + image_seq_len: int, + base_seq_len: int = 256, + max_seq_len: int = 4096, + base_shift: float = 0.5, + max_shift: float = 1.15, +) -> float: + slope = (max_shift - base_shift) / (max_seq_len - base_seq_len) + intercept = base_shift - slope * base_seq_len + return image_seq_len * slope + intercept + + +class VideoAudioScheduler: + """Composite scheduler dispatching video and audio updates.""" + + def __init__(self, video_scheduler, audio_scheduler): + self.video_scheduler = video_scheduler + self.audio_scheduler = audio_scheduler + + def step(self, noise_pred, t, latents, return_dict=False, generator=None): + video_out = self.video_scheduler.step( + noise_pred[0], + t[0], + latents[0], + return_dict=False, + generator=generator, + )[0] + audio_out = self.audio_scheduler.step( + noise_pred[1], + t[1], + latents[1], + return_dict=False, + generator=generator, + )[0] + return ((video_out, audio_out),) + + +class I2VVideoAudioScheduler: + """Update the unconditioned video frames and the full audio state.""" + + def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width): + self.video_scheduler = pipeline.scheduler + self.audio_scheduler = audio_scheduler + self._pipeline = pipeline + self._latent_num_frames = latent_num_frames + self._latent_height = latent_height + self._latent_width = latent_width + + def step(self, noise_pred, t, latents, return_dict=False, generator=None): + video_out = self._pipeline._step_video_latents_i2v( + noise_pred[0], + latents[0], + t[0], + self._latent_num_frames, + self._latent_height, + self._latent_width, + ) + audio_out = self.audio_scheduler.step( + noise_pred[1], + t[1], + latents[1], + return_dict=False, + generator=generator, + )[0] + return ((video_out, audio_out),) + + +def prepare_scheduler_stage( + pipeline: Any, + request_inputs: LTXRequestInputs, + *, + device: torch.device, + sigmas: list[float] | None, + timesteps: list[int] | None, + latent_num_frames: int, + latent_height: int, + latent_width: int, +) -> tuple[Any, Any, torch.Tensor]: + sigmas = ( + np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps) + if sigmas is None + else sigmas + ) + mu = calculate_shift( + pipeline.scheduler.config.get("max_image_seq_len", 4096), + pipeline.scheduler.config.get("base_image_seq_len", 1024), + pipeline.scheduler.config.get("max_image_seq_len", 4096), + pipeline.scheduler.config.get("base_shift", 0.95), + pipeline.scheduler.config.get("max_shift", 2.05), + ) + audio_scheduler = copy.deepcopy(pipeline.scheduler) + video_audio_scheduler = pipeline._make_video_audio_scheduler( + audio_scheduler, + latent_num_frames, + latent_height, + latent_width, + ) + retrieve_timesteps( + audio_scheduler, + request_inputs.num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + timesteps_tensor, _ = retrieve_timesteps( + pipeline.scheduler, + request_inputs.num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + return audio_scheduler, video_audio_scheduler, timesteps_tensor + + +def prepare_rope_coords_stage( + pipeline: Any, + forward_ctx: LTXForwardContext, + latents: torch.Tensor, + audio_latents: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + video_coords = pipeline.transformer.rope.prepare_video_coords( + latents.shape[0], + forward_ctx.latent_num_frames, + forward_ctx.latent_height, + forward_ctx.latent_width, + latents.device, + fps=forward_ctx.request_inputs.frame_rate, + ) + audio_coords = pipeline.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], + forward_ctx.padded_audio_num_frames, + audio_latents.device, + ) + return video_coords, audio_coords + + +def build_transformer_kwargs( + pipeline: Any, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + *, + hidden_states: torch.Tensor, + audio_hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None, + audio_encoder_attention_mask: torch.Tensor | None, + ts: torch.Tensor, +) -> dict[str, Any]: + return { + "hidden_states": hidden_states, + "audio_hidden_states": audio_hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "audio_encoder_hidden_states": audio_encoder_hidden_states, + **pipeline._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx), + "encoder_attention_mask": encoder_attention_mask, + "audio_encoder_attention_mask": audio_encoder_attention_mask, + "num_frames": forward_ctx.latent_num_frames, + "height": forward_ctx.latent_height, + "width": forward_ctx.latent_width, + "fps": forward_ctx.request_inputs.frame_rate, + "audio_num_frames": forward_ctx.padded_audio_num_frames, + "video_coords": denoise_ctx.video_coords, + "audio_coords": denoise_ctx.audio_coords, + "attention_kwargs": forward_ctx.attention_kwargs, + "return_dict": False, + } + + +def step_denoised_latents( + pipeline: Any, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + noise_pred_video: torch.Tensor, + noise_pred_audio: torch.Tensor, + timestep: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + latents = pipeline.scheduler_step_maybe_with_cfg( + (noise_pred_video, noise_pred_audio), + (timestep, timestep), + (denoise_ctx.latents, denoise_ctx.audio_latents), + do_true_cfg=pipeline.do_classifier_free_guidance, + per_request_scheduler=forward_ctx.video_audio_scheduler, + ) + return pipeline._synchronize_cfg_parallel_step_output( + latents, + do_true_cfg=pipeline.do_classifier_free_guidance, + ) + + +class LTXPhaseExecutor: + """Prepare and execute one LTX phase without owning model modules.""" + + @staticmethod + def run( + pipeline: Any, + req: DiffusionRequestBatch, + request_inputs: LTXRequestInputs, + *, + noise_scale: float, + sigmas: list[float] | None, + timesteps: list[int] | None, + attention_kwargs: dict[str, Any] | None, + image: Any | None = None, + prompt_context: LTXPromptContext | None = None, + ) -> LTXPhaseResult: + pipeline._check_forward_inputs(request_inputs, image=image) + cfg_parallel_ready = pipeline._setup_forward_runtime(request_inputs, attention_kwargs) + device = pipeline.device + if prompt_context is None: + prompt_context = pipeline._prepare_prompt_context( + prompt=request_inputs.prompt, + negative_prompt=request_inputs.negative_prompt, + prompt_embeds=request_inputs.prompt_embeds, + negative_prompt_embeds=request_inputs.negative_prompt_embeds, + prompt_attention_mask=request_inputs.prompt_attention_mask, + negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask, + num_videos_per_prompt=request_inputs.num_videos_per_prompt, + max_sequence_length=request_inputs.max_sequence_length, + ) + + latent_num_frames, latent_height, latent_width = pipeline._resolve_video_latent_dimensions(request_inputs) + latents, conditioning_mask = pipeline._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=device, + noise_scale=noise_scale, + image=image, + ) + audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins = ( + pipeline._prepare_audio_latents_stage( + request_inputs, + prompt_context, + device=device, + noise_scale=noise_scale, + ) + ) + audio_scheduler, video_audio_scheduler, timesteps_tensor = prepare_scheduler_stage( + pipeline, + request_inputs, + device=device, + sigmas=sigmas, + timesteps=timesteps, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + ) + forward_ctx = LTXForwardContext( + req=req, + request_inputs=request_inputs, + prompt_context=prompt_context, + device=device, + cfg_parallel_ready=cfg_parallel_ready, + attention_kwargs=attention_kwargs, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + latent_mel_bins=latent_mel_bins, + original_audio_num_frames=original_audio_num_frames, + padded_audio_num_frames=padded_audio_num_frames, + timesteps=timesteps_tensor, + audio_scheduler=audio_scheduler, + video_audio_scheduler=video_audio_scheduler, + ) + video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents) + denoise_ctx = LTXDenoiseContext( + latents=latents, + audio_latents=audio_latents, + video_coords=video_coords, + audio_coords=audio_coords, + conditioning_mask=conditioning_mask, + ) + denoise_ctx = pipeline._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + state = LTXDenoiseExecutor.run( + pipeline, + LTXAVState(video=denoise_ctx.latents, audio=denoise_ctx.audio_latents), + forward_ctx.timesteps, + lambda index, timestep, state: pipeline._denoise_step( + index, + timestep, + state, + forward_ctx, + denoise_ctx, + ), + ) + denoise_ctx.latents = state.video + denoise_ctx.audio_latents = state.audio + latents, audio_latents = pipeline._unpack_and_denormalize_stage( + forward_ctx, + state.video, + state.audio, + ) + return LTXPhaseResult(forward_context=forward_ctx, video=latents, audio=audio_latents) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py new file mode 100644 index 00000000000..395f8e7af73 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Guidance strategies for the LTX model family.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol + +import torch +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg + +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, +) + +if TYPE_CHECKING: + from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext + from .ltx2_latents import LTXAVState + + +class LTXGuidanceStrategy(Protocol): + """Model-version guidance behavior consumed by the shared pipeline.""" + + def validate_cfg_world_size(self, cfg_world_size: int) -> None: ... + + def prepare_denoise_context( + self, + pipeline: Any, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> LTXDenoiseContext: ... + + def timestep_kwargs( + self, + ts: torch.Tensor, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> dict[str, torch.Tensor]: ... + + def combine_cfg_noise( + self, + pipeline: Any, + positive_noise_pred: tuple[torch.Tensor, torch.Tensor], + negative_noise_pred: tuple[torch.Tensor, torch.Tensor], + true_cfg_scale: float, + cfg_normalize: bool, + context: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + def predict_noise( + self, + pipeline: Any, + index: int, + timestep: torch.Tensor, + state: LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + +class _LTXGuidanceStrategyBase: + def validate_cfg_world_size(self, cfg_world_size: int) -> None: + del cfg_world_size + + def prepare_denoise_context( + self, + pipeline: Any, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> LTXDenoiseContext: + del pipeline, forward_ctx + return denoise_ctx + + def timestep_kwargs( + self, + ts: torch.Tensor, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> dict[str, torch.Tensor]: + del forward_ctx, denoise_ctx + return {"timestep": ts} + + +def combine_velocity_via_x0( + sample: torch.Tensor, + positive_velocity: torch.Tensor, + negative_velocity: torch.Tensor, + sigma: torch.Tensor, + guidance_scale: float, +) -> torch.Tensor: + """Apply CFG in x0 space and convert the result back to velocity.""" + x0_cond = sample - positive_velocity * sigma + x0_uncond = sample - negative_velocity * sigma + x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond) + return (sample - x0_guided) / sigma + + +class LTXLegacyVelocityGuidance(_LTXGuidanceStrategyBase): + """Existing LTX2 Diffusers-style velocity-space CFG behavior.""" + + def combine_cfg_noise( + self, + pipeline: Any, + positive_noise_pred: tuple[torch.Tensor, torch.Tensor], + negative_noise_pred: tuple[torch.Tensor, torch.Tensor], + true_cfg_scale: float, + cfg_normalize: bool, + context: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + del context + video_pos, audio_pos = positive_noise_pred + video_neg, audio_neg = negative_noise_pred + video_combined = CFGParallelMixin.combine_cfg_noise( + pipeline, + video_pos, + video_neg, + true_cfg_scale, + cfg_normalize, + ) + audio_combined = CFGParallelMixin.combine_cfg_noise( + pipeline, + audio_pos, + audio_neg, + true_cfg_scale, + cfg_normalize, + ) + if pipeline.guidance_rescale and pipeline.guidance_rescale > 0: + video_combined = rescale_noise_cfg( + video_combined, + video_pos, + guidance_rescale=pipeline.guidance_rescale, + ) + audio_combined = rescale_noise_cfg( + audio_combined, + audio_pos, + guidance_rescale=pipeline.guidance_rescale, + ) + return video_combined, audio_combined + + def predict_noise( + self, + pipeline: Any, + index: int, + timestep: torch.Tensor, + state: LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> tuple[torch.Tensor, torch.Tensor]: + del index + prompt_context = forward_ctx.prompt_context + video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype) + audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype) + expanded_timestep = timestep.expand(video_input.shape[0]) + positive_kwargs = pipeline._build_transformer_kwargs( + forward_ctx, + denoise_ctx, + hidden_states=video_input, + audio_hidden_states=audio_input, + encoder_hidden_states=prompt_context.positive_connector_prompt_embeds, + audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds, + encoder_attention_mask=prompt_context.positive_connector_attention_mask, + audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask, + ts=expanded_timestep, + ) + negative_kwargs = ( + { + **positive_kwargs, + "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds, + "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds, + "encoder_attention_mask": prompt_context.negative_connector_attention_mask, + "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask, + } + if pipeline.do_classifier_free_guidance + else None + ) + return pipeline.predict_noise_maybe_with_cfg( + do_true_cfg=pipeline.do_classifier_free_guidance, + true_cfg_scale=forward_ctx.request_inputs.guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + ) + + +class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase): + """Official LTX x0-space CFG behavior currently used by LTX2.3.""" + + def validate_cfg_world_size(self, cfg_world_size: int) -> None: + if cfg_world_size not in (1, 2): + raise ValueError( + f"LTX x0 guidance supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." + ) + + def prepare_denoise_context( + self, + pipeline: Any, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> LTXDenoiseContext: + if pipeline.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready: + denoise_ctx.video_coords = denoise_ctx.video_coords.repeat( + (2,) + (1,) * (denoise_ctx.video_coords.ndim - 1) + ) + denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat( + (2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1) + ) + return denoise_ctx + + def timestep_kwargs( + self, + ts: torch.Tensor, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> dict[str, torch.Tensor]: + del forward_ctx, denoise_ctx + return {"timestep": ts, "sigma": ts} + + def combine_cfg_noise( + self, + pipeline: Any, + positive_noise_pred: tuple[torch.Tensor, torch.Tensor], + negative_noise_pred: tuple[torch.Tensor, torch.Tensor], + true_cfg_scale: float, + cfg_normalize: bool, + context: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor]: + required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma") + if any(context.get(name) is None for name in required): + raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.") + + video_pos, audio_pos = positive_noise_pred + video_neg, audio_neg = negative_noise_pred + video_combined = combine_velocity_via_x0( + context["video_latents"], + video_pos, + video_neg, + context["video_sigma"], + true_cfg_scale, + ) + audio_combined = combine_velocity_via_x0( + context["audio_latents"], + audio_pos, + audio_neg, + context["audio_sigma"], + true_cfg_scale, + ) + if cfg_normalize: + video_combined = pipeline.cfg_normalize_function(video_pos, video_combined) + audio_combined = pipeline.cfg_normalize_function(audio_pos, audio_combined) + return video_combined, audio_combined + + def predict_parallel_cfg( + self, + pipeline: Any, + true_cfg_scale: float, + positive_kwargs: dict[str, Any], + negative_kwargs: dict[str, Any], + cfg_normalize: bool = True, + output_slice: int | None = None, + **context: Any, + ) -> tuple[torch.Tensor, torch.Tensor]: + def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + if output_slice is None: + return pred + return pred[0][:, :output_slice], pred[1][:, :output_slice] + + cfg_world_size = get_classifier_free_guidance_world_size() + self.validate_cfg_world_size(cfg_world_size) + if cfg_world_size != 2: + raise ValueError(f"Parallel LTX x0 guidance requires cfg_parallel_size 2, but got {cfg_world_size}.") + + cfg_rank = get_classifier_free_guidance_rank() + branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs + local_video_pred, local_audio_pred = maybe_slice(pipeline.predict_noise(**branch_kwargs)) + cfg_group = get_cfg_group() + gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True) + gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True) + return self.combine_cfg_noise( + pipeline, + (gathered_video[0], gathered_audio[0]), + (gathered_video[1], gathered_audio[1]), + true_cfg_scale, + cfg_normalize, + context, + ) + + def predict_noise( + self, + pipeline: Any, + index: int, + timestep: torch.Tensor, + state: LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> tuple[torch.Tensor, torch.Tensor]: + prompt_context = forward_ctx.prompt_context + guidance_scale = forward_ctx.request_inputs.guidance_scale + audio_scheduler = forward_ctx.audio_scheduler + if forward_ctx.cfg_parallel_ready: + video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype) + audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype) + ts = timestep.expand(video_input.shape[0]) + positive_kwargs = pipeline._build_transformer_kwargs( + forward_ctx, + denoise_ctx, + hidden_states=video_input, + audio_hidden_states=audio_input, + encoder_hidden_states=prompt_context.positive_connector_prompt_embeds, + audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds, + encoder_attention_mask=prompt_context.positive_connector_attention_mask, + audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask, + ts=ts, + ) + negative_kwargs = { + **positive_kwargs, + "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds, + "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds, + "encoder_attention_mask": prompt_context.negative_connector_attention_mask, + "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask, + } + return self.predict_parallel_cfg( + pipeline, + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + video_latents=state.video, + audio_latents=state.audio, + video_sigma=pipeline.scheduler.sigmas[index], + audio_sigma=audio_scheduler.sigmas[index], + ) + + video_input = torch.cat([state.video] * 2) if pipeline.do_classifier_free_guidance else state.video + video_input = video_input.to(prompt_context.connector_prompt_embeds.dtype) + audio_input = torch.cat([state.audio] * 2) if pipeline.do_classifier_free_guidance else state.audio + audio_input = audio_input.to(prompt_context.connector_prompt_embeds.dtype) + ts = timestep.expand(video_input.shape[0]) + transformer_kwargs = pipeline._build_transformer_kwargs( + forward_ctx, + denoise_ctx, + hidden_states=video_input, + audio_hidden_states=audio_input, + encoder_hidden_states=prompt_context.connector_prompt_embeds, + audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds, + encoder_attention_mask=prompt_context.connector_attention_mask, + audio_encoder_attention_mask=prompt_context.connector_attention_mask, + ts=ts, + ) + with pipeline._transformer_cache_context("cond_uncond"): + video_pred, audio_pred = pipeline.transformer(**transformer_kwargs) + video_pred = video_pred.float() + audio_pred = audio_pred.float() + + if pipeline.do_classifier_free_guidance: + video_uncond, video_cond = video_pred.chunk(2) + video_pred = combine_velocity_via_x0( + state.video, + video_cond, + video_uncond, + pipeline.scheduler.sigmas[index], + guidance_scale, + ) + audio_uncond, audio_cond = audio_pred.chunk(2) + audio_pred = combine_velocity_via_x0( + state.audio, + audio_cond, + audio_uncond, + audio_scheduler.sigmas[index], + guidance_scale, + ) + return video_pred, audio_pred + + +LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance() +LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py new file mode 100644 index 00000000000..4da3edf9205 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -0,0 +1,293 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared latent layout and normalization primitives for LTX pipelines.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +from diffusers.utils.torch_utils import randn_tensor + + +@dataclass +class LTXAVState: + """Packed video and audio latents carried between denoise steps.""" + + video: torch.Tensor + audio: torch.Tensor + + +def pack_latents( + latents: torch.Tensor, + patch_size: int = 1, + patch_size_t: int = 1, +) -> torch.Tensor: + batch_size, _, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + + +def unpack_latents( + latents: torch.Tensor, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, +) -> torch.Tensor: + batch_size = latents.size(0) + latents = latents.reshape( + batch_size, + num_frames, + height, + width, + -1, + patch_size_t, + patch_size, + patch_size, + ) + return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + + +def normalize_latents( + latents: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + scaling_factor: float = 1.0, +) -> torch.Tensor: + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + return (latents - latents_mean) * scaling_factor / latents_std + + +def normalize_audio_latents( + latents: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, +) -> torch.Tensor: + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) / latents_std + + +def denormalize_latents( + latents: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, + scaling_factor: float = 1.0, +) -> torch.Tensor: + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + return latents * latents_std / scaling_factor + latents_mean + + +def denormalize_audio_latents( + latents: torch.Tensor, + latents_mean: torch.Tensor, + latents_std: torch.Tensor, +) -> torch.Tensor: + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return latents * latents_std + latents_mean + + +def create_noised_state( + latents: torch.Tensor, + noise_scale: float | torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None = None, +) -> torch.Tensor: + noise = randn_tensor( + latents.shape, + generator=generator, + device=latents.device, + dtype=latents.dtype, + ) + return noise_scale * noise + (1 - noise_scale) * latents + + +def pack_audio_latents( + latents: torch.Tensor, + patch_size: int | None = None, + patch_size_t: int | None = None, +) -> torch.Tensor: + if patch_size is not None and patch_size_t is not None: + batch_size, _, latent_length, latent_mel_bins = latents.shape + post_patch_latent_length = latent_length / patch_size_t + post_patch_mel_bins = latent_mel_bins / patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_latent_length, + patch_size_t, + post_patch_mel_bins, + patch_size, + ) + return latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + return latents.transpose(1, 2).flatten(2, 3) + + +def unpack_audio_latents( + latents: torch.Tensor, + latent_length: int, + num_mel_bins: int, + patch_size: int | None = None, + patch_size_t: int | None = None, +) -> torch.Tensor: + if patch_size is not None and patch_size_t is not None: + batch_size = latents.size(0) + latents = latents.reshape( + batch_size, + latent_length, + num_mel_bins, + -1, + patch_size_t, + patch_size, + ) + return latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + return latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) + + +def unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: + return latents[:, :num_frames] + + +def get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int: + if sp_size > 1: + audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size + return audio_latent_length + + +def resolve_video_latent_shape( + height: int, + width: int, + num_frames: int, + *, + vae_spatial_compression_ratio: int, + vae_temporal_compression_ratio: int, +) -> tuple[int, int, int]: + return ( + (num_frames - 1) // vae_temporal_compression_ratio + 1, + height // vae_spatial_compression_ratio, + width // vae_spatial_compression_ratio, + ) + + +def prepare_video_latents( + pipeline: Any, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, +) -> torch.Tensor: + if latents is not None: + if latents.ndim == 5: + latents = normalize_latents( + latents, + pipeline.vae.latents_mean, + pipeline.vae.latents_std, + pipeline.vae.config.scaling_factor, + ) + latents = pack_latents( + latents, + pipeline.transformer_spatial_patch_size, + pipeline.transformer_temporal_patch_size, + ) + if latents.ndim != 3: + raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") + return create_noised_state(latents, noise_scale, generator).to(device=device, dtype=dtype) + + num_frames, height, width = resolve_video_latent_shape( + height, + width, + num_frames, + vae_spatial_compression_ratio=pipeline.vae_spatial_compression_ratio, + vae_temporal_compression_ratio=pipeline.vae_temporal_compression_ratio, + ) + shape = (batch_size, num_channels_latents, num_frames, height, width) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return pack_latents( + latents, + pipeline.transformer_spatial_patch_size, + pipeline.transformer_temporal_patch_size, + ) + + +def prepare_audio_latents( + pipeline: Any, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, +) -> tuple[torch.Tensor, int, int]: + original_latent_length = audio_latent_length + latent_mel_bins = num_mel_bins // pipeline.audio_vae_mel_compression_ratio + sp_size = getattr(pipeline.od_config.parallel_config, "sequence_parallel_size", 1) or 1 + padded_latent_length = get_sp_padded_audio_latent_length(original_latent_length, int(sp_size)) + + if latents is not None: + if latents.ndim == 4: + latents = pack_audio_latents(latents) + if latents.ndim != 3: + raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") + latents = normalize_audio_latents( + latents, + pipeline.audio_vae.latents_mean, + pipeline.audio_vae.latents_std, + ) + latents = create_noised_state(latents, noise_scale, generator) + + if latents.shape[1] not in {original_latent_length, padded_latent_length}: + raise ValueError( + "Provided `audio_latents` has incompatible audio frame count " + f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}." + ) + if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length: + padding = torch.zeros( + latents.shape[0], + padded_latent_length - original_latent_length, + latents.shape[2], + dtype=latents.dtype, + device=latents.device, + ) + latents = torch.cat([latents, padding], dim=1) + return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length + + shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return pack_audio_latents(latents), original_latent_length, padded_latent_length diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py new file mode 100644 index 00000000000..08bf2033039 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py @@ -0,0 +1,619 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Shared runtime surface for LTX pipeline variants.""" + +from __future__ import annotations + +from collections.abc import Iterable +from contextlib import nullcontext +from typing import Any, ClassVar + +import torch +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request + +from . import ltx2_latents as latent_ops +from .ltx2_components import LTXComponentProfile, initialize_pipeline_components +from .ltx2_conditioning import LTXPromptContext, LTXTextConditioningMixin +from .ltx2_denoise import ( + LTXDenoiseContext, + LTXForwardContext, + LTXPhaseExecutor, + LTXPhaseResult, + VideoAudioScheduler, + build_transformer_kwargs, + step_denoised_latents, +) +from .ltx2_guidance import LTXGuidanceStrategy +from .ltx2_recipes import LTXOneStageRecipe +from .ltx2_request import LTXRequestInputs, LTXRequestMixin + + +def _expand_per_prompt_decode_value( + value: float | list[float], + *, + prompt_batch_size: int, + effective_batch_size: int, + field_name: str, +) -> list[float]: + if not isinstance(value, list): + return [value] * effective_batch_size + if len(value) == 1: + return value * effective_batch_size + if len(value) == effective_batch_size: + return value + if prompt_batch_size > 0 and len(value) == prompt_batch_size and effective_batch_size % prompt_batch_size == 0: + repeats = effective_batch_size // prompt_batch_size + return [item for item in value for _ in range(repeats)] + raise ValueError( + f"`{field_name}` must have length 1, prompt batch size ({prompt_batch_size}), or effective batch size" + f" ({effective_batch_size}); got {len(value)}." + ) + + +def _prepare_decode_timestep_conditioning( + *, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + prompt_batch_size: int, + effective_batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + decode_timestep_values = _expand_per_prompt_decode_value( + decode_timestep, + prompt_batch_size=prompt_batch_size, + effective_batch_size=effective_batch_size, + field_name="decode_timestep", + ) + decode_noise_scale_values = ( + decode_timestep_values + if decode_noise_scale is None + else _expand_per_prompt_decode_value( + decode_noise_scale, + prompt_batch_size=prompt_batch_size, + effective_batch_size=effective_batch_size, + field_name="decode_noise_scale", + ) + ) + return ( + torch.tensor(decode_timestep_values, device=device, dtype=dtype), + torch.tensor(decode_noise_scale_values, device=device, dtype=dtype)[:, None, None, None, None], + ) + + +class LTXPipelineRuntime( + LTXRequestMixin, + LTXTextConditioningMixin, + nn.Module, + CFGParallelMixin, + ProgressBarMixin, + SupportsComponentDiscovery, + DiffusionPipelineProfilerMixin, +): + """Shared Omni runtime for explicitly composed LTX denoise phases.""" + + component_profile: ClassVar[LTXComponentProfile] + guidance_strategy: ClassVar[LTXGuidanceStrategy] + one_stage_recipe: ClassVar[LTXOneStageRecipe] + supports_request_batch = False + connector_batches_cfg = False + distributed_video_decode = True + support_image_input = False + dummy_run_num_frames = 2 + preserve_sp_padded_audio_duration = False + reports_stage_durations = False + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: + del prefix + super().__init__() + initialize_pipeline_components(self, od_config) + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def prepare_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + return latent_ops.prepare_video_latents( + self, + batch_size, + num_channels_latents, + height, + width, + num_frames, + noise_scale, + dtype, + device, + generator, + latents, + ) + + def prepare_audio_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, int, int]: + return latent_ops.prepare_audio_latents( + self, + batch_size, + num_channels_latents, + audio_latent_length, + num_mel_bins, + noise_scale, + dtype, + device, + generator, + latents, + ) + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def guidance_rescale(self): + return self._guidance_rescale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale is not None and self._guidance_scale > 1.0 + + @property + def interrupt(self): + return self._interrupt + + def _transformer_cache_context(self, context_name: str): + cache_context = getattr(self.transformer, "cache_context", None) + if callable(cache_context): + return cache_context(context_name) + return nullcontext() + + def predict_noise(self, **kwargs): + with self._transformer_cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer(**kwargs) + return noise_pred_video.float(), noise_pred_audio.float() + + def combine_cfg_noise( + self, + positive_noise_pred, + negative_noise_pred, + true_cfg_scale, + cfg_normalize=False, + kwargs: dict[str, Any] | None = None, + **context: Any, + ): + if kwargs is not None: + context = {**kwargs, **context} + return self.guidance_strategy.combine_cfg_noise( + self, + positive_noise_pred, + negative_noise_pred, + true_cfg_scale, + cfg_normalize, + context, + ) + + def predict_noise_with_parallel_cfg(self, *args, **kwargs): + predict_parallel_cfg = getattr(self.guidance_strategy, "predict_parallel_cfg", None) + if predict_parallel_cfg is None: + raise NotImplementedError("The selected LTX guidance strategy does not implement parallel CFG.") + return predict_parallel_cfg(self, *args, **kwargs) + + def _synchronize_cfg_parallel_step_output( + self, + latents: tuple[torch.Tensor, torch.Tensor], + do_true_cfg: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1): + return latents + + # CUDA async execution otherwise permits numerical drift to accumulate + # across CFG-parallel denoise steps. + latents = tuple(tensor.contiguous() for tensor in latents) + device = next((tensor.device for tensor in latents if tensor.is_cuda), None) + if device is not None: + torch.cuda.current_stream(device).synchronize() + return latents + + def _setup_forward_runtime( + self, + request_inputs: LTXRequestInputs, + attention_kwargs: dict[str, Any] | None, + ) -> bool: + self._guidance_scale = request_inputs.guidance_scale + self._guidance_rescale = request_inputs.guidance_rescale + del attention_kwargs + self._interrupt = False + cfg_world_size = get_classifier_free_guidance_world_size() + if self.do_classifier_free_guidance: + self.guidance_strategy.validate_cfg_world_size(cfg_world_size) + return self.do_classifier_free_guidance and cfg_world_size > 1 + + def _check_forward_inputs( + self, + request_inputs: LTXRequestInputs, + image: Any | None = None, + ) -> None: + self.check_inputs( + prompt=request_inputs.prompt, + height=request_inputs.height, + width=request_inputs.width, + prompt_embeds=request_inputs.prompt_embeds, + negative_prompt_embeds=request_inputs.negative_prompt_embeds, + prompt_attention_mask=request_inputs.prompt_attention_mask, + negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask, + ) + + def _resolve_request_image( + self, + req: DiffusionRequestBatch, + image: Any | None, + request_inputs: LTXRequestInputs, + ) -> Any | None: + del req, request_inputs + return image + + def _make_output(self, output: tuple[torch.Tensor, torch.Tensor]) -> DiffusionOutput: + if self.reports_stage_durations: + return DiffusionOutput( + output=output, + stage_durations=getattr(self, "stage_durations", None), + ) + return DiffusionOutput(output=output) + + def _decode_output( + self, + *, + latents: torch.Tensor, + audio_latents: torch.Tensor, + output_type: str, + connector_prompt_embeds: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + prompt_batch_size: int, + ) -> DiffusionOutput: + if output_type == "latent": + return self._make_output((latents, audio_latents)) + + latents = latents.to(connector_prompt_embeds.dtype) + if not self.vae.config.timestep_conditioning: + timestep_decode = None + else: + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + timestep_decode, decode_noise_scale_t = _prepare_decode_timestep_conditioning( + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + prompt_batch_size=prompt_batch_size, + effective_batch_size=latents.shape[0], + device=device, + dtype=latents.dtype, + ) + latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise + + dist_initialized = torch.distributed.is_initialized() + is_output_rank = not dist_initialized or torch.distributed.get_rank() == 0 + vae_decode_needs_all_ranks = False + is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None) + if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled): + try: + # Distributed tiled decode is collective, so every rank must enter it. + vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled()) + except Exception: + pass + + should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks + if should_decode_video: + video = self.vae.decode(latents.to(self.vae.dtype), timestep_decode, return_dict=False)[0] + else: + video = torch.empty(0, device=latents.device, dtype=latents.dtype) + + if self.distributed_video_decode and not is_output_rank: + return self._make_output( + ( + torch.empty(0, device=video.device, dtype=video.dtype), + torch.empty(0, device=audio_latents.device, dtype=audio_latents.dtype), + ) + ) + + if video.numel() > 0: + video = self.video_processor.postprocess_video(video, output_type=output_type) + generated_mel = self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0] + audio = self.vocoder(generated_mel) + return self._make_output((video, audio)) + + def _prepare_video_latents_stage( + self, + request_inputs: LTXRequestInputs, + prompt_context: LTXPromptContext, + *, + device: torch.device, + noise_scale: float, + image: Any | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + latents = self.prepare_latents( + prompt_context.batch_size * request_inputs.num_videos_per_prompt, + self.transformer.config.in_channels, + request_inputs.height, + request_inputs.width, + request_inputs.num_frames, + noise_scale, + torch.float32, + device, + request_inputs.generator, + request_inputs.latents, + ) + return latents, None + + def _resolve_video_latent_dimensions(self, request_inputs: LTXRequestInputs) -> tuple[int, int, int]: + latent_num_frames, latent_height, latent_width = latent_ops.resolve_video_latent_shape( + request_inputs.height, + request_inputs.width, + request_inputs.num_frames, + vae_spatial_compression_ratio=self.vae_spatial_compression_ratio, + vae_temporal_compression_ratio=self.vae_temporal_compression_ratio, + ) + latents = request_inputs.latents + if latents is not None: + if latents.ndim == 5: + _, _, latent_num_frames, latent_height, latent_width = latents.shape + elif latents.ndim != 3: + raise ValueError( + f"Provided `latents` tensor has shape {latents.shape}, expected a packed 3D or unpacked 5D tensor." + ) + return latent_num_frames, latent_height, latent_width + + def _prepare_audio_latents_stage( + self, + request_inputs: LTXRequestInputs, + prompt_context: LTXPromptContext, + *, + device: torch.device, + noise_scale: float, + ) -> tuple[torch.Tensor, int, int, int]: + duration_s = request_inputs.num_frames / request_inputs.frame_rate + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + audio_num_frames = round(duration_s * audio_latents_per_second) + audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, request_inputs.audio_latents) + + num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64 + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + num_channels = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8 + audio_latents, original_num_frames, padded_num_frames = self.prepare_audio_latents( + prompt_context.batch_size * request_inputs.num_videos_per_prompt, + num_channels_latents=num_channels, + audio_latent_length=audio_num_frames, + num_mel_bins=num_mel_bins, + noise_scale=noise_scale, + dtype=torch.float32, + device=device, + generator=request_inputs.generator, + latents=request_inputs.audio_latents, + ) + return audio_latents, original_num_frames, padded_num_frames, latent_mel_bins + + def _resolve_audio_latent_length( + self, + requested_length: int, + audio_latents: torch.Tensor | None, + ) -> int: + if audio_latents is None or audio_latents.ndim != 4: + return requested_length + + provided_length = audio_latents.shape[2] + if not self.preserve_sp_padded_audio_duration: + return provided_length + + sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 + padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size)) + return requested_length if provided_length in {requested_length, padded_length} else provided_length + + def _make_video_audio_scheduler( + self, + audio_scheduler: Any, + latent_num_frames: int, + latent_height: int, + latent_width: int, + ) -> Any: + return VideoAudioScheduler(self.scheduler, audio_scheduler) + + def _prepare_denoise_context_for_cfg( + self, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> LTXDenoiseContext: + return self.guidance_strategy.prepare_denoise_context(self, forward_ctx, denoise_ctx) + + def _denoise_timestep_kwargs( + self, + ts: torch.Tensor, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> dict[str, torch.Tensor]: + return self.guidance_strategy.timestep_kwargs(ts, forward_ctx, denoise_ctx) + + def _build_transformer_kwargs( + self, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + *, + hidden_states: torch.Tensor, + audio_hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None, + audio_encoder_attention_mask: torch.Tensor | None, + ts: torch.Tensor, + ) -> dict[str, Any]: + return build_transformer_kwargs( + self, + forward_ctx, + denoise_ctx, + hidden_states=hidden_states, + audio_hidden_states=audio_hidden_states, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + audio_encoder_attention_mask=audio_encoder_attention_mask, + ts=ts, + ) + + def _predict_noise_for_step( + self, + index: int, + timestep: torch.Tensor, + state: latent_ops.LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self.guidance_strategy.predict_noise( + self, + index, + timestep, + state, + forward_ctx, + denoise_ctx, + ) + + def _denoise_step( + self, + index: int, + timestep: torch.Tensor, + state: latent_ops.LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, + ) -> latent_ops.LTXAVState: + denoise_ctx.latents = state.video + denoise_ctx.audio_latents = state.audio + noise_pred_video, noise_pred_audio = self._predict_noise_for_step( + index, + timestep, + state, + forward_ctx, + denoise_ctx, + ) + video, audio = step_denoised_latents( + self, + forward_ctx, + denoise_ctx, + noise_pred_video, + noise_pred_audio, + timestep, + ) + return latent_ops.LTXAVState(video=video, audio=audio) + + def _unpack_and_denormalize_stage( + self, + forward_ctx: LTXForwardContext, + latents: torch.Tensor, + audio_latents: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + latents = latent_ops.unpack_latents( + latents, + forward_ctx.latent_num_frames, + forward_ctx.latent_height, + forward_ctx.latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents = latent_ops.denormalize_latents( + latents, + self.vae.latents_mean, + self.vae.latents_std, + self.vae.config.scaling_factor, + ) + + audio_latents = latent_ops.unpad_audio_latents(audio_latents, forward_ctx.original_audio_num_frames) + audio_latents = latent_ops.denormalize_audio_latents( + audio_latents, + self.audio_vae.latents_mean, + self.audio_vae.latents_std, + ) + audio_latents = latent_ops.unpack_audio_latents( + audio_latents, + forward_ctx.original_audio_num_frames, + num_mel_bins=forward_ctx.latent_mel_bins, + ) + return latents, audio_latents + + def run_phase( + self, + req: DiffusionRequestBatch, + request_inputs: LTXRequestInputs, + *, + noise_scale: float, + sigmas: list[float] | None, + timesteps: list[int] | None, + attention_kwargs: dict[str, Any] | None, + image: Any | None = None, + prompt_context: LTXPromptContext | None = None, + ) -> LTXPhaseResult: + """Prepare and execute one phase without decoding its output.""" + return LTXPhaseExecutor.run( + self, + req, + request_inputs, + noise_scale=noise_scale, + sigmas=sigmas, + timesteps=timesteps, + attention_kwargs=attention_kwargs, + image=image, + prompt_context=prompt_context, + ) + + def decode_phase(self, phase: LTXPhaseResult) -> DiffusionOutput | list[DiffusionOutput]: + """Decode one completed phase and restore per-request outputs.""" + forward_ctx = phase.forward_context + request_inputs = forward_ctx.request_inputs + output = self._decode_output( + latents=phase.video, + audio_latents=phase.audio, + output_type=request_inputs.output_type, + connector_prompt_embeds=forward_ctx.prompt_context.connector_prompt_embeds, + generator=request_inputs.generator, + device=forward_ctx.device, + decode_timestep=request_inputs.decode_timestep, + decode_noise_scale=request_inputs.decode_noise_scale, + prompt_batch_size=forward_ctx.batch_size, + ) + if not self.supports_request_batch: + return output + return split_diffusion_output_by_request( + output, + forward_ctx.req, + num_outputs_per_prompt=forward_ctx.num_videos_per_prompt, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return AutoWeightsLoader(self).load_weights(weights) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py new file mode 100644 index 00000000000..6491b4a5f29 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""User-visible one-stage defaults for the LTX model family.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class LTXOneStageRecipe: + height: int = 512 + width: int = 768 + num_frames: int = 121 + frame_rate: float = 24.0 + num_inference_steps: int = 40 + guidance_scale: float = 4.0 + + +LTX2_ONE_STAGE_RECIPE = LTXOneStageRecipe() +LTX23_ONE_STAGE_RECIPE = LTXOneStageRecipe() diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py new file mode 100644 index 00000000000..02cab34d238 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Request normalization shared by LTX pipeline variants.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + + +@dataclass +class LTXRequestInputs: + """Resolved request values shared by all LTX execution variants.""" + + prompt: str | list[str] | None + negative_prompt: str | list[str] | None + height: int + width: int + num_frames: int + frame_rate: float + num_inference_steps: int + guidance_scale: float + guidance_rescale: float + num_videos_per_prompt: int + generator: torch.Generator | list[torch.Generator] | None + latents: torch.Tensor | None + audio_latents: torch.Tensor | None + prompt_embeds: torch.Tensor | None + negative_prompt_embeds: torch.Tensor | None + prompt_attention_mask: torch.Tensor | None + negative_prompt_attention_mask: torch.Tensor | None + decode_timestep: float | list[float] + decode_noise_scale: float | list[float] | None + output_type: str + max_sequence_length: int + + +def _unwrap_request_tensor(value: Any) -> Any: + if isinstance(value, list): + return value[0] if value else None + return value + + +def _get_prompt_field(prompt: Any, *keys: str) -> Any: + if isinstance(prompt, str): + return None + additional = prompt.get("additional_information") + for key in keys: + value = prompt.get(key) + if value is None and isinstance(additional, dict): + value = additional.get(key) + if value is not None: + return _unwrap_request_tensor(value) + return None + + +def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None: + if sampling.audio_latents is not None: + return sampling.audio_latents + return sampling.extra_args.get("audio_latents") + + +class LTXRequestMixin: + """Normalize serving requests without coupling them to a model version.""" + + supports_request_batch = False + supports_guidance_rescale = False + + def check_inputs( + self, + prompt: str | list[str] | None, + height: int, + width: int, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + ) -> None: + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` must be divisible by 32 but are {height} and {width}.") + if prompt is not None and prompt_embeds is not None: + raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.") + if prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if prompt is not None and not isinstance(prompt, (str, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got {prompt_embeds.shape} and {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when" + f" passed directly, but got {prompt_attention_mask.shape} and" + f" {negative_prompt_attention_mask.shape}." + ) + + def _resolve_request_inputs( + self, + req: DiffusionRequestBatch, + *, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None, + height: int | None, + width: int | None, + num_frames: int | None, + frame_rate: float | None, + num_inference_steps: int | None, + timesteps: list[int] | None, + guidance_scale: float, + num_videos_per_prompt: int | None, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None, + audio_latents: torch.Tensor | None, + prompt_embeds: torch.Tensor | None, + negative_prompt_embeds: torch.Tensor | None, + prompt_attention_mask: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + output_type: str, + max_sequence_length: int | None, + guidance_rescale: float = 0.0, + ) -> LTXRequestInputs: + sampling_params_list = req.sampling_params_list + sampling = sampling_params_list[0] + prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt + if all(isinstance(item, str) or item.get("negative_prompt") is None for item in req.prompts): + negative_prompt = None + elif req.prompts: + negative_prompt = [ + "" if isinstance(item, str) else (item.get("negative_prompt") or "") for item in req.prompts + ] + + height = sampling.height or height or self.one_stage_recipe.height + width = sampling.width or width or self.one_stage_recipe.width + num_frames = sampling.num_frames or num_frames or self.one_stage_recipe.num_frames + frame_rate = sampling.resolved_frame_rate or frame_rate or self.one_stage_recipe.frame_rate + num_inference_steps = ( + sampling.num_inference_steps or num_inference_steps or self.one_stage_recipe.num_inference_steps + ) + if timesteps is None: + num_inference_steps = max(int(num_inference_steps), 2) + elif len(timesteps) < 2: + raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") + + num_videos_per_prompt = ( + sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1 + ) + max_sequence_length = sampling.max_sequence_length or max_sequence_length or self.tokenizer_max_length + if sampling.guidance_scale_provided: + guidance_scale = sampling.guidance_scale + if self.supports_guidance_rescale and sampling.guidance_rescale is not None: + guidance_rescale = sampling.guidance_rescale + + if self.supports_request_batch: + if generator is None: + generator = req.collate_request_generators(num_videos_per_prompt, generator) + latents = req.collate_request_tensors("latents", latents) + audio_latents = DiffusionRequestBatch.collate_tensors( + [_get_audio_latents_from_sampling(item) for item in sampling_params_list], + "audio_latents", + audio_latents, + ) + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "prompt_attention_mask": prompt_attention_mask, + "negative_prompt_attention_mask": negative_prompt_attention_mask, + }, + field_aliases={ + "prompt_attention_mask": ("prompt_attention_mask", "attention_mask"), + "negative_prompt_attention_mask": ( + "negative_prompt_attention_mask", + "negative_attention_mask", + ), + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt_attention_mask = prompt_fields["prompt_attention_mask"] + negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"] + if prompt_embeds is not None: + prompt = None + if negative_prompt_embeds is not None: + negative_prompt = None + else: + if generator is None: + generator = sampling.generator + if generator is None and sampling.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(sampling.seed) + latents = sampling.latents if sampling.latents is not None else latents + sampling_audio_latents = _get_audio_latents_from_sampling(sampling) + if sampling_audio_latents is not None: + audio_latents = sampling_audio_latents + + request_prompt_embeds = [_get_prompt_field(item, "prompt_embeds") for item in req.prompts] + if any(value is not None for value in request_prompt_embeds): + prompt_embeds = torch.stack(request_prompt_embeds) # type: ignore[arg-type] + request_negative_embeds = [_get_prompt_field(item, "negative_prompt_embeds") for item in req.prompts] + if any(value is not None for value in request_negative_embeds): + negative_prompt_embeds = torch.stack(request_negative_embeds) # type: ignore[arg-type] + request_prompt_masks = [ + _get_prompt_field(item, "prompt_attention_mask", "attention_mask") for item in req.prompts + ] + if any(value is not None for value in request_prompt_masks): + prompt_attention_mask = torch.stack(request_prompt_masks) # type: ignore[arg-type] + request_negative_masks = [ + _get_prompt_field(item, "negative_prompt_attention_mask", "negative_attention_mask") + for item in req.prompts + ] + if any(value is not None for value in request_negative_masks): + negative_prompt_attention_mask = torch.stack(request_negative_masks) # type: ignore[arg-type] + + if sampling.decode_timestep is not None: + decode_timestep = sampling.decode_timestep + if sampling.decode_noise_scale is not None: + decode_noise_scale = sampling.decode_noise_scale + if sampling.output_type is not None: + output_type = sampling.output_type + + return LTXRequestInputs( + prompt=prompt, + negative_prompt=negative_prompt, + height=int(height), + width=int(width), + num_frames=int(num_frames), + frame_rate=float(frame_rate), + num_inference_steps=int(num_inference_steps), + guidance_scale=guidance_scale, + guidance_rescale=guidance_rescale, + num_videos_per_prompt=int(num_videos_per_prompt), + generator=generator, + latents=latents, + audio_latents=audio_latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + output_type=output_type, + max_sequence_length=int(max_sequence_length), + ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index dcf848390d4..6072c04560a 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -23,6 +23,7 @@ import torch.distributed import torch.nn as nn import torch.nn.functional as F +from cache_dit import ForwardPattern from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings, PixArtAlphaTextProjection from diffusers.utils import ( BaseOutput, @@ -44,6 +45,7 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available @@ -1449,6 +1451,14 @@ class LTX2VideoTransformer3DModel(nn.Module): The normalization layer to use. """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_0, + }, + has_separate_cfg=True, + check_forward_pattern=False, + ) + _supports_gradient_checkpointing = True _skip_layerwise_casting_patterns = ["norm"] _repeated_blocks = ["LTX2VideoTransformerBlock"] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index a409776d29d..c32a6918a5b 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -1,748 +1,71 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One-stage entry points for the LTX model family.""" + from __future__ import annotations -import copy -import inspect -import json -import os -from collections.abc import Iterable -from contextlib import nullcontext -from typing import TYPE_CHECKING, Any, ClassVar +from typing import Any, ClassVar -import numpy as np import torch -if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.base_config import QuantizationConfig -from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler -from diffusers.pipelines.ltx2 import LTX2TextConnectors -from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES -from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg, retrieve_timesteps -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from torch import nn -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration -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.cfg_parallel import CFGParallelMixin -from vllm_omni.diffusion.distributed.parallel_state import ( - get_classifier_free_guidance_world_size, -) -from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager -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.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.lora.request import LoRARequest - -from .ltx2_transformer import LTX2VideoTransformer3DModel -from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline - -logger = init_logger(__name__) - - -def load_transformer_config(model_path: str, subfolder: str = "transformer", local_files_only: bool = True) -> dict: - """Load transformer config from model directory or HF Hub.""" - if local_files_only: - config_path = os.path.join(model_path, subfolder, "config.json") - if os.path.exists(config_path): - with open(config_path) as f: - return json.load(f) - else: - try: - from huggingface_hub import hf_hub_download - - config_path = hf_hub_download( - repo_id=model_path, - filename=f"{subfolder}/config.json", - ) - with open(config_path) as f: - return json.load(f) - except Exception: - pass - return {} - - -def create_transformer_from_config( - config: dict, - quant_config: QuantizationConfig | None = None, -) -> LTX2VideoTransformer3DModel: - """Create LTX2VideoTransformer3DModel from config dict.""" - if not config and quant_config is None: - return LTX2VideoTransformer3DModel() - - signature = inspect.signature(LTX2VideoTransformer3DModel.__init__) - allowed_keys = set(signature.parameters.keys()) - kwargs = {k: v for k, v in config.items() if k in allowed_keys} - if quant_config is not None: - kwargs["quant_config"] = quant_config - - return LTX2VideoTransformer3DModel(**kwargs) - - -def get_ltx2_post_process_func( - od_config: OmniDiffusionConfig, -): - def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): - if isinstance(output, tuple) and len(output) == 2: - video, audio = output - if isinstance(audio, torch.Tensor): - audio = audio.detach().cpu() - return {"video": video, "audio": audio} - return output - - return post_process_func +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from .ltx2_components import ( + LTX2_COMPONENT_PROFILE, + LTX23_COMPONENT_PROFILE, + LTXComponentProfile, +) +from .ltx2_components import ( + get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401 +) +from .ltx2_conditioning import LTXI2VConditioningMixin +from .ltx2_guidance import ( + LTX_LEGACY_VELOCITY_GUIDANCE, + LTX_OFFICIAL_X0_GUIDANCE, +) +from .ltx2_pipeline_runtime import LTXPipelineRuntime +from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, LTXOneStageRecipe +from .ltx2_request import LTXRequestInputs -def _unwrap_request_tensor(value: Any) -> Any: - if isinstance(value, list): - return value[0] if value else None - return value - - -def _get_prompt_field(prompt: Any, key: str) -> Any: - if isinstance(prompt, str): - return None - value = prompt.get(key) - if value is None: - additional = prompt.get("additional_information") - if isinstance(additional, dict): - value = additional.get(key) - return _unwrap_request_tensor(value) - - -def calculate_shift( - image_seq_len, - base_seq_len: int = 256, - max_seq_len: int = 4096, - base_shift: float = 0.5, - max_shift: float = 1.15, -): - m = (max_shift - base_shift) / (max_seq_len - base_seq_len) - b = base_shift - m * base_seq_len - mu = image_seq_len * m + b - return mu - - -class _VideoAudioScheduler: - """Composite scheduler dispatching to video and audio schedulers.""" - - def __init__(self, video_scheduler, audio_scheduler): - self.video_scheduler = video_scheduler - self.audio_scheduler = audio_scheduler - def step(self, noise_pred, t, latents, return_dict=False, generator=None): - video_out = self.video_scheduler.step( - noise_pred[0], - t[0], - latents[0], - return_dict=False, - generator=generator, - )[0] - audio_out = self.audio_scheduler.step( - noise_pred[1], - t[1], - latents[1], - return_dict=False, - generator=generator, - )[0] - return ((video_out, audio_out),) +class LTXOneStagePipeline(LTXPipelineRuntime): + """Single execution path configured by model-version and task entries.""" + component_profile: ClassVar[LTXComponentProfile] + one_stage_recipe: ClassVar[LTXOneStageRecipe] -class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): - # Audio is diffused jointly with video; warmup must size audio tokens. - dummy_run_num_frames = 2 + supports_request_batch = True + supports_guidance_rescale = False - def __init__( + def _forward_impl( self, + req: DiffusionRequestBatch, + request_inputs: LTXRequestInputs, *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - - self.device = get_local_device() - dtype = getattr(od_config, "dtype", torch.bfloat16) - - model = od_config.model - local_files_only = os.path.exists(model) - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder - # race; prefetch the whole component set before any from_pretrained. - ltx2_subfolders = [ - "tokenizer", - "text_encoder", - "connectors", - "vae", - "audio_vae", - "vocoder", - "scheduler", - ] - prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only) - - self.tokenizer = AutoTokenizer.from_pretrained( - model, - subfolder="tokenizer", - local_files_only=local_files_only, - ) - # prefer mmap loading as default device is cuda, and the output of text encoder - # could be deterministic. - with torch.device("cpu"): - self.text_encoder = from_pretrained_with_prefetch( - Gemma3ForConditionalGeneration.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - self.connectors = from_pretrained_with_prefetch( - LTX2TextConnectors.from_pretrained, - model, - subfolder="connectors", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - - self.vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Video.from_pretrained, - model, - subfolder="vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - self.audio_vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Audio.from_pretrained, - model, - subfolder="audio_vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - self.vocoder = from_pretrained_with_prefetch( - LTX2Vocoder.from_pretrained, - model, - subfolder="vocoder", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ).to(self.device) - - 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.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( - model, subfolder="scheduler", local_files_only=local_files_only - ) - - self.vae_spatial_compression_ratio = ( - self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 - ) - self.vae_temporal_compression_ratio = ( - self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 - ) - self.audio_vae_mel_compression_ratio = ( - self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 - ) - self.audio_vae_temporal_compression_ratio = ( - self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 - ) - self.transformer_spatial_patch_size = ( - self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 - ) - self.transformer_temporal_patch_size = ( - self.transformer.config.patch_size_t if getattr(self, "transformer", None) is not None else 1 - ) - - self.audio_sampling_rate = ( - self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000 - ) - self.audio_hop_length = ( - self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160 - ) - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - tokenizer_max_length = 1024 - if getattr(self, "tokenizer", None) is not None: - tokenizer_max_length = self.tokenizer.model_max_length - if tokenizer_max_length is None or tokenizer_max_length > 100000: - encoder_config = getattr(self.text_encoder, "config", None) - config_max_len = getattr(encoder_config, "max_position_embeddings", None) - if config_max_len is None: - config_max_len = getattr(encoder_config, "max_seq_len", None) - tokenizer_max_length = config_max_len or 1024 - self.tokenizer_max_length = int(tokenizer_max_length) - - self._guidance_scale = None - self._guidance_rescale = None - self._attention_kwargs = None - self._interrupt = False - self._num_timesteps = None - self._current_timestep = None - - def _get_gemma_prompt_embeds( - self, - prompt: str | list[str], - num_videos_per_prompt: int = 1, - max_sequence_length: int = 1024, - scale_factor: int = 8, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - device = device or self.device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - if getattr(self, "tokenizer", None) is not None: - self.tokenizer.padding_side = "left" - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - - prompt = [p.strip() for p in prompt] - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids - prompt_attention_mask = text_inputs.attention_mask - text_input_ids = text_input_ids.to(device) - prompt_attention_mask = prompt_attention_mask.to(device) - - text_encoder_outputs = self.text_encoder( - input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True - ) - text_encoder_hidden_states = text_encoder_outputs.hidden_states - text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) - # `diffusers>=0.38` moved per_layer_masked_mean_norm inside the connector, - # so pack to 3D here and let the connector handle normalization. - prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype) - - _, seq_len, _ = prompt_embeds.shape - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) - prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) - - return prompt_embeds, prompt_attention_mask - - def encode_prompt( - self, - prompt: str | list[str], - negative_prompt: str | list[str] | None = None, - do_classifier_free_guidance: bool = True, - num_videos_per_prompt: int = 1, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - prompt_attention_mask: torch.Tensor | None = None, - negative_prompt_attention_mask: torch.Tensor | None = None, - max_sequence_length: int = 1024, - scale_factor: int = 8, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - device = device or self.device - - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - num_videos_per_prompt=num_videos_per_prompt, - max_sequence_length=max_sequence_length, - scale_factor=scale_factor, - device=device, - dtype=dtype, - ) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - - if prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - elif batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - - negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=negative_prompt, - num_videos_per_prompt=num_videos_per_prompt, - max_sequence_length=max_sequence_length, - scale_factor=scale_factor, - device=device, - dtype=dtype, - ) - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - def check_inputs( - self, - prompt, - height, - width, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - ): - if height % 32 != 0 or width % 32 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - elif prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if prompt_embeds is not None and prompt_attention_mask is None: - raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") - - if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: - raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: - raise ValueError( - "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when " - "passed directly, but got: `prompt_attention_mask` " - f"{prompt_attention_mask.shape} != `negative_prompt_attention_mask` " - f"{negative_prompt_attention_mask.shape}." - ) - - @staticmethod - def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = latents.shape - post_patch_num_frames = num_frames // patch_size_t - post_patch_height = height // patch_size - post_patch_width = width // patch_size - latents = latents.reshape( - batch_size, - -1, - post_patch_num_frames, - patch_size_t, - post_patch_height, - patch_size, - post_patch_width, - patch_size, + noise_scale: float, + sigmas: list[float] | None, + timesteps: list[int] | None, + attention_kwargs: dict[str, Any] | None, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + phase = self.run_phase( + req, + request_inputs, + noise_scale=noise_scale, + sigmas=sigmas, + timesteps=timesteps, + attention_kwargs=attention_kwargs, + image=image, ) - latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) - return latents - - @staticmethod - def _unpack_latents( - latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 - ) -> torch.Tensor: - batch_size = latents.size(0) - latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) - latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) - return latents - - @staticmethod - def _normalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - # Normalize latents across the channel dimension [B, C, F, H, W] - latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents = (latents - latents_mean) * scaling_factor / latents_std - return latents - - @staticmethod - def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): - latents_mean = latents_mean.to(latents.device, latents.dtype) - latents_std = latents_std.to(latents.device, latents.dtype) - return (latents - latents_mean) / latents_std - - @staticmethod - def _denormalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents = latents * latents_std / scaling_factor + latents_mean - return latents - - @staticmethod - def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): - latents_mean = latents_mean.to(latents.device, latents.dtype) - latents_std = latents_std.to(latents.device, latents.dtype) - return (latents * latents_std) + latents_mean - - @staticmethod - def _create_noised_state( - latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None - ): - noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) - noised_latents = noise_scale * noise + (1 - noise_scale) * latents - return noised_latents - - @staticmethod - def _pack_audio_latents( - latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None - ) -> torch.Tensor: - if patch_size is not None and patch_size_t is not None: - batch_size, num_channels, latent_length, latent_mel_bins = latents.shape - post_patch_latent_length = latent_length / patch_size_t - post_patch_mel_bins = latent_mel_bins / patch_size - latents = latents.reshape( - batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size - ) - latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) - else: - latents = latents.transpose(1, 2).flatten(2, 3) - return latents - - @staticmethod - def _unpack_audio_latents( - latents: torch.Tensor, - latent_length: int, - num_mel_bins: int, - patch_size: int | None = None, - patch_size_t: int | None = None, - ) -> torch.Tensor: - if patch_size is not None and patch_size_t is not None: - batch_size = latents.size(0) - latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) - latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) - else: - latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) - return latents - - @staticmethod - def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: - return latents[:, :num_frames] - - def prepare_latents( - self, - batch_size: int = 1, - num_channels_latents: int = 128, - height: int = 512, - width: int = 768, - num_frames: int = 121, - noise_scale: float = 0.0, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | None = None, - latents: torch.Tensor | None = None, - ) -> torch.Tensor: - if latents is not None: - if latents.ndim == 5: - latents = self._normalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - # latents are of shape [B, C, F, H, W], need to be packed - latents = self._pack_latents( - latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ) - if latents.ndim != 3: - raise ValueError( - f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is [batch_size, num_seq, num_features]." # noqa - ) - latents = self._create_noised_state(latents, noise_scale, generator) - return latents.to(device=device, dtype=dtype) - - height = height // self.vae_spatial_compression_ratio - width = width // self.vae_spatial_compression_ratio - num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - - shape = (batch_size, num_channels_latents, num_frames, height, width) - - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) - - return latents - - def prepare_audio_latents( - self, - batch_size: int = 1, - num_channels_latents: int = 8, - audio_latent_length: int = 1, # 1 is just a dummy value - num_mel_bins: int = 64, - noise_scale: float = 0.0, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, int, int]: - original_latent_length = audio_latent_length - padded_latent_length = original_latent_length - - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - - sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) - if sp_size > 1: - padded_latent_length += (sp_size - (original_latent_length % sp_size)) % sp_size - - if latents is not None: - if latents.ndim == 4: - # latents are of shape [B, C, L, M], need to be packed - latents = self._pack_audio_latents(latents) - if latents.ndim != 3: - raise ValueError( - f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is " - "[batch_size, num_seq, num_features] or [batch_size, num_channels, audio_length, mel_bins]." - ) - latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) - latents = self._create_noised_state(latents, noise_scale, generator) - - if latents.shape[1] not in {original_latent_length, padded_latent_length}: - raise ValueError( - "Provided `audio_latents` has incompatible audio frame count " - f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}." - ) - - if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length: - padding = torch.zeros( - latents.shape[0], - padded_latent_length - original_latent_length, - latents.shape[2], - dtype=latents.dtype, - device=latents.device, - ) - latents = torch.cat([latents, padding], dim=1) - - return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length - - shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) - - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self._pack_audio_latents(latents) - return latents, original_latent_length, padded_latent_length - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def guidance_rescale(self): - return self._guidance_rescale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale > 1.0 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def current_timestep(self): - return self._current_timestep - - @property - def attention_kwargs(self): - return self._attention_kwargs - - @property - def interrupt(self): - return self._interrupt - - def _transformer_cache_context(self, context_name: str): - cache_context = getattr(self.transformer, "cache_context", None) - if callable(cache_context): - return cache_context(context_name) - return nullcontext() - - def predict_noise(self, **kwargs): - with self._transformer_cache_context("cond_uncond"): - noise_pred_video, noise_pred_audio = self.transformer(**kwargs) - return noise_pred_video.float(), noise_pred_audio.float() - - def combine_cfg_noise(self, positive_noise_pred, negative_noise_pred, true_cfg_scale, cfg_normalize=False): - """Per-element CFG combine with guidance_rescale support.""" - (video_pos, audio_pos) = positive_noise_pred - (video_neg, audio_neg) = negative_noise_pred - video_combined = super().combine_cfg_noise(video_pos, video_neg, true_cfg_scale, cfg_normalize) - audio_combined = super().combine_cfg_noise(audio_pos, audio_neg, true_cfg_scale, cfg_normalize) - if self._guidance_rescale and self._guidance_rescale > 0: - video_combined = rescale_noise_cfg(video_combined, video_pos, guidance_rescale=self._guidance_rescale) - audio_combined = rescale_noise_cfg(audio_combined, audio_pos, guidance_rescale=self._guidance_rescale) - return (video_combined, audio_combined) - - def _synchronize_cfg_parallel_step_output( - self, - latents: tuple[torch.Tensor, torch.Tensor], - do_true_cfg: bool, - ) -> tuple[torch.Tensor, torch.Tensor]: - if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1): - return latents - - # Without this sync, CUDA async execution causes non-deterministic - # numerical drift across denoising steps in CFG parallel mode, - # producing different video outputs across runs. - latents = tuple(tensor.contiguous() for tensor in latents) - device = next((tensor.device for tensor in latents if tensor.is_cuda), None) - if device is not None: - torch.cuda.current_stream(device).synchronize() - return latents + return self.decode_phase(phase) @torch.no_grad() def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, @@ -769,433 +92,82 @@ def forward( return_dict: bool = True, attention_kwargs: dict[str, Any] | None = None, max_sequence_length: int | None = None, - ) -> DiffusionOutput: - # Extract prompt/negative_prompt from request. - # Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt - if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): - negative_prompt = None - elif req.prompts: - negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - - height = req.sampling_params.height or height or 512 - width = req.sampling_params.width or width or 768 - num_frames = req.sampling_params.num_frames or num_frames or 121 - frame_rate = req.sampling_params.resolved_frame_rate or frame_rate or 24.0 - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps or 40 - if timesteps is None: - num_inference_steps = max(int(num_inference_steps), 2) - elif len(timesteps) < 2: - raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") - num_videos_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_videos_per_prompt or 1 - ) - max_sequence_length = ( - req.sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length - ) - - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - if req.sampling_params.guidance_rescale is not None: - guidance_rescale = req.sampling_params.guidance_rescale - - if generator is None: - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) - - latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents - audio_latents = ( - req.sampling_params.audio_latents - if req.sampling_params.audio_latents is not None - else req.sampling_params.extra_args.get("audio_latents", audio_latents) - ) - - # Override with pre-computed embeddings if provided in request. - req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type] - - req_negative_prompt_embeds = [_get_prompt_field(p, "negative_prompt_embeds") for p in req.prompts] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore[arg-type] - - req_prompt_attention_masks = [ - _get_prompt_field(p, "prompt_attention_mask") or _get_prompt_field(p, "attention_mask") for p in req.prompts - ] - if any(m is not None for m in req_prompt_attention_masks): - prompt_attention_mask = torch.stack(req_prompt_attention_masks) # type: ignore[arg-type] - - req_negative_attention_masks = [ - _get_prompt_field(p, "negative_prompt_attention_mask") or _get_prompt_field(p, "negative_attention_mask") - for p in req.prompts - ] - if any(m is not None for m in req_negative_attention_masks): - negative_prompt_attention_mask = torch.stack(req_negative_attention_masks) # type: ignore[arg-type] - - if req.sampling_params.decode_timestep is not None: - decode_timestep = req.sampling_params.decode_timestep - if req.sampling_params.decode_noise_scale is not None: - decode_noise_scale = req.sampling_params.decode_noise_scale - if req.sampling_params.output_type is not None: - output_type = req.sampling_params.output_type - - self.check_inputs( + *, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + del return_dict + request_inputs = self._resolve_request_inputs( + req, prompt=prompt, + negative_prompt=negative_prompt, height=height, width=width, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._attention_kwargs = attention_kwargs - self._interrupt = False - self._current_timestep = None - - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self.device - - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, + num_frames=num_frames, + frame_rate=frame_rate, + num_inference_steps=num_inference_steps, + timesteps=timesteps, + guidance_scale=self.one_stage_recipe.guidance_scale if guidance_scale is None else guidance_scale, + guidance_rescale=guidance_rescale, num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + latents=latents, + audio_latents=audio_latents, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, prompt_attention_mask=prompt_attention_mask, negative_prompt_attention_mask=negative_prompt_attention_mask, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + output_type=output_type, max_sequence_length=max_sequence_length, - device=device, - ) - # Compute positive prompt connectors - 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 - ) - - # Compute negative prompt connectors when CFG is enabled - negative_connector_prompt_embeds = None - negative_connector_audio_prompt_embeds = None - negative_connector_attention_mask = None - if self.do_classifier_free_guidance: - ( - negative_connector_prompt_embeds, - negative_connector_audio_prompt_embeds, - negative_connector_attention_mask, - ) = self.connectors( - negative_prompt_embeds, - negative_prompt_attention_mask, - padding_side=tokenizer_padding_side, - ) - - latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - latent_height = height // self.vae_spatial_compression_ratio - latent_width = width // self.vae_spatial_compression_ratio - if latents is not None: - if latents.ndim == 5: - logger.info( - "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." # noqa - ) - _, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W] - elif latents.ndim == 3: - logger.warning( - f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be" - f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct." - ) - else: - raise ValueError( - f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." # noqa - ) - video_sequence_length = latent_num_frames * latent_height * latent_width - - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - noise_scale, - torch.float32, - device, - generator, - latents, - ) - - duration_s = num_frames / frame_rate - audio_latents_per_second = ( - self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) - ) - audio_num_frames = round(duration_s * audio_latents_per_second) - if audio_latents is not None: - if audio_latents.ndim == 4: - logger.info( - "Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." # noqa - ) - _, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M] - elif audio_latents.ndim == 3: - logger.warning( - f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims" - f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct." - ) - else: - raise ValueError( - f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." # noqa - ) - - num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - - num_channels_latents_audio = ( - self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 - ) - audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents( - batch_size * num_videos_per_prompt, - num_channels_latents=num_channels_latents_audio, - audio_latent_length=audio_num_frames, - num_mel_bins=num_mel_bins, - noise_scale=noise_scale, - dtype=torch.float32, - device=device, - generator=generator, - latents=audio_latents, - ) - - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - mu = calculate_shift( - video_sequence_length, - self.scheduler.config.get("base_image_seq_len", 1024), - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_shift", 0.95), - self.scheduler.config.get("max_shift", 2.05), - ) - audio_scheduler = copy.deepcopy(self.scheduler) - video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler) - _ = retrieve_timesteps( - audio_scheduler, - num_inference_steps, - device, - timesteps, - sigmas=sigmas, - mu=mu, - ) - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas=sigmas, - mu=mu, - ) - self._num_timesteps = len(timesteps) - - video_coords = self.transformer.rope.prepare_video_coords( - latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate - ) - audio_coords = self.transformer.audio_rope.prepare_audio_coords( - audio_latents.shape[0], padded_audio_num_frames, audio_latents.device - ) - # No coord duplication needed: mixin handles CFG via separate forward calls, - # not batch=2. Each forward gets batch=1 coords directly. - - with self.progress_bar(total=len(timesteps)) as pbar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - - latent_model_input = latents.to(prompt_embeds.dtype) - audio_latent_model_input = audio_latents.to(prompt_embeds.dtype) - timestep = t.expand(latent_model_input.shape[0]) - do_true_cfg = self.do_classifier_free_guidance - - positive_kwargs = { - "hidden_states": latent_model_input, - "audio_hidden_states": audio_latent_model_input, - "encoder_hidden_states": connector_prompt_embeds, - "audio_encoder_hidden_states": connector_audio_prompt_embeds, - "timestep": timestep, - "encoder_attention_mask": connector_attention_mask, - "audio_encoder_attention_mask": connector_attention_mask, - "num_frames": latent_num_frames, - "height": latent_height, - "width": latent_width, - "fps": frame_rate, - "audio_num_frames": padded_audio_num_frames, - "video_coords": video_coords, - "audio_coords": audio_coords, - "attention_kwargs": attention_kwargs, - "return_dict": False, - } - negative_kwargs = ( - { - **positive_kwargs, - "encoder_hidden_states": negative_connector_prompt_embeds, - "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds, - "encoder_attention_mask": negative_connector_attention_mask, - "audio_encoder_attention_mask": negative_connector_attention_mask, - } - if do_true_cfg - else None - ) - - noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) - - latents, audio_latents = self.scheduler_step_maybe_with_cfg( - (noise_pred_video, noise_pred_audio), - (t, t), - (latents, audio_latents), - do_true_cfg=do_true_cfg, - per_request_scheduler=video_audio_scheduler, - ) - latents, audio_latents = self._synchronize_cfg_parallel_step_output( - (latents, audio_latents), - do_true_cfg=do_true_cfg, - ) - - pbar.update() - - latents = self._unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - latents = self._denormalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - - audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames) - audio_latents = self._denormalize_audio_latents( - audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std - ) - audio_latents = self._unpack_audio_latents( - audio_latents, - original_audio_num_frames, - num_mel_bins=latent_mel_bins, ) + image = self._resolve_request_image(req, image, request_inputs) + forward_kwargs = { + "noise_scale": noise_scale, + "sigmas": sigmas, + "timesteps": timesteps, + "attention_kwargs": attention_kwargs, + } + if self.support_image_input: + forward_kwargs["image"] = image + return self._forward_impl( + req, + request_inputs, + **forward_kwargs, + ) + + +class LTX2Pipeline(LTXOneStagePipeline): + """LTX2 one-stage text-to-video entry.""" + + supports_guidance_rescale = True + component_profile = LTX2_COMPONENT_PROFILE + guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE + one_stage_recipe = LTX2_ONE_STAGE_RECIPE + _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules) + _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules) + _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules) + _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules) + + +class LTX23Pipeline(LTXOneStagePipeline): + """LTX2.3 one-stage text-to-video entry.""" + + connector_batches_cfg = True + preserve_sp_padded_audio_duration = True + reports_stage_durations = True + component_profile = LTX23_COMPONENT_PROFILE + guidance_strategy = LTX_OFFICIAL_X0_GUIDANCE + one_stage_recipe = LTX23_ONE_STAGE_RECIPE + _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules) + _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules) + _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules) + _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules) - if output_type == "latent": - video = latents - audio = audio_latents - else: - latents = latents.to(prompt_embeds.dtype) - - if not self.vae.config.timestep_conditioning: - timestep = None - else: - noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) - if not isinstance(decode_timestep, list): - decode_timestep = [decode_timestep] * batch_size - if decode_noise_scale is None: - decode_noise_scale = decode_timestep - elif not isinstance(decode_noise_scale, list): - decode_noise_scale = [decode_noise_scale] * batch_size - - timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) - decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ - :, None, None, None, None - ] - latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise - - latents = latents.to(self.vae.dtype) - video = self.vae.decode(latents, timestep, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - - audio_latents = audio_latents.to(self.audio_vae.dtype) - generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] - audio = self.vocoder(generated_mel_spectrograms) - - if not return_dict: - return DiffusionOutput(output=(video, audio)) - - return DiffusionOutput(output=(video, audio)) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) - - -class LTX2TwoStagesPipeline(nn.Module, SupportsComponentDiscovery): - """LTX2TwoStagesPipeline is for two stages image to video generation""" - - dummy_run_num_frames = 2 - - _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] - _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] - _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - - self.device = get_local_device() - self.dtype = getattr(od_config, "dtype", torch.bfloat16) - self.model_path = od_config.model - self.distilled = False - # User provided model path may contain '/' in the end and basename function - # will not return the expected directory name, so we need to remove it by normpath - if "distilled" in os.path.basename(os.path.normpath(self.model_path)): - self.distilled = True - else: - raise NotImplementedError(f"{self.model_path} is not supported for {self.__class__.__name__}.") - - self.pipe = LTX2Pipeline(od_config=od_config, prefix=prefix) - self.upsample_pipe = LTX2LatentUpsamplePipeline( - vae=self.pipe.vae, - od_config=od_config, - ) - - self.lora_manager = DiffusionLoRAManager( - pipeline=self.pipe, - device=self.device, - dtype=self.dtype, - max_cached_adapters=od_config.max_cpu_loras, - ) - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="pipe.transformer.", - fall_back_to_pt=True, - ), - ] - + @torch.no_grad() def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, @@ -1203,9 +175,9 @@ def forward( num_frames: int | None = None, frame_rate: float | None = None, num_inference_steps: int | None = None, + sigmas: list[float] | None = None, timesteps: list[int] | None = None, guidance_scale: float = 4.0, - guidance_rescale: float = 0.0, noise_scale: float = 0.0, num_videos_per_prompt: int | None = 1, generator: torch.Generator | list[torch.Generator] | None = None, @@ -1221,9 +193,12 @@ def forward( return_dict: bool = True, attention_kwargs: dict[str, Any] | None = None, max_sequence_length: int | None = None, - ): - video_latent, audio_latent = self.pipe( - req=req, + *, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + """Preserve the pre-refactor LTX2.3 positional argument order.""" + return super().forward( + req, prompt=prompt, negative_prompt=negative_prompt, height=height, @@ -1231,10 +206,10 @@ def forward( num_frames=num_frames, frame_rate=frame_rate, num_inference_steps=num_inference_steps, - sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None, + sigmas=sigmas, timesteps=timesteps, guidance_scale=guidance_scale, - guidance_rescale=guidance_rescale, + guidance_rescale=0.0, noise_scale=noise_scale, num_videos_per_prompt=num_videos_per_prompt, generator=generator, @@ -1246,65 +221,32 @@ def forward( negative_prompt_attention_mask=negative_prompt_attention_mask, decode_timestep=decode_timestep, decode_noise_scale=decode_noise_scale, - output_type="latent", + output_type=output_type, return_dict=return_dict, attention_kwargs=attention_kwargs, max_sequence_length=max_sequence_length, - ).output + image=image, + ) - upscaled_video_latent = self.upsample_pipe( - latents=video_latent, - output_type="latent", - return_dict=False, - )[0] - if not self.distilled: - # Load Stage 2 distilled LoRA - lora_path = f"{self.model_path}/ltx-2-19b-distilled-lora-384.safetensors" - lora_request = LoRARequest( - lora_name="stage_2_distilled", - lora_int_id=1, - lora_path=lora_path, - ) - self.lora_manager.set_active_adapter(lora_request, lora_scale=1.0) +class LTX2ImageToVideoPipeline(LTXI2VConditioningMixin, LTX2Pipeline): + """LTX2 one-stage image-to-video entry.""" - # Change scheduler to use Stage 2 distilled sigmas as is - new_scheduler = FlowMatchEulerDiscreteScheduler.from_config( - self.pipe.scheduler.config, - use_dynamic_shifting=False, - shift_terminal=None, - ) - self.pipe.scheduler = new_scheduler - # We only want to change num_inference_steps here, so no need - # to deep copy the whole request - stage_2_req = copy.copy(req) - stage_2_req.sampling_params = req.sampling_params.clone() - stage_2_req.sampling_params.num_inference_steps = 3 +class LTX23ImageToVideoPipeline(LTXI2VConditioningMixin, LTX23Pipeline): + """LTX2.3 one-stage image-to-video entry.""" - video, audio = self.pipe( - req=stage_2_req, - latents=upscaled_video_latent, - audio_latents=audio_latent, - prompt=prompt, - negative_prompt=negative_prompt, - noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], - sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, - guidance_scale=1.0, - generator=generator, - output_type="np", - return_dict=False, - ).output - return DiffusionOutput(output=(video, audio)) +class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): + """LTX2 T2V entry for FastGen DMD2-distilled models.""" - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__(od_config=od_config, prefix=prefix) + self.__init_dmd2__() -class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline): - """LTX-2 T2V pipeline for FastGen DMD2-distilled models.""" +class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline): + """LTX2 I2V entry for FastGen DMD2-distilled models.""" def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__(od_config=od_config, prefix=prefix) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py deleted file mode 100644 index f333db8ce39..00000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ /dev/null @@ -1,1294 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -""" -Fully independent LTX-2.3 pipeline for vLLM-Omni. - -This pipeline does NOT inherit from LTX2Pipeline because: -- LTX-2.3 connectors run per_token_rms_norm + per-modality video/audio - projection internally (per_modality_projections=True), - 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 -""" - -from __future__ import annotations - -import copy -import json -import os -from collections.abc import Iterable -from contextlib import nullcontext -from typing import Any, ClassVar - -import numpy as np -import torch -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 -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -from huggingface_hub import hf_hub_download -from torch import nn -from transformers import AutoTokenizer, Gemma3ForConditionalGeneration -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, - get_classifier_free_guidance_rank, - get_classifier_free_guidance_world_size, -) -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 ( - _get_prompt_field, - _VideoAudioScheduler, - calculate_shift, - create_transformer_from_config, - load_transformer_config, -) - -logger = init_logger(__name__) - -# Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0) -try: - from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE -except ImportError: - LTX2VocoderWithBWE = None - - -def _detect_vocoder_output_sample_rate(model: str) -> int | None: - """Detect the vocoder output sample rate from vocoder/config.json. - - This runs at factory time (engine process) so the rate is captured in - the post-process closure and doesn't need cross-process communication. - - Returns: - Output sample rate (e.g. 48000 for LTX-2.3 BWE vocoder) or None. - """ - vocoder_config_path = os.path.join(model, "vocoder", "config.json") - if not os.path.exists(vocoder_config_path): - try: - vocoder_config_path = hf_hub_download(model, "vocoder/config.json") - except Exception: - return None - try: - with open(vocoder_config_path) as f: - cfg = json.load(f) - return cfg.get("output_sampling_rate") - except Exception: - return None - - -def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): - """Factory for the LTX-2.3 post-process function. - - Detects the vocoder output sample rate at factory time and captures it - in the closure so that the audio_sample_rate flows through - DiffusionEngine -> OmniRequestOutput -> serving_video. - """ - output_sr = _detect_vocoder_output_sample_rate(od_config.model) - - def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): - if isinstance(output, tuple) and len(output) == 2: - video, audio = output - if isinstance(audio, torch.Tensor): - audio = audio.detach().cpu() - result: dict[str, Any] = {"video": video, "audio": audio} - if output_sr is not None: - result["audio_sample_rate"] = output_sr - return result - return output - - return post_process_func - - -class LTX23Pipeline( - nn.Module, - CFGParallelMixin, - ProgressBarMixin, - SupportsComponentDiscovery, - DiffusionPipelineProfilerMixin, -): - """Fully independent LTX-2.3 pipeline. - - Key differences from LTX2Pipeline: - - Text encoding: uses ALL 49 hidden states from Gemma-3-12B, flattened - - Connectors: uses padding_side API (not additive_mask) - - Vocoder: uses LTX2VocoderWithBWE (48kHz output) - - Transformer: passes sigma for prompt_adaln - """ - - # 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, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - self.od_config = od_config - self.device = get_local_device() - dtype = getattr(od_config, "dtype", torch.bfloat16) - model = od_config.model - local_files_only = os.path.exists(model) - - # Weight sources for transformer (loaded via AutoWeightsLoader) - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="transformer.", - fall_back_to_pt=True, - ), - ] - - # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder - # race; prefetch the whole component set before any from_pretrained. - ltx2_subfolders = [ - "tokenizer", - "text_encoder", - "connectors", - "vae", - "audio_vae", - "vocoder", - "scheduler", - ] - prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only) - - # --- Tokenizer (lightweight, stays wherever) --- - self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) - - # --- Text encoder --- - with torch.device("cpu"): - self.text_encoder = from_pretrained_with_prefetch( - Gemma3ForConditionalGeneration.from_pretrained, - model, - subfolder="text_encoder", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - - # --- Connectors (LTX-2.3 connectors include caption projection) --- - self.connectors = from_pretrained_with_prefetch( - LTX2TextConnectors.from_pretrained, - model, - subfolder="connectors", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - - # --- VAE, Audio VAE --- - self.vae = from_pretrained_with_prefetch( - DistributedAutoencoderKLLTX2Video.from_pretrained, - model, - subfolder="vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - self.audio_vae = from_pretrained_with_prefetch( - AutoencoderKLLTX2Audio.from_pretrained, - model, - subfolder="audio_vae", - prefetch_list=ltx2_subfolders, - local_files_only=local_files_only, - torch_dtype=dtype, - ) - - # --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 --- - vocoder_cls = LTX2VocoderWithBWE or LTX2Vocoder - try: - self.vocoder = vocoder_cls.from_pretrained( - model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only - ) - except (TypeError, OSError, ValueError): - self.vocoder = LTX2Vocoder.from_pretrained( - model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only - ) - - # --- Transformer: created empty, weights loaded via AutoWeightsLoader --- - 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( - model, subfolder="scheduler", local_files_only=local_files_only - ) - - # --- Derived compression ratios --- - self.vae_spatial_compression_ratio = self.vae.spatial_compression_ratio if self.vae is not None else 32 - self.vae_temporal_compression_ratio = self.vae.temporal_compression_ratio if self.vae is not None else 8 - self.audio_vae_mel_compression_ratio = self.audio_vae.mel_compression_ratio if self.audio_vae is not None else 4 - self.audio_vae_temporal_compression_ratio = ( - self.audio_vae.temporal_compression_ratio if self.audio_vae is not None else 4 - ) - self.transformer_spatial_patch_size = self.transformer.config.patch_size if self.transformer is not None else 1 - self.transformer_temporal_patch_size = ( - self.transformer.config.patch_size_t if self.transformer is not None else 1 - ) - self.audio_sampling_rate = self.audio_vae.config.sample_rate if self.audio_vae is not None else 16000 - self.audio_hop_length = self.audio_vae.config.mel_hop_length if self.audio_vae is not None else 160 - - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) - - # Tokenizer max length - tokenizer_max_length = 1024 - if self.tokenizer is not None: - tokenizer_max_length = self.tokenizer.model_max_length - if tokenizer_max_length is None or tokenizer_max_length > 100000: - encoder_config = getattr(self.text_encoder, "config", None) - config_max_len = getattr(encoder_config, "max_position_embeddings", None) - if config_max_len is None: - config_max_len = getattr(encoder_config, "max_seq_len", None) - tokenizer_max_length = config_max_len or 1024 - self.tokenizer_max_length = int(tokenizer_max_length) - - # Pipeline state - self._guidance_scale = None - self._attention_kwargs = None - self._interrupt = False - 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) - # ------------------------------------------------------------------ - - def _get_gemma_prompt_embeds( - self, - prompt: str | list[str], - num_videos_per_prompt: int = 1, - max_sequence_length: int = 1024, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - """Encode prompts using Gemma-3-12B, returning ALL 49 hidden states flattened. - - Stacks all 49 hidden states and flattens to [B, seq, hidden * 49]. The - connectors unflatten, apply per_token_rms_norm, and project internally - (same shape contract as LTX-2 since the `diffusers==0.38` connector - migration; the two differ only in the connector's internal norm path). - """ - device = device or self.device - dtype = dtype or self.text_encoder.dtype - - prompt = [prompt] if isinstance(prompt, str) else prompt - batch_size = len(prompt) - - if self.tokenizer is not None: - self.tokenizer.padding_side = "left" - if self.tokenizer.pad_token is None: - self.tokenizer.pad_token = self.tokenizer.eos_token - - prompt = [p.strip() for p in prompt] - text_inputs = self.tokenizer( - prompt, - padding="max_length", - max_length=max_sequence_length, - truncation=True, - add_special_tokens=True, - return_tensors="pt", - ) - text_input_ids = text_inputs.input_ids.to(device) - prompt_attention_mask = text_inputs.attention_mask.to(device) - - text_encoder_outputs = self.text_encoder( - input_ids=text_input_ids, - attention_mask=prompt_attention_mask, - output_hidden_states=True, - ) - - hidden_states = text_encoder_outputs.hidden_states - - # LTX-2.3: Stack ALL 49 hidden states and flatten - # [49 x (B, seq, 3840)] -> [B, seq, 3840, 49] -> [B, seq, 188160] - prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) - - _, seq_len, _ = prompt_embeds.shape - prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) - prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) - - prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) - prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) - - return prompt_embeds, prompt_attention_mask - - def encode_prompt( - self, - prompt: str | list[str], - negative_prompt: str | list[str] | None = None, - do_classifier_free_guidance: bool = True, - num_videos_per_prompt: int = 1, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - prompt_attention_mask: torch.Tensor | None = None, - negative_prompt_attention_mask: torch.Tensor | None = None, - max_sequence_length: int = 1024, - device: torch.device | None = None, - dtype: torch.dtype | None = None, - ): - device = device or self.device - - prompt = [prompt] if isinstance(prompt, str) else prompt - if prompt is not None: - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - if prompt_embeds is None: - prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=prompt, - num_videos_per_prompt=num_videos_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, - ) - - if do_classifier_free_guidance and negative_prompt_embeds is None: - negative_prompt = negative_prompt or "" - negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - - if prompt is not None and type(prompt) is not type(negative_prompt): - raise TypeError( - f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !=" - f" {type(prompt)}." - ) - if isinstance(negative_prompt, list) and batch_size != len(negative_prompt): - raise ValueError( - f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" - f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" - " the batch size of `prompt`." - ) - - negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( - prompt=negative_prompt, - num_videos_per_prompt=num_videos_per_prompt, - max_sequence_length=max_sequence_length, - device=device, - dtype=dtype, - ) - - return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask - - # ------------------------------------------------------------------ - # Latent utilities (shared with LTX2Pipeline) - # ------------------------------------------------------------------ - - @staticmethod - def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = latents.shape - post_patch_num_frames = num_frames // patch_size_t - post_patch_height = height // patch_size - post_patch_width = width // patch_size - latents = latents.reshape( - batch_size, - -1, - post_patch_num_frames, - patch_size_t, - post_patch_height, - patch_size, - post_patch_width, - patch_size, - ) - latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) - return latents - - @staticmethod - def _unpack_latents( - latents: torch.Tensor, - num_frames: int, - height: int, - width: int, - patch_size: int = 1, - patch_size_t: int = 1, - ) -> torch.Tensor: - batch_size = latents.size(0) - latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) - latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) - return latents - - @staticmethod - def _normalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents = (latents - latents_mean) * scaling_factor / latents_std - return latents - - @staticmethod - def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): - latents_mean = latents_mean.to(latents.device, latents.dtype) - latents_std = latents_std.to(latents.device, latents.dtype) - return (latents - latents_mean) / latents_std - - @staticmethod - def _denormalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents = latents * latents_std / scaling_factor + latents_mean - return latents - - @staticmethod - def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): - latents_mean = latents_mean.to(latents.device, latents.dtype) - latents_std = latents_std.to(latents.device, latents.dtype) - return (latents * latents_std) + latents_mean - - @staticmethod - def _pack_audio_latents( - latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None - ) -> torch.Tensor: - if patch_size is not None and patch_size_t is not None: - batch_size, num_channels, latent_length, latent_mel_bins = latents.shape - post_patch_latent_length = latent_length / patch_size_t - post_patch_mel_bins = latent_mel_bins / patch_size - latents = latents.reshape( - batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size - ) - latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) - else: - latents = latents.transpose(1, 2).flatten(2, 3) - return latents - - @staticmethod - def _unpack_audio_latents( - latents: torch.Tensor, - latent_length: int, - num_mel_bins: int, - patch_size: int | None = None, - patch_size_t: int | None = None, - ) -> torch.Tensor: - if patch_size is not None and patch_size_t is not None: - batch_size = latents.size(0) - latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) - latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) - else: - latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) - return latents - - @staticmethod - def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: - return latents[:, :num_frames] - - @staticmethod - def _get_sp_padded_audio_latent_length(audio_latent_length: int, sp_size: int) -> int: - if sp_size > 1: - audio_latent_length += (sp_size - (audio_latent_length % sp_size)) % sp_size - return audio_latent_length - - def _resolve_audio_latent_length(self, audio_latent_length: int, audio_latents: torch.Tensor | None) -> int: - if audio_latents is None or audio_latents.ndim != 4: - return audio_latent_length - - provided_latent_length = audio_latents.shape[2] - sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 - padded_latent_length = self._get_sp_padded_audio_latent_length(audio_latent_length, int(sp_size)) - - # Keep requested duration semantics when callers pass 4D latents that - # are already padded for SP; other 4D lengths retain shape inference. - if provided_latent_length in {audio_latent_length, padded_latent_length}: - return audio_latent_length - return provided_latent_length - - # ------------------------------------------------------------------ - # Latent preparation - # ------------------------------------------------------------------ - - def prepare_latents( - self, - batch_size: int = 1, - num_channels_latents: int = 128, - height: int = 512, - width: int = 768, - num_frames: int = 121, - noise_scale: float = 0.0, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | None = None, - latents: torch.Tensor | None = None, - ) -> torch.Tensor: - if latents is not None: - if latents.ndim == 5: - latents = self._normalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - latents = self._pack_latents( - latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ) - if latents.ndim != 3: - raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) - latents = noise_scale * noise + (1 - noise_scale) * latents - return latents.to(device=device, dtype=dtype) - - height = height // self.vae_spatial_compression_ratio - width = width // self.vae_spatial_compression_ratio - num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - shape = (batch_size, num_channels_latents, num_frames, height, width) - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) - return latents - - def prepare_audio_latents( - self, - batch_size: int = 1, - num_channels_latents: int = 8, - audio_latent_length: int = 1, - num_mel_bins: int = 64, - noise_scale: float = 0.0, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, int, int]: - original_latent_length = audio_latent_length - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - - sp_size = getattr(self.od_config.parallel_config, "sequence_parallel_size", 1) or 1 - padded_latent_length = self._get_sp_padded_audio_latent_length(original_latent_length, int(sp_size)) - - if latents is not None: - if latents.ndim == 4: - latents = self._pack_audio_latents(latents) - if latents.ndim != 3: - raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") - latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) - noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) - latents = noise_scale * noise + (1 - noise_scale) * latents - - if latents.shape[1] not in {original_latent_length, padded_latent_length}: - raise ValueError( - "Provided `audio_latents` has incompatible audio frame count " - f"{latents.shape[1]}; expected {original_latent_length} or {padded_latent_length}." - ) - - if latents.shape[1] == original_latent_length and padded_latent_length > original_latent_length: - padding = torch.zeros( - latents.shape[0], - padded_latent_length - original_latent_length, - latents.shape[2], - dtype=latents.dtype, - device=latents.device, - ) - latents = torch.cat([latents, padding], dim=1) - - return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length - - shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) - if isinstance(generator, list) and len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" - f" size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = self._pack_audio_latents(latents) - return latents, original_latent_length, padded_latent_length - - # ------------------------------------------------------------------ - # Properties - # ------------------------------------------------------------------ - - @property - def guidance_scale(self): - return self._guidance_scale - - @property - def do_classifier_free_guidance(self): - return self._guidance_scale is not None and self._guidance_scale > 1.0 - - @property - def num_timesteps(self): - return self._num_timesteps - - @property - def current_timestep(self): - return self._current_timestep - - @property - def interrupt(self): - return self._interrupt - - # ------------------------------------------------------------------ - # Input validation - # ------------------------------------------------------------------ - - def check_inputs( - self, - prompt, - height, - width, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - ): - if height % 32 != 0 or width % 32 != 0: - raise ValueError(f"`height` and `width` must be divisible by 32 but are {height} and {width}.") - if prompt is not None and prompt_embeds is not None: - raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.") - elif prompt is None and prompt_embeds is None: - raise ValueError("Provide either `prompt` or `prompt_embeds`.") - elif prompt is not None and not isinstance(prompt, (str, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if prompt_embeds is not None and prompt_attention_mask is None: - raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") - - if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: - raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: - raise ValueError( - "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when " - "passed directly, but got: `prompt_attention_mask` " - f"{prompt_attention_mask.shape} != `negative_prompt_attention_mask` " - f"{negative_prompt_attention_mask.shape}." - ) - - # ------------------------------------------------------------------ - # Cache context - # ------------------------------------------------------------------ - - def _transformer_cache_context(self, context_name: str): - cache_context = getattr(self.transformer, "cache_context", None) - if callable(cache_context): - return cache_context(context_name) - return nullcontext() - - # ------------------------------------------------------------------ - # CFG helpers - # ------------------------------------------------------------------ - - @staticmethod - def _combine_x0_space_cfg( - sample: torch.Tensor, - positive_noise_pred: torch.Tensor, - negative_noise_pred: torch.Tensor, - sigma: torch.Tensor, - guidance_scale: float, - ) -> torch.Tensor: - x0_cond = sample - positive_noise_pred * sigma - x0_uncond = sample - negative_noise_pred * sigma - x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond) - return (sample - x0_guided) / sigma - - def predict_noise(self, **kwargs): - with self._transformer_cache_context("cond_uncond"): - noise_pred_video, noise_pred_audio = self.transformer(**kwargs) - return noise_pred_video.float(), noise_pred_audio.float() - - def combine_cfg_noise( - self, - positive_noise_pred, - negative_noise_pred, - true_cfg_scale, - cfg_normalize=False, - *, - video_latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - video_sigma: torch.Tensor | None = None, - audio_sigma: torch.Tensor | None = None, - ): - if video_latents is None or audio_latents is None or video_sigma is None or audio_sigma is None: - raise ValueError("LTX23Pipeline applies CFG in x0-space and requires video/audio latents and sigmas.") - - video_pos, audio_pos = positive_noise_pred - video_neg, audio_neg = negative_noise_pred - video_combined = self._combine_x0_space_cfg( - video_latents, - video_pos, - video_neg, - video_sigma, - true_cfg_scale, - ) - audio_combined = self._combine_x0_space_cfg( - audio_latents, - audio_pos, - audio_neg, - audio_sigma, - true_cfg_scale, - ) - if cfg_normalize: - video_combined = self.cfg_normalize_function(video_pos, video_combined) - audio_combined = self.cfg_normalize_function(audio_pos, audio_combined) - return video_combined, audio_combined - - def predict_noise_with_parallel_cfg( - self, - true_cfg_scale: float, - positive_kwargs: dict[str, Any], - negative_kwargs: dict[str, Any], - cfg_normalize: bool = True, - output_slice: int | None = None, - *, - video_latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - video_sigma: torch.Tensor | None = None, - audio_sigma: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: - if output_slice is None: - return pred - return pred[0][:, :output_slice], pred[1][:, :output_slice] - - cfg_world_size = get_classifier_free_guidance_world_size() - if cfg_world_size != 2: - raise ValueError(f"LTX23Pipeline parallel CFG requires cfg_parallel_size 2, but got {cfg_world_size}.") - - cfg_group = get_cfg_group() - cfg_rank = get_classifier_free_guidance_rank() - branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs - local_video_pred, local_audio_pred = maybe_slice(self.predict_noise(**branch_kwargs)) - - gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True) - gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True) - positive_noise_pred = (gathered_video[0], gathered_audio[0]) - negative_noise_pred = (gathered_video[1], gathered_audio[1]) - - return self.combine_cfg_noise( - positive_noise_pred, - negative_noise_pred, - true_cfg_scale, - cfg_normalize, - video_latents=video_latents, - audio_latents=audio_latents, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - def _synchronize_cfg_parallel_step_output( - self, - latents: tuple[torch.Tensor, torch.Tensor], - do_true_cfg: bool, - ) -> tuple[torch.Tensor, torch.Tensor]: - if not (do_true_cfg and get_classifier_free_guidance_world_size() > 1): - return latents - - latents = tuple(tensor.contiguous() for tensor in latents) - device = next((tensor.device for tensor in latents if tensor.is_cuda), None) - if device is not None: - torch.cuda.current_stream(device).synchronize() - return latents - - # ------------------------------------------------------------------ - # Forward pass - # ------------------------------------------------------------------ - - @torch.no_grad() - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - sigmas: list[float] | None = None, - timesteps: list[int] | None = None, - guidance_scale: float = 4.0, - noise_scale: float = 0.0, - num_videos_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - prompt_attention_mask: torch.Tensor | None = None, - negative_prompt_attention_mask: torch.Tensor | None = None, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - output_type: str = "np", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int | None = None, - ) -> DiffusionOutput: - # ---- Extract from request ---- - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt - if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): - negative_prompt = None - elif req.prompts: - negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - - height = req.sampling_params.height or height or 512 - width = req.sampling_params.width or width or 768 - num_frames = req.sampling_params.num_frames or num_frames or 121 - frame_rate = req.sampling_params.resolved_frame_rate or frame_rate or 24.0 - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps or 40 - # Enforce minimum of 2 timesteps for flow matching scheduler - if timesteps is None: - num_inference_steps = max(int(num_inference_steps), 2) - elif len(timesteps) < 2: - raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") - num_videos_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_videos_per_prompt or 1 - ) - max_sequence_length = ( - req.sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length - ) - - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - - if generator is None: - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) - - latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents - audio_latents = ( - req.sampling_params.audio_latents - if req.sampling_params.audio_latents is not None - else req.sampling_params.extra_args.get("audio_latents", audio_latents) - ) - - # Override with pre-computed embeddings if provided in request - req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - prompt_embeds = torch.stack(req_prompt_embeds) - - req_negative_prompt_embeds = [_get_prompt_field(p, "negative_prompt_embeds") for p in req.prompts] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) - - req_prompt_attention_masks = [ - _get_prompt_field(p, "prompt_attention_mask") or _get_prompt_field(p, "attention_mask") for p in req.prompts - ] - if any(m is not None for m in req_prompt_attention_masks): - prompt_attention_mask = torch.stack(req_prompt_attention_masks) - - req_negative_attention_masks = [ - _get_prompt_field(p, "negative_prompt_attention_mask") or _get_prompt_field(p, "negative_attention_mask") - for p in req.prompts - ] - if any(m is not None for m in req_negative_attention_masks): - negative_prompt_attention_mask = torch.stack(req_negative_attention_masks) - - if req.sampling_params.decode_timestep is not None: - decode_timestep = req.sampling_params.decode_timestep - if req.sampling_params.decode_noise_scale is not None: - decode_noise_scale = req.sampling_params.decode_noise_scale - if req.sampling_params.output_type is not None: - output_type = req.sampling_params.output_type - - self.check_inputs( - prompt=prompt, - height=height, - width=width, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._attention_kwargs = attention_kwargs - self._interrupt = False - self._current_timestep = None - cfg_world_size = get_classifier_free_guidance_world_size() - if self.do_classifier_free_guidance and cfg_world_size not in (1, 2): - raise ValueError( - f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." - ) - cfg_parallel_ready = self.do_classifier_free_guidance and cfg_world_size > 1 - - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self.device - - # ---- Encode prompts ---- - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - max_sequence_length=max_sequence_length, - device=device, - ) - - # ---- Connectors (LTX-2.3: padding_side API) ---- - # Concatenate negative + positive embeddings BEFORE connector call, - # matching diffusers which calls connectors once with batch=2. - # This ensures batch-dependent operations produce identical results. - if self.do_classifier_free_guidance: - 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) - - 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 - ) - - positive_connector_prompt_embeds = connector_prompt_embeds - positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds - positive_connector_attention_mask = connector_attention_mask - negative_connector_prompt_embeds = None - negative_connector_audio_prompt_embeds = None - negative_connector_attention_mask = None - if self.do_classifier_free_guidance: - split_batch = batch_size * num_videos_per_prompt - negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch] - positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:] - negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch] - positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:] - negative_connector_attention_mask = connector_attention_mask[:split_batch] - positive_connector_attention_mask = connector_attention_mask[split_batch:] - - # ---- Prepare latents ---- - latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - latent_height = height // self.vae_spatial_compression_ratio - latent_width = width // self.vae_spatial_compression_ratio - if latents is not None and latents.ndim == 5: - _, _, latent_num_frames, latent_height, latent_width = latents.shape - - num_channels_latents = self.transformer.config.in_channels - latents = self.prepare_latents( - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - noise_scale, - torch.float32, - device, - generator, - latents, - ) - - duration_s = num_frames / frame_rate - audio_latents_per_second = ( - self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) - ) - audio_num_frames = round(duration_s * audio_latents_per_second) - audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, audio_latents) - - num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64 - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - num_channels_latents_audio = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8 - audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents( - batch_size * num_videos_per_prompt, - num_channels_latents=num_channels_latents_audio, - audio_latent_length=audio_num_frames, - num_mel_bins=num_mel_bins, - noise_scale=noise_scale, - dtype=torch.float32, - device=device, - generator=generator, - latents=audio_latents, - ) - - # ---- Scheduler setup ---- - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - # Use max_image_seq_len (not actual video_sequence_length) for mu calculation, - # matching diffusers' LTX2Pipeline which hardcodes this value. - mu = calculate_shift( - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_image_seq_len", 1024), - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_shift", 0.95), - self.scheduler.config.get("max_shift", 2.05), - ) - audio_scheduler = copy.deepcopy(self.scheduler) - video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_scheduler) - _ = retrieve_timesteps(audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu) - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas=sigmas, - mu=mu, - ) - self._num_timesteps = len(timesteps) - - # ---- RoPE coordinates ---- - video_coords = self.transformer.rope.prepare_video_coords( - latents.shape[0], - latent_num_frames, - latent_height, - latent_width, - latents.device, - fps=frame_rate, - ) - audio_coords = self.transformer.audio_rope.prepare_audio_coords( - audio_latents.shape[0], - padded_audio_num_frames, - audio_latents.device, - ) - - # ---- CFG: duplicate coords for single-rank batch=2 CFG ---- - # Connector outputs are already batch=2 (neg+pos concatenated before connector call) - if self.do_classifier_free_guidance and not cfg_parallel_ready: - video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) - audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) - - # ---- Denoising loop ---- - # Uses x0-space CFG (delta formulation) matching diffusers' LTX2Pipeline. - # The velocity predictions are converted to x0, guidance is applied in x0 - # space, then converted back to velocity for the scheduler step. - with self.progress_bar(total=len(timesteps)) as pbar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - - if cfg_parallel_ready: - latent_model_input = latents.to(positive_connector_prompt_embeds.dtype) - audio_latent_model_input = audio_latents.to(positive_connector_prompt_embeds.dtype) - ts = t.expand(latent_model_input.shape[0]) - positive_kwargs = { - "hidden_states": latent_model_input, - "audio_hidden_states": audio_latent_model_input, - "encoder_hidden_states": positive_connector_prompt_embeds, - "audio_encoder_hidden_states": positive_connector_audio_prompt_embeds, - "timestep": ts, - "sigma": ts, - "encoder_attention_mask": positive_connector_attention_mask, - "audio_encoder_attention_mask": positive_connector_attention_mask, - "num_frames": latent_num_frames, - "height": latent_height, - "width": latent_width, - "fps": frame_rate, - "audio_num_frames": padded_audio_num_frames, - "video_coords": video_coords, - "audio_coords": audio_coords, - "attention_kwargs": attention_kwargs, - "return_dict": False, - } - negative_kwargs = { - **positive_kwargs, - "encoder_hidden_states": negative_connector_prompt_embeds, - "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds, - "encoder_attention_mask": negative_connector_attention_mask, - "audio_encoder_attention_mask": negative_connector_attention_mask, - } - noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg( - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - video_latents=latents, - audio_latents=audio_latents, - video_sigma=self.scheduler.sigmas[i], - audio_sigma=audio_scheduler.sigmas[i], - ) - - latents, audio_latents = self.scheduler_step_maybe_with_cfg( - (noise_pred_video, noise_pred_audio), - (t, t), - (latents, audio_latents), - do_true_cfg=self.do_classifier_free_guidance, - per_request_scheduler=video_audio_scheduler, - ) - latents, audio_latents = self._synchronize_cfg_parallel_step_output( - (latents, audio_latents), - do_true_cfg=self.do_classifier_free_guidance, - ) - else: - latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents - latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype) - audio_latent_model_input = ( - torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents - ) - audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype) - ts = t.expand(latent_model_input.shape[0]) - - with self._transformer_cache_context("cond_uncond"): - noise_pred_video, noise_pred_audio = self.transformer( - hidden_states=latent_model_input, - audio_hidden_states=audio_latent_model_input, - encoder_hidden_states=connector_prompt_embeds, - audio_encoder_hidden_states=connector_audio_prompt_embeds, - timestep=ts, - sigma=ts, # LTX-2.3: sigma for prompt_adaln - encoder_attention_mask=connector_attention_mask, - audio_encoder_attention_mask=connector_attention_mask, - num_frames=latent_num_frames, - height=latent_height, - width=latent_width, - fps=frame_rate, - audio_num_frames=padded_audio_num_frames, - video_coords=video_coords, - audio_coords=audio_coords, - attention_kwargs=attention_kwargs, - return_dict=False, - ) - - noise_pred_video = noise_pred_video.float() - noise_pred_audio = noise_pred_audio.float() - - if self.do_classifier_free_guidance: - noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) - noise_pred_video = self._combine_x0_space_cfg( - latents, - noise_pred_video_cond, - noise_pred_video_uncond, - self.scheduler.sigmas[i], - guidance_scale, - ) - - noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) - noise_pred_audio = self._combine_x0_space_cfg( - audio_latents, - noise_pred_audio_cond, - noise_pred_audio_uncond, - audio_scheduler.sigmas[i], - guidance_scale, - ) - - latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0] - audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] - - pbar.update() - - # ---- Unpack and denormalize ---- - latents = self._unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - latents = self._denormalize_latents( - latents, - self.vae.latents_mean, - self.vae.latents_std, - self.vae.config.scaling_factor, - ) - - audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames) - audio_latents = self._denormalize_audio_latents( - audio_latents, - self.audio_vae.latents_mean, - self.audio_vae.latents_std, - ) - audio_latents = self._unpack_audio_latents( - audio_latents, - original_audio_num_frames, - num_mel_bins=latent_mel_bins, - ) - - # ---- Decode ---- - if output_type == "latent": - video = latents - audio = audio_latents - else: - latents = latents.to(connector_prompt_embeds.dtype) - - if not self.vae.config.timestep_conditioning: - timestep_decode = None - else: - noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) - if not isinstance(decode_timestep, list): - decode_timestep = [decode_timestep] * batch_size - if decode_noise_scale is None: - decode_noise_scale = decode_timestep - elif not isinstance(decode_noise_scale, list): - decode_noise_scale = [decode_noise_scale] * batch_size - timestep_decode = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) - decode_noise_scale_t = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ - :, None, None, None, None - ] - latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise - - latents = latents.to(self.vae.dtype) - video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] - if video.numel() > 0: - video = self.video_processor.postprocess_video(video, output_type=output_type) - - audio_latents = audio_latents.to(self.audio_vae.dtype) - generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] - - audio = self.vocoder(generated_mel_spectrograms) - - 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) - return loader.load_weights(weights) - - -class LTX23ImageToVideoPipeline(nn.Module): - """LTX-2.3 image-to-video pipeline placeholder.""" - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__() - raise NotImplementedError( - "LTX23ImageToVideoPipeline is not yet implemented. " - "Use LTX23Pipeline for single-stage text-to-video generation." - ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py deleted file mode 100644 index d30ef546400..00000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py +++ /dev/null @@ -1,18 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -"""Re-exports for LTX-2.3 I2V pipeline variants. - -The registry loads pipeline classes by (mod_folder, mod_relname, cls_name). -This module exposes the I2V class names so the registry can find them. -""" - -from .pipeline_ltx2_3 import ( - LTX23ImageToVideoPipeline, - get_ltx2_post_process_func, # noqa: F401 - loaded by registry via getattr -) - -__all__ = [ - "LTX23ImageToVideoPipeline", - "get_ltx2_post_process_func", -] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py deleted file mode 100644 index e93f4e7447f..00000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ /dev/null @@ -1,903 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import copy -import os -from collections.abc import Iterable -from typing import Any, ClassVar - -import numpy as np -import PIL.Image -import torch -import torch.nn as nn -from diffusers import FlowMatchEulerDiscreteScheduler -from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps -from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents -from diffusers.utils.torch_utils import randn_tensor -from diffusers.video_processor import VideoProcessor -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.utils import get_local_device -from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager -from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader -from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery -from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.lora.request import LoRARequest - -from .pipeline_ltx2 import ( - LTX2Pipeline, - _get_prompt_field, - calculate_shift, -) -from .pipeline_ltx2 import ( - get_ltx2_post_process_func as _get_ltx2_post_process_func, -) -from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline - -logger = init_logger(__name__) - - -def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): - return _get_ltx2_post_process_func(od_config) - - -class _I2VVideoAudioScheduler: - """Composite scheduler for I2V: uses _step_video_latents_i2v for video, standard step for audio.""" - - def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width): - self.video_scheduler = pipeline.scheduler - self.audio_scheduler = audio_scheduler - self._pipeline = pipeline - self._latent_num_frames = latent_num_frames - self._latent_height = latent_height - self._latent_width = latent_width - - def step(self, noise_pred, t, latents, return_dict=False, generator=None): - video_out = self._pipeline._step_video_latents_i2v( - noise_pred[0], - latents[0], - t[0], - self._latent_num_frames, - self._latent_height, - self._latent_width, - ) - audio_out = self.audio_scheduler.step(noise_pred[1], t[1], latents[1], return_dict=False, generator=generator)[ - 0 - ] - return ((video_out, audio_out),) - - -class LTX2ImageToVideoPipeline(LTX2Pipeline): - support_image_input = True - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__(od_config=od_config, prefix=prefix) - self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear") - - @staticmethod - def _normalize_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 - ) -> torch.Tensor: - latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) - latents = (latents - latents_mean) * scaling_factor / latents_std - return latents - - def prepare_latents( - self, - image: torch.Tensor | None = None, - batch_size: int = 1, - num_channels_latents: int = 128, - height: int = 512, - width: int = 768, - num_frames: int = 121, - noise_scale: float = 0.0, - dtype: torch.dtype | None = None, - device: torch.device | None = None, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: - height = height // self.vae_spatial_compression_ratio - width = width // self.vae_spatial_compression_ratio - num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - - shape = (batch_size, num_channels_latents, num_frames, height, width) - mask_shape = (batch_size, 1, num_frames, height, width) - - if latents is not None: - if latents.ndim == 5: - # conditioning_mask needs to the same shape as latents in two stages generation. - batch_size, _, num_frames, height, width = latents.shape - mask_shape = (batch_size, 1, num_frames, height, width) - conditioning_mask = latents.new_zeros(mask_shape) - conditioning_mask[:, :, 0] = 1.0 - - latents = self._normalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - latents = self._create_noised_state(latents, noise_scale * (1 - conditioning_mask), generator) - # latents are of shape [B, C, F, H, W], need to be packed - latents = self._pack_latents( - latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ) - else: - conditioning_mask = latents.new_zeros(mask_shape) - conditioning_mask[:, :, 0] = 1.0 - - conditioning_mask = self._pack_latents( - conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ).squeeze(-1) - - if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape: - raise ValueError( - "Provided `latents` tensor has shape" - f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." - ) - return latents.to(device=device, dtype=dtype), conditioning_mask - - if image is None: - raise ValueError("`image` must be provided when `latents` is None.") - - image_batch_size = image.shape[0] - if image_batch_size == 0: - raise ValueError("`image` batch is empty.") - if batch_size % image_batch_size != 0: - raise ValueError( - f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) " - "for image-to-video outputs." - ) - num_videos_per_prompt = batch_size // image_batch_size - - if isinstance(generator, list): - if len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective" - f" batch size of {batch_size}. Make sure the batch size matches the length of the generators." - ) - - image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)] - init_latents = [ - retrieve_latents(self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), image_generators[i], "argmax") - for i in range(image_batch_size) - ] - else: - init_latents = [ - retrieve_latents(self.vae.encode(img.unsqueeze(0).unsqueeze(2)), generator, "argmax") for img in image - ] - - init_latents = torch.cat(init_latents, dim=0).to(dtype) - if num_videos_per_prompt > 1: - init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0) - init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std) - init_latents = init_latents.repeat(1, 1, num_frames, 1, 1) - - conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype) - conditioning_mask[:, :, 0] = 1.0 - - noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) - latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask) - - conditioning_mask = self._pack_latents( - conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ).squeeze(-1) - latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) - - return latents, conditioning_mask - - def check_inputs( - self, - image, - height, - width, - prompt, - latents=None, - prompt_embeds=None, - negative_prompt_embeds=None, - prompt_attention_mask=None, - negative_prompt_attention_mask=None, - ): - if height % 32 != 0 or width % 32 != 0: - raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") - - if image is None and latents is None: - raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.") - - if prompt is not None and prompt_embeds is not None: - raise ValueError( - f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" - " only forward one of the two." - ) - if prompt is None and prompt_embeds is None: - raise ValueError( - "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." - ) - if prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): - raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - - if prompt_embeds is not None and prompt_attention_mask is None: - raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") - - if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: - raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") - - if prompt_embeds is not None and negative_prompt_embeds is not None: - if prompt_embeds.shape != negative_prompt_embeds.shape: - raise ValueError( - "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" - f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" - f" {negative_prompt_embeds.shape}." - ) - if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: - raise ValueError( - "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when" - " passed directly, but" - f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" - f" {negative_prompt_attention_mask.shape}." - ) - - def _step_video_latents_i2v( - self, - noise_pred_video: torch.Tensor, - latents: torch.Tensor, - t: torch.Tensor, - latent_num_frames: int, - latent_height: int, - latent_width: int, - ) -> torch.Tensor: - noise_pred_video = self._unpack_latents( - noise_pred_video, - latent_num_frames, - latent_height, - latent_width, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - latents_unpacked = self._unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - - noise_pred_video = noise_pred_video[:, :, 1:] - noise_latents = latents_unpacked[:, :, 1:] - pred_latents = self.scheduler.step(noise_pred_video, t, noise_latents, return_dict=False)[0] - - latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2) - latents = self._pack_latents( - latents_unpacked, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size - ) - return latents - - @torch.no_grad() - def forward( - self, - req: OmniDiffusionRequest, - image: PIL.Image.Image | torch.Tensor | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - sigmas: list[float] | None = None, - timesteps: list[int] | None = None, - guidance_scale: float = 4.0, - guidance_rescale: float = 0.0, - noise_scale: float = 0.0, - num_videos_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - prompt_attention_mask: torch.Tensor | None = None, - negative_prompt_attention_mask: torch.Tensor | None = None, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - output_type: str = "np", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int | None = None, - ) -> DiffusionOutput: - # Extract prompt/negative_prompt from request. - # Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt - if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): - negative_prompt = None - elif req.prompts: - negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - - height = req.sampling_params.height or height or 512 - width = req.sampling_params.width or width or 768 - num_frames = req.sampling_params.num_frames or num_frames or 121 - frame_rate = req.sampling_params.resolved_frame_rate or frame_rate or 24.0 - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps or 40 - if timesteps is None: - num_inference_steps = max(int(num_inference_steps), 2) - elif len(timesteps) < 2: - raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") - num_videos_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_videos_per_prompt or 1 - ) - max_sequence_length = ( - req.sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length - ) - - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale - if req.sampling_params.guidance_rescale is not None: - guidance_rescale = req.sampling_params.guidance_rescale - - if generator is None: - generator = req.sampling_params.generator - if generator is None and req.sampling_params.seed is not None: - generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) - - latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents - audio_latents = ( - req.sampling_params.audio_latents - if req.sampling_params.audio_latents is not None - else req.sampling_params.extra_args.get("audio_latents", audio_latents) - ) - - # Override with pre-computed embeddings if provided in request. - req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts] - if any(p is not None for p in req_prompt_embeds): - prompt_embeds = torch.stack(req_prompt_embeds) # type: ignore[arg-type] - - req_negative_prompt_embeds = [_get_prompt_field(p, "negative_prompt_embeds") for p in req.prompts] - if any(p is not None for p in req_negative_prompt_embeds): - negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) # type: ignore[arg-type] - - req_prompt_attention_masks = [ - _get_prompt_field(p, "prompt_attention_mask") or _get_prompt_field(p, "attention_mask") for p in req.prompts - ] - if any(m is not None for m in req_prompt_attention_masks): - prompt_attention_mask = torch.stack(req_prompt_attention_masks) # type: ignore[arg-type] - - req_negative_attention_masks = [ - _get_prompt_field(p, "negative_prompt_attention_mask") or _get_prompt_field(p, "negative_attention_mask") - for p in req.prompts - ] - if any(m is not None for m in req_negative_attention_masks): - negative_prompt_attention_mask = torch.stack(req_negative_attention_masks) # type: ignore[arg-type] - - if req.sampling_params.decode_timestep is not None: - decode_timestep = req.sampling_params.decode_timestep - if req.sampling_params.decode_noise_scale is not None: - decode_noise_scale = req.sampling_params.decode_noise_scale - if req.sampling_params.output_type is not None: - output_type = req.sampling_params.output_type - - if image is None and req.prompts: - raw_images = [] - for prompt_item in req.prompts: - if isinstance(prompt_item, str): - raw_image = None - else: - multi_modal_data = prompt_item.get("multi_modal_data") or {} - raw_image = multi_modal_data.get("image") - if raw_image is None: - additional = prompt_item.get("additional_information") or {} - raw_image = ( - additional.get("preprocessed_image") - or additional.get("pixel_values") - or additional.get("image") - ) - if isinstance(raw_image, list): - raw_image = raw_image[0] if raw_image else None - if isinstance(raw_image, str): - raw_image = PIL.Image.open(raw_image).convert("RGB") - raw_images.append(raw_image) - - if any(img is None for img in raw_images): - if latents is None: - raise ValueError("Image is required for LTX2 I2V generation.") - if len(raw_images) == 1: - image = raw_images[0] - elif raw_images: - image = raw_images - - self.check_inputs( - image=image, - height=height, - width=width, - prompt=prompt, - latents=latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - ) - - self._guidance_scale = guidance_scale - self._guidance_rescale = guidance_rescale - self._attention_kwargs = attention_kwargs - self._interrupt = False - self._current_timestep = None - - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] - - device = self.device - - ( - prompt_embeds, - prompt_attention_mask, - negative_prompt_embeds, - negative_prompt_attention_mask, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - num_videos_per_prompt=num_videos_per_prompt, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - max_sequence_length=max_sequence_length, - device=device, - ) - # Compute positive prompt connectors - 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 - ) - - # Compute negative prompt connectors when CFG is enabled - negative_connector_prompt_embeds = None - negative_connector_audio_prompt_embeds = None - negative_connector_attention_mask = None - if self.do_classifier_free_guidance: - ( - negative_connector_prompt_embeds, - negative_connector_audio_prompt_embeds, - negative_connector_attention_mask, - ) = self.connectors( - negative_prompt_embeds, - negative_prompt_attention_mask, - padding_side=tokenizer_padding_side, - ) - - latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 - latent_height = height // self.vae_spatial_compression_ratio - latent_width = width // self.vae_spatial_compression_ratio - if latents is not None: - if latents.ndim == 5: - logger.info( - "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." # noqa - ) - _, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W] - elif latents.ndim == 3: - logger.warning( - f"You have supplied packed `latents` of shape {latents.shape}, so the latent dims cannot be" - f" inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct." - ) - else: - raise ValueError( - f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." # noqa - ) - video_sequence_length = latent_num_frames * latent_height * latent_width - - if latents is None: - if isinstance(image, torch.Tensor): - if image.ndim == 3: - image = image.unsqueeze(0) - elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor): - image = torch.stack(image, dim=0) - else: - image = self.video_processor.preprocess(image, height=height, width=width) - image = image.to(device=device, dtype=prompt_embeds.dtype) - - num_channels_latents = self.transformer.config.in_channels - latents, conditioning_mask = self.prepare_latents( - image, - batch_size * num_videos_per_prompt, - num_channels_latents, - height, - width, - num_frames, - noise_scale, - torch.float32, - device, - generator, - latents, - ) - - duration_s = num_frames / frame_rate - audio_latents_per_second = ( - self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) - ) - audio_num_frames = round(duration_s * audio_latents_per_second) - if audio_latents is not None: - if audio_latents.ndim == 4: - logger.info( - "Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." # noqa - ) - _, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M] - elif audio_latents.ndim == 3: - logger.warning( - f"You have supplied packed `audio_latents` of shape {audio_latents.shape}, so the latent dims" - f" cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct." - ) - else: - raise ValueError( - f"Provided `audio_latents` tensor has shape {audio_latents.shape}, but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." # noqa - ) - - num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 - latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio - num_channels_latents_audio = ( - self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 - ) - audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents( - batch_size * num_videos_per_prompt, - num_channels_latents=num_channels_latents_audio, - audio_latent_length=audio_num_frames, - num_mel_bins=num_mel_bins, - noise_scale=noise_scale, - dtype=torch.float32, - device=device, - generator=generator, - latents=audio_latents, - ) - - sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas - mu = calculate_shift( - video_sequence_length, - self.scheduler.config.get("base_image_seq_len", 1024), - self.scheduler.config.get("max_image_seq_len", 4096), - self.scheduler.config.get("base_shift", 0.95), - self.scheduler.config.get("max_shift", 2.05), - ) - audio_scheduler = copy.deepcopy(self.scheduler) - _ = retrieve_timesteps( - audio_scheduler, - num_inference_steps, - device, - timesteps, - sigmas=sigmas, - mu=mu, - ) - timesteps, num_inference_steps = retrieve_timesteps( - self.scheduler, - num_inference_steps, - device, - timesteps, - sigmas=sigmas, - mu=mu, - ) - self._num_timesteps = len(timesteps) - - video_coords = self.transformer.rope.prepare_video_coords( - latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate - ) - audio_coords = self.transformer.audio_rope.prepare_audio_coords( - audio_latents.shape[0], padded_audio_num_frames, audio_latents.device - ) - - i2v_scheduler = _I2VVideoAudioScheduler( - pipeline=self, - audio_scheduler=audio_scheduler, - latent_num_frames=latent_num_frames, - latent_height=latent_height, - latent_width=latent_width, - ) - # No coord duplication needed: mixin handles CFG via separate forward calls. - - with self.progress_bar(total=len(timesteps)) as pbar: - for i, t in enumerate(timesteps): - if self.interrupt: - continue - - self._current_timestep = t - - latent_model_input = latents.to(prompt_embeds.dtype) - audio_latent_model_input = audio_latents.to(prompt_embeds.dtype) - timestep = t.expand(latent_model_input.shape[0]) - video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask) - do_true_cfg = self.do_classifier_free_guidance - - positive_kwargs = { - "hidden_states": latent_model_input, - "audio_hidden_states": audio_latent_model_input, - "encoder_hidden_states": connector_prompt_embeds, - "audio_encoder_hidden_states": connector_audio_prompt_embeds, - "timestep": video_timestep, - "audio_timestep": timestep, - "encoder_attention_mask": connector_attention_mask, - "audio_encoder_attention_mask": connector_attention_mask, - "num_frames": latent_num_frames, - "height": latent_height, - "width": latent_width, - "fps": frame_rate, - "audio_num_frames": padded_audio_num_frames, - "video_coords": video_coords, - "audio_coords": audio_coords, - "attention_kwargs": attention_kwargs, - "return_dict": False, - } - negative_kwargs = ( - { - **positive_kwargs, - "encoder_hidden_states": negative_connector_prompt_embeds, - "audio_encoder_hidden_states": negative_connector_audio_prompt_embeds, - "encoder_attention_mask": negative_connector_attention_mask, - "audio_encoder_attention_mask": negative_connector_attention_mask, - } - if do_true_cfg - else None - ) - - noise_pred_video, noise_pred_audio = self.predict_noise_maybe_with_cfg( - do_true_cfg=do_true_cfg, - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - ) - - latents, audio_latents = self.scheduler_step_maybe_with_cfg( - (noise_pred_video, noise_pred_audio), - (t, t), - (latents, audio_latents), - do_true_cfg=do_true_cfg, - per_request_scheduler=i2v_scheduler, - ) - latents, audio_latents = self._synchronize_cfg_parallel_step_output( - (latents, audio_latents), - do_true_cfg=do_true_cfg, - ) - - pbar.update() - - latents = self._unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - self.transformer_spatial_patch_size, - self.transformer_temporal_patch_size, - ) - latents = self._denormalize_latents( - latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor - ) - - audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames) - audio_latents = self._denormalize_audio_latents( - audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std - ) - audio_latents = self._unpack_audio_latents( - audio_latents, - original_audio_num_frames, - num_mel_bins=latent_mel_bins, - ) - - if output_type == "latent": - video = latents - audio = audio_latents - else: - latents = latents.to(prompt_embeds.dtype) - - if not self.vae.config.timestep_conditioning: - timestep = None - else: - noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) - if not isinstance(decode_timestep, list): - decode_timestep = [decode_timestep] * batch_size - if decode_noise_scale is None: - decode_noise_scale = decode_timestep - elif not isinstance(decode_noise_scale, list): - decode_noise_scale = [decode_noise_scale] * batch_size - - timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) - decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ - :, None, None, None, None - ] - latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise - - latents = latents.to(self.vae.dtype) - video = self.vae.decode(latents, timestep, return_dict=False)[0] - video = self.video_processor.postprocess_video(video, output_type=output_type) - - audio_latents = audio_latents.to(self.audio_vae.dtype) - generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] - audio = self.vocoder(generated_mel_spectrograms) - - if not return_dict: - return DiffusionOutput(output=(video, audio)) - - return DiffusionOutput(output=(video, audio)) - - -class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsComponentDiscovery): - """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation""" - - support_image_input = True - dummy_run_num_frames = 2 - - _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] - _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] - _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] - - def __init__( - self, - *, - od_config: OmniDiffusionConfig, - prefix: str = "", - ): - super().__init__() - - self.device = get_local_device() - self.dtype = getattr(od_config, "dtype", torch.bfloat16) - self.model_path = od_config.model - self.distilled = False - # User provided model path may contain '/' in the end and basename function - # will not return the expected directory name, so we need to remove it by normpath - if "distilled" in os.path.basename(os.path.normpath(self.model_path)): - self.distilled = True - else: - raise NotImplementedError(f"{self.model_path} is not supported for {self.__class__.__name__}.") - - self.pipe = LTX2ImageToVideoPipeline(od_config=od_config, prefix=prefix) - self.upsample_pipe = LTX2LatentUpsamplePipeline( - vae=self.pipe.vae, - od_config=od_config, - ) - - self.lora_manager = DiffusionLoRAManager( - pipeline=self.pipe, - device=self.device, - dtype=self.dtype, - max_cached_adapters=od_config.max_cpu_loras, - ) - - self.weights_sources = [ - DiffusersPipelineLoader.ComponentSource( - model_or_path=od_config.model, - subfolder="transformer", - revision=None, - prefix="pipe.transformer.", - fall_back_to_pt=True, - ), - ] - - @torch.no_grad() - def forward( - self, - req: OmniDiffusionRequest, - image: PIL.Image.Image | torch.Tensor | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - height: int | None = None, - width: int | None = None, - num_frames: int | None = None, - frame_rate: float | None = None, - num_inference_steps: int | None = None, - sigmas: list[float] | None = None, - timesteps: list[int] | None = None, - guidance_scale: float = 4.0, - guidance_rescale: float = 0.0, - noise_scale: float = 0.0, - num_videos_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - audio_latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - prompt_attention_mask: torch.Tensor | None = None, - negative_prompt_attention_mask: torch.Tensor | None = None, - decode_timestep: float | list[float] = 0.0, - decode_noise_scale: float | list[float] | None = None, - output_type: str = "np", - return_dict: bool = True, - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int | None = None, - ): - video_latent, audio_latent = self.pipe( - req=req, - image=image, - prompt=prompt, - negative_prompt=negative_prompt, - height=height, - width=width, - num_frames=num_frames, - frame_rate=frame_rate, - num_inference_steps=num_inference_steps, - sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None, - timesteps=timesteps, - guidance_scale=guidance_scale, - guidance_rescale=guidance_rescale, - noise_scale=noise_scale, - num_videos_per_prompt=num_videos_per_prompt, - generator=generator, - latents=latents, - audio_latents=audio_latents, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - prompt_attention_mask=prompt_attention_mask, - negative_prompt_attention_mask=negative_prompt_attention_mask, - decode_timestep=decode_timestep, - decode_noise_scale=decode_noise_scale, - output_type="latent", - return_dict=return_dict, - attention_kwargs=attention_kwargs, - max_sequence_length=max_sequence_length, - ).output - - upscaled_video_latent = self.upsample_pipe( - latents=video_latent, - output_type="latent", - return_dict=False, - )[0] - - if not self.distilled: - # Load Stage 2 distilled LoRA - lora_path = f"{self.model_path}/ltx-2-19b-distilled-lora-384.safetensors" - lora_request = LoRARequest( - lora_name="stage_2_distilled", - lora_int_id=1, - lora_path=lora_path, - ) - self.lora_manager.set_active_adapter(lora_request, lora_scale=1.0) - - # Change scheduler to use Stage 2 distilled sigmas as is - new_scheduler = FlowMatchEulerDiscreteScheduler.from_config( - self.pipe.scheduler.config, - use_dynamic_shifting=False, - shift_terminal=None, - ) - self.pipe.scheduler = new_scheduler - - stage_2_req = copy.copy(req) - stage_2_req.sampling_params = req.sampling_params.clone() - stage_2_req.sampling_params.num_inference_steps = 3 - - video, audio = self.pipe( - req=stage_2_req, - latents=upscaled_video_latent, - audio_latents=audio_latent, - prompt=prompt, - negative_prompt=negative_prompt, - noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], - sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, - guidance_scale=1.0, - generator=generator, - output_type="np", - return_dict=False, - ).output - - return DiffusionOutput(output=(video, audio)) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) - - -class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline): - """LTX-2 I2V pipeline for FastGen DMD2-distilled models.""" - - def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): - super().__init__(od_config=od_config, prefix=prefix) - self.__init_dmd2__() diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py new file mode 100644 index 00000000000..8b786b8175e --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -0,0 +1,207 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Two-stage entry points for the LTX model family.""" + +from __future__ import annotations + +import os +from dataclasses import replace +from typing import Any, ClassVar + +import torch +from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +from .ltx2_components import ( + LTX2_COMPONENT_PROFILE, +) +from .ltx2_components import ( + get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401 +) +from .ltx2_conditioning import LTXI2VConditioningMixin +from .ltx2_guidance import LTX_LEGACY_VELOCITY_GUIDANCE +from .ltx2_pipeline_runtime import LTXPipelineRuntime +from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE +from .ltx2_request import LTXRequestInputs +from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline + + +class LTX2TwoStagesPipeline(LTXPipelineRuntime): + """Legacy distilled-only LTX2 two-stage compatibility entry.""" + + component_profile = LTX2_COMPONENT_PROFILE + guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE + one_stage_recipe = LTX2_ONE_STAGE_RECIPE + supports_request_batch = False + support_image_input = False + + _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules) + _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules) + _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules) + _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules) + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + model_path = od_config.model + self.distilled = "distilled" in os.path.basename(os.path.normpath(model_path)) + if not self.distilled: + raise NotImplementedError(f"{model_path} is not supported for {self.__class__.__name__}.") + + super().__init__(od_config=od_config, prefix=prefix) + self.upsample_pipe = LTX2LatentUpsamplePipeline(vae=self.vae, od_config=od_config) + + def _run_two_stage( + self, + req: DiffusionRequestBatch, + request_inputs: LTXRequestInputs, + *, + noise_scale: float, + timesteps: list[int] | None, + attention_kwargs: dict[str, Any] | None, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + stage1 = self.run_phase( + req, + request_inputs, + noise_scale=noise_scale, + sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None, + timesteps=timesteps, + attention_kwargs=attention_kwargs, + image=image, + ) + upscaled_video_latent = self.upsample_pipe( + latents=stage1.video, + output_type="latent", + return_dict=False, + )[0] + + stage2_inputs = replace( + request_inputs, + num_inference_steps=3, + guidance_scale=1.0, + latents=upscaled_video_latent, + audio_latents=stage1.audio, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + ) + stage2 = self.run_phase( + req, + stage2_inputs, + noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + timesteps=None, + attention_kwargs=attention_kwargs, + prompt_context=stage1.forward_context.prompt_context, + ) + return self.decode_phase(stage2) + + @torch.no_grad() + def forward( + self, + req: DiffusionRequestBatch, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + frame_rate: float | None = None, + num_inference_steps: int | None = None, + timesteps: list[int] | None = None, + guidance_scale: float = 4.0, + guidance_rescale: float = 0.0, + noise_scale: float = 0.0, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "np", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + max_sequence_length: int | None = None, + *, + image: Any | None = None, + ) -> DiffusionOutput | list[DiffusionOutput]: + del return_dict + request_inputs = self._resolve_request_inputs( + req, + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + num_inference_steps=num_inference_steps, + timesteps=timesteps, + guidance_scale=( + getattr(getattr(self, "one_stage_recipe", None), "guidance_scale", 4.0) + if guidance_scale is None + else guidance_scale + ), + guidance_rescale=guidance_rescale, + num_videos_per_prompt=num_videos_per_prompt, + generator=generator, + latents=latents, + audio_latents=audio_latents, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + output_type=output_type, + max_sequence_length=max_sequence_length, + ) + if self.support_image_input: + image = self._resolve_request_image(req, image, request_inputs) + return self._run_two_stage( + req, + request_inputs, + noise_scale=noise_scale, + timesteps=timesteps, + attention_kwargs=attention_kwargs, + image=image, + ) + + +class LTX2ImageToVideoTwoStagesPipeline(LTXI2VConditioningMixin, LTX2TwoStagesPipeline): + """LTX2 two-stage image-to-video entry.""" + + def forward( + self, + req: DiffusionRequestBatch, + image: Any | None = None, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + frame_rate: float | None = None, + num_inference_steps: int | None = None, + sigmas: list[float] | None = None, + *args: Any, + **kwargs: Any, + ) -> DiffusionOutput | list[DiffusionOutput]: + """Preserve the legacy I2V sigma slot, which was ignored by this pipeline.""" + del sigmas + return super().forward( + req, + image, + prompt, + negative_prompt, + height, + width, + num_frames, + frame_rate, + num_inference_steps, + *args, + **kwargs, + ) diff --git a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py index 65da889f123..513d9333547 100644 --- a/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py +++ b/vllm_omni/diffusion/models/magi_human/pipeline_magi_human.py @@ -12,7 +12,7 @@ from collections.abc import Iterable from dataclasses import dataclass from pathlib import Path -from typing import Any, Literal +from typing import Any, ClassVar, Literal import numpy as np import torch @@ -46,12 +46,14 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import ( DiffusersPipelineLoader, ) +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.t5_encoder.t5_gemma_encoder import T5GemmaEncoderModelTP from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( DiffusionPipelineProfilerMixin, ) from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from .magi_human_dit import ( DiTModel, @@ -1679,7 +1681,11 @@ def _resolve_subdir( # =========================================================================== # Main Pipeline # =========================================================================== -class MagiHumanPipeline(nn.Module, ProgressBarMixin, DiffusionPipelineProfilerMixin): +class MagiHumanPipeline(nn.Module, ProgressBarMixin, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): + _dit_modules: ClassVar[list[str]] = ["dit", "sr_dit"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] + def __init__(self, od_config: OmniDiffusionConfig, **kwargs): super().__init__() model_path = od_config.model @@ -2155,7 +2161,7 @@ def encode_prompt( @torch.inference_mode() def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | None = None, height: int = 256, width: int = 448, diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index e5533896331..614c8268c0d 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -1,7 +1,7 @@ from __future__ import annotations from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import torch from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL @@ -9,8 +9,10 @@ from torch import nn from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm from vllm.config import VllmConfig +from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config @@ -18,8 +20,10 @@ from .rope_real import RotaryPosEmbedReal from .schedulers import FlowMatchEulerDiscreteScheduler +logger = init_logger(__name__) -class MammothModa2DiTPipeline(nn.Module): + +class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): """ MammothModa2 DiT + VAE generation stage (non-autoregressive). @@ -28,6 +32,10 @@ class MammothModa2DiTPipeline(nn.Module): """ + _dit_modules: ClassVar[list[str]] = ["gen_transformer"] + _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"] + _vae_modules: ClassVar[list[str]] = ["gen_vae"] + have_multimodal_outputs = True # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage. @@ -54,9 +62,19 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config) self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config) - llm_hidden_size = int(getattr(self.config.llm_config, "hidden_size", 0) or 0) + # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config + llm_hidden_size = 0 + text_config = self.config.get_text_config() + if text_config is None: + logger.warning("No text config; failed to infer llm_hidden_size.") + elif not hasattr(text_config, "hidden_size"): + logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.") + else: + llm_hidden_size = int(text_config.hidden_size or 0) if llm_hidden_size <= 0: - raise ValueError("Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.hidden_size") + raise ValueError( + "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size" + ) self._reinit_caption_embedder(llm_hidden_size) # Optional: image condition refiner (Q-Former) @@ -127,6 +145,43 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: dtype=dtype, ) + def _split_ar_conditions( + self, + *, + full_hidden_states: torch.Tensor, + full_token_ids: list[int], + answer_start_index: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Split AR-stage hidden states into text / image condition embeds. + + The token ids that distinguish question (text) tokens, generated visual + tokens, and multi-modal placeholder tokens are read from the model config + (``gen_vocab_start_index`` and the vision placeholder token ids), so the + caller no longer needs to pass them. Mirrors the masking the bespoke + MammothModa2 example performed via ar2dit. + """ + gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index) + visual_ids = [ + int(self.config.image_token_id), + int(self.config.video_token_id), + int(self.config.vision_start_token_id), + int(self.config.vision_end_token_id), + ] + + device = full_hidden_states.device + token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device) + positions = torch.arange(token_ids.shape[0], device=device) + questions_mask = positions < answer_start_index + answers_mask = ~questions_mask + gen_token_mask = token_ids >= gen_vocab_start_index + visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device)) + text_mask = questions_mask & ~(visual_token_mask | gen_token_mask) + image_mask = answers_mask & gen_token_mask + + text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous() + image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous() + return text_cond, image_cond + @torch.inference_mode() def forward( self, @@ -136,14 +191,35 @@ def forward( ) -> OmniOutput: runtime_addi = kwargs.get("runtime_additional_information", None) info = runtime_addi[0] - text_cond = info["text_prompt_embeds"] - image_cond = info["image_prompt_embeds"] + + # Sampling knobs are declared in vllm_omni/model_extras/mammothmodal2_preview.py + # and routed via extra_body -> sampling_params.extra_args (surfaced here as + # ``sampling_extra_args``). Fall back to runtime_additional_information for the + # legacy bespoke-example path during the transition. + extra_args_list = kwargs.get("sampling_extra_args") or [] + extra_args = extra_args_list[0] if extra_args_list else {} + text_guidance_scale = float(extra_args.get("text_guidance_scale", info["text_guidance_scale"][0])) + cfg_range_val = extra_args.get("cfg_range", info["cfg_range"]) + cfg_range = float(cfg_range_val[0]), float(cfg_range_val[1]) + num_inference_steps = int(extra_args.get("num_inference_steps", info["num_inference_steps"][0])) + negative_cond = info.get("negative_prompt_embeds") negative_attention_mask = info.get("negative_prompt_attention_mask") image_hw = info["image_height"][0], info["image_width"][0] - text_guidance_scale = info["text_guidance_scale"][0] - cfg_range = info["cfg_range"][0], info["cfg_range"][1] - num_inference_steps = info["num_inference_steps"][0] + + # Split the AR hidden states into text / image conditions. The token ids that + # drive the split are sourced from the model config (see _split_ar_conditions), + # formerly supplied by the bespoke example via additional_information. Legacy + # fallback: ar2dit may have already produced the split conditions. + if "text_prompt_embeds" in info: + text_cond = info["text_prompt_embeds"] + image_cond = info["image_prompt_embeds"] + else: + text_cond, image_cond = self._split_ar_conditions( + full_hidden_states=info["full_hidden_states"], + full_token_ids=info["full_token_ids"], + answer_start_index=int(info["answer_start_index"][0]), + ) # Move to model device/dtype. model_device = next(self.parameters()).device diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py index 0da26aff8c7..c67f8cb24bf 100644 --- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py +++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py @@ -24,7 +24,6 @@ import logging import os -from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -50,6 +49,8 @@ ) from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -58,34 +59,6 @@ logger = logging.getLogger(__name__) -@dataclass -class _ZPipelineSamplingParams: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req.sampling_params``. Unlike SimpleNamespace this fails loudly - at construction if a required field is missing.""" - - height: int - width: int - num_inference_steps: int - guidance_scale: float - generator: torch.Generator | None = None - strength: float | None = None - sigmas: list[float] | None = None - max_sequence_length: int = 512 - guidance_rescale: float | None = None - num_outputs_per_prompt: int = 1 - - -@dataclass -class _ZPipelineRequest: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req``.""" - - request_id: str - sampling_params: _ZPipelineSamplingParams - prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}]) - - class MingImagePipeline(ZImagePipeline): """Ming-flash-omni-2.0 text-to-image diffusion pipeline. @@ -95,6 +68,8 @@ class MingImagePipeline(ZImagePipeline): ships ``byt5/``) """ + supports_request_batch = False + def __init__( self, *, @@ -117,6 +92,10 @@ def __init__( self.device = self._execution_device # Ming convention alias self._dtype = dtype + # Request-scoped conditioning handed to the inherited encode_prompt override + self._pending_prompt_embeds: list[torch.Tensor] | None = None + self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None + # Ming's per-checkpoint image-gen configuration. We cannot rely on # ``od_config.hf_config.image_gen_config`` because the diffusion # stage is started with ``hf_config_name: thinker_config`` (the @@ -198,7 +177,7 @@ def __init__( self.condition_encoder.load_from_checkpoint(model_path) # Optional ByT5 glyph/text encoder. Only loaded when the checkpoint - # ``extra_args.image_gen.byte5_text`` requests will be ignored. + # ships a byt5/ subfolder; otherwise byte5_text requests are ignored byte5_dir = Path(model_path) / "byt5" if byte5_dir.exists(): self.byte5 = MingByT5Encoder.from_checkpoint(byte5_dir, device=self.device, dtype=dtype) @@ -217,13 +196,12 @@ def _resolve_byte5_texts(extra: dict, sampling_params) -> list[str]: """Resolve byte5 glyph texts. Two sources, in order of priority: - 1. ``extra["byte5_text"]`` — auto-extracted from the user prompt's - quoted spans by ``thinker2imagegen`` (already wrapped as - ``'Text "<glyph>". '`` by Ming's ``get_text_from_prompt``). - 2. ``sampling_params.extra_args["image_gen"]["byte5_text"]`` — an - explicit override for programmatic callers; raw strings without - the ``Text "..."`` wrapper are auto-wrapped here to match the - distribution ByT5 was trained on. + 1. `extra["byte5_text"]`: auto-extracted from the user prompt's quoted spans + by thinker2imagegen (already wrapped as `'Text "<glyph>". '` + by Ming's get_text_from_prompt). + 2. `sampling_params.extra_args["byte5_text"]`: + raw strings without the `Text "..."` wrapper are auto-wrapped here + to match the distribution ByT5 was trained on. """ # Source 1: auto-extracted, already wrapped. Return as-is if non-empty. raw = extra.get("byte5_text") @@ -236,7 +214,7 @@ def _resolve_byte5_texts(extra: dict, sampling_params) -> list[str]: # Source 2: explicit override — wrap raw strings so the byte5 encoder # sees the same ``Text "<glyph>". `` format Ming used during training. - raw = ((getattr(sampling_params, "extra_args", None) or {}).get("image_gen") or {}).get("byte5_text") + raw = (getattr(sampling_params, "extra_args", None) or {}).get("byte5_text") if isinstance(raw, str): raw = [raw] if isinstance(raw, list): @@ -268,23 +246,38 @@ def _encode_reference_image(self, ref, height: int, width: int) -> torch.Tensor latent = self.vae.encode(ref).latent_dist.mode() return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor + def encode_prompt(self, *args, **kwargs): # noqa: ARG002 + """Return Ming's precomputed conditioning instead of encoding text. + + NOTE: Ming has no Z-Image text_encoder; its conditioning (cap_feats, optionally ByT5-augmented) + is computed in forward and stashed on `self._pending_*` immediately before + the `super().forward` call, so we simply hand it back here. + """ + if self._pending_prompt_embeds is None: + raise RuntimeError( + "MingImagePipeline.encode_prompt called without pending " + "conditioning; it must run within MingImagePipeline.forward." + ) + + return self._pending_prompt_embeds, self._pending_negative_prompt_embeds + # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ @torch.inference_mode() - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Run one text-to-image generation request. Args: - req: Diffusion request. The cross-stage thinker hidden states + req: Single-request batch. The cross-stage thinker hidden states must be present at ``req.prompts[0]["extra"]["thinker_hidden_states"]`` as a ``[N, H]`` (or ``[1, N, H]``) tensor, placed there by ``thinker2imagegen``. Returns: - DiffusionOutput with ``.output`` set to a ``[B, 3, H, W]`` + One DiffusionOutput with ``.output`` set to a ``[B, 3, H, W]`` image tensor in ``[-1, 1]``. The vllm-omni diffusion engine's output adapter converts this to PIL/base64 downstream. """ @@ -366,21 +359,23 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: negative_cap_feats = torch.cat((negative_cap_feats, torch.zeros_like(byte5_feats)), dim=1) logger.debug("[MingImagePipeline.forward] byte5 cat'd: cap_feats=%s", tuple(cap_feats.shape)) - # Sampling knobs: extra_args.image_gen.* > sampling_params.* > MingImageGenConfig defaults. + # Sampling knobs, in priority order: + # top-level extra_args[key] > sampling_params.* attr > + # MingImageGenConfig default. Knobs live flat on extra_args sp = req.sampling_params cfg = self.image_gen_config - ig = (sp.extra_args or {}).get("image_gen") or {} + ea = sp.extra_args or {} resolved: dict[str, Any] = {} - for ig_key, sp_attr, default in ( + for ea_key, sp_attr, default in ( ("height", "height", cfg.default_height), ("width", "width", cfg.default_width), ("steps", "num_inference_steps", cfg.num_inference_steps), ("cfg", "guidance_scale", cfg.guidance_scale), ("seed", "seed", None), ): - for v in (ig.get(ig_key), getattr(sp, sp_attr), default): + for v in (ea.get(ea_key), getattr(sp, sp_attr), default): if v is not None: - resolved[ig_key] = v + resolved[ea_key] = v break height = int(resolved["height"]) @@ -393,7 +388,7 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: # ``sp.generator`` causes two problems: # (1) if the caller pre-seeded it with sp.seed (e.g. the top-level # ``seed`` key on OmniDiffusionSamplingParams), any override via - # ``extra_args.image_gen.seed`` would be silently ignored; and + # `extra_args["seed"]` would be silently ignored; and # (2) a persistent generator instance accumulates state across # requests → same-seed replays produce different outputs. if seed is not None: @@ -409,17 +404,27 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])] else: negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds] + self._pending_prompt_embeds = prompt_embeds + self._pending_negative_prompt_embeds = negative_prompt_embeds - z_sp = _ZPipelineSamplingParams( + # Build a real single-request DiffusionRequestBatch; + # the prompt text itself is a neutral placeholder here. + z_sp = OmniDiffusionSamplingParams( height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, + output_type="pt", ) - z_req = _ZPipelineRequest( - request_id=req.request_id or "ming-imagegen", - sampling_params=z_sp, + z_req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt={"prompt": ""}, + sampling_params=z_sp, + request_id=req.request_id or "ming-imagegen", + ) + ] ) # Reference image (img2img) → VAE-encoded latent published on the @@ -435,32 +440,18 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: num_inference_steps, guidance_scale, seed, - ig, + ea, None if ref_latent is None else tuple(ref_latent.shape), ) try: - output = super().forward( - z_req, - prompt=None, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - generator=generator, - output_type="pt", - return_dict=True, - ) + outputs: DiffusionOutput = super().forward(z_req) finally: set_forward_context_ref_latent(None) + # Drop request-scoped conditioning so we don't retain GPU tensors. + self._pending_prompt_embeds = None + self._pending_negative_prompt_embeds = None - if hasattr(output, "output") and output.output is not None: - raw = output.output - elif hasattr(output, "images"): - raw = output.images - else: - raw = output + raw = outputs.output if not isinstance(raw, torch.Tensor): raise RuntimeError(f"ZImagePipeline returned non-tensor output: {type(raw).__name__}") if logger.isEnabledFor(logging.DEBUG): @@ -470,7 +461,7 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: raw.float().min().item(), raw.float().max().item(), ) - return DiffusionOutput(output=raw) + return outputs # ---------------------------------------------------------------------- diff --git a/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py b/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py index abeacff2aa6..518453a171e 100644 --- a/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py +++ b/vllm_omni/diffusion/models/mistral_encoder/mistral_encoder.py @@ -17,7 +17,7 @@ import logging from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn @@ -37,6 +37,9 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding from vllm.model_executor.model_loader.weight_utils import default_weight_loader +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + logger = logging.getLogger(__name__) @@ -118,6 +121,7 @@ def __init__( num_kv_heads: int, head_dim: int, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.hidden_size = hidden_size @@ -140,6 +144,7 @@ def __init__( total_num_heads=num_heads, total_num_kv_heads=num_kv_heads, bias=False, + quant_config=quant_config, prefix=f"{prefix}.qkv_proj", ) @@ -147,6 +152,7 @@ def __init__( input_size=num_heads * head_dim, output_size=hidden_size, bias=False, + quant_config=quant_config, prefix=f"{prefix}.o_proj", ) @@ -194,18 +200,21 @@ def __init__( hidden_size: int, intermediate_size: int, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.gate_up_proj = MergedColumnParallelLinear( input_size=hidden_size, output_sizes=[intermediate_size, intermediate_size], bias=False, + quant_config=quant_config, prefix=f"{prefix}.gate_up_proj", ) self.down_proj = RowParallelLinear( input_size=intermediate_size, output_size=hidden_size, bias=False, + quant_config=quant_config, prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() @@ -227,6 +236,7 @@ def __init__( intermediate_size: int, rms_norm_eps: float, prefix: str = "", + quant_config: QuantizationConfig | None = None, ): super().__init__() self.self_attn = MistralEncoderAttention( @@ -235,11 +245,13 @@ def __init__( num_kv_heads=num_kv_heads, head_dim=head_dim, prefix=f"{prefix}.self_attn", + quant_config=quant_config, ) self.mlp = MistralEncoderMLP( hidden_size=hidden_size, intermediate_size=intermediate_size, prefix=f"{prefix}.mlp", + quant_config=quant_config, ) self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps) @@ -295,7 +307,12 @@ class MistralEncoderModel(nn.Module): parallel layers for TP but simple SDPA for attention (no PagedAttention). """ - def __init__(self, config: PretrainedConfig, prefix: str = ""): + def __init__( + self, + config: PretrainedConfig, + prefix: str = "", + quant_config: QuantizationConfig | None = None, + ): super().__init__() self._processor = None self._system_message_t2i: str | None = None @@ -339,6 +356,7 @@ def __init__(self, config: PretrainedConfig, prefix: str = ""): m = self.language_model.model m.embed_tokens = VocabParallelEmbedding(self.vocab_size, self.hidden_size) + layer_prefix_root = f"{prefix}.language_model.model.layers" if prefix else "language_model.model.layers" m.layers = nn.ModuleList( [ @@ -349,7 +367,8 @@ def __init__(self, config: PretrainedConfig, prefix: str = ""): head_dim=self.head_dim, intermediate_size=self.intermediate_size, rms_norm_eps=self.rms_norm_eps, - prefix=f"language_model.model.layers.{i}", + prefix=f"{layer_prefix_root}.{i}", + quant_config=quant_config, ) for i in range(self.num_layers) ] diff --git a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py index 4fa56ea9313..392e4286888 100644 --- a/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py +++ b/vllm_omni/diffusion/models/nextstep_1_1/pipeline_nextstep_1_1.py @@ -5,7 +5,7 @@ import os import re from collections.abc import Iterable -from typing import Literal +from typing import ClassVar, Literal import numpy as np import torch @@ -28,13 +28,14 @@ ) 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 SupportsComponentDiscovery from vllm_omni.diffusion.models.nextstep_1_1.modeling_flux_vae import AutoencoderKL from vllm_omni.diffusion.models.nextstep_1_1.modeling_nextstep import ( NextStepConfig, NextStepModel, ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -131,7 +132,7 @@ def post_process_func(images: torch.Tensor): return post_process_func -class NextStep11Pipeline(nn.Module, DiffusionPipelineProfilerMixin): +class NextStep11Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): """ NextStep-1.1 Pipeline for text-to-image generation. @@ -140,6 +141,10 @@ class NextStep11Pipeline(nn.Module, DiffusionPipelineProfilerMixin): to generate images autoregressively. """ + _dit_modules: ClassVar[list[str]] = ["model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -554,7 +559,7 @@ def decoding( @torch.no_grad() def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] | None = None, height: int | None = None, width: int | None = None, @@ -707,7 +712,8 @@ def forward( sampled_images = sampled_images.detach().cpu().to(torch.float32) return DiffusionOutput( - output=sampled_images, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + output=sampled_images, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py index 3f03563a1cc..8c2c9db88fe 100644 --- a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py +++ b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py @@ -22,6 +22,10 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -808,6 +812,36 @@ def forward( return hidden_states +class _OmniGen2SPInputBoundary(nn.Module): + """SP input-sharding boundary marker for ``OmniGen2Transformer2DModel._sp_plan``. + + Resets ``sp_original_seq_len`` on every call: OmniGen2 runs the transformer + once per CFG branch and the unconditional / conditional branches have + different joint sequence lengths (text prefix differs). The SP framework's + auto-pad records the length only on the first call, so without this reset + a stale length from an earlier branch truncates the gather of later ones. + """ + + def forward(self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + from vllm_omni.diffusion.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + if is_forward_context_available(): + ctx = get_forward_context() + ctx.sp_original_seq_len = None + ctx.sp_padding_size = 0 + return hidden_states, rotary_emb + + +class _OmniGen2SPOutputBoundary(nn.Module): + """SP output-gathering boundary marker for ``OmniGen2Transformer2DModel._sp_plan``.""" + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + class OmniGen2Transformer2DModel(nn.Module): """ OmniGen2 Transformer 2D Model. @@ -836,6 +870,17 @@ class OmniGen2Transformer2DModel(nn.Module): timestep_scale: Scale factor for timestep embeddings """ + # ``attention_mask`` is intentionally omitted from the plan: OmniGen2Attention + # accepts but never reads it. ``auto_pad=True`` requires Ulysses SP — Ring SP + # cannot use attention masks. + _sp_plan = { + "sp_input_boundary": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + }, + "sp_output_boundary": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, patch_size: int = 2, @@ -983,6 +1028,9 @@ def __init__( # Add learnable embeddings to distinguish different images self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size)) # support max 5 ref images + self.sp_input_boundary = _OmniGen2SPInputBoundary() + self.sp_output_boundary = _OmniGen2SPOutputBoundary() + def img_patch_embed_and_refine( self, hidden_states, @@ -1251,9 +1299,13 @@ def forward( hidden_states = joint_hidden_states + hidden_states, rotary_emb = self.sp_input_boundary(hidden_states, rotary_emb) + for layer in self.layers: hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb) + hidden_states = self.sp_output_boundary(hidden_states) + # 4. Output norm & projection hidden_states = self.norm_out(hidden_states, temb) diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index 83ea7ed7363..2bdd513ef77 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -40,6 +40,7 @@ ) from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -258,44 +259,42 @@ def pre_process_func( request: OmniDiffusionRequest, ) -> OmniDiffusionRequest: """Pre-process requests for OmniGen2Pipeline.""" - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if raw_image is not None: - if isinstance(raw_image, list): - images = [PIL.Image.open(img) if isinstance(img, str) else img for img in raw_image] - elif isinstance(raw_image, str): - images = [PIL.Image.open(raw_image)] - else: - images = [raw_image] - - first_raw = images[0] - if isinstance(first_raw, PIL.Image.Image): - new_h, new_w = image_processor.get_new_height_width( - first_raw, max_pixels=1024 * 1024, max_side_length=1024 - ) - if request.sampling_params.height is None: - request.sampling_params.height = new_h - if request.sampling_params.width is None: - request.sampling_params.width = new_w - - preprocessed_images = [] - for image in images: - if not ( - isinstance(image, torch.Tensor) and len(image.shape) > 1 and image.shape[1] == latent_channels - ): - image = image_processor.preprocess(image, max_pixels=1024 * 1024, max_side_length=1024) - preprocessed_images.append(image) - - prompt["additional_information"]["preprocessed_images"] = preprocessed_images - - request.prompts[i] = prompt + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if raw_image is not None: + if isinstance(raw_image, list): + images = [PIL.Image.open(img) if isinstance(img, str) else img for img in raw_image] + elif isinstance(raw_image, str): + images = [PIL.Image.open(raw_image)] + else: + images = [raw_image] + + first_raw = images[0] + if isinstance(first_raw, PIL.Image.Image): + new_h, new_w = image_processor.get_new_height_width( + first_raw, max_pixels=1024 * 1024, max_side_length=1024 + ) + if request.sampling_params.height is None: + request.sampling_params.height = new_h + if request.sampling_params.width is None: + request.sampling_params.width = new_w + + preprocessed_images = [] + for image in images: + if not (isinstance(image, torch.Tensor) and len(image.shape) > 1 and image.shape[1] == latent_channels): + image = image_processor.preprocess(image, max_pixels=1024 * 1024, max_side_length=1024) + preprocessed_images.append(image) + + prompt["additional_information"]["preprocessed_images"] = preprocessed_images + + request.prompt = prompt return request return pre_process_func @@ -1005,7 +1004,7 @@ def cfg_range(self): @torch.no_grad() def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, prompt_embeds: torch.FloatTensor | None = None, diff --git a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py index 9671e275845..62ac027ca8b 100644 --- a/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py +++ b/vllm_omni/diffusion/models/omnivoice/pipeline_omnivoice.py @@ -26,7 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.models.interface import SupportAudioOutput -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.models.omnivoice.duration import RuleDurationEstimator from vllm_omni.model_executor.models.omnivoice.omnivoice_decoder import OmniVoiceDecoder from vllm_omni.model_executor.models.omnivoice.omnivoice_generator import OmniVoiceGenerator @@ -206,7 +206,7 @@ def _encode_ref_audio(self, audio_signal: torch.Tensor, sr: int) -> torch.Tensor return tokens @torch.inference_mode() - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Generate speech audio from text, optionally with voice cloning. Accepts either a plain text prompt or a structured dict: @@ -248,7 +248,7 @@ def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: audio_field = audio_field[0] elif len(audio_field) > 1: return DiffusionOutput( - error=f"OmniVoice voice cloning supports a single reference audio; got {len(audio_field)}" + error=f"OmniVoice voice cloning supports a single reference audio; got {len(audio_field)}" # noqa: E501 ) else: audio_field = None diff --git a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py index f77c5d419da..0baad965ba7 100644 --- a/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py +++ b/vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py @@ -18,8 +18,8 @@ import inspect import json import os -from collections.abc import Callable, Iterable -from typing import Any +from collections.abc import Iterable +from typing import Any, ClassVar import numpy as np import torch @@ -39,9 +39,10 @@ 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.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = init_logger(__name__) @@ -142,7 +143,13 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): +class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery): + supports_request_batch = False + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -544,117 +551,33 @@ def current_timestep(self): def interrupt(self): return self._interrupt - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - guidance_scale: float = 5.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - num_images_per_prompt: int | None = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: torch.FloatTensor | None = None, - negative_prompt_embeds: torch.FloatTensor | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end: Callable[[int, int, dict], None] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 256, - ) -> DiffusionOutput: - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `list[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - negative_prompt (`str` or `list[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - not greater than `1`). - guidance_scale (`float`, *optional*, defaults to 1.0): - True classifier-free guidance (guidance scale) is enabled when `guidance_scale` > 1 and - `negative_prompt` is provided. - height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The height in pixels of the generated image. This is set to 1024 by default for the best results. - width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor): - The width in pixels of the generated image. This is set to 1024 by default for the best results. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - sigmas (`list[float]`, *optional*): - Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in - their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed - will be used. - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `list[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`torch.FloatTensor`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple. - joint_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`list`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.ovis_image.OvisImagePipelineOutput`] or `tuple`: - [`~pipelines.ovis_image.OvisImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When - returning a tuple, the first element is a list with the generated images. - """ + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + prompt_embeds = None + negative_prompt_embeds = None + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale - ) - generator = req.sampling_params.generator or generator + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + guidance_scale = req.sampling_params.guidance_scale + generator = req.sampling_params.generator num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + max_sequence_length = req.sampling_params.max_sequence_length or 256 + + joint_attention_kwargs: dict[str, Any] | None = None + output_type = req.sampling_params.output_type or "pil" + callback_on_step_end_tensor_inputs = ["latents"] # Steps: # 1. Check Inputs @@ -684,12 +607,7 @@ def forward( self._interrupt = False # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = len(prompt) do_classifier_free_guidance = guidance_scale > 1.0 prompt_embeds, text_ids = self.encode_prompt( @@ -755,9 +673,8 @@ def forward( latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) + stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None + return DiffusionOutput(output=image, stage_durations=stage_durations) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index cdc6de99840..986c2ef86fc 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -28,14 +28,15 @@ 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.dmd2 import DMD2PipelineMixin +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( QwenImageCFGParallelMixin, ) from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.prompt_utils import ( validate_prompt_sequence_lengths, ) @@ -43,6 +44,7 @@ normalize_min_aligned_size, ) from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request if TYPE_CHECKING: from vllm_omni.diffusion.worker.input_batch import InputBatch @@ -71,8 +73,19 @@ def get_qwen_image_post_process_func( image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2) def post_process_func( - images: torch.Tensor, + images: torch.Tensor | dict[str, Any], ): + if isinstance(images, dict) and isinstance(images.get("payload"), dict): + payload = dict(images["payload"]) + image_payload = payload.get("image") + if image_payload is None: + raise ValueError("Qwen-Image postprocess expected payload['image'] in output envelope.") + payload["image"] = image_processor.postprocess(image_payload) + metadata = images.get("metadata") or {} + return { + "payload": payload, + "metadata": metadata if isinstance(metadata, dict) else {}, + } return image_processor.postprocess(images) return post_process_func @@ -250,7 +263,14 @@ def apply_rotary_emb_qwen( return x_out.type_as(x) -class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): +class QwenImagePipeline( + nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + supports_request_batch = True + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + supports_step_execution: ClassVar[bool] = True def __init__( @@ -323,7 +343,9 @@ def __init__( ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) self.transformer = QwenImageTransformer2DModel( - od_config=od_config, quant_config=od_config.quantization_config, **transformer_kwargs + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) @@ -730,10 +752,8 @@ def _prepare_generation_context( else: guidance = None - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) + txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds) + negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds) return { "prompt_embeds": prompt_embeds, @@ -756,7 +776,7 @@ def prepare_encode( ) -> "DiffusionRequestState": """Populate *state* with encoded prompts, latents, timesteps, and CFG config.""" sampling = state.sampling - prompt, negative_prompt = self._extract_prompts(state.prompts or []) + prompt, negative_prompt = self._extract_prompts([state.prompt] if state.prompt is not None else []) ctx = self._prepare_generation_context( prompt=prompt, @@ -969,51 +989,54 @@ def post_decode( height = state.sampling.height or self.default_sample_size * self.vae_scale_factor width = state.sampling.width or self.default_sample_size * self.vae_scale_factor - output_type = kwargs.get("output_type", "pil") + output_type = kwargs.get("output_type") or state.sampling.output_type or "pil" return self._decode_latents(state.latents, height, width, output_type) - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 1.0, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 1024, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: + sampling_params_list = req.sampling_params_list + common_sampling_params = sampling_params_list[0] extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) - prompt = extracted_prompt or prompt + prompt = extracted_prompt - height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale + num_inference_steps = common_sampling_params.num_inference_steps or 50 + sigmas = common_sampling_params.sigmas + max_sequence_length = common_sampling_params.max_sequence_length or 1024 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 + ) + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": None, + "prompt_embeds_mask": None, + "negative_prompt_embeds": None, + "negative_prompt_embeds_mask": None, + }, ) + prompt_embeds = prompt_fields["prompt_embeds"] + prompt_embeds_mask = prompt_fields["prompt_embeds_mask"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + negative_prompt_embeds_mask = prompt_fields["negative_prompt_embeds_mask"] + if prompt_embeds is not None: + prompt = None + if negative_prompt_embeds is not None: + negative_prompt = None + true_cfg_scale = common_sampling_params.true_cfg_scale or 4.0 + if common_sampling_params.guidance_scale_provided: + guidance_scale = common_sampling_params.guidance_scale + else: + guidance_scale = 1.0 + + latents = req.collate_request_tensors("latents", None) + output_type = common_sampling_params.output_type or "pil" + attention_kwargs = None + callback_on_step_end_tensor_inputs = ["latents"] ctx = self._prepare_generation_context( prompt=prompt, @@ -1058,7 +1081,13 @@ def forward( ) self._current_timestep = None - return self._decode_latents(latents, height, width, output_type) + + result = self._decode_latents(latents, height, width, output_type) + return split_diffusion_output_by_request( + result, + req, + num_outputs_per_prompt=num_images_per_prompt, + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 35022e5eb2c..3d6e937c773 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -7,7 +7,7 @@ import math import os from collections.abc import Iterable -from typing import Any, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -28,7 +28,7 @@ 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 SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( QwenImageCFGParallelMixin, ) @@ -45,6 +45,7 @@ normalize_min_aligned_size, ) from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -74,57 +75,55 @@ def pre_process_func( request: OmniDiffusionRequest, ): """Pre-process requests for QwenImageEditPipeline.""" - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - # Only handles single image - if not raw_image: # None or empty list - raise ValueError("""Received no input image. This model requires one input image to run.""") - elif isinstance(raw_image, list): - if len(raw_image) > 1: - raise ValueError( - """Received multiple input images. Only a single image is supported by this model.""" - ) - else: - raw_image = raw_image[0] - - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + # Only handles single image + if not raw_image: # None or empty list + raise ValueError("""Received no input image. This model requires one input image to run.""") + elif isinstance(raw_image, list): + if len(raw_image) > 1: + raise ValueError("""Received multiple input images. Only a single image is supported by this model.""") else: - image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) - - image_size = image.size - calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1]) - height = request.sampling_params.height or calculated_height - width = request.sampling_params.width or calculated_width - - # Ensure dimensions are multiples of vae_scale_factor * 2 - height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) - - # Store calculated dimensions in request - prompt["additional_information"]["calculated_height"] = calculated_height - prompt["additional_information"]["calculated_width"] = calculated_width - request.sampling_params.height = height - request.sampling_params.width = width - - # Preprocess image - if image is not None and not ( - isinstance(image, torch.Tensor) and len(image.shape) > 1 and image.shape[1] == latent_channels - ): - image = image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image - image = image_processor.preprocess(image, calculated_height, calculated_width) - image = image.unsqueeze(2) - - # Store preprocessed image and prompt image in request - prompt["additional_information"]["preprocessed_image"] = image - prompt["additional_information"]["prompt_image"] = prompt_image - request.prompts[i] = prompt + raw_image = raw_image[0] + + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image) + else: + image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) + + image_size = image.size + calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1]) + height = request.sampling_params.height or calculated_height + width = request.sampling_params.width or calculated_width + + # Ensure dimensions are multiples of vae_scale_factor * 2 + height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) + + # Store calculated dimensions in request + prompt["additional_information"]["calculated_height"] = calculated_height + prompt["additional_information"]["calculated_width"] = calculated_width + request.sampling_params.height = height + request.sampling_params.width = width + + # Preprocess image + if image is not None and not ( + isinstance(image, torch.Tensor) and len(image.shape) > 1 and image.shape[1] == latent_channels + ): + image = image_processor.resize(image, calculated_height, calculated_width) + prompt_image = image + image = image_processor.preprocess(image, calculated_height, calculated_width) + image = image.unsqueeze(2) + + # Store preprocessed image and prompt image in request + prompt["additional_information"]["preprocessed_image"] = image + prompt["additional_information"]["prompt_image"] = prompt_image + request.prompt = prompt return request return pre_process_func @@ -221,7 +220,13 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): +class QwenImageEditPipeline( + nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -251,6 +256,7 @@ def __init__( prefetch_subfolders( model, qwen_subfolders, + local_files_only=local_files_only, ) self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( @@ -276,7 +282,11 @@ def __init__( local_files_only=local_files_only, ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) - self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) + self.transformer = QwenImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.processor = from_pretrained_with_prefetch( Qwen2VLProcessor.from_pretrained, @@ -667,30 +677,7 @@ def current_timestep(self): def interrupt(self): return self._interrupt - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 1.0, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 1024, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Forward pass for image editing.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. @@ -721,32 +708,28 @@ def forward( height = req.sampling_params.height width = req.sampling_params.width else: - # fallback to run pre-processing in pipeline (debug only) - image_size = image[0].size if isinstance(image, list) else image.size - calculated_width, calculated_height, _ = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1]) - height = height or calculated_height - width = width or calculated_width - - height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = self.image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image - image = self.image_processor.preprocess(image, calculated_height, calculated_width) - image = image.unsqueeze(2) - - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") + + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + max_sequence_length = req.sampling_params.max_sequence_length or 1024 + generator = req.sampling_params.generator + true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0 if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 1.0 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + prompt_embeds_mask = None + negative_prompt_embeds_mask = None + output_type = req.sampling_params.output_type or "pil" + attention_kwargs = None + callback_on_step_end_tensor_inputs = ["latents"] # 1. check inputs # 2. encode prompts @@ -773,16 +756,9 @@ def forward( self._current_timestep = None self._interrupt = False - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = 1 - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + has_neg_prompt = negative_prompt is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index e053f85dc10..a1163716345 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -5,7 +5,7 @@ import logging import os from collections.abc import Iterable -from typing import Any, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -27,7 +27,7 @@ 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.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.qwen_image.cfg_parallel import ( QwenImageCFGParallelMixin, ) @@ -40,6 +40,7 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.prompt_utils import ( @@ -49,6 +50,7 @@ normalize_min_aligned_size, ) from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -87,74 +89,72 @@ def pre_process_func( request: OmniDiffusionRequest, ): """Pre-process requests for QwenImageEditPlusPipeline.""" - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - # Handle single image or list of images - if raw_image is None: - continue + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + # Handle single image or list of images + if raw_image is None: + return request + + if not isinstance(raw_image, list): + raw_image = [raw_image] + if len(raw_image) > MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES: + raise ValueError( + f"Received {len(raw_image)} input images. " + f"At most {MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES} images are supported by this model." + ) + image = [ + PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im) + for im in raw_image + ] - if not isinstance(raw_image, list): - raw_image = [raw_image] - if len(raw_image) > MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES: - raise ValueError( - f"Received {len(raw_image)} input images. " - f"At most {MAX_QWEN_IMAGE_EDIT_PLUS_INPUT_IMAGES} images are supported by this model." - ) - image = [ - PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image | np.ndarray | torch.Tensor, im) - for im in raw_image - ] + # Calculate dimensions based on first image + image_size = image[0].size + calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1]) + height = request.sampling_params.height or calculated_height + width = request.sampling_params.width or calculated_width + + # Ensure dimensions are multiples of vae_scale_factor * 2 + height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) + + # Store calculated dimensions in request + prompt["additional_information"]["calculated_height"] = calculated_height + prompt["additional_information"]["calculated_width"] = calculated_width + request.sampling_params.height = height + request.sampling_params.width = width + + # Preprocess images into condition_images (for prompt encoding) and vae_images (for VAE encoding) + condition_images = [] + vae_images = [] + condition_image_sizes = [] + vae_image_sizes = [] + + for img in image: + if isinstance(img, torch.Tensor) and len(img.shape) > 1 and img.shape[1] == latent_channels: + # Already a latent tensor + continue - # Calculate dimensions based on first image - image_size = image[0].size - calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1]) - height = request.sampling_params.height or calculated_height - width = request.sampling_params.width or calculated_width - - # Ensure dimensions are multiples of vae_scale_factor * 2 - height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) - - # Store calculated dimensions in request - prompt["additional_information"]["calculated_height"] = calculated_height - prompt["additional_information"]["calculated_width"] = calculated_width - request.sampling_params.height = height - request.sampling_params.width = width - - # Preprocess images into condition_images (for prompt encoding) and vae_images (for VAE encoding) - condition_images = [] - vae_images = [] - condition_image_sizes = [] - vae_image_sizes = [] - - for img in image: - if isinstance(img, torch.Tensor) and len(img.shape) > 1 and img.shape[1] == latent_channels: - # Already a latent tensor - continue - - image_width, image_height = img.size - condition_width, condition_height = calculate_dimensions( - CONDITION_IMAGE_SIZE, image_width / image_height - ) - vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) + image_width, image_height = img.size + condition_width, condition_height = calculate_dimensions(CONDITION_IMAGE_SIZE, image_width / image_height) + vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) - condition_image_sizes.append((condition_width, condition_height)) - vae_image_sizes.append((vae_width, vae_height)) + condition_image_sizes.append((condition_width, condition_height)) + vae_image_sizes.append((vae_width, vae_height)) - condition_images.append(image_processor.resize(img, condition_height, condition_width)) - vae_images.append(image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) + condition_images.append(image_processor.resize(img, condition_height, condition_width)) + vae_images.append(image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) - # Store preprocessed images in request - prompt["additional_information"]["condition_images"] = condition_images - prompt["additional_information"]["vae_images"] = vae_images - prompt["additional_information"]["condition_image_sizes"] = condition_image_sizes - prompt["additional_information"]["vae_image_sizes"] = vae_image_sizes - request.prompts[i] = prompt + # Store preprocessed images in request + prompt["additional_information"]["condition_images"] = condition_images + prompt["additional_information"]["vae_images"] = vae_images + prompt["additional_information"]["condition_image_sizes"] = condition_image_sizes + prompt["additional_information"]["vae_image_sizes"] = vae_image_sizes + request.prompt = prompt return request return pre_process_func @@ -185,8 +185,12 @@ def post_process_func( class QwenImageEditPlusPipeline( - nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin + nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery ): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -246,7 +250,11 @@ def __init__( ).to(self.device) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) - self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) + self.transformer = QwenImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.processor = from_pretrained_with_prefetch( Qwen2VLProcessor.from_pretrained, @@ -615,30 +623,7 @@ def current_timestep(self): def interrupt(self): return self._interrupt - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - image: PIL.Image.Image | list[PIL.Image.Image] | torch.Tensor | None = None, - true_cfg_scale: float = 4.0, - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 1.0, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 1024, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Forward pass for image editing with support for multiple images.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. @@ -666,55 +651,35 @@ def forward( ): condition_images = additional_information.get("condition_images") vae_images = additional_information.get("vae_images") - condition_image_sizes = additional_information.get("condition_image_sizes") vae_image_sizes = additional_information.get("vae_image_sizes") calculated_height = additional_information.get("calculated_height") calculated_width = additional_information.get("calculated_width") - height = req.sampling_params.height - width = req.sampling_params.width + height = req.sampling_params.height or calculated_height + width = req.sampling_params.width or calculated_width + image = None else: - # fallback to run pre-processing in pipeline (debug only) - if image is None: - raise ValueError("Image is required for QwenImageEditPlusPipeline") - - if not isinstance(image, list): - image = [image] - - image_size = image[0].size - calculated_width, calculated_height = calculate_dimensions(VAE_IMAGE_SIZE, image_size[0] / image_size[1]) - height = height or calculated_height - width = width or calculated_width - - height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) + raise RuntimeError("Missing preprocess images that should have been created by the preprocess function.") - condition_images = [] - vae_images = [] - condition_image_sizes = [] - vae_image_sizes = [] - - for img in image: - image_width, image_height = img.size - condition_width, condition_height = calculate_dimensions( - CONDITION_IMAGE_SIZE, image_width / image_height - ) - vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height) - condition_image_sizes.append((condition_width, condition_height)) - vae_image_sizes.append((vae_width, vae_height)) - condition_images.append(self.image_processor.resize(img, condition_height, condition_width)) - vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2)) - - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + max_sequence_length = req.sampling_params.max_sequence_length or 1024 + generator = req.sampling_params.generator + true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0 if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 1.0 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + prompt_embeds_mask = None + negative_prompt_embeds_mask = None + output_type = req.sampling_params.output_type or "pil" + attention_kwargs = None + callback_on_step_end_tensor_inputs = ["latents"] # 1. check inputs # 2. encode prompts @@ -741,16 +706,9 @@ def forward( self._current_timestep = None self._interrupt = False - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = 1 - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + has_neg_prompt = negative_prompt is not None do_true_cfg = true_cfg_scale > 1 and has_neg_prompt self.check_cfg_parallel_validity(true_cfg_scale, has_neg_prompt) @@ -812,10 +770,8 @@ def forward( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) + txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds) + negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds) latents = self.diffuse( prompt_embeds, diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index a1fcff60d40..e2fe6797698 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -7,7 +7,7 @@ import math import os from collections.abc import Iterable -from typing import Any, cast +from typing import ClassVar, cast import numpy as np import PIL.Image @@ -25,7 +25,7 @@ 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 SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import ( AutoencoderKLQwenImage, ) @@ -35,6 +35,7 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( QwenImageTransformer2DModel, ) +from vllm_omni.diffusion.models.qwen_image.rope_utils import txt_seq_lens_from_embeds from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.prompt_utils import ( @@ -44,6 +45,7 @@ normalize_min_aligned_size, ) from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, @@ -79,63 +81,61 @@ def pre_process_func( request: OmniDiffusionRequest, ): """Pre-process requests for QwenImageLayeredPipeline.""" - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if not raw_image: # None or empty list - raise ValueError("""Received no input image. This model requires one input image to run.""") - elif isinstance(raw_image, list): - if len(raw_image) > 1: - raise ValueError( - """Received multiple input images. Only a single image is supported by this model.""" - ) - else: - raw_image = raw_image[0] - - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if not raw_image: # None or empty list + raise ValueError("""Received no input image. This model requires one input image to run.""") + elif isinstance(raw_image, list): + if len(raw_image) > 1: + raise ValueError("""Received multiple input images. Only a single image is supported by this model.""") else: - image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) + raw_image = raw_image[0] - if isinstance(image, PIL.Image.Image) and image.mode != "RGBA": - image = image.convert("RGBA") + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image) + else: + image = cast(PIL.Image.Image | torch.Tensor | np.ndarray, raw_image) - # 1. calculate dimensions - image_size = image.size - assert request.sampling_params.resolution in [640, 1024], ( - f"resolution must be either 640 or 1024, but got {request.sampling_params.resolution}" - ) - calculated_width, calculated_height = calculate_dimensions( - request.sampling_params.resolution * request.sampling_params.resolution, image_size[0] / image_size[1] - ) - height = calculated_height - width = calculated_width - - height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) - - # Store calculated dimensions in request - prompt["additional_information"]["calculated_height"] = calculated_height - prompt["additional_information"]["calculated_width"] = calculated_width - request.sampling_params.height = height - request.sampling_params.width = width - - # 2. Preprocess image - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels): - image = image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image - image = image_processor.preprocess(image, calculated_height, calculated_width) - image = image.unsqueeze(2) - # image = image.to(dtype=self.text_encoder.dtype) # do it later - - # Store preprocessed image and prompt image in request - prompt["additional_information"]["preprocessed_image"] = image - prompt["additional_information"]["prompt_image"] = prompt_image - request.prompts[i] = prompt + if isinstance(image, PIL.Image.Image) and image.mode != "RGBA": + image = image.convert("RGBA") + + # 1. calculate dimensions + image_size = image.size + assert request.sampling_params.resolution in [640, 1024], ( + f"resolution must be either 640 or 1024, but got {request.sampling_params.resolution}" + ) + calculated_width, calculated_height = calculate_dimensions( + request.sampling_params.resolution * request.sampling_params.resolution, image_size[0] / image_size[1] + ) + height = calculated_height + width = calculated_width + + height, width = normalize_min_aligned_size(height, width, vae_scale_factor * 2) + + # Store calculated dimensions in request + prompt["additional_information"]["calculated_height"] = calculated_height + prompt["additional_information"]["calculated_width"] = calculated_width + request.sampling_params.height = height + request.sampling_params.width = width + + # 2. Preprocess image + if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == latent_channels): + image = image_processor.resize(image, calculated_height, calculated_width) + prompt_image = image + image = image_processor.preprocess(image, calculated_height, calculated_width) + image = image.unsqueeze(2) + # image = image.to(dtype=self.text_encoder.dtype) # do it later + + # Store preprocessed image and prompt image in request + prompt["additional_information"]["preprocessed_image"] = image + prompt["additional_information"]["prompt_image"] = prompt_image + request.prompt = prompt return request return pre_process_func @@ -202,7 +202,13 @@ def retrieve_latents( raise AttributeError("Could not access latents of provided encoder_output") -class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin): +class QwenImageLayeredPipeline( + nn.Module, SupportImageInput, QwenImageCFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + color_format = "RGBA" def __init__( @@ -269,7 +275,11 @@ def __init__( ] transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel) - self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs) + self.transformer = QwenImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + **transformer_kwargs, + ) # Pipeline configuration & processing parameters self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8 @@ -638,31 +648,7 @@ def current_timestep(self): def interrupt(self): return self._interrupt - def forward( - self, - req: OmniDiffusionRequest, - image: PIL.Image.Image | torch.Tensor | None = None, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - true_cfg_scale: float = 4.0, - layers: int | None = 4, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float | None = None, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - prompt_embeds_mask: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds_mask: torch.Tensor | None = None, - output_type: str | None = "pil", - attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int = 1024, - resolution: int = 640, - cfg_normalize: bool = False, - use_en_prompt: bool = False, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """Forward pass for image layered.""" # 1. Get preprocessed image from request (pre-processing is done in DiffusionEngine) @@ -678,26 +664,27 @@ def forward( prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") - layers = req.sampling_params.layers if req.sampling_params.layers is not None else layers - resolution = req.sampling_params.resolution if req.sampling_params.resolution is not None else resolution - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - cfg_normalize = ( - req.sampling_params.cfg_normalize if req.sampling_params.cfg_normalize is not None else cfg_normalize - ) - use_en_prompt = ( - req.sampling_params.use_en_prompt if req.sampling_params.use_en_prompt is not None else use_en_prompt - ) - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - sigmas = req.sampling_params.sigmas or sigmas - generator = req.sampling_params.generator or generator - true_cfg_scale = req.sampling_params.true_cfg_scale or true_cfg_scale + layers = req.sampling_params.layers + max_sequence_length = req.sampling_params.max_sequence_length or 1024 + cfg_normalize = req.sampling_params.cfg_normalize + use_en_prompt = req.sampling_params.use_en_prompt + num_inference_steps = req.sampling_params.num_inference_steps or 50 + sigmas = req.sampling_params.sigmas + generator = req.sampling_params.generator + true_cfg_scale = req.sampling_params.true_cfg_scale or 4.0 + guidance_scale = None if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + prompt_embeds = None + negative_prompt_embeds = None + prompt_embeds_mask = None + negative_prompt_embeds_mask = None + output_type = req.sampling_params.output_type or "pil" + attention_kwargs = None if not isinstance(first_prompt, str) and "preprocessed_image" in ( additional_information := first_prompt.get("additional_information", {}) @@ -710,25 +697,7 @@ def forward( height = req.sampling_params.height width = req.sampling_params.width else: - # fallback to run pre-processing in pipeline (debug only) - if isinstance(image, PIL.Image.Image) and image.mode != "RGBA": - image = image.convert("RGBA") - image_size = image[0].size if isinstance(image, list) else image.size - assert resolution in [640, 1024], f"resolution must be either 640 or 1024, but got {resolution}" - calculated_width, calculated_height = calculate_dimensions( - resolution * resolution, image_size[0] / image_size[1] - ) - height = calculated_height - width = calculated_width - - height, width = normalize_min_aligned_size(height, width, self.vae_scale_factor * 2) - - if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels): - image = self.image_processor.resize(image, calculated_height, calculated_width) - prompt_image = image - image = self.image_processor.preprocess(image, calculated_height, calculated_width) - image = image.unsqueeze(2) - image = image.to(dtype=self.text_encoder.dtype) + raise RuntimeError("Missing preprocess image that should have been created by the preprocess function.") # 2. check inputs self.check_inputs( @@ -751,16 +720,9 @@ def forward( # 3. encode prompot & negative prompt if prompt is None or prompt == "" or prompt == " ": prompt = self.get_image_caption(prompt_image, use_en_prompt=use_en_prompt, device=self.device) - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = 1 - has_neg_prompt = negative_prompt is not None or ( - negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None - ) + has_neg_prompt = negative_prompt is not None if true_cfg_scale > 1 and not has_neg_prompt: logger.warning( @@ -848,10 +810,8 @@ def forward( if self.attention_kwargs is None: self._attention_kwargs = {} - txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None - negative_txt_seq_lens = ( - negative_prompt_embeds_mask.sum(dim=1).tolist() if negative_prompt_embeds_mask is not None else None - ) + txt_seq_lens = txt_seq_lens_from_embeds(prompt_embeds) + negative_txt_seq_lens = txt_seq_lens_from_embeds(negative_prompt_embeds) is_rgb = torch.tensor([0] * batch_size).to(device=self.device, dtype=torch.long) latents = self.diffuse( diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index e61b59c56a5..67e2d0db7a3 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -25,6 +25,8 @@ ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.quantization.component_config import safe_quant_config + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -48,8 +50,23 @@ logger = init_logger(__name__) -def _join_prefix(prefix: str, suffix: str) -> str: - return f"{prefix}.{suffix}" if prefix else suffix +def _normalize_qwen_image_weight_name(name: str) -> str: + name = name.removeprefix("transformer.") + if ".to_out.0." in name: + name = name.replace(".to_out.0.", ".to_out.") + return name + + +def _resolve_qwen_image_lookup_name( + name: str, + stacked_params_mapping: list[tuple[str, str, str]], +) -> tuple[str, str | None]: + lookup_name = _normalize_qwen_image_weight_name(name) + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in lookup_name or param_name in lookup_name: + continue + return lookup_name.replace(weight_name, param_name), shard_id + return lookup_name, None class ImageRopePrepare(nn.Module): @@ -444,7 +461,7 @@ def __init__( gather_output=False, return_bias=False, quant_config=quant_config, - prefix=prefix, + prefix=f"{prefix}.proj" if prefix else "proj", ) self.approximate = approximate @@ -479,7 +496,7 @@ def __init__( approximate="tanh", bias=bias, quant_config=quant_config, - prefix=_join_prefix(prefix, "net.0.proj"), + prefix=f"{prefix}.net.0", ), nn.Identity(), # placeholder for weight loading RowParallelLinear( @@ -488,7 +505,7 @@ def __init__( input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "net.2"), + prefix=f"{prefix}.net.2", ), ] @@ -532,7 +549,7 @@ def __init__( head_size=self.head_dim, total_num_heads=num_heads, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_qkv"), + prefix=f"{prefix}.to_qkv", ) self.query_num_heads = self.to_qkv.num_heads self.kv_num_heads = self.to_qkv.num_kv_heads @@ -548,7 +565,7 @@ def __init__( head_size=head_dim, total_num_heads=num_heads, quant_config=quant_config, - prefix=_join_prefix(prefix, "add_kv_proj"), + prefix=f"{prefix}.add_kv_proj", ) self.add_query_num_heads = self.add_kv_proj.num_heads self.add_kv_num_heads = self.add_kv_proj.num_kv_heads @@ -561,7 +578,7 @@ def __init__( input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_add_out"), + prefix=f"{prefix}.to_add_out", ) assert not pre_only @@ -572,7 +589,7 @@ def __init__( input_is_parallel=True, return_bias=False, quant_config=quant_config, - prefix=_join_prefix(prefix, "to_out"), + prefix=f"{prefix}.to_out.0", ) self.norm_added_q = nn.RMSNorm(head_dim, eps=eps) @@ -713,22 +730,19 @@ def __init__( self.dim = dim self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim + txt_mod_quant_config = safe_quant_config(quant_config) # Image processing modules. - # Modulation linear is kept unquantized (quant_config=None) — it - # produces shift/scale/gate values that are precision-sensitive - # (see #2728). Use column TP with gather_output=True so weights are - # sharded while downstream modulation still receives full [B, 6 * dim]. + # The re-quantized W4A16 checkpoint keeps img_mod.1 in full precision. self.img_mod = nn.Sequential( nn.SiLU(), - ColumnParallelLinear( + ReplicatedLinear( dim, 6 * dim, bias=True, - gather_output=True, return_bias=False, quant_config=None, - prefix=_join_prefix(prefix, "img_mod.1"), + prefix=f"{prefix}.img_mod.1", ), ) self.img_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) @@ -739,27 +753,27 @@ def __init__( context_pre_only=False, head_dim=attention_head_dim, quant_config=quant_config, - prefix=_join_prefix(prefix, "attn"), + prefix=f"{prefix}.attn", ) self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.img_mlp = FeedForward( dim=dim, dim_out=dim, quant_config=quant_config, - prefix=_join_prefix(prefix, "img_mlp"), + prefix=f"{prefix}.img_mlp", ) # Text processing modules. + # AutoRound keeps txt_mod.1 quantized inside transformer_blocks. self.txt_mod = nn.Sequential( nn.SiLU(), - ColumnParallelLinear( + ReplicatedLinear( dim, 6 * dim, bias=True, - gather_output=True, return_bias=False, - quant_config=None, - prefix=_join_prefix(prefix, "txt_mod.1"), + quant_config=txt_mod_quant_config, + prefix=f"{prefix}.txt_mod.1", ), ) self.txt_norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) @@ -769,7 +783,7 @@ def __init__( dim=dim, dim_out=dim, quant_config=quant_config, - prefix=_join_prefix(prefix, "txt_mlp"), + prefix=f"{prefix}.txt_mlp", ) self.zero_cond_t = zero_cond_t @@ -979,6 +993,7 @@ def __init__( self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim self.guidance_embeds = guidance_embeds + self.quant_config = quant_config if not use_layer3d_rope: self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True) @@ -1137,24 +1152,44 @@ def forward( # Check for SP auto_pad: create attention mask dynamically if padding was applied # In Ulysses mode, attention is computed on the FULL sequence (after All-to-All) hidden_states_mask = None # default - if self.parallel_config is not None and self.parallel_config.sequence_parallel_size > 1: - ctx = get_forward_context() - if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: - # Create mask for the full (padded) sequence - # valid positions = True, padding positions = False - batch_size = hidden_states.shape[0] - padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size - hidden_states_mask = torch.ones( - batch_size, - padded_seq_len, - dtype=torch.bool, - device=hidden_states.device, - ) - hidden_states_mask[:, ctx.sp_original_seq_len :] = False + ctx = get_forward_context() + if ( + self.parallel_config is not None + and self.parallel_config.sequence_parallel_size > 1 + and self.parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + # Create mask for the full (padded) sequence + # valid positions = True, padding positions = False + batch_size = hidden_states.shape[0] + padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + hidden_states_mask = torch.ones( + batch_size, + padded_seq_len, + dtype=torch.bool, + device=hidden_states.device, + ) + hidden_states_mask[:, ctx.sp_original_seq_len :] = False + if hidden_states_mask.all(): + hidden_states_mask = None + elif ( + self.parallel_config is not None + and self.parallel_config.sequence_parallel_size > 1 + and not self.parallel_config.mask_sp_padding + and ctx.sp_original_seq_len is not None + and ctx.sp_padding_size > 0 + ): + logger.warning_once( + "SP auto-padding applied %d token(s) (seq_len=%d, ulysses_degree=%d). " + "Padding tokens are not masked from attention (mask_sp_padding=False), " + "which avoids the varlen attention path but may produce minor numerical differences. " + "Set parallel_config.mask_sp_padding=True to restore strict masking.", + ctx.sp_padding_size, + ctx.sp_original_seq_len, + self.parallel_config.sequence_parallel_size, + ) - # if mask is all true, set it to None - if hidden_states_mask is not None and hidden_states_mask.all(): - hidden_states_mask = None if encoder_hidden_states_mask is not None and encoder_hidden_states_mask.all(): encoder_hidden_states_mask = None @@ -1205,22 +1240,26 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params: set[str] = set() for name, loaded_weight in weights: - original_name = name - lookup_name = name - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in original_name or param_name in original_name: - continue - lookup_name = original_name.replace(weight_name, param_name) - param = params_dict[lookup_name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if lookup_name not in params_dict and ".to_out.0." in lookup_name: - lookup_name = lookup_name.replace(".to_out.0.", ".to_out.") - param = params_dict[lookup_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) + original_name = name.removeprefix("transformer.") + lookup_name, shard_id = _resolve_qwen_image_lookup_name( + original_name, + stacked_params_mapping, + ) + + if lookup_name.endswith(".bias") and lookup_name not in params_dict: + continue + + param = params_dict.get(lookup_name) + if param is None: + logger.debug("Skipping unexpected Qwen-Image transformer weight %s", original_name) + continue + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + if shard_id is None: weight_loader(param, loaded_weight) + else: + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(original_name) loaded_params.add(lookup_name) return loaded_params diff --git a/vllm_omni/diffusion/models/qwen_image/rope_utils.py b/vllm_omni/diffusion/models/qwen_image/rope_utils.py new file mode 100644 index 00000000000..68dbccb5996 --- /dev/null +++ b/vllm_omni/diffusion/models/qwen_image/rope_utils.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + + +def txt_seq_lens_from_embeds(prompt_embeds: torch.Tensor | None) -> list[int] | None: + """Return per-row RoPE text lengths from padded encoder hidden-state width. + + Diffusers derives Qwen-Image text RoPE length from + ``encoder_hidden_states.shape[1]``, not from the count of valid (non-padding) + tokens in the attention mask. Callers that right-pad prompt embeddings for + batching must pass the padded width here so ``QwenEmbedRope`` builds a table + long enough for every text position. + """ + if prompt_embeds is None: + return None + if prompt_embeds.ndim == 2: + prompt_embeds = prompt_embeds.unsqueeze(0) + if prompt_embeds.ndim != 3: + raise ValueError(f"prompt_embeds must be 2D or 3D, got shape={tuple(prompt_embeds.shape)}") + seq_len = int(prompt_embeds.shape[1]) + batch_size = int(prompt_embeds.shape[0]) + return [seq_len] * batch_size diff --git a/vllm_omni/diffusion/models/schedulers/__init__.py b/vllm_omni/diffusion/models/schedulers/__init__.py index e683ed27203..11aedcb1d56 100644 --- a/vllm_omni/diffusion/models/schedulers/__init__.py +++ b/vllm_omni/diffusion/models/schedulers/__init__.py @@ -2,11 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from vllm_omni.diffusion.models.schedulers.scheduling_dmd2_euler import DMD2EulerScheduler +from vllm_omni.diffusion.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import ( FlowUniPCMultistepScheduler, ) __all__ = [ "DMD2EulerScheduler", + "FlowMatchEulerDiscreteScheduler", "FlowUniPCMultistepScheduler", ] diff --git a/vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py b/vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py new file mode 100644 index 00000000000..8aa7d3a0018 --- /dev/null +++ b/vllm_omni/diffusion/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -0,0 +1,660 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +# Originally from https://github.com/huggingface/diffusers/blob/v0.39.0/src/diffusers/schedulers/scheduling_flow_match_euler_discrete.py +# Copyright 2025 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Flow Match Euler scheduler copied from Diffusers v0.39.0. + +Modified by vLLM-Omni contributors for local integration and RNG compatibility. +""" + +import math +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput, is_scipy_available, logging +from diffusers.utils.torch_utils import randn_tensor + +if is_scipy_available(): + import scipy.stats + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +@dataclass +class FlowMatchEulerDiscreteSchedulerOutput(BaseOutput): + """ + Output class for the scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the + denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin): + """ + Euler scheduler. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + shift (`float`, defaults to 1.0): + The shift value for the timestep schedule. + use_dynamic_shifting (`bool`, defaults to False): + Whether to apply timestep shifting on-the-fly based on the image resolution. + base_shift (`float`, defaults to 0.5): + Value to stabilize image generation. Increasing `base_shift` reduces variation and image is more consistent + with desired output. + max_shift (`float`, defaults to 1.15): + Value change allowed to latent vectors. Increasing `max_shift` encourages more variation and image may be + more exaggerated or stylized. + base_image_seq_len (`int`, defaults to 256): + The base image sequence length. + max_image_seq_len (`int`, defaults to 4096): + The maximum image sequence length. + invert_sigmas (`bool`, defaults to False): + Whether to invert the sigmas. + shift_terminal (`float`, defaults to None): + The end value of the shifted timestep schedule. + use_karras_sigmas (`bool`, defaults to False): + Whether to use Karras sigmas for step sizes in the noise schedule during sampling. + use_exponential_sigmas (`bool`, defaults to False): + Whether to use exponential sigmas for step sizes in the noise schedule during sampling. + use_beta_sigmas (`bool`, defaults to False): + Whether to use beta sigmas for step sizes in the noise schedule during sampling. + time_shift_type (`str`, defaults to "exponential"): + The type of dynamic resolution-dependent timestep shifting to apply. Either "exponential" or "linear". + stochastic_sampling (`bool`, defaults to False): + Whether to use stochastic sampling. + """ + + _compatibles = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + use_dynamic_shifting: bool = False, + base_shift: float | None = 0.5, + max_shift: float | None = 1.15, + base_image_seq_len: int = 256, + max_image_seq_len: int = 4096, + invert_sigmas: bool = False, + shift_terminal: float = None, + use_karras_sigmas: bool = False, + use_exponential_sigmas: bool = False, + use_beta_sigmas: bool = False, + time_shift_type: Literal["exponential", "linear"] = "exponential", + stochastic_sampling: bool = False, + ): + if self.config.use_beta_sigmas and not is_scipy_available(): + raise ImportError("Make sure to install scipy if you want to use beta sigmas.") + if ( + sum( + [ + self.config.use_beta_sigmas, + self.config.use_exponential_sigmas, + self.config.use_karras_sigmas, + ] + ) + > 1 + ): + raise ValueError( + "Only one of `config.use_beta_sigmas`, `config.use_exponential_sigmas`, " + "`config.use_karras_sigmas` can be used." + ) + if time_shift_type not in {"exponential", "linear"}: + raise ValueError("`time_shift_type` must either be 'exponential' or 'linear'.") + + timesteps = np.linspace(1, num_train_timesteps, num_train_timesteps, dtype=np.float32)[::-1].copy() + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32) + + sigmas = timesteps / num_train_timesteps + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.timesteps = sigmas * num_train_timesteps + + self._step_index = None + self._begin_index = None + + self._shift = shift + + self.sigmas = sigmas.to("cpu") # to avoid too much CPU/GPU communication + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + @property + def shift(self): + """ + The value used for shifting. + """ + return self._shift + + @property + def step_index(self): + """ + The index counter for current timestep. It will increase 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self): + """ + The index for the first timestep. It should be set from pipeline with `set_begin_index` method. + """ + return self._begin_index + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index + def set_begin_index(self, begin_index: int = 0): + """ + Sets the begin index for the scheduler. This function should be run from pipeline before the inference. + + Args: + begin_index (`int`, defaults to `0`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def set_shift(self, shift: float): + """ + Sets the shift value for the scheduler. + + Args: + shift (`float`): + The shift value to be set. + """ + self._shift = shift + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: float | torch.FloatTensor, + noise: torch.FloatTensor | None = None, + ) -> torch.FloatTensor: + """ + Forward process in flow-matching + + Args: + sample (`torch.FloatTensor`): + The input sample. + timestep (`torch.FloatTensor`): + The current timestep in the diffusion chain. + noise (`torch.FloatTensor`): + The noise tensor. + + Returns: + `torch.FloatTensor`: + A scaled input sample. + """ + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to(device=sample.device, dtype=sample.dtype) + + if sample.device.type == "mps" and torch.is_floating_point(timestep): + # mps does not support float64 + schedule_timesteps = self.timesteps.to(sample.device, dtype=torch.float32) + timestep = timestep.to(sample.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(sample.device) + timestep = timestep.to(sample.device) + + # self.begin_index is None when scheduler is used for training, or pipeline does not implement set_begin_index + if self.begin_index is None: + step_indices = [self.index_for_timestep(t, schedule_timesteps) for t in timestep] + elif self.step_index is not None: + # add_noise is called after first denoising step (for inpainting) + step_indices = [self.step_index] * timestep.shape[0] + else: + # add noise is called before first denoising step to create initial latent(img2img) + step_indices = [self.begin_index] * timestep.shape[0] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(sample.shape): + sigma = sigma.unsqueeze(-1) + + sample = sigma * noise + (1.0 - sigma) * sample + + return sample + + def _sigma_to_t(self, sigma) -> float: + return sigma * self.config.num_train_timesteps + + def time_shift(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: + """ + Apply time shifting to the sigmas. + + Args: + mu (`float`): + The mu parameter for the time shift. + sigma (`float`): + The sigma parameter for the time shift. + t (`torch.Tensor`): + The input timesteps. + + Returns: + `torch.Tensor`: + The time-shifted timesteps. + """ + if self.config.time_shift_type == "exponential": + return self._time_shift_exponential(mu, sigma, t) + elif self.config.time_shift_type == "linear": + return self._time_shift_linear(mu, sigma, t) + + def stretch_shift_to_terminal(self, t: torch.Tensor) -> torch.Tensor: + r""" + Stretches and shifts the timestep schedule to ensure it terminates at the configured `shift_terminal` config + value. + + Reference: + https://github.com/Lightricks/LTX-Video/blob/a01a171f8fe3d99dce2728d60a73fecf4d4238ae/ltx_video/schedulers/rf.py#L51 + + Args: + t (`torch.Tensor`): + A tensor of timesteps to be stretched and shifted. + + Returns: + `torch.Tensor`: + A tensor of adjusted timesteps such that the final value equals `self.config.shift_terminal`. + """ + one_minus_z = 1 - t + scale_factor = one_minus_z[-1] / (1 - self.config.shift_terminal) + stretched_t = 1 - (one_minus_z / scale_factor) + return stretched_t + + def set_timesteps( + self, + num_inference_steps: int | None = None, + device: str | torch.device = None, + sigmas: list[float] | None = None, + mu: float | None = None, + timesteps: list[float] | None = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + + Args: + num_inference_steps (`int`, *optional*): + The number of diffusion steps used when generating samples with a pre-trained model. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + sigmas (`list[float]`, *optional*): + Custom values for sigmas to be used for each diffusion step. If `None`, the sigmas are computed + automatically. + mu (`float`, *optional*): + Determines the amount of shifting applied to sigmas when performing resolution-dependent timestep + shifting. + timesteps (`list[float]`, *optional*): + Custom values for timesteps to be used for each diffusion step. If `None`, the timesteps are computed + automatically. + """ + if self.config.use_dynamic_shifting and mu is None: + raise ValueError("`mu` must be passed when `use_dynamic_shifting` is set to be `True`") + + if sigmas is not None and timesteps is not None: + if len(sigmas) != len(timesteps): + raise ValueError("`sigmas` and `timesteps` should have the same length") + + if num_inference_steps is not None: + if (sigmas is not None and len(sigmas) != num_inference_steps) or ( + timesteps is not None and len(timesteps) != num_inference_steps + ): + raise ValueError( + "`sigmas` and `timesteps` should have the same length as num_inference_steps, " + "if `num_inference_steps` is provided" + ) + else: + num_inference_steps = len(sigmas) if sigmas is not None else len(timesteps) + + self.num_inference_steps = num_inference_steps + + # 1. Prepare default sigmas + is_timesteps_provided = timesteps is not None + + if is_timesteps_provided: + timesteps = np.array(timesteps).astype(np.float32) + + if sigmas is None: + if timesteps is None: + timesteps = np.linspace( + self._sigma_to_t(self.sigma_max), + self._sigma_to_t(self.sigma_min), + num_inference_steps, + ) + sigmas = timesteps / self.config.num_train_timesteps + else: + sigmas = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas) + + # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of + # "exponential" or "linear" type is applied + if self.config.use_dynamic_shifting: + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + + # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value + if self.config.shift_terminal: + sigmas = self.stretch_shift_to_terminal(sigmas) + + # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules + if self.config.use_karras_sigmas: + sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + elif self.config.use_exponential_sigmas: + sigmas = self._convert_to_exponential(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + elif self.config.use_beta_sigmas: + sigmas = self._convert_to_beta(in_sigmas=sigmas, num_inference_steps=num_inference_steps) + + # 5. Convert sigmas and timesteps to tensors and move to specified device + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device=device) + if not is_timesteps_provided: + timesteps = sigmas * self.config.num_train_timesteps + else: + timesteps = torch.from_numpy(timesteps).to(dtype=torch.float32, device=device) + + # 6. Append the terminal sigma value. + # If a model requires inverted sigma schedule for denoising but timesteps without inversion, the + # `invert_sigmas` flag can be set to `True`. This case is only required in Mochi + if self.config.invert_sigmas: + sigmas = 1.0 - sigmas + timesteps = sigmas * self.config.num_train_timesteps + sigmas = torch.cat([sigmas, torch.ones(1, device=sigmas.device)]) + else: + sigmas = torch.cat([sigmas, torch.zeros(1, device=sigmas.device)]) + + self.timesteps = timesteps + self.sigmas = sigmas + self._step_index = None + self._begin_index = None + + def index_for_timestep( + self, + timestep: float | torch.FloatTensor, + schedule_timesteps: torch.FloatTensor | None = None, + ) -> int: + """ + Get the index for the given timestep. + + Args: + timestep (`float` or `torch.FloatTensor`): + The timestep to find the index for. + schedule_timesteps (`torch.FloatTensor`, *optional*): + The schedule timesteps to validate against. If `None`, the scheduler's timesteps are used. + + Returns: + `int`: + The index of the timestep. + """ + if schedule_timesteps is None: + schedule_timesteps = self.timesteps + + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def _init_step_index(self, timestep: float | torch.FloatTensor) -> None: + if self.begin_index is None: + if isinstance(timestep, torch.Tensor): + timestep = timestep.to(self.timesteps.device) + self._step_index = self.index_for_timestep(timestep) + else: + self._step_index = self._begin_index + + def step( + self, + model_output: torch.FloatTensor, + timestep: float | torch.FloatTensor, + sample: torch.FloatTensor, + s_churn: float = 0.0, + s_tmin: float = 0.0, + s_tmax: float = float("inf"), + s_noise: float = 1.0, + generator: torch.Generator | None = None, + per_token_timesteps: torch.Tensor | None = None, + return_dict: bool = True, + ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion + process from the learned model outputs (most often the predicted noise). + + Args: + model_output (`torch.FloatTensor`): + The direct output from learned diffusion model. + timestep (`float`): + The current discrete timestep in the diffusion chain. + sample (`torch.FloatTensor`): + A current instance of a sample created by the diffusion process. + s_churn (`float`): + s_tmin (`float`): + s_tmax (`float`): + s_noise (`float`, defaults to 1.0): + Scaling factor for noise added to the sample. + generator (`torch.Generator`, *optional*): + A random number generator. + per_token_timesteps (`torch.Tensor`, *optional*): + The timesteps for each token in the sample. + return_dict (`bool`, defaults to `True`): + Whether or not to return a + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or tuple. + + Returns: + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] or `tuple`: + If return_dict is `True`, + [`~schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteSchedulerOutput`] is returned, + otherwise a tuple is returned where the first element is the sample tensor. + """ + + if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor): + raise ValueError( + ( + "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" + " `FlowMatchEulerDiscreteScheduler.step()` is not supported. Make sure to pass" + " one of the `scheduler.timesteps` as a timestep." + ), + ) + + if self.step_index is None: + self._init_step_index(timestep) + + # Upcast to avoid precision issues when computing prev_sample + sample = sample.to(torch.float32) + + if per_token_timesteps is not None: + per_token_sigmas = per_token_timesteps / self.config.num_train_timesteps + + sigmas = self.sigmas[:, None, None] + lower_mask = sigmas < per_token_sigmas[None] - 1e-6 + lower_sigmas = lower_mask * sigmas + lower_sigmas, _ = lower_sigmas.max(dim=0) + + current_sigma = per_token_sigmas[..., None] + next_sigma = lower_sigmas[..., None] + dt = current_sigma - next_sigma + else: + sigma_idx = self.step_index + sigma = self.sigmas[sigma_idx] + sigma_next = self.sigmas[sigma_idx + 1] + + current_sigma = sigma + next_sigma = sigma_next + dt = sigma_next - sigma + + if self.config.stochastic_sampling: + x0 = sample - current_sigma * model_output + noise = randn_tensor(sample.shape, generator=generator, device=sample.device, dtype=sample.dtype) + prev_sample = (1.0 - next_sigma) * x0 + next_sigma * noise + else: + prev_sample = sample + dt * model_output + + # upon completion increase step index by one + self._step_index += 1 + if per_token_timesteps is None: + # Cast sample back to model compatible dtype + prev_sample = prev_sample.to(model_output.dtype) + + if not return_dict: + return (prev_sample,) + + return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample) + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras + def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: + """ + Construct the noise schedule as proposed in [Elucidating the Design Space of Diffusion-Based Generative + Models](https://huggingface.co/papers/2206.00364). + + Args: + in_sigmas (`torch.Tensor`): + The input sigma values to be converted. + num_inference_steps (`int`): + The number of inference steps to generate the noise schedule for. + + Returns: + `torch.Tensor`: + The converted sigma values following the Karras noise schedule. + """ + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + rho = 7.0 # 7.0 is the value used in the paper + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_exponential + def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: + """ + Construct an exponential noise schedule. + + Args: + in_sigmas (`torch.Tensor`): + The input sigma values to be converted. + num_inference_steps (`int`): + The number of inference steps to generate the noise schedule for. + + Returns: + `torch.Tensor`: + The converted sigma values following an exponential schedule. + """ + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.exp(np.linspace(math.log(sigma_max), math.log(sigma_min), num_inference_steps)) + return sigmas + + # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_beta + def _convert_to_beta( + self, in_sigmas: torch.Tensor, num_inference_steps: int, alpha: float = 0.6, beta: float = 0.6 + ) -> torch.Tensor: + """ + Construct a beta noise schedule as proposed in [Beta Sampling is All You + Need](https://huggingface.co/papers/2407.12173). + + Args: + in_sigmas (`torch.Tensor`): + The input sigma values to be converted. + num_inference_steps (`int`): + The number of inference steps to generate the noise schedule for. + alpha (`float`, *optional*, defaults to `0.6`): + The alpha parameter for the beta distribution. + beta (`float`, *optional*, defaults to `0.6`): + The beta parameter for the beta distribution. + + Returns: + `torch.Tensor`: + The converted sigma values following a beta distribution schedule. + """ + + # Hack to make sure that other schedulers which copy this function don't break + # TODO: Add this logic to the other schedulers + if hasattr(self.config, "sigma_min"): + sigma_min = self.config.sigma_min + else: + sigma_min = None + + if hasattr(self.config, "sigma_max"): + sigma_max = self.config.sigma_max + else: + sigma_max = None + + sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item() + sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item() + + sigmas = np.array( + [ + sigma_min + (ppf * (sigma_max - sigma_min)) + for ppf in [ + scipy.stats.beta.ppf(timestep, alpha, beta) + for timestep in 1 - np.linspace(0, 1, num_inference_steps) + ] + ] + ) + return sigmas + + def _time_shift_exponential(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: + return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma) + + def _time_shift_linear(self, mu: float, sigma: float, t: torch.Tensor) -> torch.Tensor: + return mu / (mu + (1 / t - 1) ** sigma) + + def __len__(self) -> int: + return self.config.num_train_timesteps diff --git a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py index 7eeffdde4d4..b2852e524b1 100644 --- a/vllm_omni/diffusion/models/sd3/pipeline_sd3.py +++ b/vllm_omni/diffusion/models/sd3/pipeline_sd3.py @@ -3,6 +3,7 @@ import logging import os from collections.abc import Iterable +from typing import ClassVar import torch from diffusers.image_processor import VaeImageProcessor @@ -20,11 +21,12 @@ 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.sd3.sd3_transformer import ( SD3Transformer2DModel, ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -129,7 +131,13 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin): +class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery): + supports_request_batch = True + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "text_encoder_2", "text_encoder_3"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -428,6 +436,7 @@ def encode_prompt( prompt_2: str | list[str], prompt_3: str | list[str], prompt_embeds: torch.Tensor | None = None, + pooled_prompt_embeds: torch.Tensor | None = None, max_sequence_length: int = 256, num_images_per_prompt: int = 1, ): @@ -451,7 +460,6 @@ def encode_prompt( prompt = [prompt] if isinstance(prompt, str) else prompt - pooled_prompt_embeds = None if prompt_embeds is None: prompt_2 = prompt_2 or prompt prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2 @@ -620,46 +628,53 @@ def diffuse( return latents - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] = "", - prompt_2: str | list[str] = "", - prompt_3: str | list[str] = "", - negative_prompt: str | list[str] = "", - negative_prompt_2: str | list[str] = "", - negative_prompt_3: str | list[str] = "", - height: int | None = None, - width: int | None = None, - num_inference_steps: int = 28, - sigmas: list[float] | None = None, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - pooled_prompt_embeds: torch.Tensor | None = None, - negative_pooled_prompt_embeds: torch.Tensor | None = None, - max_sequence_length: int = 256, - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt - negative_prompt = [ - "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts - ] or negative_prompt - - height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - generator = req.sampling_params.generator or generator + sampling_params_list = req.sampling_params_list + common_sampling_params = sampling_params_list[0] + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + prompt_2 = "" + prompt_3 = "" + negative_prompt_2 = "" + negative_prompt_3 = "" + prompt_embeds = None + negative_prompt_embeds = None + + height = common_sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = common_sampling_params.width or self.default_sample_size * self.vae_scale_factor + sigmas = common_sampling_params.sigmas + max_sequence_length = common_sampling_params.max_sequence_length or 256 + num_inference_steps = common_sampling_params.num_inference_steps or 28 num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + common_sampling_params.num_outputs_per_prompt if common_sampling_params.num_outputs_per_prompt > 0 else 1 ) + generator = req.collate_request_generators(num_images_per_prompt, None) + latents = req.collate_request_tensors("latents", None) + output_type = common_sampling_params.output_type or "pil" + + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "pooled_prompt_embeds": None, + "negative_pooled_prompt_embeds": None, + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + pooled_prompt_embeds = prompt_fields["pooled_prompt_embeds"] + negative_pooled_prompt_embeds = prompt_fields["negative_pooled_prompt_embeds"] + if prompt_embeds is not None: + prompt = None + prompt_2 = None + prompt_3 = None + if negative_prompt_embeds is not None: + negative_prompt = None + negative_prompt_2 = None + negative_prompt_3 = None # 1. check inputs # 2. encode prompts # 3. prepare latents and timesteps @@ -680,13 +695,11 @@ def forward( max_sequence_length=max_sequence_length, ) - self._guidance_scale = req.sampling_params.guidance_scale + self._guidance_scale = common_sampling_params.guidance_scale self._current_timestep = None self._interrupt = False - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): + if prompt is not None: batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] @@ -696,6 +709,7 @@ def forward( prompt_2=prompt_2, prompt_3=prompt_3, prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, max_sequence_length=max_sequence_length, ) @@ -706,6 +720,7 @@ def forward( prompt_2=negative_prompt_2, prompt_3=negative_prompt_3, prompt_embeds=negative_prompt_embeds, + pooled_prompt_embeds=negative_pooled_prompt_embeds, max_sequence_length=max_sequence_length, ) @@ -736,7 +751,7 @@ def forward( ) self._current_timestep = None - if self.output_type == "latent": + if output_type == "latent": image = latents else: # Ensure the latents are the same dtype as the VAE for decode @@ -745,8 +760,13 @@ def forward( image = self.vae.decode(latents, return_dict=False)[0] - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None + return split_diffusion_output_by_request( + DiffusionOutput( + output=image, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ), + req, + num_outputs_per_prompt=num_images_per_prompt, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm_omni/diffusion/models/sd3/sd3_transformer.py b/vllm_omni/diffusion/models/sd3/sd3_transformer.py index 89f06157758..d8e03705112 100644 --- a/vllm_omni/diffusion/models/sd3/sd3_transformer.py +++ b/vllm_omni/diffusion/models/sd3/sd3_transformer.py @@ -5,6 +5,7 @@ import torch.nn.functional as F # TODO replace this with vLLM implementation +from cache_dit import ForwardPattern from diffusers.models.embeddings import CombinedTimestepTextProjEmbeddings, PatchEmbed from diffusers.models.modeling_outputs import Transformer2DModelOutput from diffusers.models.normalization import AdaLayerNormContinuous, AdaLayerNormZero, SD35AdaLayerNormZeroX @@ -19,6 +20,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import OmniDiffusionConfig logger = init_logger(__name__) @@ -386,6 +388,12 @@ class SD3Transformer2DModel(nn.Module): The Transformer model introduced in [Stable Diffusion 3](https://huggingface.co/papers/2403.03206). """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_1, + } + ) + _repeated_blocks = ["SD3TransformerBlock"] _layerwise_offload_blocks_attrs = ["transformer_blocks"] diff --git a/vllm_omni/diffusion/models/sdxl/__init__.py b/vllm_omni/diffusion/models/sdxl/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/diffusion/models/sdxl/__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/diffusion/models/sdxl/pipeline_sdxl.py b/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py new file mode 100644 index 00000000000..0bb8a48cd7c --- /dev/null +++ b/vllm_omni/diffusion/models/sdxl/pipeline_sdxl.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import logging +import os +from collections.abc import Iterable + +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.schedulers.scheduling_euler_discrete import EulerDiscreteScheduler +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import DistributedAutoencoderKL +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin +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.sdxl.sdxl_unet import SDXLUNet2DConditionModel +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +logger = logging.getLogger(__name__) + + +def get_sdxl_image_post_process_func(od_config: OmniDiffusionConfig): + if od_config.output_type == "latent": + return lambda x: x + image_processor = VaeImageProcessor(vae_scale_factor=8) + + def post_process_func(images: torch.Tensor): + if images.device.type != "cpu": + images = images.cpu() + return image_processor.postprocess(images) + + return post_process_func + + +class StableDiffusionXLPipeline( + nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery +): + _dit_modules: list[str] = ["unet"] + _encoder_modules: list[str] = ["text_encoder", "text_encoder_2"] + _vae_modules: list[str] = ["vae"] + _resident_modules: list[str] = [] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="unet", + revision=None, + prefix="unet.", + fall_back_to_pt=True, + ) + ] + + self.device = get_local_device() + model = od_config.model + local_files_only = os.path.exists(model) + dtype = od_config.dtype + + sdxl_subfolders = [ + "scheduler", + "tokenizer", + "tokenizer_2", + "text_encoder", + "text_encoder_2", + "vae", + ] + prefetch_subfolders(model, sdxl_subfolders, local_files_only=local_files_only) + + self.scheduler = EulerDiscreteScheduler.from_pretrained( + model, + subfolder="scheduler", + local_files_only=local_files_only, + torch_dtype=torch.float32, + ) + self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.tokenizer_2 = CLIPTokenizer.from_pretrained( + model, subfolder="tokenizer_2", local_files_only=local_files_only + ) + self.text_encoder = from_pretrained_with_prefetch( + CLIPTextModel.from_pretrained, + model, + subfolder="text_encoder", + prefetch_list=sdxl_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + self.text_encoder_2 = from_pretrained_with_prefetch( + CLIPTextModelWithProjection.from_pretrained, + model, + subfolder="text_encoder_2", + prefetch_list=sdxl_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ) + self.unet = SDXLUNet2DConditionModel(od_config=od_config) + self.vae = from_pretrained_with_prefetch( + DistributedAutoencoderKL.from_pretrained, + model, + subfolder="vae", + prefetch_list=sdxl_subfolders, + local_files_only=local_files_only, + torch_dtype=dtype, + ).to(self.device) + + self.vae_scale_factor = 8 + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.tokenizer_max_length = 77 + self.default_sample_size = 128 # 1024 / 8 + self.output_type = self.od_config.output_type + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler + ) + + def _get_clip_prompt_embeds( + self, + prompt: str | list[str], + tokenizer: CLIPTokenizer, + text_encoder: nn.Module, + num_images_per_prompt: int = 1, + return_pooled: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + text_inputs = tokenizer( + prompt, + padding="max_length", + max_length=self.tokenizer_max_length, + truncation=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + + text_encoder_device = next(text_encoder.parameters()).device + outputs = text_encoder(text_input_ids.to(text_encoder_device), output_hidden_states=True) + prompt_embeds = outputs.hidden_states[-2].to(dtype=self.od_config.dtype, device=self.device) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + pooled_prompt_embeds = None + if return_pooled: + pooled_prompt_embeds = outputs[0].to(dtype=self.od_config.dtype, device=self.device) + pooled_prompt_embeds = pooled_prompt_embeds.repeat(num_images_per_prompt, 1) + pooled_prompt_embeds = pooled_prompt_embeds.view(batch_size * num_images_per_prompt, -1) + + return prompt_embeds, pooled_prompt_embeds + + def encode_prompt( + self, + prompt: str | list[str], + num_images_per_prompt: int = 1, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Encode with first CLIP (ViT-L/14, 768-dim hidden states) + prompt_embeds_1, _ = self._get_clip_prompt_embeds( + prompt, + self.tokenizer, + self.text_encoder, + num_images_per_prompt, + return_pooled=False, + ) + # Encode with second CLIP (OpenCLIP ViT-bigG/14, 1280-dim hidden states + pooled) + prompt_embeds_2, pooled_prompt_embeds = self._get_clip_prompt_embeds( + prompt, + self.tokenizer_2, + self.text_encoder_2, + num_images_per_prompt, + return_pooled=True, + ) + + # Concatenate along feature dimension: 768 + 1280 = 2048 + prompt_embeds = torch.cat([prompt_embeds_1, prompt_embeds_2], dim=-1) + + return prompt_embeds, pooled_prompt_embeds + + def _get_add_time_ids( + self, + original_size: tuple[int, int], + crops_coords_top_left: tuple[int, int], + target_size: tuple[int, int], + batch_size: int, + ) -> torch.Tensor: + add_time_ids = torch.tensor( + [list(original_size) + list(crops_coords_top_left) + list(target_size)], + dtype=torch.float32, + device=self.device, + ) + return add_time_ids.repeat(batch_size, 1) + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + return latents.to(device=self.device, dtype=self.od_config.dtype) + + shape = ( + batch_size, + num_channels_latents, + height // self.vae_scale_factor, + width // self.vae_scale_factor, + ) + latents = randn_tensor(shape, generator=generator, device=self.device, dtype=self.od_config.dtype) + latents = latents * self.scheduler.init_noise_sigma + return latents + + def predict_noise(self, **kwargs): + result = self.unet(**kwargs) + return result[0] + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + def diffuse( + self, + latents: torch.Tensor, + timesteps: torch.Tensor, + prompt_embeds: torch.Tensor, + added_cond_kwargs: dict, + negative_prompt_embeds: torch.Tensor | None, + negative_added_cond_kwargs: dict | None, + do_cfg: bool, + guidance_scale: float, + ) -> torch.Tensor: + for _, t in enumerate(timesteps): + if self.interrupt: + continue + self._current_timestep = t + + # Scale model input + latent_model_input = self.scheduler.scale_model_input(latents, t) + timestep = t.expand(latent_model_input.shape[0]).to(device=self.device, dtype=self.od_config.dtype) + + positive_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": prompt_embeds, + "added_cond_kwargs": added_cond_kwargs, + "return_dict": False, + } + if do_cfg: + negative_kwargs = { + "hidden_states": latent_model_input, + "timestep": timestep, + "encoder_hidden_states": negative_prompt_embeds, + "added_cond_kwargs": negative_added_cond_kwargs, + "return_dict": False, + } + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + do_cfg, + guidance_scale, + positive_kwargs, + negative_kwargs, + cfg_normalize=False, + ) + + latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_cfg) + + return latents + + def forward( + self, + req: DiffusionRequestBatch, + prompt: str | list[str] = "", + negative_prompt: str | list[str] = "", + height: int | None = None, + width: int | None = None, + num_inference_steps: int = 50, + num_images_per_prompt: int = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> DiffusionOutput: + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + negative_prompt = [ + "" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts + ] or negative_prompt + + height = req.sampling_params.height or self.default_sample_size * self.vae_scale_factor + width = req.sampling_params.width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + generator = req.sampling_params.generator or generator + num_images_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_images_per_prompt + ) + + self._guidance_scale = req.sampling_params.guidance_scale + self._current_timestep = None + self._interrupt = False + + if isinstance(prompt, str): + batch_size = 1 + else: + batch_size = len(prompt) + + # Encode prompt + prompt_embeds, pooled_prompt_embeds = self.encode_prompt( + prompt=prompt, + num_images_per_prompt=num_images_per_prompt, + ) + + # Encode negative prompt for CFG + do_cfg = self.guidance_scale > 1.0 + negative_prompt_embeds = None + negative_pooled_prompt_embeds = None + if do_cfg: + negative_prompt_embeds, negative_pooled_prompt_embeds = self.encode_prompt( + prompt=negative_prompt, + num_images_per_prompt=num_images_per_prompt, + ) + + # Prepare added conditioning + original_size = (height, width) + target_size = (height, width) + crops_coords_top_left = (0, 0) + + add_time_ids = self._get_add_time_ids( + original_size, + crops_coords_top_left, + target_size, + batch_size * num_images_per_prompt, + ) + added_cond_kwargs = { + "text_embeds": pooled_prompt_embeds, + "time_ids": add_time_ids, + } + negative_added_cond_kwargs = None + if do_cfg: + negative_added_cond_kwargs = { + "text_embeds": negative_pooled_prompt_embeds, + "time_ids": add_time_ids, + } + + # XPU workaround: set_timesteps uses numpy internally and needs CPU tensors. + for attr in ("alphas_cumprod", "sigmas", "timesteps"): + if hasattr(self.scheduler, attr): + val = getattr(self.scheduler, attr) + if hasattr(val, "device") and val.device.type != "cpu": + setattr(self.scheduler, attr, val.cpu()) + self.scheduler.set_timesteps(num_inference_steps) + self.scheduler.sigmas = self.scheduler.sigmas.to(self.device) + self.scheduler.timesteps = self.scheduler.timesteps.to(self.device) + timesteps = self.scheduler.timesteps + + # Prepare latents (uses scheduler.init_noise_sigma which depends on set_timesteps) + latents = self.prepare_latents( + batch_size * num_images_per_prompt, + self.unet.in_channels, + height, + width, + generator, + latents, + ) + self._num_timesteps = len(timesteps) + + # Denoising loop + latents = self.diffuse( + latents=latents, + timesteps=timesteps, + prompt_embeds=prompt_embeds, + added_cond_kwargs=added_cond_kwargs, + negative_prompt_embeds=negative_prompt_embeds, + negative_added_cond_kwargs=negative_added_cond_kwargs, + do_cfg=do_cfg, + guidance_scale=self.guidance_scale, + ) + + self._current_timestep = None + + if self.output_type == "latent": + image = latents + else: + latents = latents.to(self.vae.dtype) + latents = latents / self.vae.config.scaling_factor + image = self.vae.decode(latents, return_dict=False)[0] + + return DiffusionOutput( + output=image, + 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) + return loader.load_weights(weights) diff --git a/vllm_omni/diffusion/models/sdxl/sdxl_unet.py b/vllm_omni/diffusion/models/sdxl/sdxl_unet.py new file mode 100644 index 00000000000..76382e645a6 --- /dev/null +++ b/vllm_omni/diffusion/models/sdxl/sdxl_unet.py @@ -0,0 +1,871 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) + +logger = init_logger(__name__) + + +def get_timestep_embedding( + timesteps: torch.Tensor, + embedding_dim: int, + max_period: int = 10000, +) -> torch.Tensor: + half_dim = embedding_dim // 2 + exponent = ( + -math.log(max_period) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / half_dim + ) + emb = timesteps[:, None].float() * torch.exp(exponent)[None, :] + emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1) + if embedding_dim % 2 == 1: + emb = F.pad(emb, (0, 1)) + return emb + + +class SDXLTimestepEmbedding(nn.Module): + def __init__(self, in_channels: int, time_embed_dim: int): + super().__init__() + self.linear_1 = nn.Linear(in_channels, time_embed_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + sample = self.linear_1(sample) + sample = self.act(sample) + sample = self.linear_2(sample) + return sample + + +class SDXLAddTimestepEmbedding(nn.Module): + def __init__(self, addition_time_embed_dim: int, text_embed_dim: int, time_embed_dim: int): + super().__init__() + self.linear_1 = nn.Linear(text_embed_dim + addition_time_embed_dim * 6, time_embed_dim) + self.act = nn.SiLU() + self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim) + + def forward(self, text_embeds: torch.Tensor, time_ids: torch.Tensor) -> torch.Tensor: + time_embeds = get_timestep_embedding(time_ids.flatten(), 256) + time_embeds = time_embeds.reshape(text_embeds.shape[0], -1) + add_embeds = torch.cat([text_embeds, time_embeds], dim=-1) + add_embeds = add_embeds.to(dtype=self.linear_1.weight.dtype) + add_embeds = self.linear_1(add_embeds) + add_embeds = self.act(add_embeds) + add_embeds = self.linear_2(add_embeds) + return add_embeds + + +class SDXLResnetBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + groups: int = 32, + ): + super().__init__() + self.norm1 = nn.GroupNorm(groups, in_channels, eps=1e-5) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.time_emb_proj = nn.Linear(time_embed_dim, out_channels) + self.norm2 = nn.GroupNorm(groups, out_channels, eps=1e-5) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1) + self.nonlinearity = nn.SiLU() + + self.conv_shortcut = None + if in_channels != out_channels: + self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1) + + def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor: + residual = hidden_states + + hidden_states = self.norm1(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states) + + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb) + hidden_states = hidden_states + temb[:, :, None, None] + + hidden_states = self.norm2(hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv2(hidden_states) + + if self.conv_shortcut is not None: + residual = self.conv_shortcut(residual) + + return hidden_states + residual + + +class SDXLGEGLU(nn.Module): + def __init__(self, dim_in: int, dim_out: int): + super().__init__() + self.proj = MergedColumnParallelLinear(dim_in, [dim_out, dim_out], bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = self.proj(hidden_states) + hidden_states, gate = hidden_states.chunk(2, dim=-1) + return hidden_states * F.gelu(gate) + + +class SDXLFeedForward(nn.Module): + def __init__(self, dim: int, inner_dim: int): + super().__init__() + self.geglu = SDXLGEGLU(dim, inner_dim) + self.out_proj = RowParallelLinear(inner_dim, dim, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.geglu(hidden_states) + hidden_states, _ = self.out_proj(hidden_states) + return hidden_states + + +class SDXLSelfAttention(nn.Module): + def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, prefix: str = ""): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=attention_head_dim, + total_num_heads=num_attention_heads, + bias=False, + ) + self.to_out = nn.ModuleList( + [ + RowParallelLinear(inner_dim, dim, bias=True), + ] + ) + self.attention = Attention( + num_heads=num_attention_heads, + head_size=attention_head_dim, + causal=False, + softmax_scale=1.0 / math.sqrt(attention_head_dim), + role="self", + prefix=prefix, + ) + + +class SDXLCrossAttention(nn.Module): + def __init__( + self, dim: int, cross_attention_dim: int, num_attention_heads: int, attention_head_dim: int, prefix: str = "" + ): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + self.to_q = ColumnParallelLinear(dim, inner_dim, bias=False) + self.to_k = ColumnParallelLinear(cross_attention_dim, inner_dim, bias=False) + self.to_v = ColumnParallelLinear(cross_attention_dim, inner_dim, bias=False) + self.to_out = nn.ModuleList( + [ + RowParallelLinear(inner_dim, dim, bias=True), + ] + ) + self.attention = Attention( + num_heads=num_attention_heads, + head_size=attention_head_dim, + causal=False, + softmax_scale=1.0 / math.sqrt(attention_head_dim), + role="cross", + prefix=prefix, + skip_sequence_parallel=True, + ) + + +class SDXLBasicTransformerBlock(nn.Module): + def __init__( + self, + dim: int, + num_attention_heads: int, + attention_head_dim: int, + cross_attention_dim: int, + prefix: str = "", + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.attention_head_dim = attention_head_dim + + self.norm1 = nn.LayerNorm(dim, elementwise_affine=True) + self.norm2 = nn.LayerNorm(dim, elementwise_affine=True) + self.norm3 = nn.LayerNorm(dim, elementwise_affine=True) + + self.attn1 = SDXLSelfAttention( + dim=dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + prefix=f"{prefix}.attn1", + ) + self.attn2 = SDXLCrossAttention( + dim=dim, + cross_attention_dim=cross_attention_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + prefix=f"{prefix}.attn2", + ) + + # Feed-forward + self.ff = SDXLFeedForward(dim=dim, inner_dim=dim * 4) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: None = None, + ) -> torch.Tensor: + local_num_heads = self.attn1.to_qkv.num_heads + head_dim = self.attention_head_dim + + # Self-attention + norm_hidden = self.norm1(hidden_states) + qkv, _ = self.attn1.to_qkv(norm_hidden) + q, k, v = qkv.chunk(3, dim=-1) + q = q.unflatten(-1, (local_num_heads, head_dim)) + k = k.unflatten(-1, (local_num_heads, head_dim)) + v = v.unflatten(-1, (local_num_heads, head_dim)) + attn_out = self.attn1.attention(q, k, v) + attn_out = attn_out.flatten(-2) + attn_out, _ = self.attn1.to_out[0](attn_out) + hidden_states = hidden_states + attn_out + + # Cross-attention + norm_hidden = self.norm2(hidden_states) + q, _ = self.attn2.to_q(norm_hidden) + k, _ = self.attn2.to_k(encoder_hidden_states) + v, _ = self.attn2.to_v(encoder_hidden_states) + q = q.unflatten(-1, (local_num_heads, head_dim)) + k = k.unflatten(-1, (local_num_heads, head_dim)) + v = v.unflatten(-1, (local_num_heads, head_dim)) + attn_out = self.attn2.attention(q, k, v) + attn_out = attn_out.flatten(-2) + attn_out, _ = self.attn2.to_out[0](attn_out) + hidden_states = hidden_states + attn_out + + # Feed-forward + norm_hidden = self.norm3(hidden_states) + ff_out = self.ff(norm_hidden) + hidden_states = hidden_states + ff_out + + return hidden_states + + +class SDXLTransformer2DModel(nn.Module): + def __init__( + self, + num_attention_heads: int, + attention_head_dim: int, + in_channels: int, + num_layers: int, + cross_attention_dim: int, + prefix: str = "", + ): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + + self.norm = nn.GroupNorm(32, in_channels, eps=1e-6, affine=True) + self.proj_in = nn.Linear(in_channels, inner_dim) + + self.transformer_blocks = nn.ModuleList( + [ + SDXLBasicTransformerBlock( + dim=inner_dim, + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.transformer_blocks.{i}", + ) + for i in range(num_layers) + ] + ) + + self.proj_out = nn.Linear(inner_dim, in_channels) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> torch.Tensor: + batch, channels, height, width = hidden_states.shape + residual = hidden_states + + hidden_states = self.norm(hidden_states) + hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch, height * width, channels) + hidden_states = self.proj_in(hidden_states) + + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states) + + hidden_states = self.proj_out(hidden_states) + hidden_states = hidden_states.reshape(batch, height, width, channels).permute(0, 3, 1, 2) + + return hidden_states + residual + + +class SDXLDownsample2D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.conv = nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.conv(hidden_states) + + +class SDXLUpsample2D(nn.Module): + def __init__(self, channels: int): + super().__init__() + self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest") + return self.conv(hidden_states) + + +class SDXLDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 2, + add_downsample: bool = True, + ): + super().__init__() + self.resnets = nn.ModuleList() + for i in range(num_layers): + res_in = in_channels if i == 0 else out_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + + self.downsamplers = None + if add_downsample: + self.downsamplers = nn.ModuleList([SDXLDownsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + output_states = [] + for resnet in self.resnets: + hidden_states = resnet(hidden_states, temb) + output_states.append(hidden_states) + + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + output_states.append(hidden_states) + + return hidden_states, output_states + + +class SDXLCrossAttnDownBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 2, + num_attention_heads: int = 10, + cross_attention_dim: int = 2048, + transformer_layers_per_block: int = 1, + add_downsample: bool = True, + prefix: str = "", + ): + super().__init__() + attention_head_dim = out_channels // num_attention_heads + + self.resnets = nn.ModuleList() + self.attentions = nn.ModuleList() + for i in range(num_layers): + res_in = in_channels if i == 0 else out_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + self.attentions.append( + SDXLTransformer2DModel( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.attentions.{i}", + ) + ) + + self.downsamplers = None + if add_downsample: + self.downsamplers = nn.ModuleList([SDXLDownsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + output_states = [] + for resnet, attn in zip(self.resnets, self.attentions): + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states, encoder_hidden_states) + output_states.append(hidden_states) + + if self.downsamplers is not None: + hidden_states = self.downsamplers[0](hidden_states) + output_states.append(hidden_states) + + return hidden_states, output_states + + +class SDXLUNetMidBlock2DCrossAttn(nn.Module): + def __init__( + self, + in_channels: int, + time_embed_dim: int, + num_attention_heads: int = 20, + cross_attention_dim: int = 2048, + transformer_layers_per_block: int = 10, + prefix: str = "", + ): + super().__init__() + attention_head_dim = in_channels // num_attention_heads + + self.resnets = nn.ModuleList( + [ + SDXLResnetBlock2D(in_channels, in_channels, time_embed_dim), + SDXLResnetBlock2D(in_channels, in_channels, time_embed_dim), + ] + ) + self.attentions = nn.ModuleList( + [ + SDXLTransformer2DModel( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=in_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.attentions.0", + ) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + ) -> torch.Tensor: + hidden_states = self.resnets[0](hidden_states, temb) + hidden_states = self.attentions[0](hidden_states, encoder_hidden_states) + hidden_states = self.resnets[1](hidden_states, temb) + return hidden_states + + +class SDXLUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 3, + add_upsample: bool = True, + resnet_in_channels_list: list[int] | None = None, + ): + super().__init__() + self.resnets = nn.ModuleList() + for i in range(num_layers): + if resnet_in_channels_list is not None: + res_in = resnet_in_channels_list[i] + else: + skip_channels = prev_output_channel if i == 0 else out_channels + res_in = (in_channels if i == 0 else out_channels) + skip_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + + self.upsamplers = None + if add_upsample: + self.upsamplers = nn.ModuleList([SDXLUpsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + for resnet in self.resnets: + res_hidden = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden], dim=1) + hidden_states = resnet(hidden_states, temb) + + if self.upsamplers is not None: + hidden_states = self.upsamplers[0](hidden_states) + + return hidden_states + + +class SDXLCrossAttnUpBlock2D(nn.Module): + def __init__( + self, + in_channels: int, + prev_output_channel: int, + out_channels: int, + time_embed_dim: int, + num_layers: int = 3, + num_attention_heads: int = 10, + cross_attention_dim: int = 2048, + transformer_layers_per_block: int = 1, + add_upsample: bool = True, + prefix: str = "", + resnet_in_channels_list: list[int] | None = None, + ): + super().__init__() + attention_head_dim = out_channels // num_attention_heads + + self.resnets = nn.ModuleList() + self.attentions = nn.ModuleList() + for i in range(num_layers): + if resnet_in_channels_list is not None: + res_in = resnet_in_channels_list[i] + else: + skip_channels = prev_output_channel if i == 0 else out_channels + res_in = (in_channels if i == 0 else out_channels) + skip_channels + self.resnets.append(SDXLResnetBlock2D(res_in, out_channels, time_embed_dim)) + self.attentions.append( + SDXLTransformer2DModel( + num_attention_heads=num_attention_heads, + attention_head_dim=attention_head_dim, + in_channels=out_channels, + num_layers=transformer_layers_per_block, + cross_attention_dim=cross_attention_dim, + prefix=f"{prefix}.attentions.{i}", + ) + ) + + self.upsamplers = None + if add_upsample: + self.upsamplers = nn.ModuleList([SDXLUpsample2D(out_channels)]) + + def forward( + self, + hidden_states: torch.Tensor, + temb: torch.Tensor, + encoder_hidden_states: torch.Tensor, + res_hidden_states_tuple: tuple[torch.Tensor, ...], + ) -> torch.Tensor: + for resnet, attn in zip(self.resnets, self.attentions): + res_hidden = res_hidden_states_tuple[-1] + res_hidden_states_tuple = res_hidden_states_tuple[:-1] + hidden_states = torch.cat([hidden_states, res_hidden], dim=1) + hidden_states = resnet(hidden_states, temb) + hidden_states = attn(hidden_states, encoder_hidden_states) + + if self.upsamplers is not None: + hidden_states = self.upsamplers[0](hidden_states) + + return hidden_states + + +def _build_sdxl_sp_plan() -> dict: + """Build _sp_plan for all attention modules in the SDXL UNet. + + Shards hidden_states at the first transformer block input and gathers + at proj_out for each SDXLTransformer2DModel instance. + """ + plan: dict = {} + attn_paths = [ + # down_blocks[1]: 2 attention modules + "down_blocks.1.attentions.0", + "down_blocks.1.attentions.1", + # down_blocks[2]: 2 attention modules + "down_blocks.2.attentions.0", + "down_blocks.2.attentions.1", + # mid_block: 1 attention module + "mid_block.attentions.0", + # up_blocks[0]: 3 attention modules + "up_blocks.0.attentions.0", + "up_blocks.0.attentions.1", + "up_blocks.0.attentions.2", + # up_blocks[1]: 3 attention modules + "up_blocks.1.attentions.0", + "up_blocks.1.attentions.1", + "up_blocks.1.attentions.2", + ] + for path in attn_paths: + plan[f"{path}.transformer_blocks.0"] = { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), + } + plan[f"{path}.proj_out"] = SequenceParallelOutput(gather_dim=1, expected_dims=3) + return plan + + +class SDXLUNet2DConditionModel(nn.Module): + _repeated_blocks = ["SDXLBasicTransformerBlock", "SDXLResnetBlock2D"] + _layerwise_offload_blocks_attrs = ["down_blocks", "up_blocks"] + _sp_plan = _build_sdxl_sp_plan() + + @staticmethod + def _is_shardable_block(name: str, module) -> bool: + """Match transformer blocks and ResNet blocks for HSDP sharding.""" + if not name.split(".")[-1].isdigit(): + return False + return "transformer_blocks" in name or "resnets" in name + + _hsdp_shard_conditions = [_is_shardable_block] + + def __init__(self, *, od_config: OmniDiffusionConfig): + super().__init__() + self.od_config = od_config + + # SDXL Base configuration + model_channels = 320 + time_embed_dim = model_channels * 4 # 1280 + block_out_channels = [320, 640, 1280] + layers_per_block = 2 + transformer_layers_per_block = [1, 2, 10] + cross_attention_dim = 2048 + addition_time_embed_dim = 256 + num_head_channels = 64 + + self.in_channels = 4 + self.out_channels = 4 + + # Input convolution + self.conv_in = nn.Conv2d(4, model_channels, kernel_size=3, padding=1) + + # Time embedding + self.time_proj_dim = model_channels + self.time_embedding = SDXLTimestepEmbedding(model_channels, time_embed_dim) + + # Addition embedding (for pooled text + crop/size conditioning) + self.add_embedding = SDXLAddTimestepEmbedding( + addition_time_embed_dim=addition_time_embed_dim, + text_embed_dim=1280, + time_embed_dim=time_embed_dim, + ) + + # Down blocks + self.down_blocks = nn.ModuleList() + + # Level 0: DownBlock2D (no attention) + self.down_blocks.append( + SDXLDownBlock2D( + in_channels=model_channels, + out_channels=block_out_channels[0], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block, + add_downsample=True, + ) + ) + # Level 1: CrossAttnDownBlock2D + self.down_blocks.append( + SDXLCrossAttnDownBlock2D( + in_channels=block_out_channels[0], + out_channels=block_out_channels[1], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block, + num_attention_heads=block_out_channels[1] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[1], + add_downsample=True, + prefix="down_blocks.1", + ) + ) + # Level 2: CrossAttnDownBlock2D (no downsampler) + self.down_blocks.append( + SDXLCrossAttnDownBlock2D( + in_channels=block_out_channels[1], + out_channels=block_out_channels[2], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block, + num_attention_heads=block_out_channels[2] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[2], + add_downsample=False, + prefix="down_blocks.2", + ) + ) + + # Mid block + self.mid_block = SDXLUNetMidBlock2DCrossAttn( + in_channels=block_out_channels[2], + time_embed_dim=time_embed_dim, + num_attention_heads=block_out_channels[2] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[2], + prefix="mid_block", + ) + + # Up blocks (reversed, with skip connections) + # Skip connection channels consumed from down_block_res_samples (right to left): + # down_block_res_samples channels = [320, 320, 320, 320, 640, 640, 640, 1280, 1280] + # up_blocks[0] consumes last 3: skips are [1280, 1280, 640] (consumed from right) + # up_blocks[1] consumes next 3: skips are [640, 640, 320] (consumed from right) + # up_blocks[2] consumes next 3: skips are [320, 320, 320] (consumed from right) + self.up_blocks = nn.ModuleList() + + # Level 0: CrossAttnUpBlock2D (mirrors down_blocks[2]) + # Input from mid: 1280, skips (consumed right-to-left): 1280, 1280, 640 + self.up_blocks.append( + SDXLCrossAttnUpBlock2D( + in_channels=block_out_channels[2], + prev_output_channel=block_out_channels[2], + out_channels=block_out_channels[2], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block + 1, + num_attention_heads=block_out_channels[2] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[2], + add_upsample=True, + prefix="up_blocks.0", + resnet_in_channels_list=[1280 + 1280, 1280 + 1280, 1280 + 640], + ) + ) + # Level 1: CrossAttnUpBlock2D (mirrors down_blocks[1]) + # Input from up_blocks[0]: 1280, skips (consumed right-to-left): 640, 640, 320 + self.up_blocks.append( + SDXLCrossAttnUpBlock2D( + in_channels=block_out_channels[2], + prev_output_channel=block_out_channels[1], + out_channels=block_out_channels[1], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block + 1, + num_attention_heads=block_out_channels[1] // num_head_channels, + cross_attention_dim=cross_attention_dim, + transformer_layers_per_block=transformer_layers_per_block[1], + add_upsample=True, + prefix="up_blocks.1", + resnet_in_channels_list=[1280 + 640, 640 + 640, 640 + 320], + ) + ) + # Level 2: UpBlock2D (mirrors down_blocks[0], no attention, no upsample) + # Input from up_blocks[1]: 640, skips (consumed right-to-left): 320, 320, 320 + self.up_blocks.append( + SDXLUpBlock2D( + in_channels=block_out_channels[1], + prev_output_channel=block_out_channels[0], + out_channels=block_out_channels[0], + time_embed_dim=time_embed_dim, + num_layers=layers_per_block + 1, + add_upsample=False, + resnet_in_channels_list=[640 + 320, 320 + 320, 320 + 320], + ) + ) + + # Output + self.conv_norm_out = nn.GroupNorm(32, block_out_channels[0], eps=1e-5) + self.conv_act = nn.SiLU() + self.conv_out = nn.Conv2d(block_out_channels[0], 4, kernel_size=3, padding=1) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + added_cond_kwargs: dict, + return_dict: bool = False, + ) -> tuple[torch.Tensor]: + # Time embedding + dtype = hidden_states.dtype + t_emb = get_timestep_embedding(timestep, self.time_proj_dim).to(dtype=dtype) + emb = self.time_embedding(t_emb) + + # Addition embedding (pooled text + size/crop) + text_embeds = added_cond_kwargs["text_embeds"].to(dtype=dtype) + time_ids = added_cond_kwargs["time_ids"] + add_emb = self.add_embedding(text_embeds, time_ids) + emb = emb + add_emb + + # Input conv + hidden_states = self.conv_in(hidden_states) + + # Down blocks + down_block_res_samples = [hidden_states] + for down_block in self.down_blocks: + if isinstance(down_block, SDXLCrossAttnDownBlock2D): + hidden_states, res_samples = down_block(hidden_states, emb, encoder_hidden_states) + else: + hidden_states, res_samples = down_block(hidden_states, emb) + down_block_res_samples.extend(res_samples) + + # Mid block + hidden_states = self.mid_block(hidden_states, emb, encoder_hidden_states) + + # Up blocks + for up_block in self.up_blocks: + # Determine how many skip connections this block needs + n_resnets = len(up_block.resnets) + res_samples = tuple(down_block_res_samples[-n_resnets:]) + down_block_res_samples = down_block_res_samples[:-n_resnets] + + if isinstance(up_block, SDXLCrossAttnUpBlock2D): + hidden_states = up_block(hidden_states, emb, encoder_hidden_states, res_samples) + else: + hidden_states = up_block(hidden_states, emb, res_samples) + + # Output + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states) + + return (hidden_states,) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # Self-attention QKV fusion + (".attn1.to_qkv", ".attn1.to_q", "q"), + (".attn1.to_qkv", ".attn1.to_k", "k"), + (".attn1.to_qkv", ".attn1.to_v", "v"), + ] + self.stacked_params_mapping = stacked_params_mapping + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + # Handle GEGLU: diffusers stores as ff.net.0.proj.weight/bias + # MergedColumnParallelLinear needs the two halves loaded separately + if ".ff.net.0.proj." in name: + name = name.replace(".ff.net.0.proj.", ".ff.geglu.proj.") + if name in params_dict: + param = params_dict[name] + weight_loader = param.weight_loader + half = loaded_weight.shape[0] // 2 + weight_loader(param, loaded_weight[:half], 0) + weight_loader(param, loaded_weight[half:], 1) + loaded_params.add(name) + continue + elif ".ff.net.2." in name: + name = name.replace(".ff.net.2.", ".ff.out_proj.") + + handled = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name not in params_dict: + break + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name) + handled = True + break + + if not handled: + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params diff --git a/vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py b/vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py new file mode 100644 index 00000000000..5b81b6ce44c --- /dev/null +++ b/vllm_omni/diffusion/models/sensenova_u1/fused_rmsnorm_rope.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def qk_norm_rope_kernel( + q, + k, + query, + key, + q_norm_weight, + k_norm_weight, + q_norm_hw_weight, + k_norm_hw_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + q_stride_b: tl.constexpr, + q_stride_s: tl.constexpr, + q_stride_h: tl.constexpr, + q_stride_d: tl.constexpr, + k_stride_b: tl.constexpr, + k_stride_s: tl.constexpr, + k_stride_h: tl.constexpr, + k_stride_d: tl.constexpr, + query_stride_b: tl.constexpr, + query_stride_h: tl.constexpr, + query_stride_s: tl.constexpr, + query_stride_d: tl.constexpr, + key_stride_b: tl.constexpr, + key_stride_h: tl.constexpr, + key_stride_s: tl.constexpr, + key_stride_d: tl.constexpr, + cos_t_stride_b: tl.constexpr, + cos_t_stride_s: tl.constexpr, + cos_t_stride_d: tl.constexpr, + cos_h_stride_b: tl.constexpr, + cos_h_stride_s: tl.constexpr, + cos_h_stride_d: tl.constexpr, + cos_w_stride_b: tl.constexpr, + cos_w_stride_s: tl.constexpr, + cos_w_stride_d: tl.constexpr, + head_q: tl.constexpr, + head_dim: tl.constexpr, + eps: tl.constexpr = 1e-6, +): + token_idx = tl.program_id(0) + head_pid = tl.program_id(1) + batch_idx = tl.program_id(2) + is_q = head_pid < head_q + head_idx = tl.where(is_q, head_pid, head_pid - head_q) + + offs = tl.arange(0, head_dim) + t_mask = offs < head_dim // 2 + hw_mask = ~t_mask + + in_base = tl.where( + is_q, + q + batch_idx * q_stride_b + token_idx * q_stride_s + head_idx * q_stride_h, + k + batch_idx * k_stride_b + token_idx * k_stride_s + head_idx * k_stride_h, + ) + vals = tl.load(in_base + offs * tl.where(is_q, q_stride_d, k_stride_d)).to(tl.float32) + + t_vals = tl.where(t_mask, vals, 0.0) + hw_vals = tl.where(hw_mask, vals, 0.0) + rms_t = tl.rsqrt(tl.sum(t_vals * t_vals, axis=0) / (head_dim // 2) + eps) + rms_hw = tl.rsqrt(tl.sum(hw_vals * hw_vals, axis=0) / (head_dim // 2) + eps) + rms = tl.where(t_mask, rms_t, rms_hw) + + t_weight = tl.load(tl.where(is_q, q_norm_weight, k_norm_weight) + offs, mask=t_mask, other=0.0) + hw_weight = tl.load( + tl.where(is_q, q_norm_hw_weight, k_norm_hw_weight) + (offs - head_dim // 2), + mask=hw_mask, + other=0.0, + ) + weight = tl.where(t_mask, t_weight, hw_weight) + normed = (vals * rms).to(query.dtype.element_ty) * weight + + pair_offs = tl.where( + offs < head_dim // 4, + offs + head_dim // 4, + tl.where( + offs < head_dim // 2, + offs - head_dim // 4, + tl.where( + offs < head_dim // 8 * 5, + offs + head_dim // 8, + tl.where( + offs < head_dim // 4 * 3, + offs - head_dim // 8, + tl.where(offs < head_dim // 8 * 7, offs + head_dim // 8, offs - head_dim // 8), + ), + ), + ), + ) + pair_vals = tl.load(in_base + pair_offs * tl.where(is_q, q_stride_d, k_stride_d)).to(tl.float32) + pair_rms = tl.where(pair_offs < head_dim // 2, rms_t, rms_hw) + pair_t_mask = pair_offs < head_dim // 2 + pair_hw_mask = ~pair_t_mask + pair_t_weight = tl.load( + tl.where(is_q, q_norm_weight, k_norm_weight) + pair_offs, + mask=pair_t_mask, + other=0.0, + ) + pair_hw_weight = tl.load( + tl.where(is_q, q_norm_hw_weight, k_norm_hw_weight) + (pair_offs - head_dim // 2), + mask=pair_hw_mask, + other=0.0, + ) + pair_weight = tl.where(pair_offs < head_dim // 2, pair_t_weight, pair_hw_weight) + pair_normed = (pair_vals * pair_rms).to(query.dtype.element_ty) * pair_weight + + cos_t_vals = tl.load( + cos_t + batch_idx * cos_t_stride_b + token_idx * cos_t_stride_s + offs * cos_t_stride_d, + mask=t_mask, + other=0.0, + ) + sin_t_vals = tl.load( + sin_t + batch_idx * cos_t_stride_b + token_idx * cos_t_stride_s + offs * cos_t_stride_d, + mask=t_mask, + other=0.0, + ) + hw_local = offs - head_dim // 2 + h_local = hw_local + w_local = hw_local - head_dim // 4 + h_mask = (offs >= head_dim // 2) & (offs < head_dim // 4 * 3) + w_mask = offs >= head_dim // 4 * 3 + cos_h_vals = tl.load( + cos_h + batch_idx * cos_h_stride_b + token_idx * cos_h_stride_s + h_local * cos_h_stride_d, + mask=h_mask, + other=0.0, + ) + sin_h_vals = tl.load( + sin_h + batch_idx * cos_h_stride_b + token_idx * cos_h_stride_s + h_local * cos_h_stride_d, + mask=h_mask, + other=0.0, + ) + cos_w_vals = tl.load( + cos_w + batch_idx * cos_w_stride_b + token_idx * cos_w_stride_s + w_local * cos_w_stride_d, + mask=w_mask, + other=0.0, + ) + sin_w_vals = tl.load( + sin_w + batch_idx * cos_w_stride_b + token_idx * cos_w_stride_s + w_local * cos_w_stride_d, + mask=w_mask, + other=0.0, + ) + cos_vals = cos_t_vals + cos_h_vals + cos_w_vals + sin_vals = sin_t_vals + sin_h_vals + sin_w_vals + + sign = tl.where( + (offs < head_dim // 4) + | ((offs >= head_dim // 2) & (offs < head_dim // 8 * 5)) + | ((offs >= head_dim // 4 * 3) & (offs < head_dim // 8 * 7)), + -1.0, + 1.0, + ) + out = normed * cos_vals + sign * pair_normed * sin_vals + + out_base = tl.where( + is_q, + query + batch_idx * query_stride_b + head_idx * query_stride_h + token_idx * query_stride_s, + key + batch_idx * key_stride_b + head_idx * key_stride_h + token_idx * key_stride_s, + ) + tl.store(out_base + offs * tl.where(is_q, query_stride_d, key_stride_d), out) + + +def triton_qk_norm_rope( + q, + k, + q_norm_weight, + k_norm_weight, + q_norm_hw_weight, + k_norm_hw_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + eps, +): + batch_size, seq_len, head_q, head_dim = q.shape + head_k = k.shape[2] + + query = torch.empty((batch_size, head_q, seq_len, head_dim), device=q.device, dtype=q.dtype) + key = torch.empty((batch_size, head_k, seq_len, head_dim), device=k.device, dtype=k.dtype) + grid = ( + seq_len, + head_q + head_k, + batch_size, + ) + + qk_norm_rope_kernel[grid]( + q, + k, + query, + key, + q_norm_weight, + k_norm_weight, + q_norm_hw_weight, + k_norm_hw_weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + query.stride(0), + query.stride(1), + query.stride(2), + query.stride(3), + key.stride(0), + key.stride(1), + key.stride(2), + key.stride(3), + cos_t.stride(0), + cos_t.stride(1), + cos_t.stride(2), + cos_h.stride(0), + cos_h.stride(1), + cos_h.stride(2), + cos_w.stride(0), + cos_w.stride(1), + cos_w.stride(2), + head_q=head_q, + head_dim=head_dim, + eps=eps, + ) + return query, key 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 fa7f95c66e2..67e141e6ae6 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py +++ b/vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py @@ -22,7 +22,7 @@ import os from collections.abc import Iterable from types import SimpleNamespace -from typing import ClassVar +from typing import Any, ClassVar import numpy as np import torch @@ -33,11 +33,13 @@ from vllm.logger import init_logger from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin 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 SupportsComponentDiscovery from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from .sensenova_u1_transformer import ( SenseNovaU1ForCausalLM, @@ -488,7 +490,31 @@ def _optimized_scale(positive_flat, negative_flat): # --------------------------------------------------------------------------- -class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): +class SenseNovaU1DenoisingAdapter(nn.Module): + """Denoising-only entry point used by cache backends.""" + + def __init__(self, language_model: SenseNovaU1ForCausalLM): + super().__init__() + object.__setattr__(self, "language_model", language_model) + self.do_true_cfg = True + + @property + def model(self): + return self.language_model.model + + @property + def lm_head(self): + return self.language_model.lm_head + + @property + def logits_processor(self): + return self.language_model.logits_processor + + def forward(self, *args, **kwargs): + return self.language_model(*args, **kwargs) + + +class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin): """SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni. Builds the full model graph internally: @@ -504,30 +530,6 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli support_image_input = True - # Model-specific parameters accepted via ``extra_body`` in the online - # serving API. The serving layer reads this set and forwards matching - # keys from the request payload to ``OmniDiffusionSamplingParams.extra_args``. - EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "think", - "cfg_scale", - "cfg_norm", - "timestep_shift", - "t_eps", - "img_cfg_scale", - "max_tokens", - } - ) - - # Keys from ``DiffusionOutput.custom_output`` that should be surfaced in - # the API response ``metrics`` dict. The serving layer reads this set and - # copies matching entries from ``custom_output`` into the response. - EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset( - { - "think_text", - } - ) - # CPU-offload protocol: language_model carries the denoising blocks; the # vision and FM modules are lightweight encoders pinned on GPU during the # diffusion loop. There is no separate VAE. @@ -558,6 +560,12 @@ def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): self.llm_cfg, prefix="language_model", ) + # Cache-DiT hooks pipeline.transformer(.blocks), so it must point at the + # real decoder module (exposes .blocks and real parameters). + self.transformer = self.language_model.model + # TeaCache intercepts the ForCausalLM-level denoising forward; route it + # through a dedicated adapter so it does not collide with Cache-DiT. + self.denoising_transformer = SenseNovaU1DenoisingAdapter(self.language_model) # Vision model (understanding branch) self.vision_model = NEOVisionModel(self.vis_cfg) @@ -724,7 +732,8 @@ def _t2i_predict_v( **_kw, ): B, L = z.shape[0], z.shape[1] - outputs = self.language_model( + denoising_model = self.language_model if cache_dit_skip else self.denoising_transformer + outputs = denoising_model( inputs_embeds=input_embeds, image_gen_indicators=torch.ones( (input_embeds.shape[0], input_embeds.shape[1]), dtype=torch.bool, device=input_embeds.device @@ -929,8 +938,10 @@ def _forward_text(self, p, input_images) -> DiffusionOutput: ) logger.info("Text generation: %d chars", len(text_output)) return DiffusionOutput( - output=text_output, - custom_output={"text_output": text_output}, + output={ + "payload": {"text": text_output}, + "metadata": {"text": {"text_output": text_output}}, + }, ) # ----------------------------------------------------------------------- @@ -1076,123 +1087,136 @@ def _init_noise_and_schedule(self, p): timesteps=timesteps, ) - def _denoise_step(self, image_prediction, ns, t, z, image_embeds, caches, p, step_i): - """Execute one denoising step with the appropriate CFG strategy.""" - kv_cond = caches["cond"] - idx_cond = caches["idx_cond"] - mask_cond = caches["mask_cond"] - is_it2i = "img_cond" in caches - if is_it2i: - use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0 - needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1) - has_cached_partner = use_cfg and needs_cfg - else: - has_cached_partner = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1 - - out_cond = self._t2i_predict_v( - image_embeds, - idx_cond, - mask_cond, - kv_cond, - t, - z, + def _get_cfg_kwargs(self, caches: dict, image_embeds, t, z, ns, p, branch: str, cache_dit_skip: bool = False): + required = (branch, f"idx_{branch}", f"mask_{branch}") + missing = [key for key in required if key not in caches] + if missing: + raise KeyError(f"Missing {branch} CFG cache keys: {missing}") + kwargs = dict( + input_embeds=image_embeds, + past_key_values=caches[branch], + indexes_image=caches[f"idx_{branch}"], + attn_mask=caches[f"mask_{branch}"], + t=t, + z=z, image_token_num=ns.token_h * ns.token_w, image_size=p.image_size, - cache_dit_skip=not has_cached_partner, ) + if cache_dit_skip: + kwargs["cache_dit_skip"] = True + return kwargs - if is_it2i: - return self._denoise_step_it2i(out_cond, image_embeds, ns, t, z, caches, p, step_i) - return self._denoise_step_t2i(out_cond, image_embeds, ns, t, z, caches, p, step_i) - - def _denoise_step_t2i(self, out_cond, image_embeds, ns, t, z, caches, p, step_i): - """T2I: two-way CFG (condition vs uncondition), supports cfg_zero_star.""" - if t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1: - out_uncond = self._t2i_predict_v( - image_embeds, - caches["idx_uncond"], - caches["mask_uncond"], - caches["uncond"], - t, - z, - image_token_num=ns.token_h * ns.token_w, - image_size=p.image_size, + def _denoise(self, image_prediction, ns, t, z, image_embeds, caches, p, step_i, is_it2i): + if not is_it2i: + has_cached_partner = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] and p.cfg_scale > 1 + cond_kwargs = self._get_cfg_kwargs( + caches, image_embeds, t, z, ns, p, branch="cond", cache_dit_skip=not has_cached_partner ) - if p.cfg_norm == "cfg_zero_star": - pos_flat = out_cond.view(p.batch_size, -1) - neg_flat = out_uncond.view(p.batch_size, -1) - alpha = _optimized_scale(pos_flat, neg_flat) - alpha = alpha.view(p.batch_size, *([1] * (len(out_cond.shape) - 1))).to(pos_flat.dtype) - v_pred = ( - out_cond * 0.0 - if step_i <= 0 - else (out_uncond * alpha + p.cfg_scale * (out_cond - out_uncond * alpha)) - ) - else: - v_pred = out_uncond + p.cfg_scale * (out_cond - out_uncond) - v_pred = self._apply_cfg_norm(v_pred, out_cond, p.cfg_norm) - else: - v_pred = out_cond - return v_pred - - def _denoise_step_it2i(self, out_cond, image_embeds, ns, t, z, caches, p, step_i): - """IT2I: dual CFG with ``cfg_scale`` (text) and ``img_cfg_scale`` (image).""" - use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0 - needs_cfg = not (p.cfg_scale == 1 and p.img_cfg_scale == 1) - - if not use_cfg or not needs_cfg: - return out_cond - predict_kw = dict(image_token_num=ns.token_h * ns.token_w, image_size=p.image_size) - - if p.img_cfg_scale == 1: - out_img_cond = self._t2i_predict_v( - image_embeds, - caches["idx_img_cond"], - caches["mask_img_cond"], - caches["img_cond"], - t, - z, - **predict_kw, - ) - v_pred = out_img_cond + p.cfg_scale * (out_cond - out_img_cond) - elif p.cfg_scale == p.img_cfg_scale: - out_uncond = self._t2i_predict_v( - image_embeds, - caches["idx_uncond"], - caches["mask_uncond"], - caches["uncond"], - t, - z, - **predict_kw, + in_interval = t >= p.cfg_interval[0] and t <= p.cfg_interval[1] + if not (in_interval and p.cfg_scale > 1): + return self.predict_noise(**cond_kwargs) + + uncond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="uncond") + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=p.cfg_scale, + positive_kwargs=cond_kwargs, + negative_kwargs=uncond_kwargs, + cfg_normalize=p.cfg_norm, + kwargs={"step_i": step_i, "is_it2i": is_it2i}, ) - v_pred = out_uncond + p.cfg_scale * (out_cond - out_uncond) + return noise_pred else: - out_img_cond = self._t2i_predict_v( - image_embeds, - caches["idx_img_cond"], - caches["mask_img_cond"], - caches["img_cond"], - t, - z, - cache_dit_skip=True, - **predict_kw, - ) - out_uncond = self._t2i_predict_v( - image_embeds, - caches["idx_uncond"], - caches["mask_uncond"], - caches["uncond"], - t, - z, - **predict_kw, + use_cfg = (t > p.cfg_interval[0] and t < p.cfg_interval[1]) or p.cfg_interval[0] == 0 + needs_cfg = p.cfg_scale != 1 or p.img_cfg_scale != 1 + has_cached_partner = use_cfg and needs_cfg + cond_kwargs = self._get_cfg_kwargs( + caches, image_embeds, t, z, ns, p, branch="cond", cache_dit_skip=not has_cached_partner ) + + if not use_cfg or not needs_cfg: + return self.predict_noise(**cond_kwargs) + + cfg_norm = p.cfg_norm if (p.cfg_scale > 1 or p.img_cfg_scale > 1) else None + if p.img_cfg_scale == 1: + image_cond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="img_cond") + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=p.cfg_scale, + positive_kwargs=cond_kwargs, + negative_kwargs=image_cond_kwargs, + cfg_normalize=cfg_norm, + kwargs={"is_it2i": is_it2i}, + ) + elif p.cfg_scale == p.img_cfg_scale: + uncond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="uncond") + noise_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=True, + true_cfg_scale=p.cfg_scale, + positive_kwargs=cond_kwargs, + negative_kwargs=uncond_kwargs, + cfg_normalize=cfg_norm, + kwargs={"is_it2i": is_it2i}, + ) + else: + image_cond_kwargs = self._get_cfg_kwargs( + caches, image_embeds, t, z, ns, p, branch="img_cond", cache_dit_skip=True + ) + uncond_kwargs = self._get_cfg_kwargs(caches, image_embeds, t, z, ns, p, branch="uncond") + noise_pred = self.predict_noise_with_multi_branch_cfg( + do_true_cfg=True, + true_cfg_scale={ + "cfg_scale": p.cfg_scale, + "img_cfg_scale": p.img_cfg_scale, + }, + branches_kwargs=[cond_kwargs, image_cond_kwargs, uncond_kwargs], + cfg_normalize=cfg_norm, + ) + return noise_pred + + def predict_noise(self, **kwargs): + return self._t2i_predict_v(**kwargs) + + def combine_cfg_noise(self, out_cond, out_uncond, cfg_scale, cfg_norm, kwargs: dict[str, Any] | None = None): + if not isinstance(kwargs, dict) or "is_it2i" not in kwargs: + raise ValueError("SenseNovaU1 CFG requires kwargs['is_it2i']") + if isinstance(cfg_scale, dict): + raise NotImplementedError("dict cfg_scale requires multi-branch CFG") + + out_cond = out_cond[0] + out_uncond = out_uncond[0] + batch_size = out_cond.shape[0] + + is_it2i = kwargs.get("is_it2i") + if not is_it2i and cfg_norm == "cfg_zero_star": + if not isinstance(kwargs, dict) or not isinstance(kwargs.get("step_i"), int): + raise ValueError("SenseNovaU1 T2I CFG requires kwargs['step_i']") + step_i = kwargs.get("step_i") + pos_flat = out_cond.view(batch_size, -1) + neg_flat = out_uncond.view(batch_size, -1) + alpha = _optimized_scale(pos_flat, neg_flat) + alpha = alpha.view(batch_size, *([1] * (len(out_cond.shape) - 1))).to(pos_flat.dtype) v_pred = ( - out_uncond + p.cfg_scale * (out_cond - out_img_cond) + p.img_cfg_scale * (out_img_cond - out_uncond) + out_cond * 0.0 if step_i <= 0 else (out_uncond * alpha + cfg_scale * (out_cond - out_uncond * alpha)) ) + else: + if cfg_norm == "cfg_zero_star": + logger.warning("cfg_zero_star is T2I-only; please use 'global' or 'channel' instead.") + # cfg_zero_star is T2I-only; _apply_cfg_norm treats it as a no-op for IT2I. + v_pred = out_uncond + cfg_scale * (out_cond - out_uncond) + v_pred = self._apply_cfg_norm(v_pred, out_cond, cfg_norm) + return v_pred - if p.cfg_scale > 1 or p.img_cfg_scale > 1: - v_pred = self._apply_cfg_norm(v_pred, out_cond, p.cfg_norm) + def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize): + out_cond = predictions[0] + out_img_cond = predictions[1] + out_uncond = predictions[2] + cfg_scale = true_cfg_scale["cfg_scale"] + img_cfg_scale = true_cfg_scale["img_cfg_scale"] + v_pred = out_uncond + cfg_scale * (out_cond - out_img_cond) + img_cfg_scale * (out_img_cond - out_uncond) + if cfg_scale > 1 or img_cfg_scale > 1: + v_pred = self._apply_cfg_norm(v_pred, out_cond, cfg_normalize) return v_pred @staticmethod @@ -1215,7 +1239,7 @@ def _expand_and_prepare_kv(self, kv, token_hw, batch_size): prepare_flash_kv_cache(kv, current_len=token_hw, batch_size=batch_size) @torch.inference_mode() - def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: p = self._parse_request(req) self.top_cfg.t_eps = p.t_eps @@ -1273,7 +1297,7 @@ def _forward_t2i(self, p) -> DiffusionOutput: "idx_uncond": indexes_image_uncond, "mask_uncond": {"full_attention": None}, } - return self._run_denoising_loop(ns, caches, p, think_text) + return self._run_denoising_loop(ns, caches, p, think_text, is_it2i=False) def _forward_it2i(self, p, input_images: list[Image.Image]) -> DiffusionOutput: """Image-to-image (editing) generation path with dual CFG.""" @@ -1389,9 +1413,9 @@ def _forward_it2i(self, p, input_images: list[Image.Image]) -> DiffusionOutput: if key in caches and not isinstance(caches[key], dict): self._expand_and_prepare_kv(caches[key], ns.token_h * ns.token_w, p.batch_size) - return self._run_denoising_loop(ns, caches, p, think_text) + return self._run_denoising_loop(ns, caches, p, think_text, is_it2i=True) - def _run_denoising_loop(self, ns, caches, p, think_text="") -> DiffusionOutput: + def _run_denoising_loop(self, ns, caches, p, think_text="", is_it2i: bool = False) -> DiffusionOutput: """Shared denoising loop for both T2I and IT2I.""" merge_size = self.merge_size image_prediction = ns.image_prediction @@ -1424,7 +1448,7 @@ def _run_denoising_loop(self, ns, caches, p, think_text="") -> DiffusionOutput: timestep_embeddings = timestep_embeddings + ns_emb image_embeds = image_embeds + timestep_embeddings - v_pred = self._denoise_step(image_prediction, ns, t, z, image_embeds, caches, p, step_i) + v_pred = self._denoise(image_prediction, ns, t, z, image_embeds, caches, p, step_i, is_it2i) z = z + (t_next - t) * v_pred image_prediction = _unpatchify(z, self.patch_size * merge_size, p.image_size[1], p.image_size[0]) @@ -1435,10 +1459,15 @@ def _run_denoising_loop(self, ns, caches, p, think_text="") -> DiffusionOutput: images = _to_pil(image_prediction) img = images[0] if images else None - custom = {} + metadata = {} if think_text: - custom["think_text"] = think_text - return DiffusionOutput(output=img, custom_output=custom) + metadata["text"] = {"think_text": think_text} + return DiffusionOutput( + output={ + "payload": {"image": img}, + "metadata": metadata, + } + ) # ----------------------------------------------------------------------- # Weight loading diff --git a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py index f9ff8aaf470..843358a70b6 100644 --- a/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py +++ b/vllm_omni/diffusion/models/sensenova_u1/sensenova_u1_transformer.py @@ -16,7 +16,9 @@ import torch import torch.nn as nn +from cache_dit import ForwardPattern from transformers.cache_utils import DynamicCache +from vllm.logger import init_logger from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, @@ -30,6 +32,9 @@ from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig, SensenovaCachedAdapter + +logger = init_logger(__name__) @dataclass @@ -371,6 +376,34 @@ def _project_and_rope(self, hidden_states, indexes, qkv_proj, q_norm, k_norm, q_ k = k.view(*input_shape, self.num_kv_heads, self.head_dim) v = v.view(*input_shape, self.num_kv_heads, self.head_dim) + cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)) + cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)) + cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)) + + value_states = v.transpose(1, 2) # [B, H, S, D] + try: + from .fused_rmsnorm_rope import triton_qk_norm_rope + except ImportError: + triton_qk_norm_rope = None + if triton_qk_norm_rope is not None: + query_states, key_states = triton_qk_norm_rope( + q, + k, + q_norm.weight, + k_norm.weight, + q_norm_hw.weight, + k_norm_hw.weight, + cos_t, + sin_t, + cos_h, + sin_h, + cos_w, + sin_w, + self.config.rms_norm_eps, + ) + return query_states, key_states, value_states + logger.debug("Triton fused SenseNova qk norm rope unavailable; using PyTorch fallback") + # Split head_dim into t and hw halves q_t, q_hw = q.chunk(2, dim=-1) k_t, k_hw = k.chunk(2, dim=-1) @@ -386,17 +419,13 @@ def _project_and_rope(self, hidden_states, indexes, qkv_proj, q_norm, k_norm, q_ k_h, k_w = k_hw.chunk(2, dim=-1) # RoPE for each dimension - cos_t, sin_t = self.rotary_emb(hidden_states, indexes[0].unsqueeze(0)) q_t, k_t = apply_rotary_pos_emb(q_t, k_t, cos_t, sin_t) - cos_h, sin_h = self.rotary_emb_hw(hidden_states, indexes[1].unsqueeze(0)) q_h, k_h = apply_rotary_pos_emb(q_h, k_h, cos_h, sin_h) - cos_w, sin_w = self.rotary_emb_hw(hidden_states, indexes[2].unsqueeze(0)) q_w, k_w = apply_rotary_pos_emb(q_w, k_w, cos_w, sin_w) # Reassemble: [B, H, S, head_dim] query_states = torch.cat([q_t, q_h, q_w], dim=-1) key_states = torch.cat([k_t, k_h, k_w], dim=-1) - value_states = v.transpose(1, 2) # [B, H, S, D] return query_states, key_states, value_states def forward_und(self, hidden_states, indexes, attention_mask, past_key_values=None, **kwargs): @@ -606,6 +635,14 @@ def forward( class SenseNovaU1Model(nn.Module): + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_3, + }, + has_separate_cfg=True, + cached_adapter_cls=SensenovaCachedAdapter, + ) + def __init__(self, config, quant_config=None, prefix: str = ""): super().__init__() self.config = config diff --git a/vllm_omni/diffusion/models/soulx_singer/__init__.py b/vllm_omni/diffusion/models/soulx_singer/__init__.py new file mode 100644 index 00000000000..30235d2c174 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/__init__.py @@ -0,0 +1,44 @@ +"""SoulX-Singer SVS / SVC flow-matching singing models for vLLM-Omni.""" + +__all__ = [ + "FlowMatchingAudioPipeline", + "PipelineSoulXSingerSVS", + "PipelineSoulXSingerSVC", + "get_soulxsinger_post_process_func", + "get_soulxsinger_pre_process_func", + "get_soulxsinger_svc_pre_process_func", +] + + +def __getattr__(name: str): + if name in {"FlowMatchingAudioPipeline", "get_soulxsinger_post_process_func"}: + from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( + FlowMatchingAudioPipeline, + get_soulxsinger_post_process_func, + ) + + return { + "FlowMatchingAudioPipeline": FlowMatchingAudioPipeline, + "get_soulxsinger_post_process_func": get_soulxsinger_post_process_func, + }[name] + if name in {"PipelineSoulXSingerSVC", "get_soulxsinger_svc_pre_process_func"}: + from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svc import ( + PipelineSoulXSingerSVC, + get_soulxsinger_svc_pre_process_func, + ) + + return { + "PipelineSoulXSingerSVC": PipelineSoulXSingerSVC, + "get_soulxsinger_svc_pre_process_func": get_soulxsinger_svc_pre_process_func, + }[name] + if name in {"PipelineSoulXSingerSVS", "get_soulxsinger_pre_process_func"}: + from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_svs import ( + PipelineSoulXSingerSVS, + get_soulxsinger_pre_process_func, + ) + + return { + "PipelineSoulXSingerSVS": PipelineSoulXSingerSVS, + "get_soulxsinger_pre_process_func": get_soulxsinger_pre_process_func, + }[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py b/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py new file mode 100644 index 00000000000..282ab12abdc --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/__init__.py @@ -0,0 +1,17 @@ +from .convnext import ConvNeXtV2Block +from .decoder import CFMDecoder +from .flow_matching import FlowMatchingTransformer +from .llama import DiffLlama +from .mel_transform import MelSpectrogramEncoder +from .vocoder import Vocoder +from .whisper_encoder import WhisperEncoder + +__all__ = [ + "ConvNeXtV2Block", + "FlowMatchingTransformer", + "CFMDecoder", + "DiffLlama", + "MelSpectrogramEncoder", + "Vocoder", + "WhisperEncoder", +] diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py b/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py new file mode 100644 index 00000000000..465898e8a27 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/convnext.py @@ -0,0 +1,46 @@ +import torch +import torch.nn as nn + + +class GRN(nn.Module): + def __init__(self, dim): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, 1, dim)) + self.beta = nn.Parameter(torch.zeros(1, 1, dim)) + + def forward(self, x): + Gx = torch.norm(x, p=2, dim=1, keepdim=True) + Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6) + return self.gamma * (x * Nx) + self.beta + x + + +# ref: https://github.com/SWivid/F5-TTS/blob/main/src/f5_tts/model/modules.py#L247 +class ConvNeXtV2Block(nn.Module): + def __init__( + self, + dim: int, + intermediate_dim: int, + dilation: int = 1, + ): + super().__init__() + padding = (dilation * (7 - 1)) // 2 + self.dwconv = nn.Conv1d( + dim, dim, kernel_size=7, padding=padding, groups=dim, dilation=dilation + ) # depthwise conv + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.grn = GRN(intermediate_dim) + self.pwconv2 = nn.Linear(intermediate_dim, dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = x.transpose(1, 2) # b n d -> b d n + x = self.dwconv(x) + x = x.transpose(1, 2) # b d n -> b n d + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + return residual + x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py new file mode 100644 index 00000000000..6b1fda616a9 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/decoder.py @@ -0,0 +1,24 @@ +import torch +import torch.nn as nn + +from vllm_omni.diffusion.models.soulx_singer.modules.flow_matching import FlowMatchingTransformer + + +class CFMDecoder(nn.Module): + def __init__(self, config): + super().__init__() + self.model = FlowMatchingTransformer(cfg=config, **config) + + def forward(self, mel, x_mask, decoder_inp, is_prompt): + outputs = self.model(mel, x_mask, decoder_inp, is_prompt) + + noise, x, flow_pred, final_mask, prompt_len = outputs["output"] + return noise, x, flow_pred, final_mask, prompt_len + + def reverse_diffusion(self, pt_mel, pt_decoder_inp, gt_decoder_inp, n_timesteps=32, cfg=1): + diffusion_cond = torch.cat([pt_decoder_inp, gt_decoder_inp], dim=1) + diffusion_cond_emb = self.model.cond_emb(diffusion_cond) + diffusion_prompt = pt_mel + + generated = self.model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps, cfg=cfg) + return generated diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py b/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py new file mode 100644 index 00000000000..19e87b65b64 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/flow_matching.py @@ -0,0 +1,423 @@ +# https://github.com/open-mmlab/Amphion/blob/main/models/svc/flow_matching_transformer/fmt_model.py + +# Copyright (c) 2023 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import LlamaConfig + +from vllm_omni.diffusion.models.soulx_singer.modules.llama import DiffLlama + + +class FlowMatchingTransformer(nn.Module): + def __init__( + self, + mel_dim=100, + hidden_size=1024, + num_layers=12, + num_heads=16, + cfg_drop_prob=0.2, + use_embedding=True, + cond_codebook_size=1024, + cond_scale_factor=1, + sigma=1e-5, + time_scheduler="linear", + cfg=None, + ): + super().__init__() + self.cfg = cfg + + if cfg is not None: + mel_dim = getattr(cfg, "mel_dim", mel_dim) + hidden_size = getattr(cfg, "hidden_size", hidden_size) + num_layers = getattr(cfg, "num_layers", num_layers) + num_heads = getattr(cfg, "num_heads", num_heads) + cfg_drop_prob = getattr(cfg, "cfg_drop_prob", cfg_drop_prob) + cond_codebook_size = getattr(cfg, "cond_codebook_size", cond_codebook_size) + time_scheduler = getattr(cfg, "time_scheduler", time_scheduler) + sigma = getattr(cfg, "sigma", sigma) + cond_scale_factor = getattr(cfg, "cond_scale_factor", cond_scale_factor) + + self.mel_dim = mel_dim + self.hidden_size = hidden_size + self.num_layers = num_layers + self.num_heads = num_heads + self.cfg_drop_prob = cfg_drop_prob + self.cond_codebook_size = cond_codebook_size + self.time_scheduler = time_scheduler + self.sigma = sigma + self.cond_scale_factor = cond_scale_factor + + if use_embedding: + self.cond_emb = nn.Embedding(cond_codebook_size, self.hidden_size) + else: + self.cond_emb = nn.Linear(cond_codebook_size, self.hidden_size) + + if cond_scale_factor != 1: + self.do_resampling = True + assert np.log2(cond_scale_factor).is_integer() + + up_layers = [] + for _ in range(int(np.log2(cond_scale_factor))): + up_layers.extend( + [ + nn.ConvTranspose1d(hidden_size, hidden_size, kernel_size=4, stride=2, padding=1), + nn.GELU(), + ] + ) + self.resampling_layers = nn.Sequential(*up_layers) + else: + self.do_resampling = False + + ### REPA: Use the Wav2Vec2Bert features to align. ### + self.use_repa = "repa" in cfg + self.repa_layer_index = None + if self.use_repa: + self.repa_layer_index = cfg.repa.layer_index + + self.repa_mlp_layer = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.SiLU(), + nn.Linear(hidden_size * 4, cfg.repa.output_dim), + ) + + ### CTC: Use the ASR loss ### + self.use_ctc = "ctc" in cfg + self.ctc_layer_index = None + if self.use_ctc: + self.ctc_layer_index = cfg.ctc.layer_index + + self.ctc_mlp_layer = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.SiLU(), + nn.Linear(hidden_size * 4, cfg.ctc.output_dim), + ) + + self.reset_parameters() + + # DiffLlama config must be overridden manually to make sure the + # internal RoPE head dimension is correct + llama_config = LlamaConfig( + vocab_size=0, + bos_token_id=None, + eos_token_id=None, + hidden_size=hidden_size, + num_attention_heads=num_heads, + num_hidden_layers=num_layers, + rms_norm_eps=1e-6, + output_hidden_states=False, + output_attentions=False, + use_cache=False, + ) + self.diff_estimator = DiffLlama( + mel_dim=mel_dim, hidden_size=hidden_size, num_heads=num_heads, num_layers=num_layers, config=llama_config + ) + + self.sigma = sigma + + @torch.no_grad() + def forward_diffusion(self, x, t, is_prompt=None): + """ + x: (B, T, mel_dim) + t: (B,) + """ + new_t = t + t = t.unsqueeze(-1).unsqueeze(-1) + z = torch.randn(x.shape, dtype=x.dtype, device=x.device, requires_grad=False) # (B, T, mel_dim) + + # get prompt len + if torch.rand(1) <= self.cfg_drop_prob: + prompt_len = torch.zeros(x.shape[0]).to(x) + is_prompt = torch.zeros_like(x[:, :, 0]) + else: + if is_prompt is None: + prompt_len = torch.randint(min(x.shape[1] // 4, 5), int(x.shape[1] * 0.4), (x.shape[0],)).to( + x.device + ) # (B,) + + # get is_prompt + is_prompt = torch.zeros_like(x[:, :, 0]) # (B, T) + col_indices = torch.arange(is_prompt.shape[1]).repeat(is_prompt.shape[0], 1).to(prompt_len) # (B, T) + is_prompt[col_indices < prompt_len.unsqueeze(1)] = 1 # (B, T) 1 if prompt + else: + prompt_len = is_prompt.sum(dim=1) # (B,) + + mask = torch.ones_like(x[:, :, 0]) # mask if 1, not mask if 0 + mask[is_prompt.bool()] = 0 + mask = mask[:, :, None] + + # flow matching: xt = (1 - (1 - sigma) * t) * x0 + t * x; where x0 ~ N(0, 1), x is a sample + # flow gt: x - (1 - sigma) * x0 = x - (1 - sigma) * noise + xt = ((1 - (1 - self.sigma) * t) * z + t * x) * mask + x * (1 - mask) + + return xt, z, new_t, prompt_len, mask + + def loss_t(self, x, x_mask, t, cond=None, is_prompt=None): + xt, z, new_t, prompt_len, mask = self.forward_diffusion(x, t, is_prompt) + + noise = z + + # drop all condition for cfg, so if prompt_len is 0, we also drop cond + if cond is not None: + cond = cond * torch.where( + prompt_len > 0, + torch.ones_like(prompt_len), + torch.zeros_like(prompt_len), + ).to(cond.device).unsqueeze(-1).unsqueeze(-1) + + dit_output = self.diff_estimator(xt, new_t, cond, x_mask, return_dict=True) + flow_pred = dit_output["output"] # (B, T, mel_dim) + + # final mask used for loss calculation + final_mask = mask * x_mask[..., None] # (B, T, 1) + + results = {"output": (noise, x, flow_pred, final_mask, prompt_len)} + + if self.use_repa: + repa_hidden_states = dit_output["hidden_states"][self.repa_layer_index] # (B, T, hidden_size) + + repa_pred = self.repa_mlp_layer(repa_hidden_states) # (B, T, repa_dim) + results["repa"] = repa_pred + + if self.use_ctc: + ctc_hidden_states = dit_output["hidden_states"][self.ctc_layer_index] # (B, T, hidden_size) + ctc_pred = self.ctc_mlp_layer(ctc_hidden_states) # (B, T, ctc_dim) + results["ctc"] = ctc_pred + + return results + + def compute_loss(self, x, x_mask, cond=None, is_prompt=None): + # x0: (B, T, num_quantizer) + # x_mask: (B, T) mask is 0 for padding + t = torch.rand(x.shape[0], device=x.device, requires_grad=False) + t = torch.clamp(t, 1e-5, 1.0) + # generation is harder at the start than later; use a cosine + # scheduler for timestep t when time_scheduler == "cos". + if self.time_scheduler == "cos": + t = 1 - torch.cos(t * math.pi * 0.5) + else: + pass + return self.loss_t(x, x_mask, t, cond, is_prompt) + + def reset_parameters(self): + def _reset_parameters(m): + if isinstance(m, nn.MultiheadAttention): + if m._qkv_same_embed_dim: + nn.init.normal_(m.in_proj_weight, std=0.02) + else: + nn.init.normal_(m.q_proj_weight, std=0.02) + nn.init.normal_(m.k_proj_weight, std=0.02) + nn.init.normal_(m.v_proj_weight, std=0.02) + + if m.in_proj_bias is not None: + nn.init.constant_(m.in_proj_bias, 0.0) + nn.init.constant_(m.out_proj.bias, 0.0) + if m.bias_k is not None: + nn.init.xavier_normal_(m.bias_k) + if m.bias_v is not None: + nn.init.xavier_normal_(m.bias_v) + + elif ( + isinstance(m, nn.Conv1d) + or isinstance(m, nn.ConvTranspose1d) + or isinstance(m, nn.Conv2d) + or isinstance(m, nn.ConvTranspose2d) + ): + m.weight.data.normal_(0.0, 0.02) + + elif isinstance(m, nn.Linear): + m.weight.data.normal_(mean=0.0, std=0.02) + if m.bias is not None: + m.bias.data.zero_() + + elif isinstance(m, nn.Embedding): + m.weight.data.normal_(mean=0.0, std=0.02) + if m.padding_idx is not None: + m.weight.data[m.padding_idx].zero_() + + self.apply(_reset_parameters) + + @torch.no_grad() + def reverse_diffusion( + self, + cond, + prompt, + x_mask=None, + prompt_mask=None, + n_timesteps=10, + cfg=1.0, + rescale_cfg=0.75, + ): + h = 1.0 / n_timesteps + prompt_len = prompt.shape[1] + target_len = cond.shape[1] - prompt_len + trunk_dtype = cond.dtype + + if x_mask is None: + x_mask = torch.ones(cond.shape[0], target_len, dtype=trunk_dtype, device=cond.device) + if prompt_mask is None: + prompt_mask = torch.ones(cond.shape[0], prompt_len, dtype=trunk_dtype, device=cond.device) + xt_mask = torch.cat([prompt_mask, x_mask], dim=1) + z = torch.randn( + (cond.shape[0], target_len, self.mel_dim), + dtype=trunk_dtype, + device=cond.device, + requires_grad=False, + ) + xt = z + h_tensor = torch.tensor(h, dtype=trunk_dtype, device=cond.device) + + # t from 0 to 1: x0 = z ~ N(0, 1) + for i in range(n_timesteps): + xt_input = torch.cat([prompt, xt], dim=1) + t = torch.full( + (z.shape[0],), + fill_value=(i + 0.5) * h, + dtype=trunk_dtype, + device=z.device, + ) + flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) + flow_pred = flow_pred[:, prompt_len:, :] + + # cfg + if cfg > 0: + uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask) + pos_flow_pred_std = flow_pred.std() + flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) + rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() + flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg + + xt = xt + flow_pred * h_tensor + + return xt + + @torch.no_grad() + def reverse_diffusion_v2( + self, + cond, + prompt, + x_mask=None, + prompt_mask=None, + n_timesteps=10, + cfg=1.0, + rescale_cfg=0.75, + ): + h = 1.0 / n_timesteps + prompt_len = prompt.shape[1] + target_len = cond.shape[1] - prompt_len * 2 + + if x_mask is None: + x_mask = torch.ones(cond.shape[0], target_len).to(cond.device) # (B, T) + if prompt_mask is None: + prompt_mask = torch.ones(cond.shape[0], prompt_len).to(cond.device) # (B, prompt_len) + xt_mask = torch.cat([prompt_mask, x_mask, prompt_mask], dim=1) + z = torch.randn( + (cond.shape[0], target_len, self.mel_dim), + dtype=cond.dtype, + device=cond.device, + requires_grad=False, + ) + xt = z + + # t from 0 to 1: x0 = z ~ N(0, 1) + for i in range(n_timesteps): + xt_input = torch.cat([prompt, xt, prompt], dim=1) + t = (0 + (i + 0.5) * h) * torch.ones(z.shape[0], dtype=z.dtype, device=z.device) + flow_pred = self.diff_estimator(xt_input, t, cond, xt_mask) + flow_pred = flow_pred[:, prompt_len:-prompt_len, :] + + # cfg + if cfg > 0: + uncond_flow_pred = self.diff_estimator(xt, t, torch.zeros_like(cond)[:, : xt.shape[1], :], x_mask) + pos_flow_pred_std = flow_pred.std() + flow_pred_cfg = flow_pred + cfg * (flow_pred - uncond_flow_pred) + rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() + flow_pred = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg + + dxt = flow_pred * h + xt = xt + dxt + + return xt + + def forward( + self, + x: torch.Tensor, + x_mask: torch.Tensor, + cond_code: torch.Tensor, + is_prompt: torch.Tensor | None = None, + ): + """ + Args: + x: (B, T, mel_dim) + x_mask: (B, T) + cond_code: (B, T), Note that cond_code might be not at 50Hz! + """ + T = x.shape[1] + + cond = self.cond_emb(cond_code) # (B, T, hidden_size) + if self.do_resampling: + # Align to the frame rate of Mels + cond = self.resampling_layers(cond.transpose(1, 2)).transpose(1, 2) + + # print("cond_code: {}, after resampling: {}".format(cond_code.shape, cond.shape)) + + if cond.shape[1] >= T: # Check time dimension + cond = cond[:, :T, :] + else: + padding_frames = T - cond.shape[1] + last_frame = cond[:, -1:, :] + padding = last_frame.repeat(1, padding_frames, 1) + cond = torch.cat([cond, padding], dim=1) + + return self.compute_loss(x, x_mask, cond, is_prompt) + + +if __name__ == "__main__": + model_cfg = { + "mel_dim": 128, + "hidden_size": 256, + "num_layers": 8, + "num_heads": 8, + "cfg_drop_prob": 0.2, + "use_embedding": False, + "cond_codebook_size": 256, + "cond_scale_factor": 1, + "sigma": 1e-5, + "time_scheduler": "cos", + } + + device = "cuda" + x = torch.randn(2, 100, 128).to(device) + x_mask = torch.ones(2, 100).to(device) + # cond_code = torch.randint(0, 16384, (2, 25)).to(device) + cond_code = torch.randn(2, 100, 256).to(device) + + model = FlowMatchingTransformer(cfg=model_cfg, **model_cfg).to(device) + outputs = model(x, x_mask, cond_code) + print(outputs) + + noise, x, flow_pred, final_mask, prompt_len = outputs["output"] + final_mask = final_mask.squeeze(-1) + + flow_gt = x - (1 - 1e-5) * noise + + # [B, n_frames, D] + diff_loss = F.l1_loss(flow_pred, flow_gt, reduction="none").float() * final_mask.unsqueeze(-1) + diff_loss = torch.mean(diff_loss, dim=2).sum() / final_mask.sum() + + print("diff_loss:", diff_loss.item()) + + diffusion_cond = torch.randn(2, 150, 256).to(device) + diffusion_cond_emb = model.cond_emb(diffusion_cond) + diffusion_prompt = torch.randn(2, 50, 128).to(device) + n_timesteps = 32 + + generated = model.reverse_diffusion(diffusion_cond_emb, diffusion_prompt, n_timesteps=n_timesteps) + print("generated:", generated.shape) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/llama.py b/vllm_omni/diffusion/models/soulx_singer/modules/llama.py new file mode 100644 index 00000000000..23a3377f257 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/llama.py @@ -0,0 +1,393 @@ +import math + +import torch +import torch.nn as nn +from cache_dit import ForwardPattern +from transformers import LlamaConfig, LlamaModel +from transformers.modeling_outputs import BaseModelOutputWithPast +from transformers.models.llama.modeling_llama import LlamaDecoderLayer + +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput + + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + device = x.device + dtype = x.dtype + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device, dtype=dtype) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class LlamaAdaptiveRMSNorm(nn.Module): + def __init__(self, hidden_size=1024, eps=1e-6, dim_cond=1024): + super().__init__() + self.to_weight = nn.Linear(dim_cond, hidden_size) + nn.init.zeros_(self.to_weight.weight) + nn.init.ones_(self.to_weight.bias) + self.variance_epsilon = eps + self._is_hf_initialized = True + + def forward(self, hidden_states, cond_embedding): + input_dtype = hidden_states.dtype + variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + + weight = self.to_weight(cond_embedding) + if len(weight.shape) == 2: + weight = weight.unsqueeze(1) + + return (weight * hidden_states).to(input_dtype) + + +class LlamaNARDecoderLayer(LlamaDecoderLayer): + def __init__(self, config: LlamaConfig, layer_idx: int): + super().__init__(config, layer_idx) + + self.input_layernorm = LlamaAdaptiveRMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + dim_cond=config.hidden_size, + ) + self.post_attention_layernorm = LlamaAdaptiveRMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + dim_cond=config.hidden_size, + ) + + def forward( + self, + hidden_states: torch.Tensor, + cond_embedding: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + output_attentions: bool | None = False, + cache_position: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.FloatTensor, ...]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states, cond_embedding=cond_embedding) + + # RoPE is computed by the parent model and passed in as (cos, sin). + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + cache_position=cache_position, + output_attentions=output_attentions, + **kwargs, + ) + + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm( + hidden_states, + cond_embedding=cond_embedding, + ) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +class DiffLlama(LlamaModel): + # Regional torch.compile in DiffusionModelRunner targets these blocks. + _repeated_blocks = ["LlamaNARDecoderLayer"] + _layerwise_offload_blocks_attrs = ["layers"] + + # cache-dit: Llama-style backbone with static condition embeddings injected + # before the decoder loop (like Stable Audio Open). Pattern_3 handles the + # cross-attention-style conditioning on the ``layers`` ModuleList. + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "layers": ForwardPattern.Pattern_3, + }, + ) + + # Sequence parallelism: shard mel sequence before the first decoder layer and + # gather after the output projection. Tensor parallelism still requires + # migrating attention to vLLM parallel linear layers. + _sp_plan = { + "layers.0": { + "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3), + }, + "mel_out_mlp": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + + @staticmethod + def _is_transformer_block(name: str, module: nn.Module) -> bool: + del module + return name.startswith("layers.") and name.split(".")[-1].isdigit() + + _hsdp_shard_conditions = [_is_transformer_block] + + def __init__( + self, + mel_dim=100, + hidden_size=1024, + num_heads=16, + num_layers=16, + dropout=0.1, + ffn_dropout=0.1, + attention_dropout=0.0, + config: LlamaConfig | None = None, + ): + if config is None: + config = LlamaConfig( + vocab_size=32000, + hidden_size=hidden_size, + intermediate_size=hidden_size * 4, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + max_position_embeddings=4096, + rms_norm_eps=1e-6, + attention_dropout=attention_dropout, + use_cache=False, + attn_implementation="eager", + ) + + config.use_cache = False + + super().__init__(config) + + # Keep custom decoder layers on the same attention backend. + layer_config = LlamaConfig( + vocab_size=config.vocab_size, + hidden_size=hidden_size, + intermediate_size=hidden_size * 4, + num_hidden_layers=num_layers, + num_attention_heads=num_heads, + num_key_value_heads=num_heads, + max_position_embeddings=4096, + rms_norm_eps=config.rms_norm_eps, + attention_dropout=attention_dropout, + use_cache=False, + attn_implementation="eager", + bos_token_id=None, + eos_token_id=None, + ) + + self.layers = nn.ModuleList( + [ + LlamaNARDecoderLayer( + layer_config, + layer_idx=i, + ) + for i in range(num_layers) + ] + ) + + self.norm = LlamaAdaptiveRMSNorm(hidden_size, dim_cond=hidden_size) + + self.diff_step_embedding = SinusoidalPosEmb(hidden_size) + self.diff_step_mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.SiLU(), + nn.Linear(hidden_size * 4, hidden_size), + ) + + self.cond_mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.SiLU(), + nn.Linear(hidden_size * 4, hidden_size), + ) + + self.mel_mlp = nn.Sequential( + nn.Linear(mel_dim, hidden_size * 4), + nn.SiLU(), + nn.Linear(hidden_size * 4, hidden_size), + ) + + self.mel_out_mlp = nn.Sequential( + nn.Linear(hidden_size, hidden_size * 4), + nn.SiLU(), + nn.Linear(hidden_size * 4, mel_dim), + ) + + self.embed_tokens = None + + self.post_init() + + def _prepare_decoder_attention_mask( + self, + attention_mask, + input_shape, + inputs_embeds, + past_key_values_length=0, + ): + """ + Build a 4D additive padding mask. + + This model is non-autoregressive, so no causal mask is added here. + """ + + combined_attention_mask = None + + def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None): + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = ( + mask[:, None, None, :] + .expand( + bsz, + 1, + tgt_len, + src_len, + ) + .to(dtype) + ) + + inverted_mask = 1.0 - expanded_mask + + return inverted_mask.masked_fill( + inverted_mask.to(torch.bool), + torch.finfo(dtype).min, + ) + + if attention_mask is not None: + expanded_attn_mask = _expand_mask( + attention_mask, + inputs_embeds.dtype, + tgt_len=input_shape[-1], + ).to(inputs_embeds.device) + + combined_attention_mask = ( + expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask + ) + + return combined_attention_mask + + def forward( + self, + x, + diffusion_step, + cond, + x_mask, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: list[torch.FloatTensor] | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = False, + ) -> BaseModelOutputWithPast | torch.Tensor | dict: + batch_size, seq_length, _ = x.shape + + cond_embedding = self.cond_mlp(cond) + x = self.mel_mlp(x) + + diffusion_step = self.diff_step_embedding(diffusion_step).to(device=x.device, dtype=x.dtype) + diffusion_step = self.diff_step_mlp(diffusion_step) + + x = x + cond_embedding + + inputs_embeds = x + attention_mask = x_mask + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + # KV cache is disabled because this diffusion decoder processes full sequences. + + if past_key_values is not None: + raise NotImplementedError("KV cache is disabled for this non-autoregressive DiffLlama implementation.") + + if position_ids is None: + device = input_ids.device if input_ids is not None else inputs_embeds.device + position_ids = torch.arange( + 0, + seq_length, + dtype=torch.long, + device=device, + ) + position_ids = position_ids.unsqueeze(0).expand(batch_size, -1) + else: + position_ids = position_ids.view(batch_size, seq_length).long() + + if attention_mask is None: + attention_mask = torch.ones( + (batch_size, seq_length), + dtype=torch.bool, + device=inputs_embeds.device, + ) + + attention_mask = self._prepare_decoder_attention_mask( + attention_mask, + (batch_size, seq_length), + inputs_embeds, + past_key_values_length=0, + ) + + hidden_states = inputs_embeds + + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_layer_hidden_states = [] + + cache_position = torch.arange( + 0, + seq_length, + device=hidden_states.device, + dtype=torch.long, + ) + + # LlamaAttention expects precomputed rotary embeddings. + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for idx, decoder_layer in enumerate(self.layers): + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + raise NotImplementedError + + layer_outputs = decoder_layer( + hidden_states, + cond_embedding=diffusion_step, + attention_mask=attention_mask, + position_embeddings=position_embeddings, + output_attentions=output_attentions, + cache_position=cache_position, + ) + + hidden_states = layer_outputs[0] + all_layer_hidden_states.append(hidden_states.clone()) + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states, cond_embedding=diffusion_step) + + if output_hidden_states: + all_hidden_states += (hidden_states,) + + hidden_states = self.mel_out_mlp(hidden_states) + + if return_dict: + return { + "output": hidden_states, + "hidden_states": all_layer_hidden_states, + "attentions": all_self_attns, + } + + return hidden_states diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py b/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py new file mode 100644 index 00000000000..cc767dcf307 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/mel_transform.py @@ -0,0 +1,151 @@ +import math +from typing import Any + +import numpy as np +import torch +import torch.nn as nn + +from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn + + +def dynamic_range_compression(x, c=1, clip_val=1e-5): + return np.log(np.clip(x, a_min=clip_val, a_max=None) * c) + + +def dynamic_range_decompression(x, c=1): + return np.exp(x) / c + + +def dynamic_range_compression_torch(x, c=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * c) + + +def dynamic_range_decompression_torch(x, c=1): + return torch.exp(x) / c + + +def spectral_normalize_torch(magnitudes): + output = dynamic_range_compression_torch(magnitudes) + return output + + +def spectral_de_normalize_torch(magnitudes): + output = dynamic_range_decompression_torch(magnitudes) + return output + + +class MelSpectrogram(nn.Module): + def __init__( + self, + n_fft, + num_mels, + sampling_rate, + hop_size, + win_size, + fmin, + fmax, + center=False, + ): + super().__init__() + self.n_fft = n_fft + self.hop_size = hop_size + self.win_size = win_size + self.sampling_rate = sampling_rate + self.num_mels = num_mels + self.fmin = fmin + self.fmax = fmax + self.center = center + + mel_basis = {} + hann_window = {} + + mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + mel_basis = mel.to(torch.float32) + hann_window = torch.hann_window(win_size) + + self.register_buffer("mel_basis", mel_basis) + self.register_buffer("hann_window", hann_window) + + def forward(self, y): + y = torch.nn.functional.pad( + y.unsqueeze(1), + ( + int((self.n_fft - self.hop_size) / 2), + int((self.n_fft - self.hop_size) / 2), + ), + mode="reflect", + ) + y = y.squeeze(1) + spec = torch.stft( + y, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=self.center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + spec = torch.view_as_real(spec) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(self.mel_basis, spec) + spec = spectral_normalize_torch(spec) + + return spec + + +def load_mel_spectrogram(): + return load_mel_spectrogram_from_cfg(None) + + +def _get_from_mapping(cfg: Any, key: str, default: Any = None) -> Any: + """Safely read a field from a dict/OmegaConf-like object.""" + if cfg is None: + return default + if isinstance(cfg, dict): + return cfg.get(key, default) + return getattr(cfg, key, default) + + +def load_mel_spectrogram_from_cfg(audio_cfg: Any | None = None) -> MelSpectrogram: + """Build MelSpectrogram from `audio_config`-like config. + + Expected keys (either in dict or Hydra/OmegaConf object): + - hop_size, sample_rate (or sampling_rate), n_fft, num_mels, win_size, fmin, fmax + """ + # Defaults keep current behavior. + mel_cfg: dict[str, Any] = { + "hop_size": _get_from_mapping(audio_cfg, "hop_size", 480), + "sampling_rate": _get_from_mapping( + audio_cfg, + "sampling_rate", + _get_from_mapping(audio_cfg, "sample_rate", 24000), + ), + "n_fft": _get_from_mapping(audio_cfg, "n_fft", 1920), + "num_mels": _get_from_mapping(audio_cfg, "num_mels", 128), + "win_size": _get_from_mapping(audio_cfg, "win_size", 1920), + "fmin": _get_from_mapping(audio_cfg, "fmin", 0), + "fmax": _get_from_mapping(audio_cfg, "fmax", 12000), + } + + mel_model = MelSpectrogram(**mel_cfg) + mel_model.eval() + return mel_model + + +class MelSpectrogramEncoder(nn.Module): + def __init__(self, audio_config: dict | None = None): + super().__init__() + self.model = load_mel_spectrogram_from_cfg(audio_config) + audio_config = audio_config or {} + self.mel_mean = audio_config.get("mel_mean", -4.92) + self.mel_var = audio_config.get("mel_var", 8.14) + + def forward(self, x): + x = self.model(x).transpose(1, 2) + x = (x - self.mel_mean) / math.sqrt(self.mel_var) + return x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py b/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py new file mode 100644 index 00000000000..70933dd92f9 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/note_transcription.py @@ -0,0 +1,958 @@ +import math + +import numpy as np +import packaging.version as version +import torch +import torch.nn as nn + +# ---------------------- note transcription common layers ---------------------- + + +class LayerNorm(nn.LayerNorm): + def __init__(self, nout: int, dim: int = -1, eps: float = 1e-5): + """Construct an LayerNorm object.""" + super().__init__(nout, eps=eps) + self.dim = dim + + def forward(self, x): + if self.dim == -1: + return super().forward(x) + return super().forward(x.transpose(1, -1)).transpose(1, -1) + + +class Reshape(nn.Module): + def __init__(self, *args): + super().__init__() + self.shape = args + + def forward(self, x): + return x.view(self.shape) + + +class Permute(nn.Module): + def __init__(self, *args): + super().__init__() + self.args = args + + def forward(self, x): + return x.permute(self.args) + + +class LambdaLayer(nn.Module): + def __init__(self, lambd): + super().__init__() + self.lambd = lambd + + def forward(self, x): + return self.lambd(x) + + +class ResidualBlock(nn.Module): + """ + Implements a residual block pattern: + [Norm] → [Conv1d] → [Scale] → [Activation] → [Conv1d] + repeated n times, each as a residual unit. + """ + + def __init__( + self, + channels: int, + kernel_size: int, + dilation: int, + n: int = 2, + c_multiple: int = 2, + ln_eps: float = 1e-12, + ): + super().__init__() + + self.blocks = nn.ModuleList( + [ + nn.Sequential( + LayerNorm(channels, dim=1, eps=ln_eps), + nn.Conv1d( + channels, + c_multiple * channels, + kernel_size, + dilation=dilation, + padding=(dilation * (kernel_size - 1)) // 2, + ), + LambdaLayer(lambda x: x * kernel_size**-0.5), # scaling for stability + nn.LeakyReLU(), + nn.Conv1d(c_multiple * channels, channels, kernel_size=1, dilation=dilation), + ) + for _ in range(n) + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] + for block in self.blocks: + x_res = block(x) + x = (x + x_res) * nonpadding + return x + + +class ConvBlocks(nn.Module): + """ + Decodes the expanded phoneme encoding into spectrograms. + """ + + def __init__( + self, + hidden_size, + out_dims, + dilations, + kernel_size, + layers_in_block=2, + c_multiple=2, + ln_eps=1e-5, + is_btc=True, + post_net_kernel=3, + ): + super().__init__() + self.is_btc = is_btc + + res_blocks = [ + ResidualBlock( + hidden_size, + kernel_size=kernel_size, + dilation=1, + n=layers_in_block, + c_multiple=c_multiple, + ln_eps=ln_eps, + ) + for _ in range(1) + ] # default 1 + self.res_blocks = nn.Sequential(*res_blocks) + + # Output normalization and conv + self.last_norm = LayerNorm(hidden_size, dim=1, eps=ln_eps) + self.post_net1 = nn.Conv1d(hidden_size, out_dims, kernel_size=post_net_kernel, padding=post_net_kernel // 2) + + def forward(self, x: torch.Tensor, nonpadding: torch.Tensor | None = None) -> torch.Tensor: + if self.is_btc: + x = x.transpose(1, 2) + if nonpadding is None: + nonpadding = (x.abs().sum(1) > 0).float()[:, None, :] + elif self.is_btc: + nonpadding = nonpadding.transpose(1, 2) + x = self.res_blocks(x) * nonpadding + x = self.last_norm(x) * nonpadding + x = self.post_net1(x) * nonpadding + if self.is_btc: + x = x.transpose(1, 2) + return x + + +def Linear(in_features: int, out_features: int, bias: bool = True, init_type: str = "xavier"): + m = nn.Linear(in_features, out_features, bias) + if init_type == "xavier": + nn.init.xavier_uniform_(m.weight) + elif init_type == "kaiming": + nn.init.kaiming_normal_(m.weight, mode="fan_in") + if bias: + nn.init.constant_(m.bias, 0.0) + return m + + +def Embedding(num_embeddings, embedding_dim, padding_idx=None, init_type="normal"): + m = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx) + if init_type == "normal": + nn.init.normal_(m.weight, mean=0, std=embedding_dim**-0.5) + elif init_type == "kaiming": + nn.init.kaiming_normal_(m.weight, mode="fan_in") + if padding_idx is not None: + nn.init.constant_(m.weight[padding_idx], 0) + return m + + +# ---------------------- note transcription espnet attention ---------------------- + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + """ + + def __init__(self, n_head, n_feat, dropout_rate, flash=False): + """Construct an MultiHeadedAttention object.""" + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.attn = None + self.dropout = nn.Dropout(p=dropout_rate) + self.dropout_rate = dropout_rate + self.flash = flash + + def forward_qkv(self, query, key, value): + """Transform query, key and value. + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + Returns: + torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention(self, value, scores, mask): + """Compute attention context vector. + Args: + value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + """ + n_batch = value.size(0) + if mask is not None: + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + min_value = float(np.finfo(torch.tensor(0, dtype=scores.dtype).numpy().dtype).min) + scores = scores.masked_fill(mask, min_value) + self.attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2) + else: + self.attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(self.attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward(self, query, key, value, mask): + """Compute scaled dot product attention. + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + """ + q, k, v = self.forward_qkv(query, key, value) + if version.parse(torch.__version__) >= version.parse("2.0") and self.flash: + n_batch = value.size(0) + x = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=mask.unsqueeze(1) if mask is not None else None, dropout_p=self.dropout_rate + ) + x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) + return self.linear_out(x) + else: + scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) + return self.forward_attention(v, scores, mask) + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding. + Paper: https://arxiv.org/abs/1901.02860 + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + """ + + def __init__(self, n_head, n_feat, dropout_rate): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__(n_head, n_feat, dropout_rate) + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def rel_shift(self, x, zero_triu=False): + """Compute relative positinal encoding. + Args: + x (torch.Tensor): Input tensor (batch, time, size). + zero_triu (bool): If true, return the lower triangular part of the matrix. + Returns: + torch.Tensor: Output tensor. + """ + zero_pad = torch.zeros((*x.size()[:3], 1), device=x.device, dtype=x.dtype) + x_padded = torch.cat([zero_pad, x], dim=-1) + + x_padded = x_padded.view(*x.size()[:2], x.size(3) + 1, x.size(2)) + x = x_padded[:, :, 1:].view_as(x) + + if zero_triu: + ones = torch.ones((x.size(2), x.size(3))) + x = x * torch.tril(ones, x.size(3) - x.size(2))[None, None, :, :] + + return x + + def forward(self, query, key, value, pos_emb, mask): + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + pos_emb (torch.Tensor): Positional embedding tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time2) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + matrix_bd = self.rel_shift(matrix_bd) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask) + + +# ---------------------- note transcription positional encoding ---------------------- + + +class PositionalEncoding(torch.nn.Module): + """Positional encoding. + Args: + d_model (int): Embedding dimension. + max_len (int): Maximum input length. + reverse (bool): Whether to reverse the input position. + """ + + def __init__(self, d_model, max_len=5000, reverse=False): + """Construct an PositionalEncoding object.""" + super().__init__() + self.d_model = d_model + self.reverse = reverse + self.xscale = math.sqrt(self.d_model) + self.pe = None + + self.extend_pe(torch.tensor(0.0).expand(1, max_len)) + + def extend_pe(self, x): + """Reset the positional encodings.""" + if self.pe is not None: + if self.pe.size(1) >= x.size(1): + if self.pe.dtype != x.dtype or self.pe.device != x.device: + self.pe = self.pe.to(dtype=x.dtype, device=x.device) + return + pe = torch.zeros(x.size(1), self.d_model) + if self.reverse: + position = torch.arange(x.size(1) - 1, -1, -1.0, dtype=torch.float32).unsqueeze(1) + else: + position = torch.arange(0, x.size(1), dtype=torch.float32).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, self.d_model, 2, dtype=torch.float32) * -(math.log(10000.0) / self.d_model) + ) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.pe = pe.to(device=x.device, dtype=x.dtype) + + def forward(self, x: torch.Tensor): + """Add positional encoding. + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + """ + self.extend_pe(x) + x = x * self.xscale + self.pe[:, : x.size(1)] + return x + + +class RelPositionalEncoding(PositionalEncoding): + """Relative positional encoding module. + See : Appendix B in https://arxiv.org/abs/1901.02860 + Args: + d_model (int): Embedding dimension. + max_len (int): Maximum input length. + """ + + def __init__(self, d_model, max_len=5000): + super().__init__(d_model, max_len, reverse=True) + + def forward(self, x): + """Compute positional encoding. + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + torch.Tensor: Positional embedding tensor (1, time, `*`). + """ + self.extend_pe(x) + x = x * self.xscale + pos_emb = self.pe[:, : x.size(1)] + return x, pos_emb + + +# ---------------------- note transcription conformer layers ---------------------- + + +class ConvolutionModule(nn.Module): + def __init__(self, channels: int, kernel_size: int, activation: nn.Module = nn.ReLU(), bias: bool = True): + super().__init__() + if (kernel_size - 1) % 2 != 0: + raise ValueError("kernel_size should be a odd number for 'SAME' padding") + + self.pointwise_conv1 = nn.Conv1d(channels, 2 * channels, kernel_size=1, bias=bias) + self.depthwise_conv = nn.Conv1d( + channels, channels, kernel_size=kernel_size, padding=(kernel_size - 1) // 2, groups=channels, bias=bias + ) + self.norm = nn.BatchNorm1d(channels) + self.pointwise_conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=bias) + self.activation = activation + + def forward(self, x): + # (batch, time, channels) -> (batch, channels, time) + x = x.transpose(1, 2) + + x = self.pointwise_conv1(x) + x = nn.functional.glu(x, dim=1) + + x = self.depthwise_conv(x) + x = self.norm(x) + x = self.activation(x) + + x = self.pointwise_conv2(x) + x = x.transpose(1, 2) + return x + + +class MultiLayeredConv1d(nn.Module): + """ + Multi-layered Conv1d that replaces feed-forward in Transformer block. + + Reference: + FastSpeech: Fast, Robust and Controllable Text to Speech + https://arxiv.org/pdf/1905.09263.pdf + """ + + def __init__(self, in_chans, hidden_chans, kernel_size): + super().__init__() + padding = (kernel_size - 1) // 2 + self.w_1 = nn.Conv1d(in_chans, hidden_chans, kernel_size, padding=padding) + self.w_2 = nn.Conv1d(hidden_chans, in_chans, kernel_size, padding=padding) + + def forward(self, x): + """ + Args: + x (torch.Tensor): (batch, time, in_chans) + Returns: + torch.Tensor: (batch, time, in_chans) + """ + # Conv1d expects (batch, channels, time) + x = x.transpose(1, 2) + x = torch.relu(self.w_1(x)) + x = self.w_2(x) + x = x.transpose(1, 2) + return x + + +class Swish(nn.Module): + """Swish activation function.""" + + def forward(self, x): + return x * torch.sigmoid(x) + + +class EncoderLayer(nn.Module): + """ + A single encoder layer: optional macaron FFN, self-attn, conv, FFN, norms. + Args: + size (int): Input dimension. + self_attn (nn.Module): Self-attention module. + feed_forward (nn.Module): Standard feed-forward network. + feed_forward_macaron (nn.Module): Optional, second FFN in macaron structure. + conv_module (nn.Module): Optional, convolution module. + normalize_before (bool): If True, use pre-layer norm. Else, post-layer norm. + concat_after (bool): If True, concat attn in/out, else residual. + """ + + def __init__( + self, + size, + self_attn, + feed_forward, + feed_forward_macaron, + conv_module, + normalize_before=True, + concat_after=False, + ): + super().__init__() + self.size = size + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + + self.norm_ff = LayerNorm(size) + self.norm_mha = LayerNorm(size) + + self.ff_scale = 0.5 if feed_forward_macaron is not None else 1.0 + self.norm_ff_macaron = LayerNorm(size) if feed_forward_macaron is not None else None + + if conv_module is not None: + self.norm_conv = LayerNorm(size) + self.norm_final = LayerNorm(size) + else: + self.norm_conv = None + self.norm_final = None + + self.normalize_before = normalize_before + self.concat_after = concat_after + if self.concat_after: + self.concat_linear = nn.Linear(2 * size, size) + + def forward(self, x_input, mask, cache=None) -> tuple: + """ + Forward pass for the encoder layer. + + Args: + x_input (Tensor or Tuple): Input tensor of shape (batch, time, size) + or tuple (x, pos_emb) where pos_emb is the positional embedding. + mask (Tensor): Mask tensor of shape (batch, time) or (batch, 1, time). + cache (Tensor, optional): Input cache for streaming (can be None). + + Returns: + Tuple[Tensor, Tensor]: (outputs, mask), where outputs either includes + pos_emb or not, depending on input. + """ + # Unpack possible (x, pos_emb) tuple input + if isinstance(x_input, tuple): + x, pos_emb = x_input + else: + x, pos_emb = x_input, None + + # --- Macaron Feed-Forward (first FFN) --- + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.feed_forward_macaron(x) + if not self.normalize_before: + x = self.norm_ff_macaron(x) + + # --- Self-Attention --- + residual = x + if self.normalize_before: + x = self.norm_mha(x) + + # Handle cache for streaming + if cache is not None: + assert cache.shape == (x.shape[0], x.shape[1] - 1, self.size) + x_q = x[:, -1:, :] + residual = residual[:, -1:, :] + mask = mask[:, -1:, :] if mask is not None else None + else: + x_q = x + + # Apply self-attention (with or without positional embeddings) + if pos_emb is not None: + attn_out = self.self_attn(x_q, x, x, pos_emb, mask) + else: + attn_out = self.self_attn(x_q, x, x, mask) + + if self.concat_after: + x = torch.cat((x, attn_out), dim=-1) + x = residual + self.concat_linear(x) + else: + x = residual + attn_out + if not self.normalize_before: + x = self.norm_mha(x) + + # --- Convolution Module --- + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x = residual + self.conv_module(x) + if not self.normalize_before: + x = self.norm_conv(x) + + # --- Feed-Forward (second FFN) --- + residual = x + if self.normalize_before: + x = self.norm_ff(x) + x = residual + self.ff_scale * self.feed_forward(x) + if not self.normalize_before: + x = self.norm_ff(x) + + # --- Final normalization, if present --- + if self.norm_final is not None: + x = self.norm_final(x) + + # --- Append to cache if streaming --- + if cache is not None: + x = torch.cat([cache, x], dim=1) + + # --- Final return, with or without positional embedding --- + if pos_emb is not None: + return (x, pos_emb), mask + else: + return x, mask + + +class ConformerLayers(nn.Module): + def __init__(self, hidden_size, num_layers, kernel_size=9, num_heads=4, use_last_norm=True, save_hidden=False): + super().__init__() + self.use_last_norm = use_last_norm + self.layers = nn.ModuleList() + positionwise_layer = MultiLayeredConv1d + positionwise_layer_args = (hidden_size, hidden_size * 4, 1) + self.pos_embed = RelPositionalEncoding(hidden_size) + self.encoder_layers = nn.ModuleList( + [ + EncoderLayer( + hidden_size, + RelPositionMultiHeadedAttention(num_heads, hidden_size, 0.0), + positionwise_layer(*positionwise_layer_args), + positionwise_layer(*positionwise_layer_args), + ConvolutionModule(hidden_size, kernel_size, Swish()), + ) + for _ in range(num_layers) + ] + ) + if self.use_last_norm: + self.layer_norm = nn.LayerNorm(hidden_size) + else: + self.layer_norm = nn.Linear(hidden_size, hidden_size) + self.save_hidden = save_hidden + if save_hidden: + self.hiddens = [] + + def forward(self, x, padding_mask=None): + """ + + :param x: [B, T, H] + :param padding_mask: [B, T] + :return: [B, T, H] + """ + self.hiddens = [] + nonpadding_mask = x.abs().sum(-1) > 0 + x = self.pos_embed(x) + for layer in self.encoder_layers: + x, mask = layer(x, nonpadding_mask[:, None, :]) + if self.save_hidden: + self.hiddens.append(x[0]) + x = x[0] + x = self.layer_norm(x) * nonpadding_mask.float()[:, :, None] + return x + + +# ---------------------- note transcription backbone layers ---------------------- + + +class UnetDown(nn.Module): + def __init__( + self, + hidden_size, + n_layers, + kernel_size, + down_rates, + channel_multiples=None, + is_btc=True, + constant_channels=False, + ): + super().__init__() + assert n_layers == len(down_rates) # downs, down sample rate + down_rates = [int(i) for i in down_rates] + self.n_layers = n_layers + self.hidden_size = hidden_size + self.is_btc = is_btc + channel_multiples = channel_multiples if channel_multiples is not None else down_rates + self.layers = nn.ModuleList() + self.downs = nn.ModuleList() + in_channels = hidden_size + for i in range(self.n_layers): + out_channels = int(in_channels * channel_multiples[i]) if not constant_channels else in_channels + self.layers.append( + nn.Sequential( + ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), + nn.Conv1d(in_channels, out_channels, kernel_size, padding=(kernel_size - 1) // 2), + ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), + ) + ) + self.downs.append(nn.Sequential(nn.AvgPool1d(down_rates[i]))) + in_channels = out_channels + self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6) + self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2) + + def forward(self, x, **kwargs): + # x [B, T, C] + if self.is_btc: + x = x.transpose(1, 2) # [B, C, T] + skip_xs = [] + for i in range(self.n_layers): + skip_x = self.layers[i](x) + x = self.downs[i](skip_x) + if self.is_btc: + skip_xs.append(skip_x.transpose(1, 2)) # [B, T, C] + else: + skip_xs.append(skip_x) + x = self.post_net(self.last_norm(x)) + if self.is_btc: + x = x.transpose(1, 2) + return x, skip_xs + + +class UnetMid(nn.Module): + def __init__(self, hidden_size, kernel_size, n_layers=None, in_dims=None, out_dims=None, is_btc=True, net=None): + super().__init__() + in_dims = in_dims if in_dims is not None else hidden_size + out_dims = out_dims if out_dims is not None else hidden_size + self.pre = nn.Conv1d(in_dims, hidden_size, kernel_size, padding=kernel_size // 2) + self.post = nn.Conv1d(hidden_size, out_dims, kernel_size, padding=kernel_size // 2) + self.is_btc = is_btc + if net is not None: + self.net = net + else: + self.net = ConvBlocks( + hidden_size, + out_dims=hidden_size, + dilations=None, + kernel_size=kernel_size, + layers_in_block=2, + c_multiple=2, + post_net_kernel=3, + is_btc=is_btc, + ) + + def forward(self, x, cond=None, **kwargs): + # x [B, T, C] + if self.is_btc: + x = self.pre(x.transpose(1, 2)).transpose(1, 2) + else: + x = self.pre(x) + if cond is None: + cond = 0 + x = self.net(x + cond) + if self.is_btc: + x = self.post(x.transpose(1, 2)).transpose(1, 2) + else: + x = self.post(x) + return x + + +class UnetUp(nn.Module): + def __init__( + self, + hidden_size, + n_layers, + kernel_size, + up_rates, + channel_multiples=None, + is_btc=True, + constant_channels=False, + use_skip_layer=False, + skip_scale=1.0, + ): + super().__init__() + assert n_layers == len(up_rates) # this is reversed in up module, from the output to the interface with middle + up_rates = [int(i) for i in up_rates] + self.n_layers = n_layers + self.hidden_size = hidden_size + self.is_btc = is_btc + self.skip_scale = skip_scale + channel_multiples = channel_multiples if channel_multiples is not None else up_rates + # in_channels = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size + self.in_channels_lst = ( + (np.cumprod([1] + channel_multiples) * hidden_size).astype(int) + if not constant_channels + else [hidden_size for _ in range(self.n_layers + 1)] + ) + in_channels = self.in_channels_lst[-1] + self.ups = nn.ModuleList() + self.skip_layers = nn.ModuleList() + self.layers = nn.ModuleList() + for i in range(self.n_layers - 1, -1, -1): + out_channels = self.in_channels_lst[i] if not constant_channels else in_channels + self.ups.append( + nn.Sequential( + nn.ConvTranspose1d( + in_channels, + in_channels, + kernel_size=kernel_size, + stride=up_rates[i], + padding=kernel_size // 2, + output_padding=up_rates[i] - 1, + ), + LayerNorm(in_channels, dim=1, eps=1e-6), + nn.LeakyReLU(), + ) + ) + self.layers.append( + nn.Sequential( + nn.Conv1d(in_channels * 2, out_channels, kernel_size, padding=(kernel_size - 1) // 2), + ResidualBlock(out_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5), + ) + ) + if use_skip_layer: + self.skip_layers.append( + ResidualBlock(in_channels, kernel_size, dilation=1, n=1, c_multiple=1, ln_eps=1e-5) + ) + else: + self.skip_layers.append(nn.Identity()) + + in_channels = out_channels + self.out_channels = out_channels + self.last_norm = LayerNorm(out_channels, dim=1, eps=1e-6) + self.post_net = nn.Conv1d(out_channels, out_channels, kernel_size=kernel_size, padding=kernel_size // 2) + + def forward(self, x, skips, **kwargs): + # x [B, T, C] + if self.is_btc: + x = x.transpose(1, 2) # [B, C, T] + for i in range(self.n_layers): + x = self.ups[i](x) + skip_x = ( + skips[self.n_layers - i - 1] if not self.is_btc else skips[self.n_layers - i - 1].transpose(1, 2) + ) # [B, T, C] -> [B, C, T] + skip_x = self.skip_layers[i](skip_x) * self.skip_scale + x = torch.cat((x, skip_x), dim=1) # [B, C, T] + x = self.layers[i](x) + x = self.post_net(self.last_norm(x)) + if self.is_btc: + x = x.transpose(1, 2) + return x + + +class Unet(nn.Module): + def __init__( + self, + hidden_size, + down_layers, + up_layers, + kernel_size, + updown_rates, + mid_layers=None, + channel_multiples=None, + is_btc=True, + constant_channels=False, + mid_net=None, + use_skip_layer=False, + skip_scale=1.0, + ): + super().__init__() + assert len(updown_rates) == down_layers == up_layers, f"{len(updown_rates)}, {down_layers}, {up_layers}" + if channel_multiples is not None: + assert len(channel_multiples) == len(updown_rates) + else: + channel_multiples = updown_rates + self.down = UnetDown( + hidden_size, down_layers, kernel_size, updown_rates, channel_multiples, is_btc, constant_channels + ) + down_out_dims = int(np.cumprod(channel_multiples)[-1] * hidden_size) if not constant_channels else hidden_size + self.mid = UnetMid( + hidden_size, + kernel_size, + mid_layers, + in_dims=down_out_dims, + out_dims=down_out_dims, + is_btc=is_btc, + net=mid_net, + ) + self.up = UnetUp( + hidden_size, + up_layers, + kernel_size, + updown_rates, + channel_multiples, + is_btc, + constant_channels, + use_skip_layer, + skip_scale, + ) + + def forward(self, x, mid_cond=None, **kwargs): + x, skips = self.down(x) + x = self.mid(x, mid_cond) + x = self.up(x, skips) + return x + + +class BackboneNet(nn.Module): + def __init__(self, hparams): + super().__init__() + self.hidden_size = hidden_size = hparams["hidden_size"] + updown_rates = [2, 2, 2] + channel_multiples = [1, 1, 1] + if hparams.get("updown_rates", None) is not None: + updown_rates = [int(i) for i in hparams.get("updown_rates", None).split("-")] + if hparams.get("channel_multiples", None) is not None: + channel_multiples = [float(i) for i in hparams.get("channel_multiples", None).split("-")] + assert len(updown_rates) == len(channel_multiples) + # convs + if hparams.get("bkb_net", "conv") == "conv": + self.net = Unet( + hidden_size, + down_layers=len(updown_rates), + mid_layers=hparams.get("bkb_layers", 12), + up_layers=len(updown_rates), + kernel_size=3, + updown_rates=updown_rates, + channel_multiples=channel_multiples, + is_btc=True, + constant_channels=False, + mid_net=None, + use_skip_layer=hparams.get("unet_skip_layer", False), + ) + # conformer + elif hparams.get("bkb_net", "conv") == "conformer": + mid_net = ConformerLayers( + hidden_size, + num_layers=hparams.get("bkb_layers", 12), + kernel_size=hparams.get("conformer_kernel", 9), + num_heads=4, + ) + self.net = Unet( + hidden_size, + down_layers=len(updown_rates), + up_layers=len(updown_rates), + kernel_size=3, + updown_rates=updown_rates, + channel_multiples=channel_multiples, + is_btc=True, + constant_channels=False, + mid_net=mid_net, + use_skip_layer=hparams.get("unet_skip_layer", False), + ) + + def forward(self, x): + return self.net(x) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py new file mode 100644 index 00000000000..9e020f61b5a --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/__init__.py @@ -0,0 +1,7 @@ +"""Preprocess neural network modules for SoulX-Singer.""" + +from .pipeline import SoulXPreprocessPipeline + +__all__ = [ + "SoulXPreprocessPipeline", +] diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py new file mode 100644 index 00000000000..0752c73f882 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/asr.py @@ -0,0 +1,326 @@ +"""Lyric transcription (Paraformer + Parakeet) for SoulX preprocess.""" + +import os +import re +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +from omegaconf import open_dict +from vllm.logger import init_logger + +from .utils import get_audio_duration + +logger = init_logger(__name__) + + +def _import_funasr_automodel(): + """Import FunASR without shelling out to ffmpeg at module import time.""" + import subprocess + from functools import wraps + + real_check_output = subprocess.check_output + + @wraps(real_check_output) + def _safe_check_output(cmd, *args, **kwargs): + if cmd and Path(str(cmd[0])).name == "ffmpeg": + raise FileNotFoundError("ffmpeg unavailable for FunASR import probe") + return real_check_output(cmd, *args, **kwargs) + + subprocess.check_output = _safe_check_output + try: + from funasr import AutoModel + finally: + subprocess.check_output = real_check_output + return AutoModel + + +def build_words_with_gaps(raw_words, raw_timestamps, audio_duration: float): + words, word_durs = [], [] + prev = 0.0 + for w, t in zip(raw_words, raw_timestamps): + s, e = float(t[0]), float(t[1]) + if s > prev: + words.append("<SP>") + word_durs.append(s - prev) + words.append(w) + word_durs.append(e - s) + prev = e + + if audio_duration > prev: + if not words: + words.append("<SP>") + word_durs.append(audio_duration) + elif words[-1] != "<SP>": + words.append("<SP>") + word_durs.append(audio_duration - prev) + else: + word_durs[-1] += audio_duration - prev + return words, word_durs + + +class ParaformerASR(nn.Module): + """Mandarin/Cantonese ASR via FunASR Paraformer.""" + + def __init__(self, model_path: str, device: str): + super().__init__() + from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load + + _patch_torchaudio_load() + AutoModel = _import_funasr_automodel() + self.asr = AutoModel(model=model_path, disable_update=True, device=device) + logger.info("Loaded Paraformer ASR from %s", model_path) + + def _generate(self, source: str) -> tuple[list[str], list[float], float]: + out = self.asr.generate(source, output_timestamp=True)[0] + raw_words = out["text"].replace("@", "").split(" ") + raw_timestamps = [[t[0] / 1000, t[1] / 1000] for t in out["timestamp"]] + duration = float(raw_timestamps[-1][1]) if raw_timestamps else 0.0 + return raw_words, raw_timestamps, duration + + def _maybe_refine_with_sidecar_f0( + self, + wav_path: str, + words: list[str], + word_durs: list[float], + *, + sample_rate: int, + hop_size: int, + ) -> tuple[list[str], list[float]]: + """Match upstream ``_ASRZhModel.process``: refine using ``{wav}_f0.npy`` if present.""" + f0_path = os.path.splitext(wav_path)[0] + "_f0.npy" + if not os.path.isfile(f0_path): + return words, word_durs + f0 = np.load(f0_path) + if f0.size == 0: + return words, word_durs + return refine_word_durs_with_f0( + words, + word_durs, + f0, + sample_rate=sample_rate, + hop_size=hop_size, + ) + + @torch.inference_mode() + def forward(self, audio: str | np.ndarray, sample_rate: int = 16000) -> tuple[list[str], list[float]]: + if isinstance(audio, str): + raw_words, raw_timestamps, _ = self._generate(audio) + audio_duration = get_audio_duration(audio) + words, word_durs = build_words_with_gaps(raw_words, raw_timestamps, audio_duration) + return self._maybe_refine_with_sidecar_f0( + audio, + words, + word_durs, + sample_rate=24000, + hop_size=480, + ) + + waveform = np.asarray(audio, dtype=np.float32).squeeze() + audio_duration = len(waveform) / float(sample_rate) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp_path = tmp.name + try: + sf.write(tmp_path, waveform, sample_rate) + raw_words, raw_timestamps, _ = self._generate(tmp_path) + words, word_durs = build_words_with_gaps(raw_words, raw_timestamps, audio_duration) + return self._maybe_refine_with_sidecar_f0( + tmp_path, + words, + word_durs, + sample_rate=24000, + hop_size=480, + ) + finally: + Path(tmp_path).unlink(missing_ok=True) + Path(os.path.splitext(tmp_path)[0] + "_f0.npy").unlink(missing_ok=True) + + +def _disable_nemo_cuda_graphs(model) -> None: + decoding_cfg = model.cfg.decoding + with open_dict(decoding_cfg): + if hasattr(decoding_cfg, "greedy") and decoding_cfg.greedy is not None: + decoding_cfg.greedy.use_cuda_graph_decoder = False + if hasattr(decoding_cfg, "beam") and decoding_cfg.beam is not None: + decoding_cfg.beam.allow_cuda_graphs = False + model.change_decoding_strategy(decoding_cfg) + + +class ParakeetASR(nn.Module): + """English ASR via NeMo Parakeet-TDT.""" + + def __init__(self, model_path: str, device: str): + super().__init__() + try: + import nemo.collections.asr as nemo_asr # type: ignore + except Exception as e: # pragma: no cover + raise ImportError( + "NeMo (nemo_toolkit) is required for ASR English but is not available in this Python env." + ) from e + + self.asr = nemo_asr.models.ASRModel.restore_from( + restore_path=model_path, + map_location=device, + ) + self.asr.eval() + _disable_nemo_cuda_graphs(self.asr) + logger.info("Loaded Parakeet ASR from %s", model_path) + + @staticmethod + def _clean_word(word: str) -> str: + return re.sub(r"[\?\.,:]", "", word).strip() + + @staticmethod + def _extract_word_segments(output: Any) -> list[dict[str, Any]]: + ts = getattr(output, "timestamp", None) + if not ts or not isinstance(ts, dict): + return [] + word_ts = ts.get("word") + return word_ts if isinstance(word_ts, list) else [] + + @torch.inference_mode() + def forward(self, audio: str | np.ndarray, sample_rate: int = 16000) -> tuple[list[str], list[float]]: + if isinstance(audio, str): + sources = [audio] + duration = 0.0 + else: + waveform = np.asarray(audio, dtype=np.float32).squeeze() + duration = len(waveform) / float(sample_rate) + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp_path = tmp.name + sf.write(tmp_path, waveform, sample_rate) + sources = [tmp_path] + + try: + outputs = self.asr.transcribe( + sources, + timestamps=True, + batch_size=1, + num_workers=0, + ) + finally: + if not isinstance(audio, str): + Path(sources[0]).unlink(missing_ok=True) + + output = outputs[0] if outputs else None + raw_words: list[str] = [] + raw_timestamps: list[list[float]] = [] + if output is not None: + for w in self._extract_word_segments(output): + s, e = float(w.get("start", 0.0)), float(w.get("end", 0.0)) + word = self._clean_word(str(w.get("word", ""))) + if word: + raw_words.append(word) + raw_timestamps.append([s, e]) + + if raw_timestamps: + duration = max(duration, float(raw_timestamps[-1][1])) + return build_words_with_gaps(raw_words, raw_timestamps, duration) + + +def refine_word_durs_with_f0( + words: list[str], + word_durs: list[float], + f0: np.ndarray, + *, + sample_rate: int = 24000, + hop_size: int = 480, +) -> tuple[list[str], list[float]]: + boundaries = np.cumsum([0, *[int(dur * sample_rate / hop_size) for dur in word_durs]]).tolist() + sil_tolerance = 5 + ext_tolerance = 5 + new_words: list[str] = [] + new_word_durs: list[float] = [] + if words: + new_words.append(words[0]) + new_word_durs.append(word_durs[0]) + + for i in range(1, len(words)): + word = words[i] + if word == "<SP>": + start_frame = boundaries[i] + end_frame = boundaries[i + 1] + num_frames = end_frame - start_frame + frame_idx = start_frame + unvoiced_count = 0 + while frame_idx < end_frame: + if f0[frame_idx] <= 1: + unvoiced_count += 1 + if unvoiced_count >= sil_tolerance: + frame_idx -= sil_tolerance - 1 + break + else: + unvoiced_count = 0 + frame_idx += 1 + + voice_frames = frame_idx - start_frame + if voice_frames >= int(num_frames * 0.9): + new_word_durs[-1] += word_durs[i] + elif voice_frames >= ext_tolerance: + dur = voice_frames * hop_size / sample_rate + new_word_durs[-1] += dur + new_words.append("<SP>") + new_word_durs.append(word_durs[i] - dur) + else: + new_words.append(word) + new_word_durs.append(word_durs[i]) + else: + new_words.append(word) + new_word_durs.append(word_durs[i]) + return new_words, new_word_durs + + +class LyricModel(nn.Module): + def __init__( + self, + zh_model_path: str, + en_model_path: str, + device: str = "cuda", + *, + target_sr: int = 24000, + hop_size: int = 480, + ): + super().__init__() + self.device_str = device + self.en_model_path = en_model_path + self.target_sr = target_sr + self.hop_size = hop_size + self.zh_asr = ParaformerASR(zh_model_path, device=device) + self.en_asr: nn.Module | None = None + + @torch.inference_mode() + def forward( + self, + audio: str | np.ndarray, + language: str = "Mandarin", + *, + sample_rate: int = 16000, + f0: np.ndarray | None = None, + ) -> tuple[list[str], list[float]]: + if language not in {"Mandarin", "Cantonese", "English"}: + raise ValueError(f"Unsupported language: {language}") + + if language.lower() == "english": + if self.en_asr is None: + self.en_asr = ParakeetASR(self.en_model_path, device=self.device_str) + self.add_module("en_asr", self.en_asr) + words, durs = self.en_asr(audio, sample_rate=sample_rate) + else: + # Path-based ASR loads ``{wav}_f0.npy`` inside ParaformerASR (upstream parity). + if isinstance(audio, str): + words, durs = self.zh_asr(audio) + else: + words, durs = self.zh_asr(audio, sample_rate=sample_rate) + if f0 is not None and len(f0) > 0: + words, durs = refine_word_durs_with_f0( + words, + durs, + f0, + sample_rate=self.target_sr, + hop_size=self.hop_size, + ) + return words, durs diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py new file mode 100644 index 00000000000..825c1622f4f --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/mel_grid_f0.py @@ -0,0 +1,96 @@ +"""Resample RMVPE F0 (16 kHz hop) onto the SoulX mel frame grid.""" + +from __future__ import annotations + +import math +import time +from typing import TYPE_CHECKING + +import numpy as np +from scipy.interpolate import interp1d +from vllm.logger import init_logger + +from .rmvpe import RMVPE_SAMPLE_RATE +from .utils import load_mono_audio + +if TYPE_CHECKING: + from .rmvpe import RMVPE + +logger = init_logger(__name__) + +RMVPE_HOP = 160 + + +def interpolate_f0_to_mel_grid( + f0_16k: np.ndarray, + original_length: int, + original_sr: int, + *, + target_sr: int, + hop_size: int, + max_duration: float = 300.0, +) -> np.ndarray: + batch_max_length = int(max_duration * target_sr / hop_size) + duration_in_seconds = original_length / original_sr + effective_target_length = int(duration_in_seconds * target_sr) + original_frames = math.ceil(effective_target_length / hop_size) + target_frames = min(original_frames, batch_max_length) + + t_16k = np.arange(len(f0_16k)) * (RMVPE_HOP / float(original_sr)) + t_target = np.arange(target_frames) * (hop_size / float(target_sr)) + + if len(f0_16k) > 0: + f_interp = interp1d( + t_16k, + f0_16k, + kind="linear", + bounds_error=False, + fill_value=0.0, + assume_sorted=True, + ) + f0 = f_interp(t_target) + else: + f0 = np.zeros(target_frames) + + if len(f0) != target_frames: + f0 = f0[:target_frames] if len(f0) > target_frames else np.pad(f0, (0, target_frames - len(f0))) + return f0.astype(np.float32, copy=False) + + +def extract_f0_file( + rmvpe: RMVPE, + wav_path: str, + *, + target_sr: int, + hop_size: int, + f0_path: str | None = None, + thread: float = 0.03, + max_duration: float = 300.0, + verbose: bool = False, +) -> np.ndarray: + if verbose: + logger.info("[f0] extracting from %s", wav_path) + t0 = time.perf_counter() + + audio, _ = load_mono_audio(wav_path, target_sr=RMVPE_SAMPLE_RATE) + f0_16k = rmvpe.infer_from_audio(audio, sample_rate=RMVPE_SAMPLE_RATE) + f0 = interpolate_f0_to_mel_grid( + np.asarray(f0_16k, dtype=np.float32), + original_length=audio.shape[-1], + original_sr=RMVPE_SAMPLE_RATE, + target_sr=target_sr, + hop_size=hop_size, + max_duration=max_duration, + ) + + if verbose: + voiced_ratio = float(np.mean(f0 > 0)) if len(f0) else 0.0 + logger.info( + "[f0] done: frames=%s voiced_ratio=%.3f time=%.3fs", + len(f0), + voiced_ratio, + time.perf_counter() - t0, + ) + if f0_path is not None: + np.save(f0_path, f0) + return f0 diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py new file mode 100644 index 00000000000..cf5f84e7a84 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/pipeline.py @@ -0,0 +1,431 @@ +"""SoulX preprocess pipeline.""" + +from __future__ import annotations + +import json +import tempfile +from collections.abc import Iterable +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery +from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.stack import SoulXPreprocessStack +from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import load_mono_audio, resample_mono +from vllm_omni.diffusion.models.soulx_singer.preprocess.metadata_utils import ( + SegmentMetadata, + _merge_group, + convert_metadata, +) +from vllm_omni.diffusion.models.soulx_singer.utils import ( + load_config, + load_wav, + preprocess_weight_paths, + resolve_preprocess_weights_root, + resolve_soulx_kind, +) + +logger = init_logger(__name__) + + +def build_svs_prompt_meta( + preprocess: SoulXPreprocessPipeline, + *, + prompt_audio: str | tuple[np.ndarray, int], + prompt_meta_raw: dict[str, Any], + metadata_processor, + vocal_sep: bool | None = None, +) -> dict[str, Any]: + """Build SVS prompt_meta; in-memory prompt wav is trimmed to mel2note length.""" + if isinstance(prompt_audio, str): + return metadata_processor.process(prompt_meta_raw, prompt_audio) + prompt_meta = metadata_processor.process(prompt_meta_raw, None) + prompt_vocal, prompt_sr = preprocess._extract_vocal(prompt_audio, vocal_sep=vocal_sep) + if prompt_sr != metadata_processor.sample_rate: + prompt_vocal = resample_mono( + prompt_vocal, + orig_sr=prompt_sr, + target_sr=metadata_processor.sample_rate, + ) + max_samples = prompt_meta["mel2note"].shape[1] * metadata_processor.hop_size + segment_wav = np.asarray(prompt_vocal[:max_samples], dtype=np.float32) + prompt_meta["wav"] = torch.from_numpy(segment_wav).unsqueeze(0).float().to(metadata_processor.device) + return prompt_meta + + +class SoulXPreprocessPipeline(nn.Module, SupportAudioInput, SupportsComponentDiscovery): + """Lazy-loaded preprocess stack integrated with vLLM-Omni diffusion lifecycle.""" + + support_audio_input: ClassVar[bool] = True + weights_sources: ClassVar[tuple] = () + _dit_modules: ClassVar[list[str]] = [] + _encoder_modules: ClassVar[list[str]] = [ + "stack.vocal_sep", + "stack.lyric", + "stack.rosvot", + ] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = ["stack.rmvpe"] + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + vocal_sep: bool = True, + midi_transcribe: bool = True, + max_merge_duration_ms: int = 60000, + verbose: bool = False, + extra_args: dict[str, Any] | None = None, + ) -> None: + super().__init__() + del prefix + self.od_config = od_config + extra_args = extra_args or {} + self.kind = resolve_soulx_kind(od_config.model) + hf_config = load_config(str(Path(od_config.model) / "config.yaml")) + audio_config = hf_config.audio + if self.kind == "svc" and midi_transcribe: + logger.warning("SVC model (config.json): midi_transcribe is ignored; lyric/MIDI transcription is SVS-only.") + midi_transcribe = False + + self.vocal_sep = vocal_sep + self.midi_transcribe = midi_transcribe + self.max_merge_duration_ms = max_merge_duration_ms + self.verbose = verbose + self.target_sr = int(audio_config.sample_rate) + self.hop_size = int(audio_config.hop_size) + self._metadata_processor = None + self._weights_root: Path | None = None + self.stack: SoulXPreprocessStack | None = None + self._ensure_stack(extra_args) + + def _ensure_stack(self, extra_args: dict[str, Any] | None = None) -> None: + extra_args = extra_args or {} + weights_root = resolve_preprocess_weights_root(self.od_config) + if self.stack is not None and self._weights_root == weights_root: + return + weights = preprocess_weight_paths(weights_root) + device = str(get_local_device()) + self.stack = SoulXPreprocessStack( + weights, + device, + target_sr=self.target_sr, + hop_size=self.hop_size, + verbose=self.verbose, + ) + self._weights_root = weights_root + + def forward(self, req) -> DiffusionOutput: + from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( + SOULX_PREPROCESSED_KEY, + ) + from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( + build_metadata_processor, + build_preprocess_payload, + build_warmup_payload, + is_warmup_request, + ) + from vllm_omni.inputs.data import OmniTextPrompt + + extra_args = dict(getattr(req.sampling_params, "extra_args", None) or {}) + prompt = OmniTextPrompt(prompt=req.prompts[0]) if isinstance(req.prompts[0], str) else req.prompts[0] + info = prompt.get("additional_information") or {} + if isinstance(info, dict): + extra_args.update(dict(info.get("extra_args") or {})) + + self._ensure_stack(extra_args) + if self._metadata_processor is None: + self._metadata_processor = build_metadata_processor(self.od_config) + + device = get_local_device() + if is_warmup_request(req): + payload = build_warmup_payload( + self.kind, + metadata_processor=self._metadata_processor, + device=device, + sample_rate=self.target_sr, + ) + else: + payload = build_preprocess_payload( + self.kind, + prompt=prompt, + extra_args=extra_args, + preprocess=self, + metadata_processor=self._metadata_processor, + device=device, + sample_rate=self.target_sr, + ) + + return DiffusionOutput( + output={ + "payload": {}, + "metadata": {SOULX_PREPROCESSED_KEY: payload}, + }, + to_cpu=True, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + del weights + self._ensure_stack({}) + return set() + + def _extract_vocal( + self, + audio: str | tuple[np.ndarray, int], + *, + vocal_sep: bool | None = None, + ) -> tuple[np.ndarray, int]: + if self.stack is None: + self._ensure_stack({}) + use_sep = self.vocal_sep if vocal_sep is None else vocal_sep + if use_sep: + return self.stack.extract_vocal(audio) + return load_mono_audio(audio) + + def extract_f0(self, vocal: np.ndarray, sample_rate: int) -> np.ndarray: + if self.stack is None: + self._ensure_stack({}) + return self.stack.extract_f0(vocal, sample_rate) + + def iter_svs_metadata( + self, + audio_source: str | tuple[np.ndarray, int], + *, + language: str = "Mandarin", + vocal_sep: bool | None = None, + max_merge_duration_ms: int | None = None, + ): + """Yield merged SVS metadata dicts one segment at a time.""" + if isinstance(audio_source, tuple): + vocal, sample_rate = audio_source + origin_wav_fn = "" + else: + vocal, sample_rate = self._extract_vocal(str(audio_source), vocal_sep=vocal_sep) + origin_wav_fn = str(audio_source) + + with tempfile.TemporaryDirectory(prefix="soulx_preprocess_") as work_dir: + work = Path(work_dir) + vocal_path = work / "vocal.wav" + sf.write(vocal_path, vocal, sample_rate) + + vocal_f0_path = str(vocal_path).replace(".wav", "_f0.npy") + vocal_f0 = self.stack.extract_f0(str(vocal_path), sample_rate=sample_rate, f0_path=vocal_f0_path) + + if not self.midi_transcribe: + end_ms = int(len(vocal) / sample_rate * 1000) + item = SegmentMetadata( + item_name="full", + wav_fn=str(vocal_path), + language=language, + start_time_ms=0, + end_time_ms=end_ms, + note_text=["<SP>"], + note_dur=[end_ms / 1000.0], + note_pitch=[0], + note_type=[1], + ) + yield convert_metadata(item) + return + + segments = self.stack.ensure_segmenter().forward( + vocal, + sample_rate, + vocal_f0, + base_name="vocal", + origin_wav_fn=origin_wav_fn, + verbose=self.verbose, + ) + cut_dir = work / "cut_wavs" + cut_dir.mkdir(parents=True, exist_ok=True) + lyric = self.stack.ensure_lyric() + rosvot = self.stack.ensure_rosvot() + + long_cut_dir = work / "long_cut_wavs" + max_dur = max_merge_duration_ms or self.max_merge_duration_ms + max_gap_ms = 10000 + current_group: list[Any] = [] + current_len = 0 + prev_end = -1 + + def _flush_group() -> dict[str, Any]: + merged_item = _merge_group(vocal, sample_rate, current_group, long_cut_dir) + merged_f0_path = merged_item.wav_fn.replace(".wav", "_f0.npy") + self.stack.extract_f0(merged_item.wav_fn, sample_rate=sample_rate, f0_path=merged_f0_path) + return convert_metadata(merged_item) + + for seg in segments: + seg_key = seg["item_name"] + seg_wav_path = cut_dir / f"{seg_key}.wav" + sf.write(seg_wav_path, seg["wav"], seg["sample_rate"]) + seg_f0_path = str(seg_wav_path).replace(".wav", "_f0.npy") + self.stack.extract_f0(str(seg_wav_path), sample_rate=seg["sample_rate"], f0_path=seg_f0_path) + + words, durs = lyric.forward( + str(seg_wav_path), + language, + sample_rate=seg["sample_rate"], + ) + seg_item = { + "item_name": seg_key, + "wav_fn": str(seg_wav_path), + "start_time_ms": seg["start_time_ms"], + "end_time_ms": seg["end_time_ms"], + "origin_wav_fn": origin_wav_fn or str(vocal_path), + "words": words, + "word_durs": durs, + "language": language, + } + meta_item = rosvot.transcribe(seg_item, segment_info=seg_item, verbose=self.verbose) + + start_time = int(meta_item.get("start_time_ms", seg["start_time_ms"])) + end_time = int(meta_item.get("end_time_ms", seg["end_time_ms"])) + if current_group and ( + start_time - prev_end > max_gap_ms or current_len + end_time - start_time > max_dur + ): + yield _flush_group() + current_group = [] + current_len = 0 + + current_group.append(meta_item) + current_len += end_time - start_time + prev_end = end_time + + if current_group: + yield _flush_group() + + @staticmethod + def build_svs_payload_from_paths( + *, + prompt_metadata_path: str, + target_metadata_path: str, + audio_path: str, + metadata_processor, + ) -> dict[str, Any]: + with open(prompt_metadata_path, encoding="utf-8") as f: + prompt_metadata = json.load(f) + with open(target_metadata_path, encoding="utf-8") as f: + target_metadata = json.load(f) + if not prompt_metadata: + raise ValueError("prompt_metadata is empty") + if not target_metadata: + raise ValueError("target_metadata is empty") + + prompt_meta = metadata_processor.process(prompt_metadata[0], audio_path) + return { + "kind": "svs", + "prompt_meta": prompt_meta, + "target_meta_list": target_metadata, + } + + def build_svs_payload_from_audio( + self, + *, + prompt_audio: str | tuple[np.ndarray, int], + target_audio: str | tuple[np.ndarray, int], + metadata_processor, + language: str = "Mandarin", + vocal_sep: bool | None = None, + prompt_vocal_sep: bool | None = None, + target_vocal_sep: bool | None = None, + prompt_max_merge_duration_ms: int | None = None, + target_max_merge_duration_ms: int | None = None, + ) -> dict[str, Any]: + p_sep = prompt_vocal_sep if prompt_vocal_sep is not None else vocal_sep + t_sep = target_vocal_sep if target_vocal_sep is not None else vocal_sep + p_merge = ( + prompt_max_merge_duration_ms if prompt_max_merge_duration_ms is not None else self.max_merge_duration_ms + ) + t_merge = ( + target_max_merge_duration_ms if target_max_merge_duration_ms is not None else self.max_merge_duration_ms + ) + prompt_list = list( + self.iter_svs_metadata( + prompt_audio, + language=language, + vocal_sep=p_sep, + max_merge_duration_ms=p_merge, + ) + ) + target_list = list( + self.iter_svs_metadata( + target_audio, + language=language, + vocal_sep=t_sep, + max_merge_duration_ms=t_merge, + ) + ) + if not prompt_list or not target_list: + raise ValueError("SVS preprocess produced empty metadata") + + # Prompt wav is trimmed to mel2note length, not the metadata time window. + prompt_meta = build_svs_prompt_meta( + self, + prompt_audio=prompt_audio, + prompt_meta_raw=prompt_list[0], + metadata_processor=metadata_processor, + vocal_sep=vocal_sep, + ) + + return { + "kind": "svs", + "prompt_meta": prompt_meta, + "target_meta_list": target_list, + } + + @staticmethod + def build_svc_payload_from_paths( + *, + prompt_wav_path: str, + target_wav_path: str, + prompt_f0_path: str, + target_f0_path: str, + sample_rate: int, + device: torch.device | str, + ) -> dict[str, Any]: + prompt_wav = load_wav(prompt_wav_path, sample_rate).to(device) + target_wav = load_wav(target_wav_path, sample_rate).to(device) + prompt_f0 = torch.from_numpy(np.load(prompt_f0_path)).unsqueeze(0).to(device) + target_f0 = torch.from_numpy(np.load(target_f0_path)).unsqueeze(0).to(device) + return { + "kind": "svc", + "prompt_wav": prompt_wav, + "target_wav": target_wav, + "prompt_f0": prompt_f0, + "target_f0": target_f0, + } + + def build_svc_payload_from_audio( + self, + *, + prompt_audio: str | tuple[np.ndarray, int], + target_audio: str | tuple[np.ndarray, int], + sample_rate: int, + device: torch.device | str, + vocal_sep: bool | None = None, + ) -> dict[str, Any]: + prompt_vocal, prompt_sr = self._extract_vocal(prompt_audio, vocal_sep=vocal_sep) + target_vocal, target_sr = self._extract_vocal(target_audio, vocal_sep=vocal_sep) + + if prompt_sr != sample_rate: + prompt_vocal = resample_mono(prompt_vocal, orig_sr=prompt_sr, target_sr=sample_rate) + if target_sr != sample_rate: + target_vocal = resample_mono(target_vocal, orig_sr=target_sr, target_sr=sample_rate) + + prompt_f0 = self.extract_f0(prompt_vocal, sample_rate) + target_f0 = self.extract_f0(target_vocal, sample_rate) + + return { + "kind": "svc", + "prompt_wav": torch.from_numpy(prompt_vocal).unsqueeze(0).to(device), + "target_wav": torch.from_numpy(target_vocal).unsqueeze(0).to(device), + "prompt_f0": torch.from_numpy(prompt_f0).unsqueeze(0).to(device), + "target_f0": torch.from_numpy(target_f0).unsqueeze(0).to(device), + } diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py new file mode 100644 index 00000000000..50ae4b2fcac --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rmvpe.py @@ -0,0 +1,593 @@ +"""RMVPE F0 extractor. + +All model code is inlined (originally from RickyL-2000/ROSVOT +``modules/pe/rmvpe/``) with two minor substitutions: + +* ``librosa.filters.mel`` → ``torchaudio.functional.melscale_fbanks`` + (already used elsewhere in vllm-omni); +* ``librosa.sequence.viterbi`` → hand-written numpy viterbi (non-default path). +* ``pyworld`` (deprecated) → removed. +""" + +from __future__ import annotations + +import math +from typing import ClassVar + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torchaudio.functional import melscale_fbanks +from torchaudio.transforms import Resample +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery + +logger = init_logger(__name__) + +RMVPE_SAMPLE_RATE = 16000 +_N_CLASS = 360 +_N_MELS = 128 +_MEL_FMIN = 30 +_MEL_FMAX = 8000 +_WINDOW_LENGTH = 1024 +_CONST = 1997.3794084376191 + + +class ConvBlockRes(nn.Module): + def __init__(self, in_channels: int, out_channels: int, momentum: float = 0.01): + super().__init__() + self.conv = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), + nn.BatchNorm2d(out_channels, momentum=momentum), + nn.ReLU(), + nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False), + nn.BatchNorm2d(out_channels, momentum=momentum), + nn.ReLU(), + ) + if in_channels != out_channels: + self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1)) + self._has_shortcut = True + else: + self._has_shortcut = False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self._has_shortcut: + return self.conv(x) + self.shortcut(x) + return self.conv(x) + x + + +class ResEncoderBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, kernel_size, n_blocks: int = 1, momentum: float = 0.01): + super().__init__() + self.n_blocks = n_blocks + self.conv = nn.ModuleList() + self.conv.append(ConvBlockRes(in_channels, out_channels, momentum)) + for _ in range(n_blocks - 1): + self.conv.append(ConvBlockRes(out_channels, out_channels, momentum)) + self.kernel_size = kernel_size + if self.kernel_size is not None: + self.pool = nn.AvgPool2d(kernel_size=kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + for i in range(self.n_blocks): + x = self.conv[i](x) + if self.kernel_size is not None: + return x, self.pool(x) + return x + + +class ResDecoderBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, stride, n_blocks: int = 1, momentum: float = 0.01): + super().__init__() + out_padding = (0, 1) if stride == (1, 2) else (1, 1) + self.n_blocks = n_blocks + self.conv1 = nn.Sequential( + nn.ConvTranspose2d( + in_channels, + out_channels, + kernel_size=(3, 3), + stride=stride, + padding=(1, 1), + output_padding=out_padding, + bias=False, + ), + nn.BatchNorm2d(out_channels, momentum=momentum), + nn.ReLU(), + ) + self.conv2 = nn.ModuleList() + self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum)) + for _ in range(n_blocks - 1): + self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum)) + + def forward(self, x: torch.Tensor, concat_tensor: torch.Tensor) -> torch.Tensor: + x = self.conv1(x) + x = torch.cat((x, concat_tensor), dim=1) + for i in range(self.n_blocks): + x = self.conv2[i](x) + return x + + +class Encoder(nn.Module): + def __init__( + self, + in_channels: int, + in_size: int, + n_encoders: int, + kernel_size, + n_blocks: int, + out_channels: int = 16, + momentum: float = 0.01, + ): + super().__init__() + self.n_encoders = n_encoders + self.bn = nn.BatchNorm2d(in_channels, momentum=momentum) + self.layers = nn.ModuleList() + self.latent_channels: list[list[int]] = [] + for _ in range(self.n_encoders): + self.layers.append(ResEncoderBlock(in_channels, out_channels, kernel_size, n_blocks, momentum=momentum)) + self.latent_channels.append([out_channels, in_size]) + in_channels = out_channels + out_channels *= 2 + in_size //= 2 + self.out_size = in_size + self.out_channel = out_channels + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, list[torch.Tensor]]: + concat_tensors: list[torch.Tensor] = [] + x = self.bn(x) + for i in range(self.n_encoders): + _c, x = self.layers[i](x) + concat_tensors.append(_c) + return x, concat_tensors + + +class Intermediate(nn.Module): + def __init__(self, in_channels: int, out_channels: int, n_inters: int, n_blocks: int, momentum: float = 0.01): + super().__init__() + self.n_inters = n_inters + self.layers = nn.ModuleList() + self.layers.append(ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)) + for _ in range(n_inters - 1): + self.layers.append(ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for i in range(self.n_inters): + x = self.layers[i](x) + return x + + +class Decoder(nn.Module): + def __init__(self, in_channels: int, n_decoders: int, stride, n_blocks: int, momentum: float = 0.01): + super().__init__() + self.layers = nn.ModuleList() + self.n_decoders = n_decoders + for _ in range(self.n_decoders): + out_channels = in_channels // 2 + self.layers.append(ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)) + in_channels = out_channels + + def forward(self, x: torch.Tensor, concat_tensors: list[torch.Tensor]) -> torch.Tensor: + for i in range(self.n_decoders): + x = self.layers[i](x, concat_tensors[-1 - i]) + return x + + +class TimbreFilter(nn.Module): + def __init__(self, latent_rep_channels: list[list[int]]): + super().__init__() + self.layers = nn.ModuleList() + for ch in latent_rep_channels: + self.layers.append(ConvBlockRes(ch[0], ch[0])) + + def forward(self, x_tensors: list[torch.Tensor]) -> list[torch.Tensor]: + return [layer(t) for layer, t in zip(self.layers, x_tensors)] + + +class DeepUnet0(nn.Module): + def __init__( + self, + kernel_size, + n_blocks: int, + en_de_layers: int = 5, + inter_layers: int = 4, + in_channels: int = 1, + en_out_channels: int = 16, + ): + super().__init__() + self.encoder = Encoder(in_channels, _N_MELS, en_de_layers, kernel_size, n_blocks, en_out_channels) + self.intermediate = Intermediate( + self.encoder.out_channel // 2, self.encoder.out_channel, inter_layers, n_blocks + ) + self.tf = TimbreFilter(self.encoder.latent_channels) + self.decoder = Decoder(self.encoder.out_channel, en_de_layers, kernel_size, n_blocks) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, concat_tensors = self.encoder(x) + x = self.intermediate(x) + x = self.decoder(x, concat_tensors) + return x + + +class BiGRU(nn.Module): + def __init__(self, input_features: int, hidden_features: int, num_layers: int): + super().__init__() + self.gru = nn.GRU(input_features, hidden_features, num_layers=num_layers, batch_first=True, bidirectional=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.gru(x)[0] + + +class E2E0(nn.Module): + """E2E0 pitch estimator network (DeepUnet0 + CNN + BiGRU + classifier).""" + + def __init__( + self, + n_blocks: int, + n_gru: int, + kernel_size, + en_de_layers: int = 5, + inter_layers: int = 4, + in_channels: int = 1, + en_out_channels: int = 16, + ): + super().__init__() + self.unet = DeepUnet0(kernel_size, n_blocks, en_de_layers, inter_layers, in_channels, en_out_channels) + self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1)) + if n_gru: + self.fc = nn.Sequential( + BiGRU(3 * _N_MELS, 256, n_gru), + nn.Linear(512, _N_CLASS), + nn.Dropout(0.25), + nn.Sigmoid(), + ) + else: + self.fc = nn.Sequential( + nn.Linear(3 * _N_MELS, _N_CLASS), + nn.Dropout(0.25), + nn.Sigmoid(), + ) + + def forward(self, mel: torch.Tensor) -> torch.Tensor: + mel = mel.transpose(-1, -2).unsqueeze(1) + x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2) + x = self.fc(x) + return x + + +class MelSpectrogram(nn.Module): + """Mel-spectrogram extractor matching ROSVOT's MelSpectrogram. + + Log-mel spectrogram via torch.stft + adaptive filterbank. + """ + + def __init__( + self, + n_mel_channels: int = _N_MELS, + sampling_rate: int = RMVPE_SAMPLE_RATE, + win_length: int = _WINDOW_LENGTH, + hop_length: int = 160, + n_fft: int | None = None, + mel_fmin: float = _MEL_FMIN, + mel_fmax: float = _MEL_FMAX, + clamp: float = 1e-5, + ): + super().__init__() + n_fft = win_length if n_fft is None else n_fft + self.hann_window: dict[str, torch.Tensor] = {} + mel_basis = melscale_fbanks( + n_freqs=n_fft // 2 + 1, + f_min=float(mel_fmin), + f_max=float(mel_fmax) if mel_fmax is not None else 8000.0, + n_mels=n_mel_channels, + sample_rate=int(sampling_rate), + norm="slaney", + mel_scale="htk", + ).T # (n_freqs, n_mels) → (n_mels, n_freqs) to match librosa/layernorm convention + self.register_buffer("mel_basis", mel_basis) # (n_mels, n_freqs) + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + self.sampling_rate = sampling_rate + self.n_mel_channels = n_mel_channels + self.clamp = clamp + + def forward(self, audio: torch.Tensor, keyshift: float = 0, speed: float = 1, center: bool = True) -> torch.Tensor: + factor = 2 ** (keyshift / 12) + n_fft_new = int(np.round(self.n_fft * factor)) + win_length_new = int(np.round(self.win_length * factor)) + hop_length_new = int(np.round(self.hop_length * speed)) + + key = str(keyshift) + "_" + str(audio.device) + if key not in self.hann_window: + self.hann_window[key] = torch.hann_window(win_length_new).to(audio.device) + + if center: + pad_left = win_length_new // 2 + pad_right = (win_length_new + 1) // 2 + audio = F.pad(audio, (pad_left, pad_right)) + + fft = torch.stft( + audio, + n_fft=n_fft_new, + hop_length=hop_length_new, + win_length=win_length_new, + window=self.hann_window[key], + center=False, + return_complex=True, + ) + magnitude = fft.abs() + + if keyshift != 0: + size = self.n_fft // 2 + 1 + if magnitude.size(1) < size: + magnitude = F.pad(magnitude, (0, 0, 0, size - magnitude.size(1))) + magnitude = magnitude[:, :size, :] * self.win_length / win_length_new + + # register_buffer confuses type checkers; cast to Tensor explicitly + mel_basis: torch.Tensor = self.mel_basis # type: ignore[assignment] + mel_output = torch.matmul(mel_basis, magnitude) + log_mel = torch.log(torch.clamp(mel_output, min=self.clamp)) + return log_mel + + +def to_local_average_f0(hidden: torch.Tensor, center: torch.Tensor | None = None, thread: float = 0.03) -> np.ndarray: + """Decode hidden representation to F0 via local weighted average.""" + idx = torch.arange(_N_CLASS, device=hidden.device)[None, None, :] + idx_cents = idx * 20 + _CONST + if center is None: + center = torch.argmax(hidden, dim=2, keepdim=True) + start = torch.clip(center - 4, min=0) + end = torch.clip(center + 5, max=_N_CLASS) + idx_mask = (idx >= start) & (idx < end) + weights = hidden * idx_mask + product_sum = torch.sum(weights * idx_cents, dim=2) + weight_sum = torch.sum(weights, dim=2) + cents = product_sum / (weight_sum + (weight_sum == 0)) + f0 = 10 * 2 ** (cents / 1200) + uv = hidden.max(dim=2)[0] < thread + f0 = f0 * ~uv + return f0.cpu().numpy() + + +def numpy_viterbi(prob: np.ndarray, transition: np.ndarray) -> np.ndarray: + """Viterbi decoding on log-probabilities (numpy). + + Args: + prob: Emission probabilities, shape (n_states, n_obs). + transition: Transition matrix, shape (n_states, n_states), row-stochastic. + + Returns: + path: Most likely state sequence, shape (n_obs,). + """ + n_states, n_obs = prob.shape + log_emit = np.log(prob + 1e-12) + log_trans = np.log(transition + 1e-12) + + delta = log_emit[:, 0] + np.zeros(n_states) # initial: log(pi) omitted (uniform) + psi = np.zeros((n_obs, n_states), dtype=np.int64) + + for t in range(1, n_obs): + scores = delta[:, None] + log_trans # (n_states, n_states) + psi[t] = np.argmax(scores, axis=0) + delta = scores[psi[t], np.arange(n_states)] + log_emit[:, t] + + # Backtrack + path = np.zeros(n_obs, dtype=np.int64) + path[-1] = np.argmax(delta) + for t in range(n_obs - 2, -1, -1): + path[t] = psi[t + 1, path[t + 1]] + return path + + +def to_viterbi_f0(hidden: torch.Tensor, thread: float = 0.03) -> np.ndarray: + """Decode hidden representation to F0 via Viterbi. + + Uses a hand-written numpy Viterbi decoder (replacing + ``librosa.sequence.viterbi``) to avoid the librosa dependency. + """ + if not hasattr(to_viterbi_f0, "_transition"): + xx, yy = np.meshgrid(range(_N_CLASS), range(_N_CLASS)) + transition = np.maximum(30 - abs(xx - yy), 0) + transition = transition / transition.sum(axis=1, keepdims=True) + to_viterbi_f0._transition = transition # type: ignore[attr-defined] + + prob = hidden.squeeze(0).cpu().numpy() + prob = prob.T + prob = prob / prob.sum(axis=0, keepdims=True) + + path = numpy_viterbi(prob, to_viterbi_f0._transition) # type: ignore[attr-defined] + center = torch.from_numpy(path).unsqueeze(0).unsqueeze(-1).to(hidden.device) + return to_local_average_f0(hidden, center=center, thread=thread) + + +def resample_align_curve( + points: np.ndarray, original_timestep: float, target_timestep: float, align_length: int = -1 +) -> np.ndarray: + """Align a curve from one time grid to another via linear interpolation.""" + t_max = (len(points) - 1) * original_timestep + curve_interp = np.interp( + np.arange(0, t_max, target_timestep), + original_timestep * np.arange(len(points)), + points, + ).astype(points.dtype) + if align_length > 0: + if len(curve_interp) > align_length: + curve_interp = curve_interp[:align_length] + elif len(curve_interp) < align_length: + curve_interp = np.concatenate( + [curve_interp, np.full(align_length - len(curve_interp), fill_value=curve_interp[-1])], + axis=0, + ) + return curve_interp + + +class RMVPE(nn.Module, SupportAudioInput, SupportsComponentDiscovery): + """RMVPE F0 extractor (SoulX-Singer). + + Wraps the E2E0 model with mel-spectrogram front-end and decoding + post-processing. + """ + + support_audio_input: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = [] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = ["."] + _layerwise_offload_blocks_attrs: ClassVar[list[str]] = [] + + def __init__(self, model_path: str | None = None, hop_length: int = 160, device: str | None = None): + super().__init__() + self.resample_kernel: dict[str, Resample] = {} + if device is None: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + self.device = torch.device(device) + + self.model = E2E0(4, 1, (2, 2)).eval() + if model_path is not None: + ckpt = torch.load(model_path, map_location=self.device) + state_dict = ckpt if not isinstance(ckpt, dict) else ckpt.get("model", ckpt) + self.model.load_state_dict(state_dict, strict=False) + logger.info("Loaded RMVPE weights from %s", model_path) + self.model = self.model.to(self.device) + + self.mel_extractor = MelSpectrogram( + _N_MELS, + RMVPE_SAMPLE_RATE, + _WINDOW_LENGTH, + hop_length, + None, + _MEL_FMIN, + _MEL_FMAX, + ).to(self.device) + self.hop_length = hop_length + + # ---- internal helpers ---- + + @torch.no_grad() + def mel2hidden(self, mel: torch.Tensor) -> torch.Tensor: + n_frames = mel.shape[-1] + pad_len = 32 * ((n_frames - 1) // 32 + 1) - n_frames + mel = F.pad(mel, (0, pad_len), mode="constant") + hidden = self.model(mel) + return hidden[:, :n_frames] + + def decode(self, hidden: torch.Tensor, thread: float = 0.03, use_viterbi: bool = False) -> np.ndarray: + if use_viterbi: + return to_viterbi_f0(hidden, thread=thread) + return to_local_average_f0(hidden, thread=thread) + + def postprocess(self, f0: np.ndarray, fmin: float = 50, fmax: float = 1000, min_gap: int = 2) -> np.ndarray: + f0[f0 < fmin] = 0 + f0[f0 > fmax] = 0 + # Eliminate glitch: if positive f0 values span fewer than min_gap frames + # between zeros, zero them out. + for idx in range(f0.shape[0] - min_gap - 1): + if f0[idx] == 0 and f0[idx + min_gap + 1] == 0 and np.sum(f0[idx : idx + min_gap + 2]) > 0: + f0[idx : idx + min_gap + 2] = 0 + return f0 + + # ---- single audio inference ---- + + def infer_from_audio( + self, audio: np.ndarray, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False + ) -> np.ndarray: + audio_t = torch.from_numpy(audio).float().unsqueeze(0).to(self.device) + if sample_rate == RMVPE_SAMPLE_RATE: + audio_res = audio_t + else: + key_str = str(sample_rate) + if key_str not in self.resample_kernel: + self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128) + self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device) + audio_res = self.resample_kernel[key_str](audio_t) + mel = self.mel_extractor(audio_res, center=True) + hidden = self.mel2hidden(mel) + f0 = self.decode(hidden, thread=thread, use_viterbi=use_viterbi).squeeze(0) + return f0 + + def get_pitch( + self, + waveform, + sample_rate: int, + hop_size: int, + length: int, + interp_uv: bool = False, + fmin: float = 50, + fmax: float = 1000, + ) -> tuple[np.ndarray, np.ndarray]: + f0 = self.infer_from_audio(waveform, sample_rate=sample_rate) + f0 = self.postprocess(f0, fmin, fmax) + uv = f0 == 0 + time_step = hop_size / sample_rate + f0_res = resample_align_curve(f0, 0.01, time_step, length) + uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5 + if not interp_uv: + f0_res[uv_res] = 0 + return f0_res, uv_res + + # ---- batch audio inference ---- + + def infer_from_audio_batch( + self, audios, sample_rate: int = 16000, thread: float = 0.03, use_viterbi: bool = False + ) -> list[np.ndarray]: + sizes: list[int] | None = None + if isinstance(audios, list): + audios = [torch.from_numpy(a).float() for a in audios] + sizes = [math.ceil((a.shape[0] + 1) / self.hop_length) for a in audios] + audios = nn.utils.rnn.pad_sequence(audios, batch_first=True, padding_value=0.0).to(self.device) + elif isinstance(audios, torch.Tensor): + if audios.device != self.device: + audios = audios.to(self.device) + else: + raise NotImplementedError(f"Unsupported audio type: {type(audios)}") + + if sample_rate == RMVPE_SAMPLE_RATE: + audios_res = audios + else: + key_str = str(sample_rate) + if key_str not in self.resample_kernel: + self.resample_kernel[key_str] = Resample(sample_rate, RMVPE_SAMPLE_RATE, lowpass_filter_width=128) + self.resample_kernel[key_str] = self.resample_kernel[key_str].to(self.device) + audios_res = self.resample_kernel[key_str](audios) + + mels = self.mel_extractor(audios_res, center=True) + hiddens = self.mel2hidden(mels) + f0 = self.decode(hiddens, thread=thread, use_viterbi=use_viterbi) + + f0s: list[np.ndarray] = [] + for i in range(f0.shape[0]): + f0s.append(f0[i, : sizes[i]] if sizes is not None else f0[i, :]) + return f0s + + def get_pitch_batch( + self, + waveforms, + sample_rate: int, + hop_size: int, + lengths: list[int], + interp_uv: bool = False, + fmin: float = 50, + fmax: float = 1000, + ) -> tuple[list[np.ndarray], list[np.ndarray]]: + f0s = self.infer_from_audio_batch(waveforms, sample_rate=sample_rate) + f0s_res: list[np.ndarray] = [] + uvs_res: list[np.ndarray] = [] + for idx, f0 in enumerate(f0s): + f0 = self.postprocess(f0, fmin, fmax, min_gap=6) + uv = f0 == 0 + length = lengths[idx] + time_step = hop_size / sample_rate + f0_res = resample_align_curve(f0, 0.01, time_step, length) + uv_res = resample_align_curve(uv.astype(np.float32), 0.01, time_step, length) > 0.5 + if not interp_uv: + f0_res[uv_res] = 0 + f0s_res.append(f0_res) + uvs_res.append(uv_res) + return f0s_res, uvs_res + + def release_cuda(self): + self.model = self.model.cpu() + self.mel_extractor = self.mel_extractor.cpu() diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py new file mode 100644 index 00000000000..0333b521e34 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py @@ -0,0 +1,619 @@ +"""ROSVOT note transcription.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery +from vllm_omni.diffusion.models.soulx_singer.modules.mel_transform import spectral_normalize_torch +from vllm_omni.diffusion.models.soulx_singer.modules.note_transcription import ( + BackboneNet, + ConvBlocks, + Embedding, +) +from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import ( + boundary2Interval, + denorm_f0, + f0_to_coarse, + load_model_ckpt, + load_mono_audio, + load_rosvot_config, + norm_interp_f0, + pad_or_cut_xd, + resample_mono, +) +from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn + +logger = init_logger(__name__) + + +def regulate_real_note_itv( + note_itv: np.ndarray, + note_bd: np.ndarray, + word_bd: np.ndarray, + word_durs: np.ndarray, + hop_size: int, + audio_sample_rate: int, +) -> tuple[np.ndarray, np.ndarray]: + assert note_itv.shape[0] == np.sum(note_bd) + 1 + assert np.sum(word_bd) <= np.sum(note_bd) + assert word_durs.shape[0] == np.sum(word_bd) + 1, f"{word_durs.shape[0]} {np.sum(word_bd) + 1}" + word_bd = np.cumsum(word_bd) * word_bd + word_itv = np.zeros((word_durs.shape[0], 2)) + word_offsets = np.cumsum(word_durs) + note2words = np.zeros(note_itv.shape[0], dtype=int) + for idx in range(len(word_offsets) - 1): + word_itv[idx, 1] = word_itv[idx + 1, 0] = word_offsets[idx] + word_itv[-1, 1] = word_offsets[-1] + note_itv_secs = note_itv * hop_size / audio_sample_rate + for idx, itv in enumerate(note_itv): + start_idx, end_idx = itv + if word_bd[start_idx] > 0: + word_dur_idx = word_bd[start_idx] + note_itv_secs[idx, 0] = word_itv[word_dur_idx, 0] + note2words[idx] = word_dur_idx + if word_bd[end_idx] > 0: + word_dur_idx = word_bd[end_idx] - 1 + note_itv_secs[idx, 1] = word_itv[word_dur_idx, 1] + note2words[idx] = word_dur_idx + note2words += 1 + return note_itv_secs, note2words + + +def regulate_ill_slur( + notes: np.ndarray, note_itv: np.ndarray, note2words: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + res_note2words: list[int] = [] + res_note_itv: list[list[float]] = [] + res_notes: list[int] = [] + note_idx = 0 + note_idx_end = 0 + while note_idx <= len(notes) - 1: + while note_idx <= note_idx_end < len(notes) and note2words[note_idx] == note2words[note_idx_end]: + note_idx_end += 1 + res_note2words.append(note2words[note_idx]) + res_note_itv.append(note_itv[note_idx].tolist()) + res_notes.append(notes[note_idx]) + for idx in range(note_idx + 1, note_idx_end): + if notes[idx] == notes[idx - 1]: + res_note_itv[-1][1] = note_itv[idx][1] + else: + res_note_itv.append(note_itv[idx].tolist()) + res_note2words.append(note2words[idx]) + res_notes.append(notes[idx]) + note_idx = note_idx_end + return ( + np.array(res_notes, dtype=notes.dtype), + np.array(res_note_itv, dtype=note_itv.dtype), + np.array(res_note2words, dtype=note2words.dtype), + ) + + +def regulate_boundary(bd_logits, threshold, min_gap=18, ref_bd=None, ref_bd_min_gap=8, non_padding=None): + # this doesn't preserve gradient + device = bd_logits.device + bd_logits = torch.sigmoid(bd_logits).data.cpu() + # bd_logits[0] = bd_logits[-1] = 1e-5 # avoid itv invalid problem + bd = (bd_logits > threshold).long() + bd_res = torch.zeros_like(bd).long() + for i in range(bd.shape[0]): + bd_i = bd[i] + last_bd_idx = -1 + start = -1 + for j in range(bd_i.shape[0]): + if bd_i[j] == 1: + if 0 <= start < j: + continue + elif start < 0: + start = j + else: + if 0 <= start < j: + if j - 1 > start: + bd_idx = start + int(torch.argmax(bd_logits[i, start:j]).item()) + else: + bd_idx = start + if bd_idx - last_bd_idx < min_gap and last_bd_idx > 0: + bd_idx = round((bd_idx + last_bd_idx) / 2) + bd_res[i, last_bd_idx] = 0 + bd_res[i, bd_idx] = 1 + last_bd_idx = bd_idx + start = -1 + + # assert ref_bd_min_gap <= min_gap // 2 + if ref_bd is not None and ref_bd_min_gap > 0: + ref = ref_bd.data.cpu() + for i in range(bd_res.shape[0]): + ref_bd_i = ref[i] + ref_bd_i_js = [] + for j in range(ref_bd_i.shape[0]): + if ref_bd_i[j] == 1: + ref_bd_i_js.append(j) + seg_sum = torch.sum(bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap]) + if seg_sum == 0: + bd_res[i, j] = 1 + elif seg_sum == 1 and bd_res[i, j] != 1: + bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap] = ref_bd_i[ + max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap + ] + elif seg_sum > 1: + for k in range(1, ref_bd_min_gap + 1): + if bd_res[i, max(0, j - k)] == 1 and ref_bd_i[max(0, j - k)] != 1: + bd_res[i, max(0, j - k)] = 0 + break + if ( + bd_res[i, min(bd_res.shape[1] - 1, j + k)] == 1 + and ref_bd_i[min(bd_res.shape[1] - 1, j + k)] != 1 + ): + bd_res[i, min(bd_res.shape[1] - 1, j + k)] = 0 + break + bd_res[i, j] = 1 + # final check + assert torch.sum(bd_res[i, ref_bd_i_js]) == len(ref_bd_i_js), ( + f"{torch.sum(bd_res[i, ref_bd_i_js])} {len(ref_bd_i_js)}" + ) + + bd_res = bd_res.to(device) + + # force valid begin and end + bd_res[:, 0] = 0 + if non_padding is not None: + for i in range(bd_res.shape[0]): + bd_res[i, sum(non_padding[i]) - 1 :] = 0 + else: + bd_res[:, -1] = 0 + + return bd_res + + +def _align_word(word_durs: list[float], mel_len: int, hop_size: int, audio_sample_rate: int) -> np.ndarray: + mel2word = np.zeros([mel_len], int) + start_time = 0.0 + for i_word, wd in enumerate(word_durs): + start_frame = int(start_time * audio_sample_rate / hop_size + 0.5) + end_frame = int((start_time + wd) * audio_sample_rate / hop_size + 0.5) + mel2word[start_frame:end_frame] = i_word + 1 + start_time += wd + return mel2word + + +class PitchDecoder(nn.Module): + def __init__(self, hparams): + super().__init__() + self.hidden_size = hidden_size = hparams["hidden_size"] + self.dropout = hparams.get("dropout", 0.0) + self.note_bd_out = nn.Linear(hidden_size, 1) + self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0)) + + # note prediction + self.pitch_attn_num_head = hparams.get("pitch_attn_num_head", 1) + self.multihead_dot_attn = nn.Linear(hidden_size, self.pitch_attn_num_head) + self.post = ConvBlocks( + hidden_size, + out_dims=hidden_size, + dilations=None, + kernel_size=3, + layers_in_block=1, + c_multiple=1, + post_net_kernel=3, + ) + self.pitch_out = nn.Linear(hidden_size, hparams.get("note_num", 100) + 4) + self.note_num = hparams.get("note_num", 100) + self.note_start = hparams.get("note_start", 30) + self.pitch_temperature = max(1e-7, hparams.get("note_pitch_temperature", 1.0)) + + def forward(self, feat, note_bd, train=True): + bsz, T, _ = feat.shape + + attn = torch.sigmoid(self.multihead_dot_attn(feat)) # [B, T, C] -> [B, T, num_head] + attn_feat = feat.unsqueeze(3) * attn.unsqueeze(2) # [B, T, C, 1] x [B, T, 1, num_head] -> [B, T, C, num_head] + attn_feat = torch.mean(attn_feat, dim=-1) # [B, T, C, num_head] -> [B, T, C] + mel2note = torch.cumsum(note_bd, 1) + note_length = torch.max(torch.sum(note_bd, dim=1)).item() + 1 # max length + note_lengths = torch.sum(note_bd, dim=1) + 1 # [B] + + attn = torch.mean(attn, dim=-1, keepdim=True) # [B, T, num_head] -> [B, T, 1] + denom = mel2note.new_zeros(bsz, note_length, dtype=attn.dtype).scatter_add_( + dim=1, index=mel2note, src=attn.squeeze(-1) + ) # [B, T] -> [B, note_length] count the note frames of each note (with padding excluded) + frame2note = mel2note.unsqueeze(-1).repeat(1, 1, self.hidden_size) # [B, T] -> [B, T, C], with padding included + note_aggregate = frame2note.new_zeros(bsz, note_length, self.hidden_size, dtype=attn_feat.dtype).scatter_add_( + dim=1, index=frame2note, src=attn_feat + ) # [B, T, C] -> [B, note_length, C] + note_aggregate = note_aggregate / (denom.unsqueeze(-1) + 1e-5) + note_logits = self.post(note_aggregate) + note_logits = self.pitch_out(note_logits) / self.pitch_temperature + # note_logits = torch.clamp(note_logits, min=-16., max=16.) # don't know need it or not + + note_pred = torch.softmax(note_logits, dim=-1) # [B, note_length, note_num] + note_pred = torch.argmax(note_pred, dim=-1) # [B, note_length] + # for some reason, note idx maybe 130 (why?) + note_pred[note_pred > self.note_num] = 0 + note_pred[note_pred < self.note_start] = 0 + + return note_lengths, note_logits, note_pred + + +class MidiExtractor(nn.Module): + def __init__(self, hparams): + super().__init__() + self.hparams = hparams.copy() + self.hidden_size = hidden_size = hparams["hidden_size"] + self.note_bd_threshold = hparams.get("note_bd_threshold", 0.5) + self.note_bd_min_gap = round( + hparams.get("note_bd_min_gap", 100) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"] + ) + self.note_bd_ref_min_gap = round( + hparams.get("note_bd_ref_min_gap", 50) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"] + ) + + self.mel_proj = nn.Conv1d(hparams["use_mel_bins"], hidden_size, kernel_size=3, padding=1) + self.mel_encoder = ConvBlocks( + hidden_size, + out_dims=hidden_size, + dilations=None, + kernel_size=3, + layers_in_block=2, + c_multiple=1, + post_net_kernel=3, + ) + self.use_pitch = hparams.get("use_pitch_embed", True) + if self.use_pitch: + self.pitch_embed = Embedding(300, hidden_size, 0, "kaiming") + self.uv_embed = Embedding(3, hidden_size, 0, "kaiming") + self.use_wbd = hparams.get("use_wbd", True) + if self.use_wbd: + self.word_bd_embed = Embedding(3, hidden_size, 0, "kaiming") + self.cond_encoder = ConvBlocks( + hidden_size, + out_dims=hidden_size, + dilations=None, + kernel_size=3, + layers_in_block=1, + c_multiple=1, + post_net_kernel=3, + ) + + # backbone + self.net = BackboneNet(hparams) + + # note bd prediction + self.note_bd_out = nn.Linear(hidden_size, 1) + self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0)) + + # note prediction + self.pitch_decoder = PitchDecoder(hparams) + + def run_encoder(self, mel, word_bd=None, pitch=None, uv=None): + mel_embed = self.mel_proj(mel.transpose(1, 2)).transpose(1, 2) + mel_embed = self.mel_encoder(mel_embed) + pitch_embed = word_bd_embed = 0 + if self.use_pitch and pitch is not None and uv is not None: + pitch_embed = self.pitch_embed(pitch) + self.uv_embed(uv) # [B, T, C] + if self.use_wbd and word_bd is not None: + word_bd_embed = self.word_bd_embed(word_bd) + feat = self.cond_encoder(mel_embed + pitch_embed + word_bd_embed) + return feat + + def forward(self, mel, word_bd=None, note_bd=None, pitch=None, uv=None, non_padding=None, train=True): + ret = {} + bsz, T, _ = mel.shape + + feat = self.run_encoder(mel, word_bd, pitch, uv) + feat = self.net(feat) # [B, T, C] + + # note bd prediction + # dropout has been dropped (inference mode) + note_bd_logits = self.note_bd_out(feat).squeeze(-1) / self.note_bd_temperature + note_bd_logits = torch.clamp(note_bd_logits, min=-16.0, max=16.0) + ret["note_bd_logits"] = note_bd_logits # [B, T] + if note_bd is None or not train: + note_bd = regulate_boundary( + note_bd_logits, + self.note_bd_threshold, + self.note_bd_min_gap, + word_bd, + self.note_bd_ref_min_gap, + non_padding, + ) + ret["note_bd_pred"] = note_bd # [B, T] + + # note pitch prediction + note_lengths, note_logits, note_pred = self.pitch_decoder(feat, note_bd, train) + ret["note_lengths"], ret["note_logits"], ret["note_pred"] = note_lengths, note_logits, note_pred + + return ret + + +class MelNet(nn.Module): + def __init__(self, hparams, device="cpu") -> None: + super().__init__() + self.n_fft = hparams["fft_size"] + self.num_mels = hparams["audio_num_mel_bins"] + self.sampling_rate = hparams["audio_sample_rate"] + self.hop_size = hparams["hop_size"] + self.win_size = hparams["win_size"] + self.fmin = hparams["fmin"] + self.fmax = hparams["fmax"] + self.device = device + + mel = librosa_mel_fn( + sr=self.sampling_rate, + n_fft=self.n_fft, + n_mels=self.num_mels, + fmin=self.fmin, + fmax=self.fmax, + ) + self.mel_basis = mel.float().to(self.device) + self.hann_window = torch.hann_window(self.win_size).to(self.device) + + def to(self, device, **kwagrs): + super().to(device=device, **kwagrs) + self.mel_basis = self.mel_basis.to(device) + self.hann_window = self.hann_window.to(device) + self.device = device + + def forward(self, y, center=False, complex=False): + if isinstance(y, np.ndarray): + y = torch.FloatTensor(y) + if len(y.shape) == 1: + y = y.unsqueeze(0) + y = y.clamp(min=-1.0, max=1.0).to(self.device) + + pad_length = math.ceil(y.shape[1] / self.hop_size) * self.hop_size - y.shape[1] + y = torch.nn.functional.pad( + y.unsqueeze(1), + [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2 + pad_length)], + mode="reflect", + ) + y = y.squeeze(1) + + spec = torch.stft( + y, + self.n_fft, + hop_length=self.hop_size, + win_length=self.win_size, + window=self.hann_window, + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + if not complex: + spec = torch.view_as_real(spec) + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) # [B, n_fft, T] + spec = torch.matmul(self.mel_basis, spec) + spec = spectral_normalize_torch(spec) + spec = spec.transpose(1, 2) # [B, T, n_fft] + else: + B, C, T, _ = spec.shape + spec = spec.transpose(1, 2) # [B, T, n_fft, 2] + return spec + + +class RosvotModel(nn.Module, SupportAudioInput, SupportsComponentDiscovery): + support_audio_input: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = [] + _encoder_modules: ClassVar[list[str]] = ["."] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["midi.net", "midi.pitch_decoder", "midi.cond_encoder"] + + def __init__( + self, + rosvot_ckpt: str | Path, + *, + config_path: str | Path = "", + pe: nn.Module | None = None, + the: float = 0.85, + verbose: bool = False, + rosvot_source_dir: str | Path | None = None, + ): + super().__init__() + self.verbose = verbose + ckpt = Path(rosvot_ckpt) + resolved_config = Path(config_path) if config_path else ckpt.with_name("config.yaml") + self.hparams = load_rosvot_config(resolved_config, hparams_str=f"note_bd_threshold={the}") + if verbose: + logger.info("ROSVOT config: %s", resolved_config) + + self.midi = MidiExtractor(self.hparams) + self.mel_net = MelNet(self.hparams) + self.pe = pe if pe is not None and self.hparams.get("use_pitch_embed", False) else None + self._checkpoint_path = str(ckpt) + self.load_checkpoint(str(ckpt), verbose=verbose) + self.eval() + + def load_checkpoint(self, checkpoint_path: str | None = None, *, verbose: bool = False) -> None: + load_model_ckpt(self.midi, checkpoint_path or self._checkpoint_path, verbose=verbose) + + @torch.no_grad() + def forward(self, wav: torch.Tensor, word_durs: list[float]) -> dict[str, Any]: + hparams = self.hparams + if wav.ndim == 1: + wav = wav.unsqueeze(0) + + mel_len = (wav.shape[-1] + hparams["hop_size"] - 1) // hparams["hop_size"] + min_word_dur = hparams.get("min_word_dur", 20) / 1000 + wd_raw = list(word_durs) + word_durs_filtered: list[float] = [] + for i, wd in enumerate(wd_raw): + if wd < min_word_dur: + if i == 0 and len(wd_raw) > 1: + wd_raw[i + 1] += wd + elif word_durs_filtered: + word_durs_filtered[-1] += wd + else: + word_durs_filtered.append(wd) + + mel2word = _align_word(word_durs_filtered, mel_len, hparams["hop_size"], hparams["audio_sample_rate"]) + if mel2word.size > 0 and mel2word[0] == 0: + mel2word = mel2word + 1 + + real_len = min(mel_len, int(np.sum(mel2word > 0))) + T = math.ceil(min(real_len, hparams["max_frames"]) / hparams["frames_multiple"]) * hparams["frames_multiple"] + device = wav.device + self.mel_net.to(device) + wav_t = pad_or_cut_xd(wav.float(), T * hparams["hop_size"], 1) + + pitch_coarse = uv_t = None + if self.pe is not None: + f0s, uvs = self.pe.get_pitch_batch( + wav_t, + sample_rate=hparams["audio_sample_rate"], + hop_size=hparams["hop_size"], + lengths=[real_len], + fmax=hparams["f0_max"], + fmin=hparams["f0_min"], + ) + f0_1d, uv_1d = norm_interp_f0(f0s[0][:T]) + f0_t = pad_or_cut_xd(torch.as_tensor(f0_1d, device=device, dtype=torch.float32), T, 0).unsqueeze(0) + uv_t = pad_or_cut_xd(torch.as_tensor(uv_1d, device=device, dtype=torch.float32), T, 0).long().unsqueeze(0) + pitch_coarse = f0_to_coarse(denorm_f0(f0_t, uv_t)).to(device) + + mel = pad_or_cut_xd(self.mel_net(wav_t)[0], T, dim=0).unsqueeze(0) + mel_nonpadding_mask = torch.zeros(1, T, device=device) + mel_nonpadding_mask[:, :real_len] = 1.0 + mel = (mel.transpose(1, 2) * mel_nonpadding_mask.unsqueeze(1)).transpose(1, 2) + mel_nonpadding = mel.abs().sum(-1) > 0 + + mel2word_t = pad_or_cut_xd(torch.as_tensor(mel2word, device=device, dtype=torch.long), T, 0) + word_bd = torch.zeros_like(mel2word_t) + word_bd[1:] = (mel2word_t[1:] != mel2word_t[:-1]).long() + word_bd[real_len:] = 0 + word_bd = word_bd.unsqueeze(0) + + outputs = self.midi( + mel=mel[:, :, : hparams.get("use_mel_bins", 80)], + word_bd=word_bd, + pitch=pitch_coarse, + uv=uv_t, + non_padding=mel_nonpadding, + ) + outputs["word_durs_filtered"] = word_durs_filtered + outputs["real_len"] = real_len + outputs["word_bd"] = word_bd + return outputs + + @staticmethod + def _load_wav(wav_src: str | np.ndarray, sample_rate: int, *, src_sample_rate: int | None = None) -> np.ndarray: + if isinstance(wav_src, str): + wav, _ = load_mono_audio(wav_src, target_sr=sample_rate) + return wav + wav = np.asarray(wav_src, dtype=np.float32) + if src_sample_rate is not None and src_sample_rate != sample_rate: + wav = resample_mono(wav, orig_sr=src_sample_rate, target_sr=sample_rate) + return wav + + @staticmethod + def _normalize_note2words(note2words: list[int]) -> list[int]: + if not note2words: + return [] + out = [note2words[0]] + for idx in range(1, len(note2words)): + out.append(max(note2words[idx], out[-1])) + return out + + @staticmethod + def _build_ep_types(note2words: list[int], align_words: list[str]) -> list[int]: + ep_types: list[int] = [] + prev = -1 + for i, w in zip(note2words, align_words): + ep_types.append(1 if w == "<SP>" else (2 if i != prev else 3)) + prev = i + return ep_types + + @torch.no_grad() + def transcribe( + self, + item: dict[str, Any], + *, + segment_info: dict[str, Any] | None = None, + verbose: bool = False, + ) -> dict[str, Any]: + if "word_durs" not in item: + raise ValueError('item must contain "word_durs" from lyric transcription') + + if item.get("wav_fn"): + wav = self._load_wav(item["wav_fn"], self.hparams["audio_sample_rate"]) + elif item.get("wav") is not None: + wav = self._load_wav( + item["wav"], + self.hparams["audio_sample_rate"], + src_sample_rate=item.get("sample_rate"), + ) + else: + raise ValueError('item must contain "wav_fn" or "wav"') + + device = next(self.parameters()).device + outputs = self(torch.from_numpy(wav).float().to(device), list(item["word_durs"])) + real_len = int(outputs["real_len"]) + word_durs_filtered = outputs["word_durs_filtered"] + word_bd = outputs["word_bd"] + item_name = item.get("item_name", "") + + note_lengths = outputs["note_lengths"].detach().cpu().numpy() + note_bd_pred = outputs["note_bd_pred"][0].detach().cpu().numpy()[:real_len] + note_pred = outputs["note_pred"][0].detach().cpu().numpy()[: note_lengths[0]] + + if note_pred.shape == (0,): + rosvot_out = {"item_name": item_name, "pitches": [], "note_durs": [], "note2words": None} + else: + note_itv_pred = boundary2Interval(note_bd_pred) + word_bd_for_reg = word_bd[0].detach().cpu().numpy()[:real_len] + hop = self.hparams["hop_size"] + sr = self.hparams["audio_sample_rate"] + + if self.hparams.get("infer_regulate_real_note_itv", True): + try: + note_itv_pred_secs, note2words = regulate_real_note_itv( + note_itv_pred, + note_bd_pred, + word_bd_for_reg, + np.array(word_durs_filtered), + hop, + sr, + ) + note_pred, note_itv_pred_secs, note2words = regulate_ill_slur( + note_pred, note_itv_pred_secs, note2words + ) + except Exception: + if verbose: + logger.exception("ROSVOT postprocess failed") + note_itv_pred_secs = note_itv_pred * hop / sr + note2words = None + else: + note_itv_pred_secs = note_itv_pred * hop / sr + note2words = None + + rosvot_out = { + "item_name": item_name, + "pitches": note_pred.tolist(), + "note_durs": [float(itv[1] - itv[0]) for itv in note_itv_pred_secs], + "note2words": note2words.tolist() if note2words is not None else None, + } + + note2words_raw = rosvot_out.get("note2words") or [] + align_words = [item["words"][idx - 1] for idx in note2words_raw if 0 < idx <= len(item["words"])] + ep_types = self._build_ep_types(self._normalize_note2words(note2words_raw), align_words) if align_words else [] + seg = segment_info or item + + return { + "item_name": seg.get("item_name", item_name), + "wav_fn": seg.get("wav_fn", item.get("wav_fn", "")), + "origin_wav_fn": seg.get("origin_wav_fn", item.get("origin_wav_fn", "")), + "start_time_ms": seg.get("start_time_ms", ""), + "end_time_ms": seg.get("end_time_ms", ""), + "language": item.get("language", ""), + "note_text": align_words, + "note_dur": rosvot_out.get("note_durs", []), + "note_type": ep_types, + "note_pitch": rosvot_out.get("pitches", []), + } diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py new file mode 100644 index 00000000000..0140182ab5c --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/segmenter.py @@ -0,0 +1,267 @@ +"""Rule-based vocal segmentation from F0 contour (no neural net).""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import ClassVar + +import numpy as np +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class VocalSegmentConfig: + hop_ms: int = 20 + smooth_ms: int = 200 + start_ms: int = 120 + end_ms: int = 200 + prepad_ms: int = 80 + postpad_ms: int = 120 + min_len_ms: int = 1000 + max_len_ms: int = 20000 + short_seg_merge_gap_ms: int = 8000 + small_gap_ms: int = 500 + lookback_ms: int = 200 + lookahead_ms: int = 100 + + +def _moving_average(x: np.ndarray, win: int) -> np.ndarray: + if win <= 1: + return x + kernel = np.ones(win, dtype=np.float32) / float(win) + return np.convolve(x, kernel, mode="same") + + +def _merge_short_segments( + segments_ms: list[list[int]], + *, + min_len_ms: int, + max_len_ms: int, + short_seg_merge_gap_ms: int, + small_gap_ms: int, +) -> list[list[int]]: + if not segments_ms: + return [] + + merged: list[list[int]] = [] + cur_start, cur_end = segments_ms[0] + for next_start, next_end in segments_ms[1:]: + cur_len = cur_end - cur_start + gap_ms = next_start - cur_end + merged_len = next_end - cur_start + should_merge = (cur_len < min_len_ms and gap_ms < short_seg_merge_gap_ms) or ( + gap_ms < small_gap_ms and merged_len < max_len_ms + ) + if should_merge: + cur_end = next_end + continue + if (cur_end - cur_start) >= min_len_ms: + merged.append([cur_start, cur_end]) + cur_start, cur_end = next_start, next_end + + if (cur_end - cur_start) >= min_len_ms: + merged.append([cur_start, cur_end]) + + if not merged: + return segments_ms + + return merged + + +def _voiced_to_segments( + voiced: np.ndarray, + *, + hop_ms: int, + smooth_ms: int, + start_ms: int, + end_ms: int, + prepad_ms: int, + postpad_ms: int, + max_len_ms: int, +) -> list[list[int]]: + smooth_frames = max(1, int(round(smooth_ms / hop_ms))) + active = _moving_average(voiced.astype(np.float32), smooth_frames) >= 0.5 + + segments: list[list[int]] = [] + start_idx = None + start_frames = max(1, int(round(start_ms / hop_ms))) + end_frames = max(1, int(round(end_ms / hop_ms))) + prepad_frames = max(0, int(round(prepad_ms / hop_ms))) + postpad_frames = max(0, int(round(postpad_ms / hop_ms))) + active_count = 0 + inactive_count = 0 + + for i, flag in enumerate(active): + if flag: + active_count += 1 + inactive_count = 0 + else: + inactive_count += 1 + active_count = 0 + + if start_idx is None: + if active_count >= start_frames: + start_idx = max(0, i - start_frames + 1 - prepad_frames) + else: + if inactive_count >= end_frames: + end_idx = min(len(active) - 1, i - end_frames + 1 + postpad_frames) + start_ms_val = start_idx * hop_ms + end_ms_val = end_idx * hop_ms + hop_ms + if end_ms_val > start_ms_val: + segments.append([int(start_ms_val), int(end_ms_val)]) + start_idx = None + + if start_idx is not None: + start_ms_val = start_idx * hop_ms + end_idx = min(len(active) - 1, len(active) - 1 + postpad_frames) + end_ms_val = end_idx * hop_ms + hop_ms + if end_ms_val > start_ms_val: + segments.append([int(start_ms_val), int(end_ms_val)]) + + def _split_segment(seg: list[int]) -> list[list[int]]: + start_ms_val, end_ms_val = seg + start_frame = int(start_ms_val // hop_ms) + end_frame = int((end_ms_val - 1) // hop_ms) + end_frame = max(start_frame, min(end_frame, len(active) - 1)) + + best_start = None + best_len = 0 + cur_start = None + cur_len = 0 + for idx in range(start_frame, end_frame + 1): + if not active[idx]: + if cur_start is None: + cur_start = idx + cur_len = 1 + else: + cur_len += 1 + else: + if cur_start is not None and cur_len > best_len: + best_start, best_len = cur_start, cur_len + cur_start = None + cur_len = 0 + if cur_start is not None and cur_len > best_len: + best_start, best_len = cur_start, cur_len + + if best_start is None: + split_frame = (start_frame + end_frame) // 2 + else: + split_frame = best_start + best_len // 2 + + split_ms = split_frame * hop_ms + if split_ms <= start_ms_val: + split_ms = start_ms_val + hop_ms + if split_ms >= end_ms_val: + split_ms = end_ms_val - hop_ms + if split_ms <= start_ms_val or split_ms >= end_ms_val: + return [seg] + return [[start_ms_val, int(split_ms)], [int(split_ms), end_ms_val]] + + queue = segments[:] + segments = [] + while queue: + seg = queue.pop(0) + if (seg[1] - seg[0]) <= max_len_ms: + segments.append(seg) + continue + parts = _split_segment(seg) + if len(parts) == 1: + segments.append(seg) + else: + queue = parts + queue + return segments + + +class VocalSegmenter(nn.Module, SupportAudioInput, SupportsComponentDiscovery): + support_audio_input: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = [] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = ["."] + _layerwise_offload_blocks_attrs: ClassVar[list[str]] = [] + + def __init__(self, config: VocalSegmentConfig | None = None, *, verbose: bool = False): + super().__init__() + self.config = config or VocalSegmentConfig() + self.verbose = verbose + + def forward( + self, + audio: np.ndarray, + sample_rate: int, + f0: np.ndarray, + *, + base_name: str = "vocal", + origin_wav_fn: str = "", + verbose: bool | None = None, + ) -> list[dict]: + verbose = self.verbose if verbose is None else verbose + if verbose: + logger.info("[vocal detection] segmenting %s", base_name) + t0 = time.time() + + cfg = self.config + y = np.asarray(audio, dtype=np.float32).squeeze() + sr = int(sample_rate) + + segments_ms = _voiced_to_segments( + f0 > 0, + hop_ms=cfg.hop_ms, + smooth_ms=cfg.smooth_ms, + start_ms=cfg.start_ms, + end_ms=cfg.end_ms, + prepad_ms=cfg.prepad_ms, + postpad_ms=cfg.postpad_ms, + max_len_ms=cfg.max_len_ms, + ) + segments_ms = _merge_short_segments( + segments_ms, + min_len_ms=cfg.min_len_ms, + max_len_ms=cfg.max_len_ms, + short_seg_merge_gap_ms=cfg.short_seg_merge_gap_ms, + small_gap_ms=cfg.small_gap_ms, + ) + + adjusted_segments: list[list[int]] = [] + prev_end = 0 + for start_ms, end_ms in segments_ms: + start_ms = max(0, start_ms - cfg.lookback_ms) + end_ms = min(end_ms + cfg.lookahead_ms, int(y.shape[0] / sr * 1000)) + if start_ms < prev_end and adjusted_segments: + adjusted_segments[-1][1] = start_ms + adjusted_segments.append([start_ms, end_ms]) + prev_end = end_ms + + segment_infos = [] + for idx, (start_ms, end_ms) in enumerate(adjusted_segments): + if end_ms - start_ms > cfg.max_len_ms: + start_ms = end_ms - cfg.max_len_ms + key = f"{base_name}_{idx}" + start_sample = int(start_ms * sr / 1000) + end_sample = int(end_ms * sr / 1000) + segment_infos.append( + { + "item_name": key, + "wav": y[start_sample:end_sample], + "sample_rate": sr, + "wav_fn": key, + "start_time_ms": int(start_sample * 1000 / sr), + "end_time_ms": int(end_sample * 1000 / sr), + "origin_wav_fn": origin_wav_fn, + "duration": int((end_sample - start_sample) * 1000 / sr), + } + ) + + if verbose: + logger.info( + "[vocal detection] done n_segments=%s time=%.3fs", + len(segment_infos), + time.time() - t0, + ) + return segment_infos diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py new file mode 100644 index 00000000000..e29b52a2963 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/stack.py @@ -0,0 +1,145 @@ +"""Lazy-loaded SoulX preprocess model tree.""" + +import os +import tempfile +from pathlib import Path +from typing import ClassVar + +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery + +from .asr import LyricModel +from .mel_grid_f0 import extract_f0_file +from .rmvpe import RMVPE +from .rosvot import RosvotModel +from .segmenter import VocalSegmenter +from .utils import resample_mono +from .vocal_sep import MelBandRoformerSeparator + +logger = init_logger(__name__) + + +class SoulXPreprocessStack(nn.Module, SupportsComponentDiscovery): + """Unified nn.Module tree for SoulX-Singer audio preprocess.""" + + _dit_modules: ClassVar[list[str]] = [] + _encoder_modules: ClassVar[list[str]] = ["vocal_sep", "lyric", "rosvot"] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = ["rmvpe"] + + def __init__( + self, + weights: dict[str, str], + device: str, + *, + target_sr: int = 24000, + hop_size: int = 480, + verbose: bool = False, + ) -> None: + super().__init__() + self.weights = weights + self.device_str = device + self.target_sr = target_sr + self.hop_size = hop_size + self.verbose = verbose + self._rmvpe: nn.Module | None = None + self._vocal_sep: nn.Module | None = None + self._segmenter: nn.Module | None = None + self._lyric: nn.Module | None = None + self._rosvot: nn.Module | None = None + + @property + def device(self) -> torch.device: + return torch.device(self.device_str) + + def ensure_rmvpe(self): + if self._rmvpe is None: + if self.verbose: + logger.info("[rmvpe] loading shared pitch model") + self._rmvpe = RMVPE(self.weights["rmvpe"], device=self.device_str) + self.add_module("rmvpe", self._rmvpe) + return self._rmvpe + + def ensure_vocal_sep(self): + if self._vocal_sep is None: + if self.verbose: + logger.info("[vocal extraction] loading model") + self._vocal_sep = MelBandRoformerSeparator( + self.weights["sep_config"], + self.weights["sep_ckpt"], + use_half=False, + ).to(self.device) + self.add_module("vocal_sep", self._vocal_sep) + if self.verbose: + logger.info("[vocal extraction] ready on %s, sr=%s", self.device_str, self._vocal_sep.sample_rate) + return self._vocal_sep + + def ensure_segmenter(self): + if self._segmenter is None: + self._segmenter = VocalSegmenter(verbose=self.verbose) + self.add_module("segmenter", self._segmenter) + return self._segmenter + + def ensure_lyric(self): + if self._lyric is None: + if self.verbose: + logger.info("[lyric transcription] loading ASR") + self._lyric = LyricModel( + self.weights["asr_zh"], + self.weights["asr_en"], + device=self.device_str, + target_sr=self.target_sr, + hop_size=self.hop_size, + ) + self.add_module("lyric", self._lyric) + return self._lyric + + def ensure_rosvot(self): + if self._rosvot is None: + if self.verbose: + logger.info("[note transcription] loading ROSVOT") + rosvot_src = os.environ.get("ROSVOT_SOURCE_DIR") + self._rosvot = RosvotModel( + self.weights["rosvot"], + pe=self.ensure_rmvpe(), + verbose=self.verbose, + rosvot_source_dir=rosvot_src, + ).to(self.device) + self.add_module("rosvot", self._rosvot) + return self._rosvot + + def extract_f0(self, vocal, sample_rate: int, *, f0_path: str | None = None) -> np.ndarray: + """Extract F0 on the mel grid; ndarray inputs are spooled to a temp wav.""" + rmvpe = self.ensure_rmvpe() + kwargs = dict( + target_sr=self.target_sr, + hop_size=self.hop_size, + f0_path=f0_path, + verbose=self.verbose, + ) + if isinstance(vocal, str): + return extract_f0_file(rmvpe, vocal, **kwargs) + tmp_path: str | None = None + try: + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + tmp_path = tmp.name + sf.write(tmp_path, np.asarray(vocal, dtype=np.float32).squeeze(), int(sample_rate)) + return extract_f0_file(rmvpe, tmp_path, **kwargs) + finally: + if tmp_path is not None: + Path(tmp_path).unlink(missing_ok=True) + + def extract_vocal(self, audio: str | tuple[np.ndarray, int]) -> tuple[np.ndarray, int]: + sep = self.ensure_vocal_sep() + if isinstance(audio, tuple): + mix, sr = audio + mix = np.asarray(mix, dtype=np.float32).squeeze() + if sr != sep.sample_rate: + mix = resample_mono(mix, orig_sr=sr, target_sr=sep.sample_rate) + return sep.separate_mono(mix, self.device), sep.sample_rate + return sep.separate_path(str(audio), self.device) diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py new file mode 100644 index 00000000000..f135fd55388 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/utils.py @@ -0,0 +1,260 @@ +"""Shared helpers for SoulX preprocess (I/O, config, checkpoints, pitch).""" + +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf +import torch +import torch.nn as nn +import torch.nn.functional as F +from omegaconf import OmegaConf +from vllm.logger import init_logger +from vllm.multimodal.audio import AudioResampler +from vllm.multimodal.media.audio import load_audio + +logger = init_logger(__name__) + + +# --- audio I/O --- + + +def resample_mono( + audio: np.ndarray, + *, + orig_sr: int, + target_sr: int, +) -> np.ndarray: + """Resample mono float waveform to ``target_sr``.""" + audio = np.asarray(audio, dtype=np.float32).squeeze() + if orig_sr == target_sr: + return audio + resampler = AudioResampler(target_sr=float(target_sr)) + return np.asarray( + resampler.resample(audio, orig_sr=float(orig_sr)), + dtype=np.float32, + ) + + +def get_audio_duration(path: str) -> float: + """Return audio duration in seconds for a file on disk.""" + return float(sf.info(path).duration) + + +def load_mono_audio( + source: str | tuple[np.ndarray, int], + *, + target_sr: int | None = None, +) -> tuple[np.ndarray, int]: + if isinstance(source, tuple): + audio, sr = source + audio = np.asarray(audio, dtype=np.float32).squeeze() + if target_sr is not None and sr != target_sr: + audio = resample_mono(audio, orig_sr=sr, target_sr=target_sr) + sr = target_sr + return audio, sr + + if target_sr is None: + audio, sr = load_audio(str(source), sr=None, mono=True) + else: + audio, sr = load_audio(str(source), sr=target_sr, mono=True) + return np.asarray(audio, dtype=np.float32).squeeze(), int(sr) + + +# --- tensor padding --- + + +def pad_or_cut_xd(values: torch.Tensor, tgt_len: int, dim: int = -1, pad_value=0) -> torch.Tensor: + ndim = values.ndim + if ndim == 1: + if values.shape[0] < tgt_len: + return F.pad(values, [0, tgt_len - values.shape[0]], value=pad_value) + return values[:tgt_len] + if ndim == 2: + if dim in (0, -2): + src_len = values.shape[0] + if src_len < tgt_len: + return F.pad(values, [0, 0, 0, tgt_len - src_len], value=pad_value) + return values[:tgt_len] + src_len = values.shape[1] + if src_len < tgt_len: + return F.pad(values, [0, tgt_len - src_len], value=pad_value) + return values[:, :tgt_len] + raise NotImplementedError(f"pad_or_cut_xd supports 1D/2D tensors, got ndim={ndim}") + + +# --- YAML config (OmegaConf) --- + + +def _apply_hparams_str(config: dict[str, Any], hparams_str: str) -> None: + if not hparams_str: + return + for item in hparams_str.split(","): + key, raw_value = item.split("=", maxsplit=1) + raw_value = raw_value.strip().strip("'\"") + node: dict[str, Any] = config + parts = key.split(".") + for part in parts[:-1]: + node = node.setdefault(part, {}) + leaf = parts[-1] + current = node.get(leaf) + if raw_value in {"True", "False"}: + node[leaf] = raw_value == "True" + elif isinstance(current, bool): + node[leaf] = raw_value.lower() == "true" + elif isinstance(current, int): + node[leaf] = int(raw_value) + elif isinstance(current, float): + node[leaf] = float(raw_value) + else: + node[leaf] = type(current)(raw_value) if current is not None else raw_value + + +def load_yaml_config( + config: str | Path, + *, + hparams_str: str = "", +) -> dict[str, Any]: + """Load YAML with optional ``base_config`` chain; returns a plain dict.""" + config_path = Path(config).expanduser() + if not config_path.is_file(): + raise FileNotFoundError(f"Config not found: {config_path}") + + merged = OmegaConf.create({}) + pending = [config_path.resolve()] + loaded: set[str] = set() + + while pending: + path = pending.pop(0) + path_key = str(path) + if path_key in loaded: + continue + loaded.add(path_key) + + current = OmegaConf.load(path_key) + base_configs = current.get("base_config") + if base_configs is not None: + entries = base_configs if OmegaConf.is_list(base_configs) else [base_configs] + for entry in entries: + base_path = Path(entry) + if not base_path.is_absolute(): + base_path = (path.parent / entry).resolve() + if not base_path.is_file(): + logger.warning("Skipping missing base config: %s", base_path) + continue + pending.insert(0, base_path) + merged = OmegaConf.merge(current, merged) + + result = OmegaConf.to_container(merged, resolve=True) + if not isinstance(result, dict): + raise TypeError(f"Expected dict config, got {type(result)}") + _apply_hparams_str(result, hparams_str) + return result + + +def load_rosvot_config(config: str | Path, *, hparams_str: str = "") -> dict[str, Any]: + return load_yaml_config(config, hparams_str=hparams_str) + + +# --- checkpoint loading --- + + +def load_model_ckpt( + model: nn.Module, + checkpoint_path: str, + *, + model_name: str = "model", + verbose: bool = True, +) -> None: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint: + raise ValueError(f"Unexpected ROSVOT checkpoint format: {checkpoint_path}") + + state_dict = checkpoint["state_dict"] + if state_dict and "." in next(iter(state_dict)): + prefix = f"{model_name}." + state_dict = {k[len(prefix) :]: v for k, v in state_dict.items() if k.startswith(prefix)} + elif model_name in state_dict and isinstance(state_dict[model_name], dict): + state_dict = state_dict[model_name] + + model.load_state_dict(state_dict, strict=True) + if verbose: + logger.info("Loaded ROSVOT weights from %s", checkpoint_path) + + +# --- pitch helpers (ROSVOT / RMVPE) --- + + +def f0_to_coarse(f0, f0_bin=256, f0_max=900.0, f0_min=50.0): + f0_mel_min = 1127 * np.log(1 + f0_min / 700) + f0_mel_max = 1127 * np.log(1 + f0_max / 700) + is_torch = isinstance(f0, torch.Tensor) + f0_mel = 1127 * (1 + f0 / 700).log() if is_torch else 1127 * np.log(1 + f0 / 700) + f0_mel[f0_mel > 0] = (f0_mel[f0_mel > 0] - f0_mel_min) * (f0_bin - 2) / (f0_mel_max - f0_mel_min) + 1 + f0_mel[f0_mel <= 1] = 1 + f0_mel[f0_mel > f0_bin - 1] = f0_bin - 1 + return (f0_mel + 0.5).long() if is_torch else np.rint(f0_mel).astype(int) + + +def norm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100): + is_torch = isinstance(f0, torch.Tensor) + if pitch_norm == "standard": + f0 = (f0 - f0_mean) / f0_std + if pitch_norm == "log": + f0 = torch.log2(f0 + 1e-8) if is_torch else np.log2(f0 + 1e-8) + if uv is not None: + f0[uv > 0] = 0 + return f0 + + +def norm_interp_f0(f0, pitch_norm="log", f0_mean=None, f0_std=None): + is_torch = isinstance(f0, torch.Tensor) + if is_torch: + device = f0.device + f0 = f0.detach().cpu().numpy() + uv = f0 == 0 + f0 = norm_f0(f0, uv, pitch_norm, f0_mean, f0_std) + if sum(uv) == len(f0): + f0[uv] = 0 + elif sum(uv) > 0: + f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv]) + if is_torch: + uv_t = torch.FloatTensor(uv) + f0_t = torch.FloatTensor(f0) + return f0_t.to(device), uv_t.to(device) + return f0, uv + + +def denorm_f0(f0, uv, pitch_norm="log", f0_mean=400, f0_std=100, pitch_padding=None, min=50, max=900): + is_torch = isinstance(f0, torch.Tensor) + if pitch_norm == "standard": + f0 = f0 * f0_std + f0_mean + if pitch_norm == "log": + f0 = 2**f0 + f0 = f0.clamp(min=min, max=max) if is_torch else np.clip(f0, a_min=min, a_max=max) + if uv is not None: + f0[uv > 0] = 0 + if pitch_padding is not None: + f0[pitch_padding] = 0 + return f0 + + +def boundary2Interval(bd): + is_torch = isinstance(bd, torch.Tensor) + if is_torch: + device = bd.device + bd = bd.detach().cpu().numpy() + ret = np.zeros(shape=(bd.sum() + 1, 2), dtype=int) + ret_idx = 0 + ret[0, 0] = 0 + for i, u in enumerate(bd): + if i == 0: + continue + if u == 1: + ret[ret_idx, 1] = i + ret[ret_idx + 1, 0] = i + ret_idx += 1 + ret[-1, 1] = bd.shape[0] - 1 + if is_torch: + return torch.LongTensor(ret).to(device) + return ret diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py new file mode 100644 index 00000000000..b62902d3800 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/vocal_sep.py @@ -0,0 +1,207 @@ +"""BS-RoFormer vocal separation adapter (external pip package + SoulX chunking).""" + +from __future__ import annotations + +from pathlib import Path +from typing import ClassVar + +import numpy as np +import torch +import torch.nn as nn +import yaml +from omegaconf import DictConfig, OmegaConf +from vllm.logger import init_logger +from vllm.multimodal.media.audio import load_audio + +from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery +from vllm_omni.diffusion.offloader.layerwise_backend import LayerWiseOffloadBackend + +logger = init_logger(__name__) + + +class MelBandRoformerSeparator(nn.Module, SupportAudioInput, SupportsComponentDiscovery): + support_audio_input: ClassVar[bool] = True + _dit_modules: ClassVar[list[str]] = [] + _encoder_modules: ClassVar[list[str]] = ["roformer"] + _vae_modules: ClassVar[list[str]] = [] + _resident_modules: ClassVar[list[str]] = [] + _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["roformer.layers"] + + def __init__( + self, + config_path: str | Path, + checkpoint_path: str | Path | None = None, + *, + chunk_length_sec: float = 5, + use_half: bool = True, + ): + super().__init__() + try: + from bs_roformer import MelBandRoformer as ExtModel + except ImportError as e: + raise ImportError( + 'Install BS-RoFormer: pip install "vllm-omni[soulx-preprocess]" or pip install BS-RoFormer' + ) from e + + cfg = self._load_config(Path(config_path)) + model_cfg = self._normalize_cfg(OmegaConf.to_container(cfg.model, resolve=True)) + self.roformer = ExtModel(**model_cfg) + self.config = cfg + self.sample_rate = int(cfg.audio.sample_rate) + + LayerWiseOffloadBackend.set_blocks_attr_names(self.roformer, ["layers"]) + dtype = torch.float16 if use_half else torch.float32 + self.roformer = self.roformer.to(dtype=dtype) + + chunk = int(chunk_length_sec * self.sample_rate) + self.config.inference = getattr(self.config, "inference", {}) + inf = self.config.inference + inf.chunk_size = chunk + inf.num_overlap = inf.get("num_overlap", 4) + inf.use_amp = False + inf.normalize = inf.get("normalize", True) + + fade = chunk // 10 + w = torch.ones(chunk) + w[:fade] = torch.linspace(0, 1, fade) + w[-fade:] = torch.linspace(1, 0, fade) + self.register_buffer("_window", w.to(dtype=dtype), persistent=False) + + if checkpoint_path: + self.load_checkpoint(str(checkpoint_path), use_half=use_half) + else: + self.roformer.eval() + + @staticmethod + def _normalize_cfg(cfg: dict) -> dict: + tuple_keys = {"multi_stft_resolutions_window_sizes", "freqs_per_bands", "multi_stft_window_sizes"} + + def _fix(v): + if isinstance(v, list): + return tuple(_fix(x) for x in v) + if isinstance(v, dict): + return {k: _fix(val) for k, val in v.items()} + return v + + for k in tuple_keys: + if k in cfg: + cfg[k] = _fix(cfg[k]) + return cfg + + @staticmethod + def _load_config(path: Path) -> DictConfig: + class Loader(yaml.SafeLoader): + pass + + Loader.add_constructor( + "tag:yaml.org,2002:python/tuple", + lambda loader, node: tuple(loader.construct_sequence(node)), + ) + with open(path) as f: + return OmegaConf.create(yaml.load(f, Loader=Loader)) + + def load_checkpoint(self, path: str, *, use_half: bool = True) -> None: + self.roformer.load_state_dict(torch.load(path, map_location="cpu"), strict=False) + if use_half: + self.roformer.half() + self.roformer.eval() + + @property + def chunk_size(self) -> int: + return int(self.config.inference.chunk_size) + + @property + def num_overlap(self) -> int: + return int(self.config.inference.num_overlap) + + @property + def use_amp(self) -> bool: + return bool(self.config.inference.use_amp) + + @property + def normalize_input(self) -> bool: + return bool(self.config.inference.normalize) + + @property + def num_channels(self) -> int: + return int(self.config.audio.get("num_channels", 1)) + + def _prepare_mix(self, mix: np.ndarray) -> tuple[np.ndarray, tuple[float, float] | None]: + if mix.ndim == 1: + mix = mix[np.newaxis, :] + if self.num_channels == 2 and mix.shape[0] == 1: + mix = np.repeat(mix, 2, axis=0) + if not self.normalize_input: + return mix, None + mono = mix.mean(0) + return (mix - mono.mean()) / mono.std(), (float(mono.mean()), float(mono.std())) + + @staticmethod + def _select_vocal(recon: torch.Tensor) -> torch.Tensor: + if recon.ndim == 4: + return recon[0, 0] + if recon.ndim == 3: + return recon[0] + if recon.ndim == 2: + return recon + raise ValueError(recon.shape) + + @torch.inference_mode() + def forward(self, mix: torch.Tensor) -> torch.Tensor: + if mix.ndim == 1: + mix = mix.unsqueeze(0) + if self.num_channels == 2 and mix.shape[0] == 1: + mix = mix.repeat(2, 1) + + model_dtype = next(self.roformer.parameters()).dtype + compute_dtype = torch.float32 if model_dtype == torch.float16 else model_dtype + if mix.dtype != compute_dtype: + mix = mix.to(dtype=compute_dtype) + + chunk_size = self.chunk_size + step = chunk_size // self.num_overlap + fade = chunk_size // 10 + border = chunk_size - step + n_samples = mix.shape[-1] + window = self._window.to(mix.device) + + if n_samples > 2 * border and border > 0: + mix = nn.functional.pad(mix, (border, border), mode="reflect") + + out = torch.zeros_like(mix) + weight = torch.zeros(mix.shape[-1], device=mix.device, dtype=mix.dtype) + pos = 0 + with torch.cuda.amp.autocast(enabled=self.use_amp): + while pos < mix.shape[-1]: + part = mix[:, pos : pos + chunk_size] + seg_len = part.shape[-1] + pad_mode = "reflect" if seg_len > chunk_size // 2 else "constant" + part = nn.functional.pad(part, (0, chunk_size - seg_len), mode=pad_mode, value=0) + stem = self._select_vocal(self.roformer(part.unsqueeze(0))) + win = window.clone() + if pos == 0: + win[:fade] = 1 + elif pos + step >= n_samples: + win[-fade:] = 1 + out[:, pos : pos + seg_len] += stem[:, :seg_len] * win[:seg_len] + weight[pos : pos + seg_len] += win[:seg_len] + pos += step + + out = out / weight + if n_samples > 2 * border and border > 0: + out = out[..., border:-border] + return out + + @torch.inference_mode() + def separate_mono(self, mix: np.ndarray, device: torch.device) -> np.ndarray: + m, norm = self._prepare_mix(mix) + target_dtype = next(self.roformer.parameters()).dtype + v = self.forward(torch.as_tensor(m, dtype=target_dtype, device=device)) + if norm: + v = v * norm[1] + norm[0] + v = v.mean(0) if v.ndim == 2 else v + return v.float().cpu().numpy().astype(np.float32) + + def separate_path(self, path: str, device: torch.device) -> tuple[np.ndarray, int]: + mix, _ = load_audio(path, sr=self.sample_rate, mono=False) + return self.separate_mono(np.asarray(mix, dtype=np.float32), device), self.sample_rate diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py new file mode 100644 index 00000000000..ca3ef852edb --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/vocoder.py @@ -0,0 +1,983 @@ +import accelerate +import numpy as np +import scipy +import torch +from omegaconf import DictConfig +from torch import nn, view_as_complex, view_as_real +from torch.nn.utils import remove_weight_norm, weight_norm +from torchaudio.functional.functional import _hz_to_mel, _mel_to_hz + + +def _deep_update_dict(base: dict, override: dict | None) -> dict: + """Recursively merge `override` into `base` (dict only).""" + if not override: + return base + for k, v in override.items(): + if isinstance(v, dict) and isinstance(base.get(k), dict): + base[k] = _deep_update_dict(base[k], v) + else: + base[k] = v + return base + + +def _get_vocoder_default_cfg() -> dict: + # Defaults keep current behavior. + return { + "preprocess": { + "hop_size": 480, + "sample_rate": 24000, + "max_length": 36000, + "n_fft": 1920, + "num_mels": 128, + "win_size": 1920, + "fmin": 0, + "fmax": 12000, + "mel_var": 8.14, + "mel_mean": -4.92, + "load_phone": False, + "load_chromagram": False, + }, + "model": { + "vocos": { + "input_channels": 128, + "dim": 1024, + "intermediate_dim": 4096, + "num_layers": 30, + "n_fft": 1920, + "hop_size": 480, + "padding": "same", + }, + "period_gan": { + "max_downsample_channels": 1024, + "channels": 64, + "channel_increasing_factor": 2, + }, + "spec_gan": { + "stft_params": { + "fft_sizes": [128, 256, 512, 1024, 2048], + "hop_sizes": [32, 64, 128, 256, 512], + "win_lengths": [128, 256, 512, 1024, 2048], + "window": "hann_window", + }, + "in_channels": 1, + "out_channels": 1, + "channels": 64, + "kernel_sizes": [5, 3], + "max_downsample_channels": 1024, + "down_scales": [2, 2, 2], + "use_weight_norm": True, + "use_complex": False, + }, + }, + "loss": { + "mel_loss": {"sample_rate": 24000}, + "disc_loss_weight": 1.0, + "mel_loss_weight": 10.0, + "adv_loss_weight": 2.0, + "fm_loss_weight": 2.0, + "spec_fm_loss_weight": 1.0, + }, + } + + +def safe_log(x: torch.Tensor, clip_val: float = 1e-7) -> torch.Tensor: + """ + Computes the element-wise logarithm of the input tensor with clipping to avoid near-zero values. + + Args: + x (Tensor): Input tensor. + clip_val (float, optional): Minimum value to clip the input tensor. Defaults to 1e-7. + + Returns: + Tensor: Element-wise logarithm of the input tensor with clipping applied. + """ + return torch.log(torch.clip(x, min=clip_val)) + + +def symlog(x: torch.Tensor) -> torch.Tensor: + return torch.sign(x) * torch.log1p(x.abs()) + + +def symexp(x: torch.Tensor) -> torch.Tensor: + return torch.sign(x) * (torch.exp(x.abs()) - 1) + + +class STFT(nn.Module): + def __init__( + self, + n_fft: int, + hop_length: int, + win_length: int, + center=True, + ): + super().__init__() + self.center = center + 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) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # x: (B, T * hop_length) + + if not self.center: + pad = self.win_length - self.hop_length + x = torch.nn.functional.pad(x, (pad // 2, pad // 2), mode="reflect") + + stft_spec = torch.stft( + x, + self.n_fft, + hop_length=self.hop_length, + win_length=self.win_length, + window=self.window, + center=self.center, + return_complex=False, + ) # (B, n_fft // 2 + 1, T, 2) + + rea = stft_spec[:, :, :, 0] # (B, n_fft // 2 + 1, T, 2) + imag = stft_spec[:, :, :, 1] # (B, n_fft // 2 + 1, T, 2) + + log_mag = torch.log( + torch.abs(torch.sqrt(torch.pow(rea, 2) + torch.pow(imag, 2))) + 1e-5 + ) # (B, n_fft // 2 + 1, T) + phase = torch.atan2(imag, rea) # (B, n_fft // 2 + 1, T) + + return log_mag, phase + + +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) + + def forward(self, spec: torch.Tensor) -> torch.Tensor: + """ + 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. + + Returns: + Tensor: Reconstructed time-domain signal of shape (B, L), where L is the length of the output signal. + """ + 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'.") + + assert spec.dim() == 3, "Expected a 3D tensor as input" + B, N, 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, pad:-pad] + + # 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()[pad:-pad] + + # Normalize + assert (window_envelope > 1e-11).all() + y = y / window_envelope + + return y + + +class MDCT(nn.Module): + """ + Modified Discrete Cosine Transform (MDCT) module. + + Args: + frame_len (int): Length of the MDCT frame. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + """ + + def __init__(self, frame_len: int, padding: str = "same"): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.frame_len = frame_len + N = frame_len // 2 + n0 = (N + 1) / 2 + window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() + self.register_buffer("window", window) + + pre_twiddle = torch.exp(-1j * torch.pi * torch.arange(frame_len) / frame_len) + post_twiddle = torch.exp(-1j * torch.pi * n0 * (torch.arange(N) + 0.5) / N) + # view_as_real: NCCL Backend does not support ComplexFloat data type + # https://github.com/pytorch/pytorch/issues/71613 + self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) + self.register_buffer("post_twiddle", view_as_real(post_twiddle)) + + def forward(self, audio: torch.Tensor) -> torch.Tensor: + """ + Apply the Modified Discrete Cosine Transform (MDCT) to the input audio. + + Args: + audio (Tensor): Input audio waveform of shape (B, T), where B is the batch size + and T is the length of the audio. + + Returns: + Tensor: MDCT coefficients of shape (B, L, N), where L is the number of output frames + and N is the number of frequency bins. + """ + if self.padding == "center": + audio = torch.nn.functional.pad(audio, (self.frame_len // 2, self.frame_len // 2)) + elif self.padding == "same": + # hop_length is 1/2 frame_len + audio = torch.nn.functional.pad(audio, (self.frame_len // 4, self.frame_len // 4)) + else: + raise ValueError("Padding must be 'center' or 'same'.") + + x = audio.unfold(-1, self.frame_len, self.frame_len // 2) + N = self.frame_len // 2 + x = x * self.window.expand(x.shape) + x = torch.fft.fft(x * view_as_complex(self.pre_twiddle).expand(x.shape), dim=-1)[..., :N] + res = x * view_as_complex(self.post_twiddle).expand(x.shape) * np.sqrt(1 / N) + return torch.real(res) * np.sqrt(2) + + +class IMDCT(nn.Module): + """ + Inverse Modified Discrete Cosine Transform (IMDCT) module. + + Args: + frame_len (int): Length of the MDCT frame. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + """ + + def __init__(self, frame_len: int, padding: str = "same"): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.frame_len = frame_len + N = frame_len // 2 + n0 = (N + 1) / 2 + window = torch.from_numpy(scipy.signal.cosine(frame_len)).float() + self.register_buffer("window", window) + + pre_twiddle = torch.exp(1j * torch.pi * n0 * torch.arange(N * 2) / N) + post_twiddle = torch.exp(1j * torch.pi * (torch.arange(N * 2) + n0) / (N * 2)) + self.register_buffer("pre_twiddle", view_as_real(pre_twiddle)) + self.register_buffer("post_twiddle", view_as_real(post_twiddle)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Apply the Inverse Modified Discrete Cosine Transform (IMDCT) to the input MDCT coefficients. + + Args: + x (Tensor): Input MDCT coefficients of shape (B, L, N), where B is the batch size, + L is the number of frames, and N is the number of frequency bins. + + Returns: + Tensor: Reconstructed audio waveform of shape (B, T), where T is the length of the audio. + """ + B, L, N = x.shape + Y = torch.zeros((B, L, N * 2), dtype=x.dtype, device=x.device) + Y[..., :N] = x + Y[..., N:] = -1 * torch.conj(torch.flip(x, dims=(-1,))) + y = torch.fft.ifft(Y * view_as_complex(self.pre_twiddle).expand(Y.shape), dim=-1) + y = torch.real(y * view_as_complex(self.post_twiddle).expand(y.shape)) * np.sqrt(N) * np.sqrt(2) + result = y * self.window.expand(y.shape) + output_size = (1, (L + 1) * N) + audio = torch.nn.functional.fold( + result.transpose(1, 2), + output_size=output_size, + kernel_size=(1, self.frame_len), + stride=(1, self.frame_len // 2), + )[:, 0, 0, :] + + if self.padding == "center": + pad = self.frame_len // 2 + elif self.padding == "same": + pad = self.frame_len // 4 + else: + raise ValueError("Padding must be 'center' or 'same'.") + + audio = audio[:, pad:-pad] + return audio + + +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) -> torch.Tensor: + """ + 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: + Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. + """ + x = self.out(x).transpose(1, 2) + mag, p = x.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 + + # Always compute complex values in float32 to avoid ComplexHalf warning (then cast audio back if needed) + orig_dtype = mag.dtype + mag, x, y = mag.float(), x.float(), y.float() + S = mag * (x + 1j * y) + audio = self.istft(S) + if orig_dtype != torch.float32: + audio = audio.to(orig_dtype) + return audio + + +class IMDCTSymExpHead(FourierHead): + """ + IMDCT Head module for predicting MDCT coefficients with symmetric exponential function + + Args: + dim (int): Hidden dimension of the model. + mdct_frame_len (int): Length of the MDCT frame. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + sample_rate (int, optional): The sample rate of the audio. If provided, the last layer will be initialized + based on perceptual scaling. Defaults to None. + clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False. + """ + + def __init__( + self, + dim: int, + mdct_frame_len: int, + padding: str = "same", + sample_rate: int | None = None, + clip_audio: bool = False, + ): + super().__init__() + out_dim = mdct_frame_len // 2 + self.out = nn.Linear(dim, out_dim) + self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) + self.clip_audio = clip_audio + + if sample_rate is not None: + # optionally init the last layer following mel-scale + m_max = _hz_to_mel(sample_rate // 2) + m_pts = torch.linspace(0, m_max, out_dim) + f_pts = _mel_to_hz(m_pts) + scale = 1 - (f_pts / f_pts.max()) + + with torch.no_grad(): + self.out.weight.mul_(scale.view(-1, 1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the IMDCTSymExpHead 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: + Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. + """ + x = self.out(x) + x = symexp(x) + x = torch.clip(x, min=-1e2, max=1e2) # safeguard to prevent excessively large magnitudes + audio = self.imdct(x) + if self.clip_audio: + audio = torch.clip(x, min=-1.0, max=1.0) + + return audio + + +class IMDCTCosHead(FourierHead): + """ + IMDCT Head module for predicting MDCT coefficients with parametrizing MDCT = exp(m) · cos(p) + + Args: + dim (int): Hidden dimension of the model. + mdct_frame_len (int): Length of the MDCT frame. + padding (str, optional): Type of padding. Options are "center" or "same". Defaults to "same". + clip_audio (bool, optional): If True, clip audio to [-1.0, 1.0]. Defaults to False. + """ + + def __init__( + self, + dim: int, + mdct_frame_len: int, + padding: str = "same", + clip_audio: bool = False, + ): + super().__init__() + self.clip_audio = clip_audio + self.out = nn.Linear(dim, mdct_frame_len) + self.imdct = IMDCT(frame_len=mdct_frame_len, padding=padding) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the IMDCTCosHead 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: + Tensor: Reconstructed time-domain audio signal of shape (B, T), where T is the length of the output signal. + """ + x = self.out(x) + m, p = x.chunk(2, dim=2) + m = torch.exp(m).clip(max=1e2) # safeguard to prevent excessively large magnitudes + audio = self.imdct(m * torch.cos(p)) + if self.clip_audio: + audio = torch.clip(x, min=-1.0, max=1.0) + return audio + + +class ConvNeXtBlock(nn.Module): + """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. + + Args: + dim (int): Number of input channels. + intermediate_dim (int): Dimensionality of the intermediate layer. + layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. + Defaults to None. + adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. + None means non-conditional LayerNorm. Defaults to None. + """ + + def __init__( + self, + dim: int, + intermediate_dim: int, + layer_scale_init_value: float, + adanorm_num_embeddings: int | None = None, + ): + super().__init__() + self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv + self.adanorm = adanorm_num_embeddings is not None + if adanorm_num_embeddings: + self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) + else: + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) # pointwise/1x1 convs, implemented with linear layers + self.act = nn.GELU() + self.pwconv2 = nn.Linear(intermediate_dim, dim) + self.gamma = ( + nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) + if layer_scale_init_value > 0 + else None + ) + + def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x = x.transpose(1, 2) # (B, C, T) -> (B, T, C) + if self.adanorm: + assert cond_embedding_id is not None + x = self.norm(x, cond_embedding_id) + else: + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.transpose(1, 2) # (B, T, C) -> (B, C, T) + + x = residual + x + return x + + +class AdaLayerNorm(nn.Module): + """ + Adaptive Layer Normalization module with learnable embeddings per `num_embeddings` classes + + Args: + num_embeddings (int): Number of embeddings. + embedding_dim (int): Dimension of the embeddings. + """ + + def __init__(self, num_embeddings: int, embedding_dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.dim = embedding_dim + self.scale = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim) + self.shift = nn.Embedding(num_embeddings=num_embeddings, embedding_dim=embedding_dim) + torch.nn.init.ones_(self.scale.weight) + torch.nn.init.zeros_(self.shift.weight) + + def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor) -> torch.Tensor: + scale = self.scale(cond_embedding_id) + shift = self.shift(cond_embedding_id) + x = nn.functional.layer_norm(x, (self.dim,), eps=self.eps) + x = x * scale + shift + return x + + +class ResBlock1(nn.Module): + """ + ResBlock adapted from HiFi-GAN V1 (https://github.com/jik876/hifi-gan) with dilated 1D convolutions, + but without upsampling layers. + + Args: + dim (int): Number of input channels. + kernel_size (int, optional): Size of the convolutional kernel. Defaults to 3. + dilation (tuple[int], optional): Dilation factors for the dilated convolutions. + Defaults to (1, 3, 5). + lrelu_slope (float, optional): Negative slope of the LeakyReLU activation function. + Defaults to 0.1. + layer_scale_init_value (float, optional): Initial value for the layer scale. None means no scaling. + Defaults to None. + """ + + def __init__( + self, + dim: int, + kernel_size: int = 3, + dilation: tuple[int, int, int] = (1, 3, 5), + lrelu_slope: float = 0.1, + layer_scale_init_value: float | None = None, + ): + super().__init__() + self.lrelu_slope = lrelu_slope + self.convs1 = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=dilation[0], + padding=self.get_padding(kernel_size, dilation[0]), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=dilation[1], + padding=self.get_padding(kernel_size, dilation[1]), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=dilation[2], + padding=self.get_padding(kernel_size, dilation[2]), + ) + ), + ] + ) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=1, + padding=self.get_padding(kernel_size, 1), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=1, + padding=self.get_padding(kernel_size, 1), + ) + ), + weight_norm( + nn.Conv1d( + dim, + dim, + kernel_size, + 1, + dilation=1, + padding=self.get_padding(kernel_size, 1), + ) + ), + ] + ) + + self.gamma = nn.ParameterList( + [ + ( + nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) + if layer_scale_init_value is not None + else None + ), + ( + nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) + if layer_scale_init_value is not None + else None + ), + ( + nn.Parameter(layer_scale_init_value * torch.ones(dim, 1), requires_grad=True) + if layer_scale_init_value is not None + else None + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for c1, c2, gamma in zip(self.convs1, self.convs2, self.gamma): + xt = torch.nn.functional.leaky_relu(x, negative_slope=self.lrelu_slope) + xt = c1(xt) + xt = torch.nn.functional.leaky_relu(xt, negative_slope=self.lrelu_slope) + xt = c2(xt) + if gamma is not None: + xt = gamma * xt + x = xt + x + return x + + def remove_weight_norm(self): + for conv in self.convs1: + remove_weight_norm(conv) + for conv in self.convs2: + remove_weight_norm(conv) + + @staticmethod + def get_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +class Backbone(nn.Module): + """Base class for the generator's backbone. It preserves the same temporal resolution across all layers.""" + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + """ + Args: + x (Tensor): Input tensor of shape (B, C, L), where B is the batch size, + C denotes output features, and L is the sequence length. + + Returns: + Tensor: Output of shape (B, L, H), where B is the batch size, L is the sequence length, + and H denotes the model dimension. + """ + raise NotImplementedError("Subclasses must implement the forward method.") + + +class VocosBackbone(Backbone): + """ + Vocos backbone module built with ConvNeXt blocks. Supports additional conditioning with Adaptive Layer Normalization + + Args: + input_channels (int): Number of input features channels. + dim (int): Hidden dimension of the model. + intermediate_dim (int): Intermediate dimension used in ConvNeXtBlock. + num_layers (int): Number of ConvNeXtBlock layers. + layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to `1 / num_layers`. + adanorm_num_embeddings (int, optional): Number of embeddings for AdaLayerNorm. + None means non-conditional model. Defaults to None. + """ + + def __init__( + self, + input_channels: int, + dim: int, + intermediate_dim: int, + num_layers: int, + layer_scale_init_value: float | None = None, + adanorm_num_embeddings: int | None = None, + ): + super().__init__() + self.input_channels = input_channels + self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3) + self.adanorm = adanorm_num_embeddings is not None + if adanorm_num_embeddings: + self.norm = AdaLayerNorm(adanorm_num_embeddings, dim, eps=1e-6) + else: + self.norm = nn.LayerNorm(dim, eps=1e-6) + layer_scale_init_value = layer_scale_init_value or 1 / num_layers + self.convnext = nn.ModuleList( + [ + ConvNeXtBlock( + dim=dim, + intermediate_dim=intermediate_dim, + layer_scale_init_value=layer_scale_init_value, + adanorm_num_embeddings=adanorm_num_embeddings, + ) + for _ in range(num_layers) + ] + ) + self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Conv1d | nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + bandwidth_id = kwargs.get("bandwidth_id", None) + x = self.embed(x) + if self.adanorm: + assert bandwidth_id is not None + x = self.norm(x.transpose(1, 2), cond_embedding_id=bandwidth_id) + else: + x = self.norm(x.transpose(1, 2)) + x = x.transpose(1, 2) + for conv_block in self.convnext: + x = conv_block(x, cond_embedding_id=bandwidth_id) + x = self.final_layer_norm(x.transpose(1, 2)) + return x + + +class VocosResNetBackbone(Backbone): + """ + Vocos backbone module built with ResBlocks. + + Args: + input_channels (int): Number of input features channels. + dim (int): Hidden dimension of the model. + num_blocks (int): Number of ResBlock1 blocks. + layer_scale_init_value (float, optional): Initial value for layer scaling. Defaults to None. + """ + + def __init__( + self, + input_channels, + dim, + num_blocks, + layer_scale_init_value=None, + ): + super().__init__() + self.input_channels = input_channels + self.embed = weight_norm(nn.Conv1d(input_channels, dim, kernel_size=3, padding=1)) + layer_scale_init_value = layer_scale_init_value or 1 / num_blocks / 3 + self.resnet = nn.Sequential( + *[ResBlock1(dim=dim, layer_scale_init_value=layer_scale_init_value) for _ in range(num_blocks)] + ) + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + x = self.embed(x) + x = self.resnet(x) + x = x.transpose(1, 2) + return x + + +class Vocos(nn.Module): + def __init__( + self, + input_channels: int = 256, + dim: int = 384, + intermediate_dim: int = 1152, + num_layers: int = 8, + n_fft: int = 800, + hop_size: int = 200, + padding: str = "same", + adanorm_num_embeddings=None, + cfg=None, + ): + super().__init__() + + input_channels = cfg.input_channels if cfg is not None and hasattr(cfg, "input_channels") else input_channels + dim = cfg.dim if cfg is not None and hasattr(cfg, "dim") else dim + intermediate_dim = ( + cfg.intermediate_dim if cfg is not None and hasattr(cfg, "intermediate_dim") else intermediate_dim + ) + num_layers = cfg.num_layers if cfg is not None and hasattr(cfg, "num_layers") else num_layers + adanorm_num_embeddings = ( + cfg.adanorm_num_embeddings + if cfg is not None and hasattr(cfg, "adanorm_num_embeddings") + else adanorm_num_embeddings + ) + n_fft = cfg.n_fft if cfg is not None and hasattr(cfg, "n_fft") else n_fft + hop_size = cfg.hop_size if cfg is not None and hasattr(cfg, "hop_size") else hop_size + padding = cfg.padding if cfg is not None and hasattr(cfg, "padding") else padding + + self.backbone = VocosBackbone( + input_channels=input_channels, + dim=dim, + intermediate_dim=intermediate_dim, + num_layers=num_layers, + adanorm_num_embeddings=adanorm_num_embeddings, + ) + self.head = ISTFTHead(dim, n_fft, hop_size, padding) + + def forward(self, x): + x = self.backbone(x) + x = self.head(x) + + return x[:, None, :] + + +class JsonHParams: + def __init__(self, **kwargs): + for k, v in kwargs.items(): + if isinstance(v, dict): + v = JsonHParams(**v) + self[k] = v + + def keys(self): + return self.__dict__.keys() + + def items(self): + return self.__dict__.items() + + def values(self): + return self.__dict__.values() + + def __len__(self): + return len(self.__dict__) + + def __getitem__(self, key): + return getattr(self, key) + + def __setitem__(self, key, value): + return setattr(self, key, value) + + def __contains__(self, key): + return key in self.__dict__ + + def __repr__(self): + return self.__dict__.__repr__() + + +def build_vocoder_model(cfg): + vocoder_model = Vocos(cfg=cfg.model.vocos) + vocoder_model.eval() + return vocoder_model + + +def load_checkpoint(build_model_func, cfg, ckpt_path): + model = build_model_func(cfg) + + if ckpt_path is not None: + accelerate.load_checkpoint_and_dispatch(model, ckpt_path) + return model + + +def load_vocos_model( + ckpt_path: str | None = None, + config: DictConfig = None, + device: str = "cuda", +): + """Load Vocos vocoder. + + Args: + config: DictConfig, config for vocoder. + ckpt_path: str | None = None, path to checkpoint. + device: str, device to load model. Note: accelerate dispatch handles device placement. + """ + + merged_cfg = _deep_update_dict(_get_vocoder_default_cfg(), config) + vocoder_cfg_obj = JsonHParams(**merged_cfg) + + vocoder_model = load_checkpoint(build_vocoder_model, vocoder_cfg_obj, ckpt_path) + vocoder_model.eval() + + for param in vocoder_model.parameters(): + param.requires_grad = False + + return vocoder_model + + +class Vocoder(nn.Module): + def __init__(self, vocoder: dict | None = None, ckpt_path: str | None = None): + super().__init__() + vocoder = vocoder or {} + + model_cfg = vocoder.get("model_cfg") or vocoder.get("cfg") + + self.model = load_vocos_model( + ckpt_path=ckpt_path, + config=model_cfg, + ) + + def forward(self, x): + x = self.model(x) + return x diff --git a/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py b/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py new file mode 100644 index 00000000000..379dda3a5f7 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/modules/whisper_encoder.py @@ -0,0 +1,96 @@ +"""Frozen Whisper encoder wrapper (wav -> encoder embeddings).""" + +import torch +import torchaudio +from transformers import WhisperFeatureExtractor, WhisperModel + +WHISPER_MEL_FRAMES = 3000 # 3000 frames at 16000 Hz + + +class WhisperEncoder: + """Auxiliary content encoder for SVC. + + Kept in FP32 for stable HF Whisper inference under vLLM-Omni's default + BF16 diffusion trunk (same idea as Stable Audio's FP32 VAE). Outputs are + cast to the trunk dtype before fusion with the CFM conditioner. + """ + + def __init__(self, device: str | torch.device | None = None) -> None: + self.fe = WhisperFeatureExtractor.from_pretrained("openai/whisper-base") + self.model = WhisperModel.from_pretrained("openai/whisper-base") + self.device = torch.device( + device or ("cuda" if torch.cuda.is_available() else "cpu"), + ) + self.model = self.model.to(device=self.device, dtype=torch.float32) + self.model.eval() + for param in self.model.parameters(): + param.requires_grad_(False) + + @property + def dtype(self) -> torch.dtype: + return torch.float32 + + def float(self) -> "WhisperEncoder": + self.model = self.model.float() + return self + + def encode( + self, + wav: torch.Tensor, + sr: int, + *, + output_dtype: torch.dtype | None = None, + ) -> torch.Tensor: + wav = ( + torchaudio.functional.resample(wav, orig_freq=sr, new_freq=self.fe.sampling_rate) + if sr != self.fe.sampling_rate + else wav + ) + wav_np = wav.detach().float().cpu().numpy().astype("float32", copy=False) + + inputs = self.fe( + wav_np, + sampling_rate=self.fe.sampling_rate, + return_tensors="pt", + padding=False, + truncation=False, + return_attention_mask=True, + ) + + input_features = inputs.input_features + num_frames = input_features.shape[-1] + if num_frames < WHISPER_MEL_FRAMES: + pad = WHISPER_MEL_FRAMES - num_frames + input_features = torch.nn.functional.pad(input_features, (0, pad)) + else: + input_features = input_features[..., :WHISPER_MEL_FRAMES] + + input_features = input_features.to(device=wav.device, dtype=self.dtype) + if self.model.device != wav.device: + self.model = self.model.to(wav.device) + + attention_mask = inputs.attention_mask.to(wav.device) if inputs.attention_mask is not None else None + + with torch.inference_mode(): + encoder_out = self.model.encoder(input_features).last_hidden_state + + if attention_mask is not None: + valid_mel_frames = attention_mask.sum(dim=1) + valid_enc_frames = (valid_mel_frames + 1) // 2 + max_valid_enc_frames = min(int(valid_enc_frames.max().item()), encoder_out.shape[1]) + encoder_out = encoder_out[:, :max_valid_enc_frames, :] + valid_len = min(int(valid_enc_frames[0].item()), max_valid_enc_frames) + if valid_len < max_valid_enc_frames: + encoder_out[0, valid_len:, :] = 0 + + if output_dtype is not None and encoder_out.dtype != output_dtype: + encoder_out = encoder_out.to(dtype=output_dtype) + return encoder_out + + +if __name__ == "__main__": + torch.manual_seed(0) + audio = torch.randn(1, 24000 * 25).float().to("cuda") + encoder = WhisperEncoder() + whisper_encoder_out = encoder.encode(audio, sr=24000) + print(whisper_encoder_out.shape) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py new file mode 100644 index 00000000000..a48b5236dd8 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_base.py @@ -0,0 +1,377 @@ +import os +import time +from collections.abc import Callable +from contextlib import contextmanager +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.nn as nn +from vllm.logger import init_logger +from vllm.utils.torch_utils import set_default_torch_dtype + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin, _unwrap, _wrap +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.interface import ( + SupportAudioInput, + SupportAudioOutput, + SupportsComponentDiscovery, +) +from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import consume_payload +from vllm_omni.diffusion.models.soulx_singer.utils import _patch_torchaudio_load, load_config +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + +_SOULX_PROFILER_TARGETS = [ + "forward", +] + + +def convert_soulx_audio_output_to_numpy(audio: torch.Tensor | dict[str, Any]) -> np.ndarray | dict[str, Any]: + if isinstance(audio, dict) and isinstance(audio.get("payload"), dict): + payload = dict(audio["payload"]) + audio_payload = payload.get("audio") + if isinstance(audio_payload, torch.Tensor): + payload["audio"] = audio_payload.detach().cpu().float().numpy() + return { + "payload": payload, + "metadata": audio.get("metadata") or {}, + } + if isinstance(audio, torch.Tensor): + return audio.detach().cpu().float().numpy() + return audio + + +class FlowMatchingAudioPipeline( + nn.Module, + CFGParallelMixin, + SupportAudioInput, + SupportAudioOutput, + SupportsComponentDiscovery, + DiffusionPipelineProfilerMixin, +): + """SoulX flow-matching pipeline with CFG / CFG-parallel support.""" + + support_audio_input: ClassVar[bool] = True + support_audio_output: ClassVar[bool] = True + audio_sample_rate: ClassVar[int] = 24000 + _DEFAULT_RESCALE_CFG: ClassVar[float] = 0.75 + _dit_modules: ClassVar[list[str]] = ["cfm_decoder.model.diff_estimator"] + _encoder_modules: ClassVar[list[str]] = ["mel", "f0_encoder"] + _vae_modules: ClassVar[list[str]] = ["vocoder"] + _resident_modules: ClassVar[list[str]] = ["cfm_decoder.model.cond_emb"] + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.device = get_local_device() + + # 1. parse model path and yaml config + model_dir = od_config.model + assert model_dir is not None, "model_dir is not provided" + + if not os.path.exists(model_dir): + raise FileNotFoundError( + f"""Model directory {model_dir} does not exist. User should refer to + `examples/offline_inference/text_to_speech/README.md` for usage details.""" + ) + + # Empty: skip DiffusersPipelineLoader.get_all_weights() → no _prepare_weights / file glob / + # safetensors or pt iterator from the framework. Load checkpoints in __init__ instead. + self.weights_sources = () + + self.model_path = model_dir + + config_path = os.path.join(model_dir, "config.yaml") + hf_config = load_config(config_path) + + self.audio_config = hf_config.audio + self.encoder_config = hf_config.model.encoder + self.flow_matching_config = hf_config.model.flow_matching + + self.mel_dim = self.audio_config.num_mels + + self.f0_bin = self.encoder_config.f0_bin + self.f0_dim = self.encoder_config.f0_dim + self.text_dim = self.encoder_config.text_dim + self.vocab_size = self.encoder_config.vocab_size + self.pitch_dim = self.encoder_config.pitch_dim + self.type_dim = self.encoder_config.type_dim + + _patch_torchaudio_load() + + def _mel_to_audio(self, generated_mel: torch.Tensor, *, squeeze: bool = False) -> torch.Tensor: + audio = self.vocoder(generated_mel.transpose(1, 2)[0:1, ...].float()).float() + return audio.squeeze() if squeeze else audio + + def _mel_from_wav(self, wav: torch.Tensor) -> torch.Tensor: + if wav.dtype != torch.float32: + wav = wav.float() + return self.mel(wav) + + def _load_soulx_checkpoint(self, checkpoint_name: str) -> None: + weight_path = os.path.join(self.model_path, checkpoint_name) + if not os.path.isfile(weight_path): + raise FileNotFoundError( + f"Model weights not found at {weight_path}. Please check the pretrained model path." + ) + state = torch.load(weight_path, map_location=self.device) + self.mel.float() + self.load_state_dict(state["state_dict"], strict=True) + self._finalize_loaded_dtypes() + logger.info("Loaded model weights from %s", weight_path) + + def _setup_soulx_profiler(self, *, extra_targets: list[str] | None = None) -> None: + """Enable stage timing when ``od_config.enable_diffusion_pipeline_profiler`` is set.""" + targets = list(_SOULX_PROFILER_TARGETS) + if extra_targets: + targets.extend(extra_targets) + self.setup_diffusion_pipeline_profiler( + profiler_targets=targets, + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, + ) + + @contextmanager + def _stage_timer(self, name: str): + if not getattr(self, "enable_diffusion_pipeline_profiler", False): + yield + return + if current_omni_platform.is_available(): + current_omni_platform.synchronize() + start = time.perf_counter() + try: + yield + finally: + if current_omni_platform.is_available(): + current_omni_platform.synchronize() + duration = time.perf_counter() - start + with self._profiler_lock: + self._stage_durations[name] = self._stage_durations.get(name, 0.0) + duration + + def _profiler_stage_durations(self) -> dict[str, float] | None: + if getattr(self, "enable_diffusion_pipeline_profiler", False): + return self.stage_durations + return None + + @staticmethod + def _build_fp32_audio_modules(audio_config) -> tuple[nn.Module, nn.Module]: + """Mel/vocoder must stay FP32 even when ``od_config.dtype`` is FP16/BF16.""" + from vllm_omni.diffusion.models.soulx_singer.modules import MelSpectrogramEncoder, Vocoder + + with set_default_torch_dtype(torch.float32): + mel = MelSpectrogramEncoder(audio_config) + vocoder = Vocoder() + return mel, vocoder + + def _finalize_loaded_dtypes(self) -> None: + """Keep mel/vocoder in FP32; cast DiT trunk to ``od_config.dtype`` after load.""" + self.mel.float() + self.vocoder.float() + trunk_dtype = self.od_config.dtype + if trunk_dtype not in (torch.float16, torch.bfloat16): + return + trunk_modules = ( + "cfm_decoder", + "f0_encoder", + "preflow", + "note_text_encoder", + "note_pitch_encoder", + "note_type_encoder", + ) + for name in trunk_modules: + module = getattr(self, name, None) + if module is not None: + module.to(dtype=trunk_dtype) + whisper = getattr(self, "whisper_encoder", None) + if whisper is not None: + whisper.float() + + @property + def transformer(self) -> nn.Module: + """Alias expected by ``CFGParallelMixin`` default hooks.""" + return self.cfm_decoder.model.diff_estimator + + @property + def diffusion_trunk_dtype(self) -> torch.dtype: + """DiT trunk dtype from ``diff_estimator`` (not ``cond_emb``, which is registered first).""" + return next(self.cfm_decoder.model.diff_estimator.parameters()).dtype + + def _to_trunk_dtype(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: + """Cast tensors to ``diffusion_trunk_dtype`` at the CFM/DiT boundary.""" + trunk_dtype = self.diffusion_trunk_dtype + return tuple(t if t.dtype == trunk_dtype else t.to(dtype=trunk_dtype) for t in tensors) + + def _resolve_diffusion_generator(self, sampling_params: Any) -> torch.Generator | None: + """Build a device-local RNG for CFM noise from ``generator`` or ``seed``.""" + generator = getattr(sampling_params, "generator", None) + if generator is not None: + return generator + seed = getattr(sampling_params, "seed", None) + if seed is None: + return None + return torch.Generator(device=self.device).manual_seed(int(seed)) + + def _prepare_cfm_loop_state( + self, + prompt: torch.Tensor, + cond: torch.Tensor, + n_timesteps: int, + *, + generator: torch.Generator | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Cast FP32 islands to trunk dtype and build CFM loop tensors/masks.""" + trunk_dtype = self.diffusion_trunk_dtype + prompt, cond = self._to_trunk_dtype(prompt, cond) + cond_emb = self.cfm_decoder.model.cond_emb(cond) + if cond_emb.dtype != trunk_dtype: + cond_emb = self._to_trunk_dtype(cond_emb)[0] + device = cond_emb.device + + h = 1.0 / n_timesteps + h_tensor = torch.tensor(h, device=device, dtype=trunk_dtype) + prompt_len = prompt.shape[1] + target_len = cond_emb.shape[1] - prompt_len + + x_mask = torch.ones(cond_emb.shape[0], target_len, device=device, dtype=trunk_dtype) + prompt_mask = torch.ones(cond_emb.shape[0], prompt_len, device=device, dtype=trunk_dtype) + xt_mask = torch.cat([prompt_mask, x_mask], dim=1) + xt = torch.randn( + cond_emb.shape[0], + target_len, + self.mel_dim, + dtype=trunk_dtype, + device=device, + generator=generator, + ) + return cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len + + def predict_noise( + self, + *, + x: torch.Tensor, + diffusion_step: torch.Tensor, + cond: torch.Tensor, + x_mask: torch.Tensor, + prompt_len: int | None = None, + **_: Any, + ) -> torch.Tensor: + """Forward through ``diff_estimator``, optionally slice off prompt frames.""" + x, diffusion_step, cond, x_mask = self._to_trunk_dtype(x, diffusion_step, cond, x_mask) + flow_pred = self.cfm_decoder.model.diff_estimator(x, diffusion_step, cond, x_mask) + if prompt_len is not None: + flow_pred = flow_pred[:, prompt_len:, :] + return flow_pred + + def combine_cfg_noise( + self, + positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], + negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], + true_cfg_scale: float, + cfg_normalize: bool = False, + rescale_cfg: float | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, ...]: + """SoulX CFG: pos + scale * (pos - neg) with std rescaling.""" + del cfg_normalize # SoulX uses rescale_cfg instead of cfg_normalize_function. + if rescale_cfg is None: + rescale_cfg = getattr(self, "rescale_cfg", self._DEFAULT_RESCALE_CFG) + + pos = _unwrap(_wrap(positive_noise_pred)) + neg = _unwrap(_wrap(negative_noise_pred)) + + pos_flow_pred_std = pos.std() + flow_pred_cfg = pos + true_cfg_scale * (pos - neg) + rescale_flow_pred = flow_pred_cfg * pos_flow_pred_std / flow_pred_cfg.std() + combined = rescale_cfg * rescale_flow_pred + (1 - rescale_cfg) * flow_pred_cfg + return combined.to(dtype=pos.dtype) + + def _run_flow_matching_loop( + self, + prompt: torch.Tensor, + cond: torch.Tensor, + n_timesteps: int, + cfg: float, + *, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Euler flow-matching loop aligned with ``FlowMatchingTransformer.reverse_diffusion``.""" + cond_emb, prompt, xt, x_mask, xt_mask, h_tensor, prompt_len = self._prepare_cfm_loop_state( + prompt, cond, n_timesteps, generator=generator + ) + trunk_dtype = self.diffusion_trunk_dtype + step_scale = 1.0 / n_timesteps + + for i in range(n_timesteps): + xt_input = torch.cat([prompt, xt], dim=1) + diffusion_step = torch.full( + (xt.shape[0],), + fill_value=(i + 0.5) * step_scale, + device=xt.device, + dtype=trunk_dtype, + ) + + flow_pred = self.predict_noise_maybe_with_cfg( + do_true_cfg=(cfg > 0), + true_cfg_scale=cfg, + positive_kwargs={ + "x": xt_input, + "diffusion_step": diffusion_step, + "cond": cond_emb, + "x_mask": xt_mask, + "prompt_len": prompt_len, + }, + negative_kwargs={ + "x": xt, + "diffusion_step": diffusion_step, + # CFG null cond: zero out target frames only (prompt cond stays in pos branch). + "cond": torch.zeros_like(cond_emb)[:, : xt.shape[1], :], + "x_mask": x_mask, + }, + cfg_normalize=False, + ) + + xt = xt + flow_pred * h_tensor + + return xt + + def _forward_batch_from_request( + self, + req: OmniDiffusionRequest, + *, + kind: str, + metadata_key: str, + infer_batch_fn: Callable[..., tuple[torch.Tensor, int]], + prepare_extra_args: Callable[[dict[str, Any], Any], dict[str, Any]] | None = None, + ) -> DiffusionOutput: + sampling_params = req.sampling_params + extra_args = dict(getattr(sampling_params, "extra_args", None) or {}) + if prepare_extra_args is not None: + extra_args = prepare_extra_args(extra_args, sampling_params) + sampling_params.extra_args = extra_args + + num_inference_steps = sampling_params.num_inference_steps or 32 + guidance_scale = sampling_params.guidance_scale or 3.0 + generator = self._resolve_diffusion_generator(sampling_params) + + with self._stage_timer("consume_payload"): + payload = consume_payload(req, kind, self.device) + + with self._stage_timer("infer_segments"): + audio, pitch_shift = infer_batch_fn( + payload, + extra_args=extra_args, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + ) + + return DiffusionOutput( + output={ + "payload": {"audio": audio}, + "metadata": {"audio": {metadata_key: pitch_shift}}, + }, + stage_durations=self._profiler_stage_durations() or {}, + ) diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py new file mode 100644 index 00000000000..85fed9ee755 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svc.py @@ -0,0 +1,411 @@ +"""SoulX-Singer SVC (voice conversion) pipeline implementation.""" + +import os +from collections.abc import Iterable +from typing import Any, ClassVar + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger +from vllm.utils.torch_utils import set_default_torch_dtype + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.soulx_singer.modules import ( + CFMDecoder, + WhisperEncoder, +) +from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( + FlowMatchingAudioPipeline, + convert_soulx_audio_output_to_numpy, +) +from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( + SOULX_PREPROCESSED_KEY, + get_soulx_preprocessed_payload, + has_precomputed, +) +from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( + attach_preprocess_for_diffusion_request, + is_warmup_request, + resolve_preprocess_audio, +) +from vllm_omni.diffusion.models.soulx_singer.utils import ( + f0_to_coarse, + load_config, + resolve_pitch_shift, + validate_soulx_extra_args, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +logger = init_logger(__name__) + + +def get_soulxsinger_svc_pre_process_func(od_config: OmniDiffusionConfig): + hf_config = load_config(os.path.join(od_config.model, "config.yaml")) + sample_rate = hf_config.audio.sample_rate + device = get_local_device() + _pipeline = None + + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + nonlocal _pipeline + + extra_args = validate_soulx_extra_args( + "svc", + dict(getattr(request.sampling_params, "extra_args", None) or {}), + ) + + # Inline build: when no warmup/no precomputed paths/no IPC payload, + # build the preprocess payload directly from audio file paths. + if not (is_warmup_request(request) or has_precomputed(extra_args, "svc")): + prompt = request.prompt + if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type] + prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type] + if prompt_audio is not None and target_audio is not None: + if _pipeline is None: + from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import ( + SoulXPreprocessPipeline, + ) + + _pipeline = SoulXPreprocessPipeline( + od_config=od_config, + vocal_sep=bool(extra_args.get("vocal_sep", False)), + verbose=bool(extra_args.get("preprocess_verbose", False)), + extra_args=extra_args, + ) + payload = _pipeline.build_svc_payload_from_audio( + prompt_audio=prompt_audio, + target_audio=target_audio, + sample_rate=sample_rate, + device=device, + vocal_sep=extra_args.get("vocal_sep"), + ) + payload.setdefault("kind", "svc") + prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment] + + return attach_preprocess_for_diffusion_request( + request, + kind="svc", + sample_rate=sample_rate, + device=device, + ) + + return pre_process_func + + +def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): + def post_process_func(audio: torch.Tensor): + return convert_soulx_audio_output_to_numpy(audio) + + return post_process_func + + +class PipelineSoulXSingerSVC(FlowMatchingAudioPipeline): + """SVC pipeline for the SoulX-Singer model.""" + + _encoder_modules: ClassVar[list[str]] = [ + "whisper_encoder", + "mel", + "f0_encoder", + ] + + EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( + { + "prompt_wav_path", + "target_wav_path", + "prompt_f0_path", + "target_f0_path", + "prompt_audio", + "target_audio", + "vocal_sep", + "preprocess_weights_dir", + "preprocess_verbose", + "auto_shift", + "pitch_shift", + } + ) + + EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"pitch_shift"}) + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__(od_config=od_config, prefix=prefix) + + self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim) + + self.cfm_decoder = CFMDecoder(self.flow_matching_config) + + self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config) + with set_default_torch_dtype(torch.float32): + self.whisper_encoder = WhisperEncoder(device=self.device) + + self._setup_soulx_profiler() + + def _encode_condition(self, *, whisper_features: torch.Tensor, f0_coarse: torch.Tensor) -> torch.Tensor: + cond = whisper_features + self.f0_encoder(f0_coarse) + return self._to_trunk_dtype(cond)[0] + + def _encode_prompt_whisper_feature(self, prompt_wav: torch.Tensor) -> torch.Tensor: + trunk_dtype = self.f0_encoder.weight.dtype + with self._stage_timer("whisper"): + return self.whisper_encoder.encode( + prompt_wav, + sr=self.audio_config.sample_rate, + output_dtype=trunk_dtype, + ) + + _encode_target_whisper_feature = _encode_prompt_whisper_feature + + def _infer_segment( + self, + *, + prompt_mel: torch.Tensor, + prompt_wav: torch.Tensor, + target_wav: torch.Tensor, + prompt_f0: torch.Tensor, + target_f0: torch.Tensor, + pitch_shift: int, + num_inference_steps: int, + guidance_scale: float, + prompt_feature: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + """Single-segment SVC inference aligned with ``SoulXSingerSVC.infer_segment``.""" + len_prompt_mel = prompt_mel.shape[1] + prompt_f0 = F.pad(prompt_f0, (0, 0, 0, max(0, len_prompt_mel - prompt_f0.shape[1])))[:, :len_prompt_mel] + + f0_coarse_prompt = f0_to_coarse(prompt_f0) + f0_coarse_target = f0_to_coarse(target_f0, f0_shift=int(pitch_shift * 5)) + f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1) + + trunk_dtype = self.f0_encoder.weight.dtype + if prompt_feature is None: + prompt_feature = self._encode_prompt_whisper_feature(prompt_wav) + with self._stage_timer("whisper"): + target_feature = self.whisper_encoder.encode( + target_wav, + sr=self.audio_config.sample_rate, + output_dtype=trunk_dtype, + ) + prompt_feature = F.pad( + prompt_feature, + (0, 0, 0, max(0, f0_coarse_prompt.shape[1] - prompt_feature.shape[1])), + )[:, : f0_coarse_prompt.shape[1], :] + target_feature = F.pad( + target_feature, + (0, 0, 0, max(0, f0_coarse_target.shape[1] - target_feature.shape[1])), + )[:, : f0_coarse_target.shape[1], :] + + whisper_features = torch.cat([prompt_feature, target_feature], dim=1) + with self._stage_timer("cond_encode"): + cond = self._encode_condition(whisper_features=whisper_features, f0_coarse=f0_coarse) + + with self._stage_timer("cfm"): + generated_mel = self._run_flow_matching_loop( + prompt=prompt_mel, + cond=cond, + n_timesteps=num_inference_steps, + cfg=guidance_scale, + generator=generator, + ) + + with self._stage_timer("vocoder"): + generated_audio = self._mel_to_audio(generated_mel, squeeze=True) + target_len = target_wav.shape[-1] + if generated_audio.shape[-1] > target_len: + generated_audio = generated_audio[:target_len] + elif generated_audio.shape[-1] < target_len: + generated_audio = F.pad(generated_audio, (0, target_len - generated_audio.shape[-1])) + return generated_audio + + def _build_vocal_segments( + self, + f0: torch.Tensor, + *, + hop_size: int, + sample_rate: int, + uv_frames_th: int = 10, + min_duration_sec: float = 15.0, + max_duration_sec: float = 30.0, + ignore_silent: bool = True, + ) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]: + """Build vocal segments based on F0 contour. + + Mirrors upstream ``SoulXSingerSVC.build_vocal_segments``. + Returns (overlap_segments, segments) tuples of (start_sec, end_sec). + """ + f0_np = ( + f0.detach().float().cpu().numpy().squeeze() if isinstance(f0, torch.Tensor) else np.asarray(f0).squeeze() + ) + total_frames = int(f0_np.shape[0]) + if total_frames == 0: + return [], [] + + f0_rate = sample_rate // hop_size + min_frames = max(1, int(round(min_duration_sec * f0_rate))) + max_frames = max(1, int(round(max_duration_sec * f0_rate))) + + split_points = [0] + + def append_split_point(point: int) -> None: + point = int(max(0, min(point, total_frames))) + while point - split_points[-1] > max_frames: + split_points.append(split_points[-1] + max_frames) + if point > split_points[-1]: + split_points.append(point) + + idx = 0 + while idx < total_frames: + if f0_np[idx] == 0: + run_start = idx + while idx < total_frames and f0_np[idx] == 0: + idx += 1 + run_end = idx + if (run_end - run_start) >= uv_frames_th: + split_point = max(run_end - 5, (run_start + run_end) // 2) + append_split_point(split_point) + else: + idx += 1 + append_split_point(total_frames) + + segments: list[tuple[int, int]] = [] + overlap_segments: list[tuple[int, int]] = [] + num_overlaps = 1 + + def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps) -> None: + segments.append((split_points[start_idx], split_points[end_idx])) + overlap_start_idx = start_idx + if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames: + overlap_start_idx = start_idx - overlaps + overlap_segments.append((split_points[overlap_start_idx], split_points[end_idx])) + + seg_start, seg_end = 0, 1 + while seg_start < len(split_points) - 1: + while seg_end < len(split_points) and (split_points[seg_end] - split_points[seg_start]) < min_frames: + seg_end += 1 + if seg_end >= len(split_points): + append_segment(seg_start, len(split_points) - 1) + break + append_segment(seg_start, seg_end) + seg_start = seg_end + seg_end = seg_start + 1 + + if ignore_silent: + filtered_idx = [] + for i, (ov_start, ov_end) in enumerate(overlap_segments): + voice_ratio = np.sum(f0_np[ov_start:ov_end] > 0) / max(1, ov_end - ov_start) + voiced_frames = np.sum(f0_np[ov_start:ov_end] > 0) + if voice_ratio > 0.05 and voiced_frames >= 10: + filtered_idx.append(i) + overlap_segments = [overlap_segments[i] for i in filtered_idx] + segments = [segments[i] for i in filtered_idx] + + # convert indices to seconds + overlap_sec = [(s / f0_rate, e / f0_rate) for s, e in overlap_segments] + seg_sec = [(s / f0_rate, e / f0_rate) for s, e in segments] + return overlap_sec, seg_sec + + def infer_svc_batch( + self, + payload: dict[str, Any], + *, + extra_args: dict[str, Any], + num_inference_steps: int, + guidance_scale: float, + generator: torch.Generator | None = None, + ) -> tuple[torch.Tensor, int]: + """Batch SVC with chunked inference for long audio (matches upstream).""" + prompt_wav = payload["prompt_wav"] + target_wav = payload["target_wav"] + prompt_f0 = payload["prompt_f0"] + target_f0 = payload["target_f0"] + + pitch_shift = resolve_pitch_shift( + auto_shift=bool(extra_args.get("auto_shift", True)), + manual_shift=int(extra_args.get("pitch_shift", 0)), + prompt_f0=prompt_f0, + target_f0=target_f0, + ) + + prompt_mel = self._mel_from_wav(prompt_wav) + + # Long audio: chunk by vocal segments (same as upstream) + max_sec = 30.0 * self.audio_config.sample_rate + if target_wav.shape[-1] < max_sec: + # Short audio: one-shot + generated_audio = self._infer_segment( + prompt_mel=prompt_mel, + prompt_wav=prompt_wav, + target_wav=target_wav, + prompt_f0=prompt_f0, + target_f0=target_f0, + pitch_shift=pitch_shift, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + ) + return generated_audio.unsqueeze(0), pitch_shift + + # Build vocal segments + overlap_segments, segments = self._build_vocal_segments( + target_f0, + hop_size=self.audio_config.hop_size, + sample_rate=self.audio_config.sample_rate, + ) + if len(segments) == 0: + segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)] + overlap_segments = [(0.0, target_wav.shape[-1] / self.audio_config.sample_rate)] + + f0_rate = self.audio_config.sample_rate // self.audio_config.hop_size + generated_audio = torch.zeros_like(target_wav) + + for idx in range(len(segments)): + overlap_start_sec, overlap_end_sec = overlap_segments[idx] + seg_start_sec, seg_end_sec = segments[idx] + + wav_start = int(round(overlap_start_sec * self.audio_config.sample_rate)) + wav_end = int(round(overlap_end_sec * self.audio_config.sample_rate)) + f0_start = int(round(overlap_start_sec * f0_rate)) + f0_end = int(round(overlap_end_sec * f0_rate)) + + wav_start = max(0, min(wav_start, target_wav.shape[-1])) + wav_end = max(wav_start, min(wav_end, target_wav.shape[-1])) + f0_start = max(0, min(f0_start, target_f0.shape[-1])) + f0_end = max(f0_start, min(f0_end, target_f0.shape[-1])) + + segment_gt_wav = target_wav[:, wav_start:wav_end] + segment_gt_f0 = target_f0[:, f0_start:f0_end] + + segment_generated = self._infer_segment( + prompt_mel=prompt_mel, + prompt_wav=prompt_wav, + target_wav=segment_gt_wav, + prompt_f0=prompt_f0, + target_f0=segment_gt_f0, + pitch_shift=pitch_shift, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + generator=generator, + ) + + seg_start = int(round(seg_start_sec * self.audio_config.sample_rate)) + seg_end = int(round(seg_end_sec * self.audio_config.sample_rate)) + segment_generated = segment_generated[seg_start - wav_start : seg_end - wav_start] + generated_audio[:, seg_start:seg_end] = segment_generated + + return generated_audio.unsqueeze(0), pitch_shift + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + return self._forward_batch_from_request( + req, + kind="svc", + metadata_key="pitch_shift", + infer_batch_fn=self.infer_svc_batch, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + del weights + self._load_soulx_checkpoint("model-svc.pt") + return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py new file mode 100644 index 00000000000..fb2c086f22e --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/pipeline_soulx_singer_svs.py @@ -0,0 +1,355 @@ +"""SoulX-Singer SVS (score-driven) pipeline implementation.""" + +from collections.abc import Iterable +from typing import Any, ClassVar + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.soulx_singer.modules import ( + CFMDecoder, + ConvNeXtV2Block, +) +from vllm_omni.diffusion.models.soulx_singer.pipeline_soulx_singer_base import ( + FlowMatchingAudioPipeline, + convert_soulx_audio_output_to_numpy, +) +from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( + SOULX_PREPROCESSED_KEY, + get_soulx_preprocessed_payload, + has_precomputed, +) +from vllm_omni.diffusion.models.soulx_singer.preprocess.pre_process import ( + attach_preprocess_for_diffusion_request, + build_metadata_processor, + is_warmup_request, + normalize_svs_control_extra_args, + resolve_preprocess_audio, +) +from vllm_omni.diffusion.models.soulx_singer.utils import ( + f0_to_coarse, + resolve_pitch_shift, + validate_soulx_extra_args, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + +logger = init_logger(__name__) + + +def get_soulxsinger_pre_process_func(od_config: OmniDiffusionConfig): + """Validate/load SVS preprocess payload for single-stage or stage-1 DiT.""" + metadata_processor = build_metadata_processor(od_config) + _pipeline = None + + def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: + nonlocal _pipeline + + extra_args = validate_soulx_extra_args( + "svs", + dict(getattr(request.sampling_params, "extra_args", None) or {}), + ) + + # Inline build: when no warmup/no precomputed paths/no IPC payload, + # build the preprocess payload directly from audio file paths. + if not (is_warmup_request(request) or has_precomputed(extra_args, "svs")): + prompt = request.prompt + if not isinstance(prompt, str) and not get_soulx_preprocessed_payload(prompt): # type: ignore[arg-type] + prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) # type: ignore[arg-type] + if prompt_audio is not None and target_audio is not None: + if _pipeline is None: + from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import ( + SoulXPreprocessPipeline, + ) + + _pipeline = SoulXPreprocessPipeline( + od_config=od_config, + vocal_sep=bool(extra_args.get("vocal_sep", False)), + verbose=bool(extra_args.get("preprocess_verbose", False)), + extra_args=extra_args, + ) + payload = _pipeline.build_svs_payload_from_audio( + prompt_audio=prompt_audio, + target_audio=target_audio, + metadata_processor=metadata_processor, + language=str(extra_args.get("language", "Mandarin")), + vocal_sep=extra_args.get("vocal_sep"), + prompt_vocal_sep=extra_args.get("prompt_vocal_sep"), + target_vocal_sep=extra_args.get("target_vocal_sep"), + prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"), + target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"), + ) + payload.setdefault("kind", "svs") + prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload # type: ignore[union-attr, assignment] + + return attach_preprocess_for_diffusion_request( + request, + kind="svs", + metadata_processor=metadata_processor, + ) + + return pre_process_func + + +def get_soulxsinger_post_process_func(od_config: OmniDiffusionConfig): + """Convert pipeline audio tensor output for offline consumers.""" + + def post_process_func(audio: torch.Tensor): + return convert_soulx_audio_output_to_numpy(audio) + + return post_process_func + + +def _expand_states(h: torch.Tensor, mel2token: torch.Tensor) -> torch.Tensor: + if mel2token.max() > h.size(1) - 1: + logger.warning( + "mel2token.max() (%s) is greater than h.size(1) - 1 (%s); clamping.", + mel2token.max(), + h.size(1) - 1, + ) + mel2token = torch.clamp(mel2token, 0, h.size(1) - 1) + mel2token_ = mel2token[..., None].repeat([1, 1, h.shape[-1]]) + return torch.gather(h, 1, mel2token_) + + +class PipelineSoulXSingerSVS(FlowMatchingAudioPipeline): + """Pipeline for the SoulX-Singer model (SoulX-Singer).""" + + _encoder_modules: ClassVar[list[str]] = [ + "mel", + "f0_encoder", + "preflow", + "note_text_encoder", + "note_pitch_encoder", + "note_type_encoder", + ] + + EXTRA_BODY_PARAMS: ClassVar[frozenset[str]] = frozenset( + { + "prompt_metadata_path", + "target_metadata_path", + "audio_path", + "prompt_audio", + "target_audio", + "language", + "vocal_sep", + "max_merge_duration", + "preprocess_weights_dir", + "preprocess_verbose", + "control", + "auto_shift", + "pitch_shift", + } + ) + + EXTRA_OUTPUT_PARAMS: ClassVar[frozenset[str]] = frozenset({"f0_shift"}) + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__(od_config=od_config, prefix=prefix) + self.metadata_processor = build_metadata_processor(od_config) + + self.f0_encoder = nn.Embedding(self.f0_bin, self.f0_dim) + self.preflow = nn.Sequential( + *[ + ConvNeXtV2Block( + self.text_dim, + self.text_dim * 2, + ) + for _ in range(self.encoder_config.num_layers) + ] + ) + + self.cfm_decoder = CFMDecoder(self.flow_matching_config) + + self.mel, self.vocoder = self._build_fp32_audio_modules(self.audio_config) + + self.note_text_encoder = nn.Embedding(self.vocab_size, self.text_dim) + self.note_pitch_encoder = nn.Embedding(256, self.pitch_dim) + self.note_type_encoder = nn.Embedding(256, self.type_dim) + + self._setup_soulx_profiler() + + def _encode_condition( + self, + *, + note_pitch: torch.Tensor, + note_type: torch.Tensor, + note_text: torch.Tensor, + mel2note: torch.Tensor, + f0_coarse: torch.Tensor, + ) -> torch.Tensor: + features = ( + self.note_pitch_encoder(note_pitch) + self.note_type_encoder(note_type) + self.note_text_encoder(note_text) + ) + features = self.preflow(features) + features = _expand_states(features, mel2note) + features = features + self.f0_encoder(f0_coarse) + return self._to_trunk_dtype(features)[0] + + def _infer_svs_segment( + self, + prompt_meta: dict, + target_meta: dict, + *, + pitch_shift: int, + num_inference_steps: int, + guidance_scale: float, + prompt_mel: torch.Tensor | None = None, + generator: torch.Generator | None = None, + ) -> torch.Tensor: + prompt_note_text = prompt_meta["phoneme"] + prompt_mel2note = prompt_meta["mel2note"] + prompt_note_type = prompt_meta["note_type"] + prompt_note_pitch = prompt_meta["note_pitch"] + prompt_f0 = prompt_meta["f0"] + + target_note_text = target_meta["phoneme"] + target_mel2note = target_meta["mel2note"] + target_note_type = target_meta["note_type"] + target_note_pitch = target_meta["note_pitch"] + target_f0 = target_meta["f0"] + + if target_f0 is None or prompt_f0 is None: + target_f0 = torch.zeros_like(target_mel2note).float() + prompt_f0 = torch.zeros_like(prompt_mel2note).float() + if target_note_pitch is None or prompt_note_pitch is None: + target_note_pitch = torch.zeros_like(target_note_type).int() + prompt_note_pitch = torch.zeros_like(prompt_note_type).int() + + len_prompt_note = prompt_note_pitch.shape[1] + len_prompt_mel = prompt_f0.shape[1] if prompt_f0 is not None else prompt_mel2note.shape[1] + + note_pitch = torch.cat([prompt_note_pitch, target_note_pitch], dim=1) + note_text = torch.cat([prompt_note_text, target_note_text], dim=1) + note_type = torch.cat([prompt_note_type, target_note_type], dim=1) + # Target note indices follow prompt notes in the concatenated score. + mel2note = torch.cat([prompt_mel2note, target_mel2note + len_prompt_note], dim=1) + + # pitch_shift is semitones; each coarse F0 bin is 20 cents (×5 bins per semitone). + f0_coarse_prompt = f0_to_coarse(prompt_f0) + f0_coarse_target = f0_to_coarse(target_f0, f0_shift=pitch_shift * 5) + f0_coarse = torch.cat([f0_coarse_prompt, f0_coarse_target], dim=1) + + note_pitch = note_pitch.clone() + note_pitch[note_pitch > 0] = note_pitch[note_pitch > 0] + pitch_shift + note_pitch = torch.clamp(note_pitch, 0, 255) + + if prompt_mel is None: + prompt_wav = prompt_meta["wav"] + with self._stage_timer("mel"): + prompt_mel = self._mel_from_wav(prompt_wav) + + if prompt_mel.shape[1] > len_prompt_mel: + prompt_mel = prompt_mel[:, :len_prompt_mel, :] + elif prompt_mel.shape[1] < len_prompt_mel: + logger.warning( + "prompt_mel length %s is shorter than metadata frames %s; padding mel.", + prompt_mel.shape[1], + len_prompt_mel, + ) + prompt_mel = F.pad(prompt_mel, (0, 0, 0, len_prompt_mel - prompt_mel.shape[1])) + + with self._stage_timer("cond_encode"): + cond = self._encode_condition( + note_pitch=note_pitch, + note_type=note_type, + note_text=note_text, + mel2note=mel2note, + f0_coarse=f0_coarse, + ) + + with self._stage_timer("cfm"): + generated_mel = self._run_flow_matching_loop( + prompt_mel, + cond, + num_inference_steps, + guidance_scale, + generator=generator, + ) + with self._stage_timer("vocoder"): + return self._mel_to_audio(generated_mel) + + def infer_svs_batch( + self, + payload: dict[str, Any], + *, + extra_args: dict[str, Any], + num_inference_steps: int, + guidance_scale: float, + generator: torch.Generator | None = None, + ) -> tuple[torch.Tensor, int]: + """Pure batch (non-streaming): run diffusion for each target segment in order and concat.""" + prompt_meta_raw = payload.get("prompt_meta") + target_meta_list = payload["target_meta_list"] + control = str(extra_args.get("control", "score")) + + prompt_meta = None + prompt_mel = None + if prompt_meta_raw: + prompt_meta = self._ensure_processed_meta(prompt_meta_raw) + prompt_meta = self._apply_control_mode(prompt_meta, control) + if prompt_meta and prompt_meta.get("wav") is not None: + prompt_mel = self._mel_from_wav(prompt_meta["wav"]) + + pieces: list[torch.Tensor] = [] + pitch_shift = 0 + + for idx, target_raw in enumerate(target_meta_list): + target_meta = self._ensure_processed_meta(target_raw) + target_meta = self._apply_control_mode(target_meta, control) + + if idx == 0: + pitch_shift = resolve_pitch_shift( + auto_shift=bool(extra_args.get("auto_shift", True)), + manual_shift=int(extra_args.get("pitch_shift", 0)), + prompt_f0=prompt_meta["f0"] if prompt_meta else None, + target_f0=target_meta["f0"], + prompt_note_pitch=prompt_meta["note_pitch"] if prompt_meta else None, + target_note_pitch=target_meta["note_pitch"], + ) + + seg = self._infer_svs_segment( + prompt_meta or {}, + target_meta, + pitch_shift=pitch_shift, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + prompt_mel=prompt_mel, + generator=generator, + ) + pieces.append(seg.squeeze()) + + if not pieces: + return torch.zeros(1, 0, device=self.device, dtype=torch.float32), pitch_shift + + full = torch.cat(pieces, dim=-1) + return full.unsqueeze(0), pitch_shift + + @staticmethod + def _apply_control_mode(meta: dict, control: str) -> dict: + meta["note_pitch"] = meta["note_pitch"] if control == "score" else None + meta["f0"] = meta["f0"] if control == "melody" else None + return meta + + def _ensure_processed_meta(self, meta: dict) -> dict: + if isinstance(meta.get("phoneme"), torch.Tensor): + return meta + return self.metadata_processor.process(meta, None) + + @torch.inference_mode() + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + return self._forward_batch_from_request( + req, + kind="svs", + metadata_key="f0_shift", + infer_batch_fn=self.infer_svs_batch, + prepare_extra_args=lambda extra_args, _sampling: normalize_svs_control_extra_args(extra_args), + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + del weights + self._load_soulx_checkpoint("model.pt") + return {name for name, _ in self.named_parameters()} diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py new file mode 100644 index 00000000000..aed87f7578c --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/__init__.py @@ -0,0 +1 @@ +"""Native SoulX-Singer preprocess helpers.""" diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py new file mode 100644 index 00000000000..5f54d47c9d7 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/g2p.py @@ -0,0 +1,86 @@ +import re + +_EN_WORD_RE = re.compile(r"^[A-Za-z]+(?:'[A-Za-z]+)*$") +_ZH_WORD_RE = re.compile(r"[\u4e00-\u9fff]") + +EN_FLAG = "en_" +YUE_FLAG = "yue_" +ZH_FLAG = "zh_" + +_g2p_zh = None +_g2p_en = None + + +def _ensure_g2p_backends() -> None: + global _g2p_zh, _g2p_en + if _g2p_zh is None: + from g2pM import G2pM + + _g2p_zh = G2pM() + if _g2p_en is None: + from g2p_en import G2p as G2pE + + _g2p_en = G2pE() + + +def is_chinese_char(word: str) -> bool: + if len(word) != 1: + return False + return bool(_ZH_WORD_RE.fullmatch(word)) + + +def is_english_word(word: str) -> bool: + if not word: + return False + return bool(_EN_WORD_RE.fullmatch(word)) + + +def g2p_cantonese(sent): + import ToJyutping + + return ToJyutping.get_jyutping_list(sent) + + +def g2p_mandarin(sent): + _ensure_g2p_backends() + return _g2p_zh(sent, tone=True, char_split=False) + + +def g2p_english(word): + _ensure_g2p_backends() + return _g2p_en(word) + + +def g2p_transform(words, lang): + zh_words = [] + transformed_words = [0] * len(words) + + for idx, w in enumerate(words): + if w == "<SP>": + transformed_words[idx] = w + continue + + w = w.replace("?", "").replace(".", "").replace("!", "").replace(",", "") + + if is_chinese_char(w): + zh_words.append([idx, w]) + else: + if is_english_word(w): + w = EN_FLAG + "-".join(g2p_english(w.lower())) + else: + w = "<SP>" + transformed_words[idx] = w + + sent = "".join([k[1] for k in zh_words]) + + if len(sent) > 0: + if lang == "Cantonese": + g2pm_rst = g2p_cantonese(sent) + g2pm_rst = [YUE_FLAG + k[1] for k in g2pm_rst] + else: + g2pm_rst = g2p_mandarin(sent) + g2pm_rst = [ZH_FLAG + k for k in g2pm_rst] + for p, w in zip([k[0] for k in zh_words], g2pm_rst): + transformed_words[p] = w + + return transformed_words diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py new file mode 100644 index 00000000000..8a861112186 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/metadata_utils.py @@ -0,0 +1,137 @@ +"""Metadata merge helpers for SoulX-Singer preprocess.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import soundfile as sf + + +@dataclass +class SegmentMetadata: + item_name: str + wav_fn: str + language: str + start_time_ms: int + end_time_ms: int + note_text: list[str] + note_dur: list[float] + note_pitch: list[int] + note_type: list[int] + origin_wav_fn: str | None = None + wav: np.ndarray | None = None + + def get(self, name: str, default: Any = None) -> Any: + return getattr(self, name, default) + + +def _merge_group( + audio: np.ndarray, + sample_rate: int, + segments: list[SegmentMetadata | dict[str, Any]], + output_dir: Path | None, + *, + end_extension_ms: int = 0, +) -> SegmentMetadata: + if not segments: + raise ValueError("segments must not be empty") + + words: list[str] = [] + durs: list[float] = [] + pitches: list[int] = [] + types: list[int] = [] + + for i, seg in enumerate(segments): + if i > 0: + prev_seg = segments[i - 1] + gap_ms = seg.get("start_time_ms", 0) - prev_seg.get("end_time_ms", 0) + if gap_ms > 0: + words.append("<SP>") + durs.append(gap_ms / 1000.0) + pitches.append(0) + types.append(1) + + words.extend(seg.get("note_text", [])) + durs.extend(seg.get("note_dur", [])) + pitches.extend(seg.get("note_pitch", [])) + types.extend(seg.get("note_type", [])) + + if end_extension_ms > 0: + words.append("<SP>") + durs.append(end_extension_ms / 1000.0) + pitches.append(0) + types.append(1) + + merged_words, merged_durs, merged_pitches, merged_types = [], [], [], [] + for word, dur, pitch, note_type in zip(words, durs, pitches, types): + if merged_words and word == "<SP>" and merged_words[-1] == "<SP>": + merged_durs[-1] += dur + else: + merged_words.append(word) + merged_durs.append(dur) + merged_pitches.append(pitch) + merged_types.append(note_type) + + languages = [s.get("language", "Mandarin") for s in segments if s.get("language")] + language = max(languages, key=languages.count) if languages else "Mandarin" + + start_ms = int(segments[0].get("start_time_ms", 0)) + end_ms = int(segments[-1].get("end_time_ms", 0)) + end_extension_ms + start_sample = int(start_ms * sample_rate // 1000) + end_sample = int(end_ms * sample_rate // 1000) + + first_item_name = segments[0].get("item_name", "segment") + song_prefix = "_".join(str(first_item_name).split("_")[:-1]) + item_name = f"{song_prefix}_{start_ms}_{end_ms}" + + wav_fn = "" + segment_audio = audio[start_sample:end_sample] + if output_dir is not None: + output_dir.mkdir(parents=True, exist_ok=True) + wav_path = output_dir / f"{item_name}.wav" + sf.write(wav_path, segment_audio, sample_rate) + wav_fn = str(wav_path) + + return SegmentMetadata( + item_name=item_name, + wav_fn=wav_fn, + language=language, + start_time_ms=start_ms, + end_time_ms=end_ms, + note_text=merged_words, + note_dur=merged_durs, + note_pitch=merged_pitches, + note_type=merged_types, + origin_wav_fn=segments[0].get("origin_wav_fn", ""), + wav=segment_audio if output_dir is None else None, + ) + + +def convert_metadata(item: SegmentMetadata, f0: np.ndarray | None = None) -> dict[str, Any]: + """Convert merged segment metadata into SoulX JSON metadata format. + + Matches upstream ``convert_metadata``: F0 always comes from ``*_f0.npy`` written + by ``extract_f0_file`` on the merged segment wav, not from an in-memory array. + """ + from vllm_omni.diffusion.models.soulx_singer.preprocess.g2p import g2p_transform + + if item.wav_fn: + f0_path = str(Path(item.wav_fn).with_suffix("")) + "_f0.npy" + f0_arr = np.load(f0_path) + elif f0 is not None: + f0_arr = f0 + else: + raise ValueError("convert_metadata requires item.wav_fn with sidecar f0 or an explicit f0 array") + + return { + "index": item.item_name, + "language": item.language, + "time": [item.start_time_ms, item.end_time_ms], + "duration": " ".join(f"{d:.2f}" for d in item.note_dur), + "text": " ".join(item.note_text), + "phoneme": " ".join(g2p_transform(item.note_text, item.language)), + "note_pitch": " ".join(map(str, item.note_pitch)), + "note_type": " ".join(map(str, item.note_type)), + "f0": " ".join(f"{x:.1f}" for x in f0_arr), + } diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py new file mode 100644 index 00000000000..ff4e4b7ca5a --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/payload.py @@ -0,0 +1,106 @@ +"""SoulX-Singer payload schema (guidance/precomputed tables, dummy builders).""" + +from __future__ import annotations + +from typing import Any + +import torch + +SOULX_PRECOMPUTED_KEYS_BY_KIND: dict[str, tuple[str, ...]] = { + "svs": ("prompt_metadata_path", "target_metadata_path", "audio_path"), + "svc": ("prompt_wav_path", "target_wav_path", "prompt_f0_path", "target_f0_path"), +} + + +# Key under which the SoulX preprocess payload is stored in +# prompt["additional_information"] or diffusion output metadata. +SOULX_PREPROCESSED_KEY = "soulx_preprocessed" + + +def get_soulx_preprocessed_payload(prompt: dict[str, Any]) -> dict[str, Any] | None: + """Extract attached payload from a diffusion request prompt dict.""" + additional = prompt.get("additional_information") or {} + payload = additional.get(SOULX_PREPROCESSED_KEY) + return payload if isinstance(payload, dict) else None + + +def relocate_tensors(payload: dict[str, Any], device: torch.device | str) -> dict[str, Any]: + """Recursively move all tensors inside the payload to the target device. + + Called after decode when the payload is about to be consumed by the + diffusion model (which may live on a different device / rank). + """ + if isinstance(payload, torch.Tensor): + return payload.to(device, non_blocking=True) + if isinstance(payload, dict): + return {key: relocate_tensors(value, device) for key, value in payload.items()} + if isinstance(payload, list): + return [relocate_tensors(item, device) for item in payload] + return payload + + +def consume_payload( + req: Any, + expected_kind: str, + device: torch.device | str, +) -> dict[str, Any]: + """Read the payload attached by pre_process_func, validate the kind, + and relocate tensors to the compute device. + """ + prompts = getattr(req, "prompts", None) or [] + if not prompts: + raise ValueError(f"SoulX-Singer {expected_kind} request has no prompts.") + prompt = prompts[0] + if isinstance(prompt, str): + raise ValueError(f"SoulX-Singer {expected_kind} forward requires pre_process_func output on the prompt.") + payload = get_soulx_preprocessed_payload(prompt) + if payload is None or payload.get("kind") != expected_kind: + raise ValueError( + f"SoulX-Singer {expected_kind} forward requires additional_information['soulx_preprocessed'] " + "produced by pre_process_func." + ) + return relocate_tensors(payload, device) + + +def has_precomputed(extra_args: dict[str, Any], kind: str) -> bool: + keys = SOULX_PRECOMPUTED_KEYS_BY_KIND[kind] + return all(extra_args.get(key) for key in keys) + + +def build_dummy_payload(kind: str, device: torch.device) -> dict[str, Any]: + mel_frames = 4 + hop_size = 480 + n_fft = 1920 + wav_samples = max(mel_frames * hop_size, n_fft) + if kind == "svc": + voiced_f0 = torch.full((1, mel_frames), 100.0, device=device, dtype=torch.float64) + return { + "kind": "svc", + "prompt_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), + "target_wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), + "prompt_f0": voiced_f0.clone(), + "target_f0": voiced_f0.clone(), + } + return { + "kind": "svs", + "prompt_meta": { + "phoneme": torch.zeros(1, 2, device=device, dtype=torch.long), + "note_pitch": torch.zeros(1, mel_frames, device=device, dtype=torch.long), + "note_type": torch.zeros(1, mel_frames, device=device, dtype=torch.long), + "mel2note": torch.zeros(1, mel_frames, device=device, dtype=torch.long), + "f0": torch.zeros(1, mel_frames, device=device, dtype=torch.float32), + "wav": torch.zeros(1, wav_samples, device=device, dtype=torch.float32), + }, + "target_meta_list": [ + { + "text": "dummy", + "phoneme": " ".join(["<SP>"] * mel_frames), + "f0": " ".join(["100"] * mel_frames), + "language": "Mandarin", + "note_pitch": " ".join(["0"] * mel_frames), + "note_type": " ".join(["1"] * mel_frames), + "duration": " ".join(["1"] * mel_frames), + "time": [0, max(int(wav_samples / 24000 * 1000), 100)], + } + ], + } diff --git a/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py b/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py new file mode 100644 index 00000000000..a8784fcba04 --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/preprocess/pre_process.py @@ -0,0 +1,223 @@ +"""Shared ``pre_process_func`` helpers for SoulX-Singer pipelines (batch / full-payload only).""" + +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.pipeline import SoulXPreprocessPipeline +from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( + SOULX_PRECOMPUTED_KEYS_BY_KIND, + SOULX_PREPROCESSED_KEY, + build_dummy_payload, + get_soulx_preprocessed_payload, + has_precomputed, +) +from vllm_omni.diffusion.models.soulx_singer.utils import ( + MetadataProcessor, + load_config, + resolve_phoneset_path, + validate_soulx_extra_args, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniTextPrompt + +logger = init_logger(__name__) + + +def is_warmup_request(request: OmniDiffusionRequest) -> bool: + return request.is_dummy_run() + + +def resolve_preprocess_audio(prompt: dict[str, Any], extra_args: dict[str, Any]) -> tuple[Any, Any]: + mm = prompt.get("multi_modal_data") or {} + prompt_audio = mm.get("audio") if isinstance(mm, dict) else None + if prompt_audio is None: + prompt_audio = extra_args.get("prompt_audio") + return prompt_audio, extra_args.get("target_audio") + + +def _preprocess_inputs_missing_error(kind: str) -> ValueError: + return ValueError( + f"SoulX-Singer {kind} requires precomputed paths " + f"{list(SOULX_PRECOMPUTED_KEYS_BY_KIND[kind])}, an attached soulx_preprocessed payload " + f"(multi-stage preprocess stage), or run with pipeline soulxsinger_{kind}." + ) + + +def build_precomputed_payload( + kind: str, + extra_args: dict[str, Any], + *, + metadata_processor, + sample_rate: int | None, + device, +) -> dict[str, Any]: + if kind == "svs": + return SoulXPreprocessPipeline.build_svs_payload_from_paths( + prompt_metadata_path=str(extra_args["prompt_metadata_path"]), + target_metadata_path=str(extra_args["target_metadata_path"]), + audio_path=str(extra_args["audio_path"]), + metadata_processor=metadata_processor, + ) + if sample_rate is None or device is None: + raise ValueError("svc precomputed payload requires sample_rate and device") + return SoulXPreprocessPipeline.build_svc_payload_from_paths( # type: ignore + prompt_wav_path=str(extra_args["prompt_wav_path"]), + target_wav_path=str(extra_args["target_wav_path"]), + prompt_f0_path=str(extra_args["prompt_f0_path"]), + target_f0_path=str(extra_args["target_f0_path"]), + sample_rate=int(sample_rate), + device=device, + ) + + +def build_warmup_payload( + kind: str, + *, + metadata_processor, + device, + sample_rate: int | None, +) -> dict[str, Any]: + if kind == "svs": + if metadata_processor is None: + raise ValueError("SVS warmup requires metadata_processor") + payload = build_dummy_payload("svs", torch.device("cpu")) + dummy_prompt = payload["prompt_meta"] + processed = metadata_processor.process(dict(payload["target_meta_list"][0])) + payload["prompt_meta"] = { + key: value.clone() if isinstance(value, torch.Tensor) else value for key, value in processed.items() + } + if isinstance(dummy_prompt.get("wav"), torch.Tensor): + payload["prompt_meta"]["wav"] = dummy_prompt["wav"].clone() + return payload + if device is None or sample_rate is None: + raise ValueError("svc warmup requires device and sample_rate") + return build_dummy_payload("svc", device) + + +def build_preprocess_payload( + kind: str, + *, + prompt: dict[str, Any], + extra_args: dict[str, Any], + preprocess: nn.Module, + metadata_processor, + device, + sample_rate: int, +) -> dict[str, Any]: + if has_precomputed(extra_args, kind): + return build_precomputed_payload( + kind, + extra_args, + metadata_processor=metadata_processor, + sample_rate=sample_rate, + device=device, + ) + + prompt_audio, target_audio = resolve_preprocess_audio(prompt, extra_args) + if prompt_audio is None or target_audio is None: + raise _preprocess_inputs_missing_error(kind) + + if kind == "svc": + return preprocess.build_svc_payload_from_audio( # type: ignore + prompt_audio=prompt_audio, + target_audio=target_audio, + sample_rate=sample_rate, + device=device, + vocal_sep=extra_args.get("vocal_sep"), + ) + + return preprocess.build_svs_payload_from_audio( # type: ignore + prompt_audio=prompt_audio, + target_audio=target_audio, + metadata_processor=metadata_processor, + language=str(extra_args.get("language", "Mandarin")), + vocal_sep=extra_args.get("vocal_sep"), + prompt_vocal_sep=extra_args.get("prompt_vocal_sep"), + target_vocal_sep=extra_args.get("target_vocal_sep"), + prompt_max_merge_duration_ms=extra_args.get("prompt_max_merge_duration"), + target_max_merge_duration_ms=extra_args.get("target_max_merge_duration"), + ) + + +def attach_preprocess_for_diffusion_request( + request: OmniDiffusionRequest, + *, + kind: str, + metadata_processor=None, + sample_rate: int | None = None, + device=None, +) -> OmniDiffusionRequest: + """Resolve preprocess payload for stage-1 diffusion (warmup / IPC / precomputed paths).""" + + extra_args = validate_soulx_extra_args( + kind, + dict(getattr(request.sampling_params, "extra_args", None) or {}), + ) + prompt = request.prompt + prompt = OmniTextPrompt(prompt=prompt) if isinstance(prompt, str) else prompt + + if is_warmup_request(request): + request.sampling_params.num_inference_steps = 1 + payload = build_warmup_payload( + kind, + metadata_processor=metadata_processor, + device=device, + sample_rate=sample_rate, + ) + elif get_soulx_preprocessed_payload(prompt): + request.prompt = prompt + if kind == "svs": + extra_args = normalize_svs_control_extra_args(extra_args) + request.sampling_params.extra_args = extra_args + return request + elif has_precomputed(extra_args, kind): + payload = build_precomputed_payload( + kind, + extra_args, + metadata_processor=metadata_processor, + sample_rate=sample_rate, + device=device, + ) + else: + raise _preprocess_inputs_missing_error(kind) + + if payload.get("kind") != kind: + raise ValueError(f"Invalid {kind} preprocess payload kind: {payload.get('kind')}") + prompt.setdefault("additional_information", {})[SOULX_PREPROCESSED_KEY] = payload + request.prompt = prompt + + if kind == "svs": + extra_args = normalize_svs_control_extra_args(extra_args) + request.sampling_params.extra_args = extra_args + return request + + +def normalize_svs_control_extra_args(extra_args: dict[str, Any]) -> dict[str, Any]: + control = extra_args.get("control") + if control is None: + control = "score" + logger.warning("control is not provided, using 'score' as default") + elif control not in ("score", "melody"): + raise ValueError(f"Invalid control: {control}. Must be one of: ['score', 'melody']") + extra_args["control"] = control + return extra_args + + +def build_metadata_processor(od_config: OmniDiffusionConfig): + model_dir = od_config.model + config_path = Path(model_dir) / "config.yaml" + hf_config = load_config(str(config_path)) + audio_config = hf_config.audio + + return MetadataProcessor( + hop_size=audio_config.hop_size, + sample_rate=audio_config.sample_rate, + phoneset_path=resolve_phoneset_path(model_dir), + device=str(get_local_device()), + ) diff --git a/vllm_omni/diffusion/models/soulx_singer/utils.py b/vllm_omni/diffusion/models/soulx_singer/utils.py new file mode 100644 index 00000000000..4f201c4077f --- /dev/null +++ b/vllm_omni/diffusion/models/soulx_singer/utils.py @@ -0,0 +1,565 @@ +import json +from pathlib import Path +from typing import Any, Literal + +import numpy as np +import torch +import torchaudio +from huggingface_hub import snapshot_download +from omegaconf import DictConfig, OmegaConf +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import OmniDiffusionConfig + +logger = init_logger(__name__) + +_MODEL_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer" +_PREPROCESS_WEIGHTS_REPO = "Soul-AILab/SoulX-Singer-Preprocess" + + +def _patch_torchaudio_load() -> None: + """Borrowed from MOSS-TTS-Nano. Patch torchaudio.load to use soundfile + if torchcodec is unavailable. + """ + try: + import torchaudio + + torchaudio # noqa + import torchcodec # noqa: F401 + + return + except Exception: + pass + + import soundfile as sf + + def _soundfile_load(path, frame_offset=0, num_frames=-1, normalize=True, channels_first=True, format=None): + data, sr = sf.read(str(path), dtype="float32", always_2d=True) + if frame_offset > 0: + data = data[frame_offset:] + if num_frames > 0: + data = data[:num_frames] + waveform = torch.from_numpy(data) + if channels_first: + waveform = waveform.T + return waveform, sr + + def _soundfile_save(path, src, sample_rate, channels_first=True, **kwargs): + wav = src.detach().cpu().float().numpy() + if channels_first and wav.ndim == 2: + wav = wav.T + sf.write(str(path), wav, sample_rate) + + try: + import torchaudio + + torchaudio.load = _soundfile_load + torchaudio.save = _soundfile_save + logger.info("Patched torchaudio.load/save to use soundfile (torchcodec unavailable)") + except Exception as e: + logger.warning("Could not patch torchaudio: %s", e) + + +def load_config(config_path: str | Path) -> DictConfig: + """ + Load a configuration file and optionally merges it with a base configuration. + + Args: + config_path (Path): Path to the configuration file. + """ + # Load the initial configuration from the given path + config = OmegaConf.load(str(config_path)) + + # Check if there is a base configuration specified and merge if necessary + if config.get("base_config", None) is not None: + base_config = OmegaConf.load(str(config.get("base_config"))) + config = OmegaConf.merge(base_config, config) + + return config + + +SoulXKind = Literal["svs", "svc"] + +_SVS_ARCHITECTURE = "SoulXSingerPipeline" +_SVC_ARCHITECTURE = "SoulXSingerSVCPipeline" +_ARCHITECTURE_TO_KIND: dict[str, SoulXKind] = { + _SVS_ARCHITECTURE: "svs", + _SVC_ARCHITECTURE: "svc", +} + + +def load_model_config_json(model_dir: str | Path) -> dict[str, Any]: + """Load ``config.json`` from a SoulX-Singer model view directory.""" + cfg_path = Path(model_dir) / "config.json" + if not cfg_path.is_file(): + raise FileNotFoundError( + f"SoulX-Singer config.json not found at {cfg_path}. " + "Create one with architectures ['SoulXSingerPipeline'] (SVS) or " + "['SoulXSingerSVCPipeline'] (SVC); see examples/offline_inference/text_to_speech/README.md." + ) + with cfg_path.open(encoding="utf-8") as f: + return json.load(f) + + +def resolve_soulx_kind(model_dir: str | Path) -> SoulXKind: + """Resolve SVS vs SVC from ``config.json`` ``architectures[0]``.""" + cfg = load_model_config_json(model_dir) + archs = cfg.get("architectures") or [] + if not archs: + raise ValueError( + f"SoulX-Singer config.json at {Path(model_dir) / 'config.json'} is missing 'architectures'. " + "Set architectures to ['SoulXSingerPipeline'] for SVS or ['SoulXSingerSVCPipeline'] for SVC." + ) + arch = archs[0] + kind = _ARCHITECTURE_TO_KIND.get(arch) + if kind is None: + raise ValueError( + f"Unknown SoulX-Singer architecture {arch!r} in config.json. Expected one of {list(_ARCHITECTURE_TO_KIND)}." + ) + return kind + + +def validate_soulx_extra_args(kind: SoulXKind, extra_args: dict[str, Any] | None) -> dict[str, Any]: + """Validate request ``extra_args`` against the model mode from config.json.""" + from vllm_omni.diffusion.models.soulx_singer.preprocess.payload import ( + SOULX_PRECOMPUTED_KEYS_BY_KIND, + has_precomputed, + ) + + extra_args = dict(extra_args or {}) + declared = extra_args.pop("kind", None) + if declared is not None and declared != kind: + logger.warning( + "Ignoring extra_args kind=%r; model config.json declares %r.", + declared, + kind, + ) + + other_kind: SoulXKind = "svc" if kind == "svs" else "svs" + if has_precomputed(extra_args, other_kind): + raise ValueError( + f"SoulX-Singer config.json declares {kind!r}, but request supplies " + f"{other_kind} precomputed paths {list(SOULX_PRECOMPUTED_KEYS_BY_KIND[other_kind])}." + ) + + if kind == "svc": + for key in ("language", "control"): + if extra_args.get(key) is not None: + logger.warning( + "Ignoring SVS-only extra_arg %r=%r on an SVC model (config.json).", + key, + extra_args.pop(key), + ) + for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svs"]: + if extra_args.get(key): + raise ValueError(f"SoulX-Singer SVC model (config.json) does not accept SVS precomputed key {key!r}.") + else: + for key in SOULX_PRECOMPUTED_KEYS_BY_KIND["svc"]: + if extra_args.get(key): + raise ValueError(f"SoulX-Singer SVS model (config.json) does not accept SVC precomputed key {key!r}.") + + return extra_args + + +# ---------------- utils for model weights ---------------- + + +def resolve_preprocess_weights_root(od_config: OmniDiffusionConfig) -> Path: + """Locate preprocess weights on disk or download from Hugging Face.""" + + # default local dir + if Path(od_config.model).parent.is_dir(): + local_dir = Path(od_config.model).parent / "SoulX-Singer-Preprocess" + else: + local_dir = Path(__file__).parents[5] / "pretrained" / "SoulX-Singer-Preprocess" + + local_dir = local_dir.expanduser() + if local_dir.is_dir(): + logger.info("Using SoulX preprocess weights from %s", local_dir) + return local_dir + + logger.info( + "SoulX preprocess weights not found locally; downloading %s", + _PREPROCESS_WEIGHTS_REPO, + ) + downloaded = snapshot_download( + _PREPROCESS_WEIGHTS_REPO, + allow_patterns=["*"], + local_dir=local_dir, + ) + return Path(downloaded) + + +def preprocess_weight_paths(weights_root: Path) -> dict[str, str]: + """Map upstream relative paths to absolute paths under ``weights_root``.""" + root = weights_root.resolve() + return { + "rmvpe": str(root / "rmvpe/rmvpe.pt"), + "sep_ckpt": str(root / "mel-band-roformer-karaoke/mel_band_roformer_karaoke_becruily.ckpt"), + "sep_config": str(root / "mel-band-roformer-karaoke/config_karaoke_becruily.yaml"), + "asr_zh": str(root / "speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch"), + "asr_en": str(root / "parakeet-tdt-0.6b-v2/parakeet-tdt-0.6b-v2.nemo"), + "rosvot": str(root / "rosvot/rosvot/model.pt"), + } + + +def resolve_phoneset_path(model_dir: str) -> str: + """Resolve phoneme vocabulary; must be present in model directory.""" + candidates = ( + Path(model_dir) / "phoneme" / "phone_set.json", + Path(model_dir) / "phone_set.json", + ) + for path in candidates: + if path.is_file(): + return str(path) + raise FileNotFoundError( + "SoulX-Singer phoneset not found. Expected one of: " + f"{[str(p) for p in candidates]}. " + "Copy phone_set.json into the model directory. " + "See `examples/offline_inference/text_to_speech/README.md` for instructions." + ) + + +# ---------------- utils for data processing ---------------- + + +def load_wav(wav_path: str, sample_rate: int) -> torch.Tensor: + """Load wav file and resample to target sample rate. + + Args: + wav_path (str): Path to wav file. + sample_rate (int): Target sample rate. + + Returns: + torch.Tensor: Waveform tensor with shape (1, T). + """ + _patch_torchaudio_load() + waveform, sr = torchaudio.load(wav_path) + + if sr != sample_rate: + waveform = torchaudio.functional.resample(waveform, sr, sample_rate) + + if len(waveform.shape) > 1 and waveform.shape[0] > 1: + waveform = torch.mean(waveform, dim=0, keepdim=True) + + return waveform + + +def f0_to_coarse(f0, f0_bin=361, f0_min=32.7031956625, f0_shift=0): + """ + Convert continuous F0 values to discrete F0 bins (SIL and C1 - B6, 361 bins). + args: + f0: continuous F0 values + f0_bin: number of F0 bins + f0_min: minimum F0 value + f0_shift: shift value for F0 bins + returns: + f0_coarse: discrete F0 bins + """ + is_torch = isinstance(f0, torch.Tensor) + uv_mask = f0 <= 0 + + if is_torch: + f0_min_tensor = f0.new_tensor(f0_min) + f0_safe = torch.maximum(f0, f0_min_tensor) + f0_cents = 1200 * torch.log2(f0_safe / f0_min_tensor) + else: + f0_safe = np.maximum(f0, f0_min) + f0_cents = 1200 * np.log2(f0_safe / f0_min) + + f0_coarse = (f0_cents / 20) + 1 + + if is_torch: + f0_coarse = torch.round(f0_coarse).long() + f0_coarse = torch.clamp(f0_coarse, min=1, max=f0_bin - 1) + else: + f0_coarse = np.rint(f0_coarse).astype(int) + f0_coarse = np.clip(f0_coarse, 1, f0_bin - 1) + + f0_coarse[uv_mask] = 0 + + if f0_shift != 0: + if is_torch: + voiced = f0_coarse > 0 + if voiced.any(): + shifted = f0_coarse[voiced] + f0_shift + f0_coarse[voiced] = torch.clamp(shifted, 1, f0_bin - 1) + else: + voiced = f0_coarse > 0 + if np.any(voiced): + shifted = f0_coarse[voiced] + f0_shift + f0_coarse[voiced] = np.clip(shifted, 1, f0_bin - 1) + + return f0_coarse + + +def resolve_pitch_shift( + *, + auto_shift: bool, + manual_shift: int, + prompt_f0: torch.Tensor | None = None, + target_f0: torch.Tensor | None = None, + prompt_note_pitch: torch.Tensor | None = None, + target_note_pitch: torch.Tensor | None = None, +) -> int: + """Resolve semitone shift for auto_shift; ``f0_to_coarse(..., f0_shift=shift * 5)`` consumes it.""" + if not auto_shift or manual_shift != 0: + return int(manual_shift) + + if prompt_note_pitch is not None and target_note_pitch is not None: + target_pitched = target_note_pitch[target_note_pitch >= 1] + prompt_pitched = prompt_note_pitch[prompt_note_pitch >= 1] + if target_pitched.numel() > 0 and prompt_pitched.numel() > 0: + shift = torch.round(prompt_pitched.median() - target_pitched.median()) + if torch.isfinite(shift): + return int(shift.item()) + + if prompt_f0 is not None and target_f0 is not None: + target_voiced = target_f0[target_f0 > 0] + prompt_voiced = prompt_f0[prompt_f0 > 0] + if target_voiced.numel() > 0 and prompt_voiced.numel() > 0: + shift = torch.round(torch.log2(prompt_voiced.median() / target_voiced.median()) * 1200 / 100) + if torch.isfinite(shift): + return int(shift.item()) + + return 0 + + +class MetadataProcessor: + """Data processor for SoulX-Singer""" + + def __init__( + self, + hop_size: int, + sample_rate: int, + phoneset_path: str, + device: str = "cuda", + ): + """Initialize data processor. + + Args: + hop_size (int): Hop size in samples. + sample_rate (int): Sample rate in Hz. + phoneset_path (str): Path to phoneme set JSON file. + device (str): Device to use for tensor operations. + """ + self.hop_size = hop_size + self.sample_rate = sample_rate + self.device = device + self.load_phoneme_id_map(phoneset_path) + + def load_phoneme_id_map(self, phoneset_path: str): + with open(phoneset_path, encoding="utf-8") as f: + phoneset = json.load(f) + self.phone2idx = {ph: idx for idx, ph in enumerate(phoneset)} + + def merge_phoneme(self, meta): + merged_items = [] + + duration = [float(x) for x in meta["duration"].split()] + phoneme = [str(x).replace("<AP>", "<SP>") for x in meta["phoneme"].split()] + note_pitch = [int(x) for x in meta["note_pitch"].split()] + note_type = [int(x) if phoneme[i] != "<SP>" else 1 for i, x in enumerate(meta["note_type"].split())] + + for i in range(len(phoneme)): + if ( + i > 0 + and phoneme[i] == phoneme[i - 1] == "<SP>" + and note_type[i] == note_type[i - 1] + and note_pitch[i] == note_pitch[i - 1] + ): + merged_items[-1][1] += duration[i] + else: + merged_items.append([phoneme[i], duration[i], note_pitch[i], note_type[i]]) + + meta["phoneme"] = [x[0] for x in merged_items] + meta["duration"] = [x[1] for x in merged_items] + meta["note_pitch"] = [x[2] for x in merged_items] + meta["note_type"] = [x[3] for x in merged_items] + + return meta + + def preprocess( + self, + note_duration: list[float], + phonemes: list[str], + note_pitch: list[int], + note_type: list[int], + ): + """ + Insert <BOW> and <EOW> for each note. + Get aligned indices for each frame. + + Args: + note_duration: Duration of each note in seconds + phonemes: Phoneme sequence for each note + note_pitch: Pitch value for each note + note_type: Type value for each note + + """ + sample_rate = self.sample_rate + hop_size = self.hop_size + duration = sum(note_duration) * sample_rate / hop_size + mel2note = torch.zeros(int(duration), dtype=torch.long, device=self.device) + + ph_locations = [] # idx at mel scale and length + new_phonemes = [] + dur_sum = 0 + + note2origin = [] + + for ph_idx in range(len(phonemes)): + dur = int(np.round(dur_sum * sample_rate / hop_size)) + dur = min(dur, len(mel2note) - 1) + new_phonemes.append("<BOW>") + note2origin.append(ph_idx) + if phonemes[ph_idx][:3] == "en_": + en_phs = ["en_" + x for x in phonemes[ph_idx][3:].split("-")] + [ + "<SEP>" + ] # <sep> between en words in one note + ph_locations.append([dur, max(1, len(en_phs))]) + new_phonemes.extend(en_phs) + note2origin.extend([ph_idx] * len(en_phs)) + else: + ph_locations.append([dur, 1]) + new_phonemes.append(phonemes[ph_idx]) + note2origin.append(ph_idx) + new_phonemes.append("<EOW>") + note2origin.append(ph_idx) + dur_sum += note_duration[ph_idx] + + ph_idx = 1 + for idx, (i, j) in enumerate(ph_locations): + next_phoneme_start = ph_locations[idx + 1][0] if idx < len(ph_locations) - 1 else len(mel2note) + if i >= len(mel2note) or i + j > len(mel2note): + break + if i < len(mel2note) and mel2note[i] > 0: + logger.warning(f"warning: overlap of {idx}: {mel2note[i]}") + while i < len(mel2note) and mel2note[i] > 0: + i += 1 + mel2note[i] = ph_idx + k = i + 1 + while k + j < next_phoneme_start: + mel2note[k : k + j] = torch.arange(ph_idx, ph_idx + j, device=self.device) + 1 + k += j + mel2note[next_phoneme_start - 1] = ph_idx + j + 1 + ph_idx += j + 2 # <BOW> + ph repeats + <EOW> + + new_phonemes = ["<PAD>"] + new_phonemes + new_note_pitch = [0] + [note_pitch[k] for k in note2origin] + new_note_type = [1] + [note_type[k] for k in note2origin] + + return { + "phoneme": torch.tensor([self.phone2idx[x] for x in new_phonemes], device=self.device).unsqueeze(0), + "note_pitch": torch.tensor(new_note_pitch, device=self.device).unsqueeze(0), + "note_type": torch.tensor(new_note_type, device=self.device).unsqueeze(0), + "mel2note": mel2note.unsqueeze(0), + } + + def process(self, meta: dict, wav_path: str | None = None) -> dict[str, torch.Tensor | None]: + meta = self.merge_phoneme(meta) + + item = self.preprocess( + meta["duration"], + meta["phoneme"], + meta["note_pitch"], + meta["note_type"], + ) + + f0 = [float(x) for x in meta.get("f0", "").split()] + min_frame = min(item["mel2note"].shape[1], len(f0)) if len(f0) > 0 else item["mel2note"].shape[1] + item["f0"] = torch.tensor(f0, device=self.device)[:min_frame].unsqueeze(0).float() if len(f0) > 0 else None + item["mel2note"] = item["mel2note"][:, :min_frame] + + if wav_path is not None: + waveform = load_wav(wav_path, self.sample_rate) + item["wav"] = waveform.to(self.device)[:, : min_frame * self.hop_size] + + return item + + +def build_vocal_segments( + f0, + *, + f0_rate: int = 50, + ignore_silent_frames_thresh: int = 5, + min_duration_sec_per_segment: float = 5.0, + max_duration_sec_per_segment: float = 30.0, + num_overlaps: int = 1, + ignore_silent_frames: bool = True, +) -> tuple[list[tuple[float, float]], list[tuple[float, float]]]: + """Build vocal segments from an F0 contour for chunked SVC inference.""" + if isinstance(f0, torch.Tensor): + f0_np = f0.detach().float().cpu().numpy() + else: + f0_np = np.asarray(f0, dtype=np.float32) + f0_np = np.squeeze(f0_np) + + total_frames = int(f0_np.shape[0]) + if total_frames == 0: + return [], [] + + min_frames = max(1, int(round(min_duration_sec_per_segment * f0_rate))) + max_frames = max(1, int(round(max_duration_sec_per_segment * f0_rate))) + + split_points = [0] + + def append_split_point(point: int): + point = int(max(0, min(point, total_frames))) + while point - split_points[-1] > max_frames: + split_points.append(split_points[-1] + max_frames) + if point > split_points[-1]: + split_points.append(point) + + idx = 0 + while idx < total_frames: + if f0_np[idx] == 0: + run_start = idx + while idx < total_frames and f0_np[idx] == 0: + idx += 1 + run_end = idx + if (run_end - run_start) >= ignore_silent_frames_thresh: + split_point = max(run_end - 5, (run_start + run_end) // 2) + append_split_point(split_point) + else: + idx += 1 + append_split_point(total_frames) + + segments: list[tuple[float, float]] = [] + overlap_segments: list[tuple[float, float]] = [] + + def append_segment(start_idx: int, end_idx: int, overlaps: int = num_overlaps): + segments.append((split_points[start_idx] / f0_rate, split_points[end_idx] / f0_rate)) + overlap_start_idx = start_idx + if start_idx > 0 and (split_points[end_idx] - split_points[start_idx - overlaps]) <= max_frames: + overlap_start_idx = start_idx - overlaps + overlap_segments.append((split_points[overlap_start_idx] / f0_rate, split_points[end_idx] / f0_rate)) + + segment_start, segment_end = 0, 1 + while segment_start < len(split_points) - 1: + while ( + segment_end < len(split_points) and (split_points[segment_end] - split_points[segment_start]) < min_frames + ): + segment_end += 1 + + if segment_end >= len(split_points): + append_segment(segment_start, len(split_points) - 1, overlaps=num_overlaps) + break + append_segment(segment_start, segment_end, overlaps=num_overlaps) + segment_start = segment_end + segment_end = segment_start + 1 + + if ignore_silent_frames: + filtered_idx = [] + for i, seg in enumerate(overlap_segments): + start_frame = int(seg[0] * f0_rate) + end_frame = int(seg[1] * f0_rate) + seg_frames = end_frame - start_frame + voice_frames = np.sum(f0_np[start_frame:end_frame] > 0) + if voice_frames / seg_frames > 0.05 and voice_frames >= 10: + filtered_idx.append(i) + + overlap_segments = [overlap_segments[i] for i in filtered_idx] + segments = [segments[i] for i in filtered_idx] + + return overlap_segments, segments diff --git a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py index f59215e0c23..41f12f65496 100644 --- a/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py +++ b/vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py @@ -29,14 +29,14 @@ 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 SupportAudioOutput +from vllm_omni.diffusion.models.interface import SupportAudioOutput, SupportsComponentDiscovery from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import ( StableAudioDiTModel, StableAudioSchedulerWrapper, ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch logger = init_logger(__name__) @@ -65,7 +65,7 @@ def post_process_func( return post_process_func -class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfilerMixin): +class StableAudioPipeline(nn.Module, SupportAudioOutput, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin): """ Pipeline for text-to-audio generation using Stable Audio Open. @@ -77,6 +77,8 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil prefix: Weight prefix for loading (default: "") """ + supports_request_batch = False + # Picked up by ``supports_audio_output`` in the diffusion engine so the # default stage metadata reports ``final_output_type="audio"`` and the # ``multimodal_output`` payload includes the sample rate (mirrors the @@ -84,6 +86,11 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil support_audio_output: ClassVar[bool] = True audio_sample_rate: ClassVar[int] = 44100 + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = ["projection_model"] + def __init__( self, *, @@ -378,37 +385,21 @@ def prepare_latents( latents = latents * self.scheduler.init_noise_sigma return latents - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - negative_prompt: str | list[str] | None = None, - audio_end_in_s: float | None = None, - audio_start_in_s: float = 0.0, - guidance_scale: float = 7.0, - num_waveforms_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.Tensor | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - output_type: str = "np", - ) -> DiffusionOutput: + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: """ Generate audio from text prompt. Args: - req: OmniDiffusionRequest containing generation parameters - prompt: Text prompt for audio generation - negative_prompt: Negative prompt for CFG - audio_end_in_s: Audio end time in seconds (max ~47s for stable-audio-open-1.0) - audio_start_in_s: Audio start time in seconds - guidance_scale: CFG scale - num_waveforms_per_prompt: Number of audio outputs per prompt - generator: Random generator for reproducibility - latents: Pre-generated latents - prompt_embeds: Pre-computed prompt embeddings - negative_prompt_embeds: Pre-computed negative prompt embeddings - output_type: Output format ("np", "pt", or "latent") + req: OmniDiffusionRequest containing generation parameters. + The `req.sampling_params.extra_args` can include the following keys: + - audio_start_in_s (`float`, *optional*, defaults to 0.0): + Start time of the audio in seconds. + - audio_end_in_s (`float`, *optional*): + End time of the audio in seconds. + - num_waveforms_per_prompt (`int`, *optional*, defaults to 1): + Number of audio outputs per prompt. + - output_type (`str`, *optional*, defaults to "np"): + Output format ("np", "pt", or "latent"). Returns: DiffusionOutput containing generated audio @@ -416,28 +407,33 @@ def forward( # Extract from request # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - num_inference_steps = req.sampling_params.num_inference_steps - if num_inference_steps is None: - num_inference_steps = 100 # Default steps - num_inference_steps = req.sampling_params.num_inference_steps - if num_inference_steps is None: - num_inference_steps = 50 # Default steps + + prompt_embeds = None + negative_prompt_embeds = None + + num_inference_steps = req.sampling_params.num_inference_steps or 50 if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 7.0 - if generator is None: - generator = req.sampling_params.generator + generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) + latents = req.sampling_params.latents # Get audio duration from request extra params or defaults - audio_start_in_s = req.sampling_params.extra_args.get("audio_start_in_s", audio_start_in_s) - audio_end_in_s = req.sampling_params.extra_args.get("audio_end_in_s", audio_end_in_s) + audio_start_in_s: float = req.sampling_params.extra_args.get("audio_start_in_s", 0.0) + audio_end_in_s: float | None = req.sampling_params.extra_args.get("audio_end_in_s", None) + num_waveforms_per_prompt: int = req.sampling_params.extra_args.get( + "num_waveforms_per_prompt", req.sampling_params.num_outputs_per_prompt or 1 + ) + output_type = req.sampling_params.output_type or "np" # Calculate audio length downsample_ratio = self.vae.hop_length @@ -467,12 +463,7 @@ def forward( ) # Determine batch size - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = prompt_embeds.shape[0] + batch_size = len(prompt) device = self.device do_classifier_free_guidance = guidance_scale > 1.0 @@ -493,7 +484,7 @@ def forward( audio_start_in_s, audio_end_in_s, device, - do_classifier_free_guidance and (negative_prompt is not None or negative_prompt_embeds is not None), + do_classifier_free_guidance and negative_prompt is not None, batch_size, ) @@ -508,7 +499,7 @@ def forward( ) # Handle CFG without negative prompt - if do_classifier_free_guidance and negative_prompt_embeds is None and negative_prompt is None: + if do_classifier_free_guidance and negative_prompt is None: negative_text_audio_duration_embeds = torch.zeros_like(text_audio_duration_embeds) text_audio_duration_embeds = torch.cat( [negative_text_audio_duration_embeds, text_audio_duration_embeds], @@ -600,9 +591,8 @@ def forward( # Trim to requested length audio = audio[:, :, waveform_start:waveform_end] - return DiffusionOutput( - output=audio, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) + stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None + return DiffusionOutput(output=audio, stage_durations=stage_durations) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights using AutoWeightsLoader for vLLM integration.""" diff --git a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py index 83a669d92aa..c7bc9b2d4f6 100644 --- a/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py +++ b/vllm_omni/diffusion/models/stable_audio/stable_audio_transformer.py @@ -9,12 +9,14 @@ import torch import torch.nn as nn +from cache_dit import ForwardPattern from diffusers.models.modeling_outputs import Transformer2DModelOutput from vllm.logger import init_logger from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module from vllm_omni.diffusion.layers.fourier import GaussianFourierProjection @@ -464,6 +466,12 @@ class StableAudioDiTModel(nn.Module): - Output: [B, out_channels, L] """ + _cache_dit_adapter_config = CacheDiTAdapterConfig( + block_forward_patterns={ + "transformer_blocks": ForwardPattern.Pattern_3, + }, + ) + _repeated_blocks = ["StableAudioDiTBlock"] _layerwise_offload_blocks_attrs = ["transformer_blocks"] _hsdp_shard_conditions = [is_transformer_block_module] diff --git a/vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py b/vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py index eca4267fa20..f7f62f8cec9 100644 --- a/vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py +++ b/vllm_omni/diffusion/models/t5_encoder/t5_gemma_encoder.py @@ -157,7 +157,7 @@ def __init__(self, config: PretrainedConfig) -> None: num_kv_heads=config.num_key_value_heads, head_dim=config.head_dim, max_position_embeddings=config.max_position_embeddings, - rope_theta=config.rope_theta, + rope_theta=getattr(config, "rope_theta", 10000.0), ) self.mlp = T5GemmaMLP( hidden_size=config.hidden_size, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 85a58596050..d3726e08760 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -8,7 +8,7 @@ import os import time from collections.abc import Iterable -from typing import Any, cast +from typing import Any, ClassVar, cast import PIL.Image import torch @@ -28,6 +28,7 @@ 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.dmd2 import DMD2PipelineMixin +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler @@ -35,6 +36,7 @@ from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -185,7 +187,7 @@ def post_process_func( ): if output_type == "latent": return video - custom_output = {} + video_metadata = {} if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): video, multiplier = interpolate_video_tensor( video, @@ -193,10 +195,10 @@ def post_process_func( scale=sampling_params.frame_interpolation_scale, model_path=sampling_params.frame_interpolation_model_path, ) - custom_output["video_fps_multiplier"] = multiplier + video_metadata["video_fps_multiplier"] = multiplier return { - "video": video_processor.postprocess_video(video, output_type=output_type), - "custom_output": custom_output, + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {"video": video_metadata} if video_metadata else {}, } return post_process_func @@ -209,56 +211,66 @@ def get_wan22_pre_process_func( import numpy as np def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if raw_image is None: - continue - - if not isinstance(raw_image, (str, PIL.Image.Image)): - raise TypeError( - f"""Unsupported image format {raw_image.__class__}.""", - """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""", - ) - image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image - - # Calculate dimensions based on aspect ratio if not provided - if request.sampling_params.height is None or request.sampling_params.width is None: - # Default max area for 720P - max_area = 720 * 1280 - aspect_ratio = image.height / image.width - - # Calculate dimensions maintaining aspect ratio - mod_value = 16 # Must be divisible by 16 - height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value - width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value - - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width - - # Resize image to target dimensions - image = image.resize( - (request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None - PIL.Image.Resampling.LANCZOS, + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if raw_image is None: + request.prompt = prompt + return request + + if not isinstance(raw_image, (str, PIL.Image.Image)): + raise TypeError( + f"""Unsupported image format {raw_image.__class__}.""", + """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""", ) - prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above + image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image + + # Calculate dimensions based on aspect ratio if not provided + if request.sampling_params.height is None or request.sampling_params.width is None: + # Default max area for 720P + max_area = 720 * 1280 + aspect_ratio = image.height / image.width + + # Calculate dimensions maintaining aspect ratio + mod_value = 16 # Must be divisible by 16 + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + # Resize image to target dimensions + image = image.resize( + (request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None + PIL.Image.Resampling.LANCZOS, + ) + prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above - request.prompts[i] = prompt + request.prompt = prompt return request return pre_process_func class Wan22Pipeline( - nn.Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin + nn.Module, + PipelineParallelMixin, + CFGParallelMixin, + ProgressBarMixin, + DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, ): + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -528,38 +540,27 @@ def diffuse( return latents - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | None = None, - negative_prompt: str | None = None, - height: int = 480, - width: int = 832, - num_inference_steps: int = 40, - guidance_scale: float | tuple[float, float] = 4.0, - frame_num: int = 81, - output_type: str | None = "np", - generator: torch.Generator | list[torch.Generator] | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - attention_kwargs: dict | None = None, - **kwargs, - ) -> DiffusionOutput: - # Get parameters from request or arguments + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + prompt: str | None = None + negative_prompt: str | None = None + prompt_embeds: torch.Tensor | None = None + negative_prompt_embeds: torch.Tensor | None = None if len(req.prompts) > 1: raise ValueError( """This model only supports a single prompt, not a batched request.""", """Please pass in a single prompt object or string, or a single-item list.""", ) - if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list - prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") - negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") - if prompt is None and prompt_embeds is None: - raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.") + if len(req.prompts) == 1: + first_prompt = req.prompts[0] + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num + if not prompt: + raise ValueError("Prompt is required for Wan2.2 generation.") + + height = req.sampling_params.height or 480 + width = req.sampling_params.width or 832 + num_frames = req.sampling_params.num_frames or 81 # Ensure dimensions are compatible with VAE and patch size # For expand_timesteps mode, we need latent dims to be even (divisible by patch_size) @@ -567,11 +568,16 @@ def forward( mod_value = self.vae_scale_factor_spatial * patch_size[1] # 16*2=32 for TI2V, 8*2=16 for I2V height = (height // mod_value) * mod_value width = (width // mod_value) * mod_value - num_steps = req.sampling_params.num_inference_steps or num_inference_steps + num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps # Respect per-request guidance_scale when explicitly provided. if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 4.0 + + output_type = req.sampling_params.output_type or "np" + attention_kwargs: dict | None = None guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] guidance_high = ( @@ -622,8 +628,7 @@ def forward( dtype = self.text_encoder.dtype # Seed / generator - if generator is None: - generator = req.sampling_params.generator + generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) @@ -632,24 +637,16 @@ def forward( current_omni_platform.synchronize() _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - if prompt_embeds is None: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) - else: - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - if negative_prompt_embeds is not None: - negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) - elif guidance_low > 1.0 or guidance_high > 1.0: - raise ValueError( - "negative_prompt_embeds must be provided when prompt_embeds are given and guidance > 1." - ) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, + num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, + max_sequence_length=req.sampling_params.max_sequence_length or 512, + device=device, + dtype=dtype, + ) + if DEBUG_PERF: current_omni_platform.synchronize() _t_text_enc_ms = (time.perf_counter() - _t_text_enc_start) * 1000 diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 68eb3b11bc0..8b042b59568 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -7,7 +7,7 @@ import os import time from collections.abc import Iterable -from typing import Any, cast +from typing import Any, ClassVar, cast import numpy as np import PIL.Image @@ -28,7 +28,7 @@ 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.dmd2 import DMD2PipelineMixin -from vllm_omni.diffusion.models.interface import SupportImageInput +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.utils import _load_json from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( @@ -43,6 +43,7 @@ from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -64,7 +65,7 @@ def post_process_func( ): if output_type == "latent": return video - custom_output = {} + video_metadata = {} if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): video, multiplier = interpolate_video_tensor( video, @@ -72,10 +73,10 @@ def post_process_func( scale=sampling_params.frame_interpolation_scale, model_path=sampling_params.frame_interpolation_model_path, ) - custom_output["video_fps_multiplier"] = multiplier + video_metadata["video_fps_multiplier"] = multiplier return { - "video": video_processor.postprocess_video(video, output_type=output_type), - "custom_output": custom_output, + "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, + "metadata": {"video": video_metadata} if video_metadata else {}, } return post_process_func @@ -87,50 +88,50 @@ def get_wan22_i2v_pre_process_func( """Pre-process function for I2V: load and resize input image.""" def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if raw_image is None: - raise ValueError( - """No image is provided. This model requires an image to run.""", - """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""", - ) - if not isinstance(raw_image, (str, PIL.Image.Image)): - raise TypeError( - f"""Unsupported image format {raw_image.__class__}.""", - """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""", - ) - image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image - - # Calculate dimensions based on aspect ratio if not provided - if request.sampling_params.height is None or request.sampling_params.width is None: - # Default max area for 480P - max_area = 480 * 832 - aspect_ratio = image.height / image.width - - # Calculate dimensions maintaining aspect ratio - mod_value = 16 # Must be divisible by 16 - height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value - width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value - - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width - - # Resize image to target dimensions - image = image.resize( - (request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None - PIL.Image.Resampling.LANCZOS, + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if raw_image is None: + raise ValueError( + """No image is provided. This model requires an image to run.""", + """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""", ) - prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above + if not isinstance(raw_image, (str, PIL.Image.Image)): + raise TypeError( + f"""Unsupported image format {raw_image.__class__}.""", + """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""", + ) + image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image + + # Calculate dimensions based on aspect ratio if not provided + if request.sampling_params.height is None or request.sampling_params.width is None: + # Default max area for 480P + max_area = 480 * 832 + aspect_ratio = image.height / image.width + + # Calculate dimensions maintaining aspect ratio + mod_value = 16 # Must be divisible by 16 + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + # Resize image to target dimensions + image = image.resize( + (request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None + PIL.Image.Resampling.LANCZOS, + ) + prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above - request.prompts[i] = prompt + request.prompt = prompt return request return pre_process_func @@ -143,6 +144,7 @@ class Wan22I2VPipeline( CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin, + SupportsComponentDiscovery, ): """ Wan2.2 Image-to-Video Pipeline. @@ -151,6 +153,10 @@ class Wan22I2VPipeline( Wan2.2-style I2V (with expand_timesteps for TI2V-5B). """ + _dit_modules: ClassVar[list[str]] = ["transformer", "transformer_2"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder", "image_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -423,66 +429,59 @@ def _create_transformer(self, config: dict) -> WanTransformer3DModel: quant_config = getattr(self.od_config, "quantization_config", None) return create_transformer_from_config(config, quant_config=quant_config) - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | None = None, - negative_prompt: str | None = None, - image: PIL.Image.Image | torch.Tensor | None = None, - height: int = 480, - width: int = 832, - num_inference_steps: int = 40, - guidance_scale: float | tuple[float, float] = 5.0, - frame_num: int = 81, - output_type: str | None = "np", - generator: torch.Generator | list[torch.Generator] | None = None, - prompt_embeds: torch.Tensor | None = None, - negative_prompt_embeds: torch.Tensor | None = None, - image_embeds: torch.Tensor | None = None, - last_image: PIL.Image.Image | torch.Tensor | None = None, - attention_kwargs: dict | None = None, - **kwargs, - ) -> DiffusionOutput: - # Get parameters from request or arguments + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: + prompt: str | None = None + negative_prompt: str | None = None + prompt_embeds: torch.Tensor | None = None + negative_prompt_embeds: torch.Tensor | None = None + image_embeds: torch.Tensor | None = None if len(req.prompts) > 1: raise ValueError( """This model only supports a single prompt, not a batched request.""", """Please pass in a single prompt object or string, or a single-item list.""", ) - if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list - prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt") - negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt") - if prompt is None and prompt_embeds is None: - raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.") + if len(req.prompts) == 1: + first_prompt = req.prompts[0] + prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + negative_prompt = None if isinstance(first_prompt, str) else first_prompt.get("negative_prompt") + + if not prompt: + raise ValueError("Prompt is required for Wan2.2 generation.") # Get image from request - if image is None: - multi_modal_data = ( - req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None - ) - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if raw_image is None: - raise ValueError("Image is required for I2V generation.") - if isinstance(raw_image, list): - if len(raw_image) > 1: - logger.warning( - """Received a list of image. Only a single image is supported by this model.""" - """Taking only the first image for now.""" - ) - raw_image = raw_image[0] - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image) - else: - image = cast(PIL.Image.Image | torch.Tensor, raw_image) + multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if raw_image is None: + raise ValueError("Image is required for I2V generation.") + if isinstance(raw_image, list): + if len(raw_image) > 1: + logger.warning( + """Received a list of image. Only a single image is supported by this model.""" + """Taking only the first image for now.""" + ) + raw_image = raw_image[0] + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image) + else: + image = cast(PIL.Image.Image | torch.Tensor, raw_image) + + last_image: PIL.Image.Image | torch.Tensor | None = None + if multi_modal_data is not None: + last_image = multi_modal_data.get("last_image", None) - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_frames = req.sampling_params.num_frames or frame_num - num_steps = req.sampling_params.num_inference_steps or num_inference_steps + height = req.sampling_params.height or 480 + width = req.sampling_params.width or 832 + num_frames = req.sampling_params.num_frames or 81 + num_steps = 40 if req.sampling_params.num_inference_steps is None else req.sampling_params.num_inference_steps # Respect per-request guidance_scale when explicitly provided. if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale + else: + guidance_scale = 5.0 + + output_type = req.sampling_params.output_type or "np" + attention_kwargs: dict | None = None # Handle guidance scales guidance_low = guidance_scale if isinstance(guidance_scale, (int, float)) else guidance_scale[0] @@ -527,8 +526,7 @@ def forward( dtype = self.transformer.dtype # Generator setup - if generator is None: - generator = req.sampling_params.generator + generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed) @@ -538,20 +536,15 @@ def forward( _t_pipeline_start = time.perf_counter() _t_text_enc_start = _t_pipeline_start - if prompt_embeds is None: - prompt_embeds, negative_prompt_embeds = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, - num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, - device=device, - dtype=dtype, - ) - else: - prompt_embeds = prompt_embeds.to(device=device, dtype=dtype) - if negative_prompt_embeds is not None: - negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype) + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, + num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, + max_sequence_length=req.sampling_params.max_sequence_length or 512, + device=device, + dtype=dtype, + ) if DEBUG_PERF: current_omni_platform.synchronize() 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 1bd32bb4275..47b8e02d158 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 @@ -42,6 +42,7 @@ ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -282,61 +283,60 @@ def get_wan22_s2v_pre_process_func( """ def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - # -- Reference image -- - raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None - if raw_image is None: - raise ValueError( - "No reference image provided. S2V requires a reference image. " - 'Set `"multi_modal_data": {"image": <path or PIL.Image>, ...}`' - ) - if isinstance(raw_image, str): - image = PIL.Image.open(raw_image).convert("RGB") - elif isinstance(raw_image, PIL.Image.Image): - image = raw_image - else: - raise TypeError(f"Unsupported image type {type(raw_image)}") - - # -- Audio -- - raw_audio = multi_modal_data.get("audio", None) if multi_modal_data is not None else None - if raw_audio is None: - raise ValueError( - "No audio provided. S2V requires an audio file path. " - 'Set `"multi_modal_data": {"audio": "<path>", ...}`' - ) - - # -- Compute target size -- - max_area = 720 * 1280 - if request.sampling_params.height is not None and request.sampling_params.width is not None: - height, width = request.sampling_params.height, request.sampling_params.width - else: - ref_h, ref_w = image.height, image.width - height, width = _get_size_less_than_area(ref_h, ref_w, target_area=max_area) - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width - - # Resize + center-crop reference image to target size - resize_op = transforms.Resize(min(height, width)) - crop_op = transforms.CenterCrop((height, width)) - ref_pil = crop_op(resize_op(image)) - - prompt["multi_modal_data"]["image"] = ref_pil - prompt["additional_information"]["audio_path"] = raw_audio - prompt["additional_information"]["pose_video"] = ( - multi_modal_data.get("pose_video", None) if multi_modal_data is not None else None + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + # -- Reference image -- + raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None + if raw_image is None: + raise ValueError( + "No reference image provided. S2V requires a reference image. " + 'Set `"multi_modal_data": {"image": <path or PIL.Image>, ...}`' ) - prompt["additional_information"]["init_first_frame"] = ( - multi_modal_data.get("init_first_frame", False) if multi_modal_data is not None else False + if isinstance(raw_image, str): + image = PIL.Image.open(raw_image).convert("RGB") + elif isinstance(raw_image, PIL.Image.Image): + image = raw_image + else: + raise TypeError(f"Unsupported image type {type(raw_image)}") + + # -- Audio -- + raw_audio = multi_modal_data.get("audio", None) if multi_modal_data is not None else None + if raw_audio is None: + raise ValueError( + 'No audio provided. S2V requires an audio file path. Set `"multi_modal_data": {"audio": "<path>", ...}`' ) - request.prompts[i] = prompt + + # -- Compute target size -- + max_area = 720 * 1280 + if request.sampling_params.height is not None and request.sampling_params.width is not None: + height, width = request.sampling_params.height, request.sampling_params.width + else: + ref_h, ref_w = image.height, image.width + height, width = _get_size_less_than_area(ref_h, ref_w, target_area=max_area) + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + # Resize + center-crop reference image to target size + resize_op = transforms.Resize(min(height, width)) + crop_op = transforms.CenterCrop((height, width)) + ref_pil = crop_op(resize_op(image)) + + prompt["multi_modal_data"]["image"] = ref_pil + prompt["additional_information"]["audio_path"] = raw_audio + prompt["additional_information"]["pose_video"] = ( + multi_modal_data.get("pose_video", None) if multi_modal_data is not None else None + ) + prompt["additional_information"]["init_first_frame"] = ( + multi_modal_data.get("init_first_frame", False) if multi_modal_data is not None else False + ) + request.prompt = prompt return request @@ -606,13 +606,17 @@ def _init_original_format(self, model_path: str, dtype: torch.dtype) -> None: t5_checkpoint = os.path.join(model_path, "models_t5_umt5-xxl-enc-bf16.pth") self.text_encoder = UMT5EncoderModel(_WAN_UMT5_CONFIG) - self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype).to(self.device) + _cpu_offload = self.od_config.enable_cpu_offload or self.od_config.enable_layerwise_offload + self.text_encoder = _load_wan_t5_as_umt5(self.text_encoder, t5_checkpoint, dtype=dtype) + if not _cpu_offload: + self.text_encoder = self.text_encoder.to(self.device) # -- VAE (original Wan2.1 VAE, loaded via diffusers from_single_file) -- vae_pth = os.path.join(model_path, "Wan2.1_VAE.pth") self.vae = DistributedAutoencoderKLWan.from_single_file(vae_pth, torch_dtype=dtype) self.vae.init_distributed() - self.vae = self.vae.to(self.device) + if not _cpu_offload: + self.vae = self.vae.to(self.device) # -- Audio encoder (wav2vec2) -- wav2vec_path = os.path.join(model_path, "wav2vec2-large-xlsr-53-english") @@ -1050,7 +1054,7 @@ def diffuse( def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | None = None, negative_prompt: str | None = None, image: PIL.Image.Image | None = None, @@ -1284,7 +1288,7 @@ def forward( ) # ---- Decode this clip ---- - if self.od_config.enable_cpu_offload: + if self.od_config.enable_cpu_offload and not getattr(self.od_config.parallel_config, "use_hsdp", False): self.transformer.to("cpu") current_omni_platform.empty_cache() diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index 5ba2c6c690f..4173b86011b 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -36,6 +36,7 @@ ) from vllm_omni.diffusion.models.wan2_2.wan2_2_vace_transformer import WanVACETransformer3DModel from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -96,68 +97,66 @@ def get_wan22_vace_pre_process_func(od_config: OmniDiffusionConfig): import numpy as np def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: - for i, prompt in enumerate(request.prompts): - multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None - if isinstance(prompt, str): - prompt = OmniTextPrompt(prompt=prompt) - if "additional_information" not in prompt: - prompt["additional_information"] = {} - - if not multi_modal_data: - request.prompts[i] = prompt - continue - - # Handle reference images for R2V - # "image" is the standard key from online serving (SupportImageInput convention) - # "reference_images" is the offline API key for backwards compatibility - ref_images = multi_modal_data.get("image") or multi_modal_data.get("reference_images") - if ref_images is not None: - if isinstance(ref_images, str): - ref_images = [PIL.Image.open(ref_images).convert("RGB")] - elif isinstance(ref_images, PIL.Image.Image): - ref_images = [ref_images] - elif isinstance(ref_images, list): - ref_images = [ - PIL.Image.open(img).convert("RGB") if isinstance(img, str) else img for img in ref_images - ] - - # Calculate dimensions from first reference image if not provided - if request.sampling_params.height is None or request.sampling_params.width is None: - first_img = ref_images[0] - max_area = 480 * 832 # VACE default is 480p - aspect_ratio = first_img.height / first_img.width - mod_value = 16 - height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value - width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value - - if request.sampling_params.height is None: - request.sampling_params.height = height - if request.sampling_params.width is None: - request.sampling_params.width = width - - prompt["additional_information"]["reference_images"] = ref_images - - # Handle source video for V2V / MV2V - source_video = multi_modal_data.get("video") - if source_video is not None: - if isinstance(source_video, list) and len(source_video) > 0: - if isinstance(source_video[0], str): - source_video = [PIL.Image.open(f).convert("RGB") for f in source_video] - prompt["additional_information"]["source_video"] = source_video - - # Handle mask for MV2V / inpainting - mask = multi_modal_data.get("mask") - if mask is not None: - if isinstance(mask, list) and len(mask) > 0: - if isinstance(mask[0], str): - mask = [PIL.Image.open(m).convert("L") for m in mask] - elif isinstance(mask, str): - mask = [PIL.Image.open(mask).convert("L")] - elif isinstance(mask, PIL.Image.Image): - mask = [mask] - prompt["additional_information"]["mask"] = mask - - request.prompts[i] = prompt + prompt = request.prompt + multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None + if isinstance(prompt, str): + prompt = OmniTextPrompt(prompt=prompt) + if "additional_information" not in prompt: + prompt["additional_information"] = {} + + if not multi_modal_data: + request.prompt = prompt + return request + + # Handle reference images for R2V + # "image" is the standard key from online serving (SupportImageInput convention) + # "reference_images" is the offline API key for backwards compatibility + ref_images = multi_modal_data.get("image") or multi_modal_data.get("reference_images") + if ref_images is not None: + if isinstance(ref_images, str): + ref_images = [PIL.Image.open(ref_images).convert("RGB")] + elif isinstance(ref_images, PIL.Image.Image): + ref_images = [ref_images] + elif isinstance(ref_images, list): + ref_images = [PIL.Image.open(img).convert("RGB") if isinstance(img, str) else img for img in ref_images] + + # Calculate dimensions from first reference image if not provided + if request.sampling_params.height is None or request.sampling_params.width is None: + first_img = ref_images[0] + max_area = 480 * 832 # VACE default is 480p + aspect_ratio = first_img.height / first_img.width + mod_value = 16 + height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value + width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value + + if request.sampling_params.height is None: + request.sampling_params.height = height + if request.sampling_params.width is None: + request.sampling_params.width = width + + prompt["additional_information"]["reference_images"] = ref_images + + # Handle source video for V2V / MV2V + source_video = multi_modal_data.get("video") + if source_video is not None: + if isinstance(source_video, list) and len(source_video) > 0: + if isinstance(source_video[0], str): + source_video = [PIL.Image.open(f).convert("RGB") for f in source_video] + prompt["additional_information"]["source_video"] = source_video + + # Handle mask for MV2V / inpainting + mask = multi_modal_data.get("mask") + if mask is not None: + if isinstance(mask, list) and len(mask) > 0: + if isinstance(mask[0], str): + mask = [PIL.Image.open(m).convert("L") for m in mask] + elif isinstance(mask, str): + mask = [PIL.Image.open(mask).convert("L")] + elif isinstance(mask, PIL.Image.Image): + mask = [mask] + prompt["additional_information"]["mask"] = mask + + request.prompt = prompt return request return pre_process_func @@ -198,6 +197,7 @@ def diffuse( attention_kwargs: dict[str, object], vace_context: torch.Tensor | None, vace_context_scale: float, + boundary_timestep: float | None = None, ) -> torch.Tensor: if attention_kwargs is None: attention_kwargs = {} @@ -206,6 +206,11 @@ def diffuse( self._current_timestep = t set_forward_context_denoise_step_idx(step_idx) + if boundary_timestep is not None and t < boundary_timestep and self.transformer_2 is not None: + current_model = self.transformer_2 + else: + current_model = self.transformer if self.transformer is not None else self.transformer_2 + latent_model_input = latents.to(dtype) timestep = t.expand(latents.shape[0]) @@ -219,6 +224,7 @@ def diffuse( "vace_context": vace_context, "vace_context_scale": vace_context_scale, "return_dict": False, + "current_model": current_model, } negative_kwargs = ( { @@ -229,6 +235,7 @@ def diffuse( "vace_context": vace_context, "vace_context_scale": vace_context_scale, "return_dict": False, + "current_model": current_model, } if do_true_cfg else None @@ -461,7 +468,7 @@ def prepare_masks( def forward( self, - req: OmniDiffusionRequest, + req: DiffusionRequestBatch, prompt: str | None = None, negative_prompt: str | None = None, height: int = 480, @@ -598,7 +605,10 @@ def forward( ) num_reference_images = 0 - if self.transformer.vace_patch_embedding is not None: + # Either expert may be the one loaded (boundary_ratio 0.0/1.0 skips one), + # so read it from whichever transformer exists. + active_transformer = self.transformer if self.transformer is not None else self.transformer_2 + if active_transformer.vace_patch_embedding is not None: video, mask, ref_images_processed = self.preprocess_conditions( video=source_video, mask=source_mask, @@ -640,6 +650,19 @@ def forward( timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) + # MoE boundary: only the two-expert checkpoints (Wan2.2 A14B) have a + # transformer_2 to switch to. For single-expert VACE (e.g. Wan2.1) this + # stays None and diffuse() uses the same transformer for every step. + boundary_timestep = None + if self.transformer_2 is not None: + boundary_ratio = ( + self.boundary_ratio if self.boundary_ratio is not None else req.sampling_params.boundary_ratio + ) + if boundary_ratio is None: + boundary_ratio = 0.875 + logger.warning_once("boundary_ratio is required for Wan2.2 VACE generation. using default value 0.875") + boundary_timestep = boundary_ratio * self.scheduler.config.num_train_timesteps + latents = self.diffuse( latents=latents, timesteps=timesteps, @@ -650,6 +673,7 @@ def forward( attention_kwargs=attention_kwargs, vace_context=vace_context, vace_context_scale=vace_context_scale, + boundary_timestep=boundary_timestep, ) self._current_timestep = None 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 8fef41f74e5..d1b74d020ab 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 @@ -31,6 +31,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed from vllm_omni.platforms import current_omni_platform @@ -1058,6 +1059,32 @@ def _sample_indices(num_frames, stride, expand_ratio, c): return indices +# --------------------------------------------------------------------------- +# SP helper modules +# --------------------------------------------------------------------------- + + +class WanS2VBlocksPrepare(nn.Module): + """Pass-through module providing SP split hook point before transformer blocks. + + Returns (hidden_states, freqs) unchanged. The SP hook shards both outputs + via split_output=True along the sequence dimension. + """ + + def forward(self, hidden_states: torch.Tensor, freqs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return hidden_states, freqs + + +class WanS2VBlocksOutputGather(nn.Module): + """Pass-through module providing SP gather hook point after transformer blocks. + + Returns hidden_states unchanged. The SP gather hook collects from all ranks. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return hidden_states + + # --------------------------------------------------------------------------- # Main model # --------------------------------------------------------------------------- @@ -1087,6 +1114,17 @@ def _is_transformer_block(name, mod): _hsdp_shard_conditions = [_is_transformer_block] + _sp_plan = { + # Shard both hidden_states and freqs via blocks_prepare split_output + # hidden_states: [B, seq, dim], freqs: [B, seq, 1, head_dim//2] (complex) + "blocks_prepare": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + 1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True), + }, + # Gather hidden_states after transformer blocks + "blocks_output_gather": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } + def __init__( self, cond_dim=0, @@ -1157,6 +1195,10 @@ def __init__( ] ) + # SP hook modules (pass-through; hooks shard/gather at these boundaries) + self.blocks_prepare = WanS2VBlocksPrepare() + self.blocks_output_gather = WanS2VBlocksOutputGather() + # Output head self.norm_out = FP32LayerNorm(dim, eps, elementwise_affine=False) self.proj_out = nn.Linear(dim, math.prod(patch_size) * out_dim) @@ -1396,8 +1438,21 @@ def inject_motion( ] return x, seq_lens, rope_embs, mask_input + def _is_sp_active(self) -> bool: + from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available + + if not is_forward_context_available(): + return False + return get_forward_context().sp_active + def after_transformer_block(self, block_idx, hidden_states): if block_idx in self.audio_injector.injected_block_id.keys(): + sp_active = self._is_sp_active() + if sp_active: + from vllm_omni.diffusion.distributed.sp_sharding import sp_gather, sp_shard + + hidden_states = sp_gather(hidden_states, dim=1, validate=False) + audio_attn_id = self.audio_injector.injected_block_id[block_idx] audio_emb = self.merged_audio_emb num_frames = audio_emb.shape[1] @@ -1424,6 +1479,9 @@ def after_transformer_block(self, block_idx, hidden_states): residual_out = rearrange(residual_out, "(b t) n c -> b (t n) c", t=num_frames) hidden_states[:, : self.original_seq_len].add_(residual_out) + if sp_active: + hidden_states = sp_shard(hidden_states, dim=1, validate=False) + return hidden_states # ------------------------------------------------------------------ @@ -1441,18 +1499,27 @@ def encode_audio(self, audio_input: torch.Tensor, motion_frames: list[int]) -> d dict with 'audio_emb' (and optionally 'audio_emb_global' when enable_adain is True). """ - 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) - result = {} - if self.enable_adain: - audio_emb_global, audio_emb = audio_emb_res - result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :] - else: - audio_emb = audio_emb_res - result["audio_emb"] = audio_emb[:, motion_frames[1] :, :] + # Under HSDP, this method is called outside the FSDP forward hook. + # Unshard root-managed params so casual_audio_encoder can run. + is_fsdp = hasattr(self, "unshard") and hasattr(self, "reshard") + if is_fsdp: + self.unshard() + try: + 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) + result = {} + if self.enable_adain: + audio_emb_global, audio_emb = audio_emb_res + result["audio_emb_global"] = audio_emb_global[:, motion_frames[1] :] + else: + audio_emb = audio_emb_res + result["audio_emb"] = audio_emb[:, motion_frames[1] :, :] + finally: + if is_fsdp: + self.reshard() return result def forward( @@ -1574,6 +1641,21 @@ def forward( context_lens = None + # SP: shard hidden_states and freqs before entering blocks + x, self.pre_compute_freqs = self.blocks_prepare(x, self.pre_compute_freqs) + + # SP: adjust segment boundary for the local shard + if self._is_sp_active(): + from vllm_omni.diffusion.distributed.parallel_state import ( + get_sequence_parallel_rank, + ) + + sp_rank = get_sequence_parallel_rank() + chunk_size = x.shape[1] + global_start = sp_rank * chunk_size + local_seg_boundary = max(0, min(e0[1] - global_start, chunk_size)) + e0 = [e0[0], local_seg_boundary] + # Transformer blocks kwargs = dict( e=e0, @@ -1588,6 +1670,9 @@ def forward( x = block(x, **kwargs) x = self.after_transformer_block(idx, x) + # SP: gather hidden_states from all ranks + x = self.blocks_output_gather(x) + # Output norm, projection & unpatchify x = x[:, : self.original_seq_len] shift, scale = self.output_scale_shift_prepare(e) 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 cda01e08375..d61d1c82b30 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -40,7 +40,7 @@ SequenceParallelInput, SequenceParallelOutput, ) -from vllm_omni.diffusion.forward_context import get_forward_context +from vllm_omni.diffusion.forward_context import build_local_sp_padding_mask, get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan @@ -406,6 +406,8 @@ def __init__( ) self.dropout = nn.Dropout(dropout) + self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) + # Unified attention layer self.attn = Attention( num_heads=self.num_heads, @@ -441,7 +443,6 @@ def forward( # Apply rotary embeddings if rotary_emb is not None: - self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) freqs_cos, freqs_sin = rotary_emb query = self.rotary_embedding(query, freqs_cos, freqs_sin) key = self.rotary_embedding(key, freqs_cos, freqs_sin) @@ -798,7 +799,7 @@ def _is_transformer_block(name: str, module) -> bool: # The _sp_plan specifies sharding/gathering at module boundaries: # - rope: Split both RoPE outputs (freqs_cos, freqs_sin) via split_output=True # - timestep_proj_prepare: Split timestep_proj for TI2V models (4D tensor) - # - blocks.0: Split hidden_states input at the first transformer block + # - _sp_shard_point: Split hidden_states before CacheDiT-wrapped transformer blocks # - proj_out: Gather outputs after the final projection layer # # Note: _sp_plan corresponds to diffusers' _cp_plan (Context Parallelism) @@ -821,10 +822,10 @@ def _is_transformer_block(name: str, module) -> bool: split_dim=1, expected_dims=4, split_output=True, auto_pad=True ), # [B, seq, 6, dim] }, - # Shard hidden_states at first transformer block input - # (after patch_embedding + flatten + transpose) - "blocks.0": { - "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True), # [B, seq, dim] + # Shard hidden_states before entering transformer blocks. This keeps + # CacheDiT block wrappers aligned with the local SP shard. + "_sp_shard_point": { + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), }, # Shard output scale/shift for TI2V (3D); T2V outputs 2D and skips sharding "output_scale_shift_prepare": { @@ -855,7 +856,6 @@ def __init__( quant_config: QuantizationConfig | None = None, ): super().__init__() - # Store config for compatibility self.config = type( "Config", @@ -937,6 +937,7 @@ def __init__( # SP helper modules self.timestep_proj_prepare = TimestepProjPrepare() + self._sp_shard_point = nn.Identity() if is_pipeline_last_stage(): self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim) else: @@ -980,10 +981,11 @@ def forward( self._cached_rope_emb = rotary_emb if is_pipeline_first_stage(): - # Patch embedding and flatten to sequence - # (hidden_states is sharded at blocks.0 input by _sp_plan) + # Patch embedding and flatten to sequence. SP sharding happens at + # _sp_shard_point so downstream block wrappers see local tensors. hidden_states = self.patch_embedding(hidden_states) hidden_states = hidden_states.flatten(2).transpose(1, 2) + hidden_states = self._sp_shard_point(hidden_states) else: if intermediate_tensors is None: raise RuntimeError("intermediate_tensors must be provided for non-first PP stages") @@ -1016,18 +1018,12 @@ def forward( parallel_config is not None and parallel_config.sequence_parallel_size > 1 and parallel_config.mask_sp_padding - and ctx.sp_original_seq_len is not None - and ctx.sp_padding_size > 0 ): - batch_size = hidden_states.shape[0] - padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size - hidden_states_mask = torch.ones( - batch_size, - padded_seq_len, - dtype=torch.bool, - device=hidden_states.device, + hidden_states_mask = build_local_sp_padding_mask( + hidden_states.shape[0], + hidden_states.shape[1], + hidden_states.device, ) - hidden_states_mask[:, ctx.sp_original_seq_len :] = False if hidden_states_mask.all(): hidden_states_mask = None elif ( diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index 871f0df862f..a13c04c99e7 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py @@ -81,9 +81,7 @@ def forward( class WanVACETransformer3DModel(WanTransformer3DModel): """VACE-extended WAN Transformer with conditioning blocks for video editing.""" - # TODO: `vace_blocks` are not layerwise-offloaded yet. The current offloader only - # supports a single block group (`blocks`); extend it to support both - # `vace_blocks` and `blocks`. + _layerwise_offload_blocks_attrs = ["vace_blocks", "blocks"] # Shard hidden_states before VACE blocks (replaces parent's blocks.0) _sp_plan = { diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 388b6316ca8..b77d2265cf3 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -19,12 +19,12 @@ import json import os from collections.abc import Callable, Iterable -from typing import Any +from typing import Any, ClassVar import PIL.Image import torch import torch.nn as nn -from diffusers.image_processor import PipelineImageInput, VaeImageProcessor +from diffusers.image_processor import VaeImageProcessor from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from diffusers.utils import logging from diffusers.utils.torch_utils import randn_tensor @@ -36,12 +36,13 @@ 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 prefetch_subfolders +from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery from vllm_omni.diffusion.models.utils import create_transformers_model from vllm_omni.diffusion.models.z_image.z_image_transformer import ( ZImageTransformer2DModel, ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin -from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -160,7 +161,13 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): +class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery): + supports_request_batch = False + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, @@ -403,124 +410,21 @@ def num_timesteps(self): def interrupt(self): return self._interrupt - def forward( - self, - req: OmniDiffusionRequest, - prompt: str | list[str] | None = None, - image: PipelineImageInput = None, - strength: float = 0.6, - height: int = 1024, - width: int = 1024, - num_inference_steps: int = 50, - sigmas: list[float] | None = None, - guidance_scale: float = 5.0, - cfg_normalization: bool = False, - cfg_truncation: float = 1.0, - negative_prompt: str | list[str] | None = None, - num_images_per_prompt: int = 1, - generator: torch.Generator | list[torch.Generator] | None = None, - latents: torch.FloatTensor | None = None, - prompt_embeds: list[torch.FloatTensor] | None = None, - negative_prompt_embeds: list[torch.FloatTensor] | None = None, - output_type: str | None = "pil", - return_dict: bool = True, - joint_attention_kwargs: dict[str, Any] | None = None, - callback_on_step_end: Callable[[int, int, dict], None] | None = None, - callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, - ) -> DiffusionOutput: - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `list[str]`, *optional*): - The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. - instead. - image (`PipelineImageInput`, *optional*): - The image to use for img2img generation. If provided, the pipeline - will perform img2img instead of text-to-image. - strength (`float`, *optional*, defaults to 0.6): - Indicates extent to transform the reference `image`. Must be between 0 and 1. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the - expense of slower inference. - sigmas (`list[float]`, *optional*): - Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in - their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed - will be used. - guidance_scale (`float`, *optional*, defaults to 5.0): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 0`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - cfg_normalization (`bool`, *optional*, defaults to False): - Whether to apply configuration normalization. - cfg_truncation (`float`, *optional*, defaults to 1.0): - The truncation value for configuration. - negative_prompt (`str` or `list[str]`, *optional*): - The prompt or prompts not to guide the image generation. If not defined, one has to pass - `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is - less than or equal to `0`). - num_images_per_prompt (`int`, *optional*, defaults to 1): - The number of images to generate per prompt. - generator (`torch.Generator` or `list[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - latents (`torch.FloatTensor`, *optional*): - Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image - generation. Can be used to tweak the same generation with different prompts. If not provided, a latents - tensor will be generated by sampling using the supplied random `generator`. - prompt_embeds (`list[torch.FloatTensor]`, *optional*): - Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not - provided, text embeddings will be generated from `prompt` input argument. - negative_prompt_embeds (`list[torch.FloatTensor]`, *optional*): - Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt - weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input - argument. - output_type (`str`, *optional*, defaults to `"pil"`): - The output format of the generate image. Choose between - [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. - return_dict (`bool`, *optional*, defaults to `True`): - Whether or not to return a [`~pipelines.stable_diffusion.ZImagePipelineOutput`] instead of a plain - tuple. - joint_attention_kwargs (`dict`, *optional*): - A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under - `self.processor` in - [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). - callback_on_step_end (`Callable`, *optional*): - A function that calls at the end of each denoising steps during the inference. The function is called - with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, - callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by - `callback_on_step_end_tensor_inputs`. - callback_on_step_end_tensor_inputs (`list`, *optional*): - The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list - will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the - `._callback_tensor_inputs` attribute of your pipeline class. - max_sequence_length (`int`, *optional*, defaults to 512): - Maximum sequence length to use with the `prompt`. - - Examples: - - Returns: - [`~pipelines.z_image.ZImagePipelineOutput`] or `tuple`: [`~pipelines.z_image.ZImagePipelineOutput`] if - `return_dict` is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the - generated images. - """ + def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") # TODO: May be some data formatting operations on the API side. Hack for now. - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] - # Handle img2img: extract image from request - if image is None and req.prompts: + prompt_embeds = None + negative_prompt_embeds = None + + image = None + if req.prompts: if len(req.prompts) > 1: logger.warning( "This model only supports a single prompt for img2img, not a batched request. " @@ -535,10 +439,8 @@ def forward( else: image = PIL.Image.open(raw_image) if isinstance(raw_image, str) else raw_image - # strength is currently only applicable for Z-Image I2I; other pipelines ignore this parameter explicit_strength = req.sampling_params.strength is not None - if explicit_strength: - strength = req.sampling_params.strength + strength = req.sampling_params.strength if explicit_strength else 0.6 if explicit_strength and image is None: logger.warning( "strength parameter (%.2f) is only applicable for image-to-image (I2I) generation. " @@ -549,20 +451,24 @@ def forward( if image is not None and strength is not None and (strength < 0 or strength > 1): raise ValueError(f"The value of strength should be in [0.0, 1.0] but is {strength}") - height = req.sampling_params.height or height - width = req.sampling_params.width or width - num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps + height = req.sampling_params.height or 1024 + width = req.sampling_params.width or 1024 + num_inference_steps = req.sampling_params.num_inference_steps or 50 generator = req.sampling_params.generator - sigmas = req.sampling_params.sigmas or sigmas - max_sequence_length = req.sampling_params.max_sequence_length or max_sequence_length - guidance_scale = ( - req.sampling_params.guidance_scale if req.sampling_params.guidance_rescale is not None else guidance_scale - ) + sigmas = req.sampling_params.sigmas + max_sequence_length = req.sampling_params.max_sequence_length or 512 + guidance_scale = req.sampling_params.guidance_scale num_images_per_prompt = ( - req.sampling_params.num_outputs_per_prompt - if req.sampling_params.num_outputs_per_prompt > 0 - else num_images_per_prompt + req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else 1 ) + latents = req.sampling_params.latents + + cfg_normalization = req.sampling_params.cfg_normalize + cfg_truncation = req.sampling_params.extra_args.get("cfg_truncation", 1.0) + joint_attention_kwargs: dict[str, Any] | None = None + callback_on_step_end: Callable[[int, int, dict], None] | None = None + callback_on_step_end_tensor_inputs = ["latents"] + output_type = req.sampling_params.output_type or "pil" vae_scale = self.vae_scale_factor * 2 if height % vae_scale != 0: @@ -584,33 +490,20 @@ def forward( self._cfg_normalization = cfg_normalization self._cfg_truncation = cfg_truncation # 2. Define call parameters - if prompt is not None and isinstance(prompt, str): - batch_size = 1 - elif prompt is not None and isinstance(prompt, list): - batch_size = len(prompt) - else: - batch_size = len(prompt_embeds) + batch_size = len(prompt) - # If prompt_embeds is provided and prompt is None, skip encoding - if prompt_embeds is not None and prompt is None: - if self.do_classifier_free_guidance and negative_prompt_embeds is None: - raise ValueError( - "When `prompt_embeds` is provided without `prompt`, " - "`negative_prompt_embeds` must also be provided for classifier-free guidance." - ) - else: - ( - prompt_embeds, - negative_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - device=device, - max_sequence_length=max_sequence_length, - ) + ( + prompt_embeds, + negative_prompt_embeds, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + device=device, + max_sequence_length=max_sequence_length, + ) # 4. Prepare latent variables num_channels_latents = self.transformer.in_channels @@ -823,9 +716,8 @@ def forward( image = self.vae.decode(latents, return_dict=False)[0] # image = self.image_processor.postprocess(image, output_type=output_type) - return DiffusionOutput( - output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None - ) + stage_durations = self.stage_durations if hasattr(self, "stage_durations") else None + return DiffusionOutput(output=image, stage_durations=stage_durations) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index 320f9bf3849..17d4ca2fd30 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -22,6 +22,38 @@ class PipelineModules: resident_modules: list[nn.Module] = field(default_factory=list) resident_names: list[str] = field(default_factory=list) + def outermost_dits(self) -> tuple[list[str], list[nn.Module]]: + """Discovered DiTs that are not nested inside another discovered DiT. + + Some pipelines declare a DiT and one of its own submodules as separate + DiTs (e.g. Cosmos3 lists both ``transformer`` and the nested + ``transformer.language_model``) so that offloading can treat them as + independent rings. HSDP instead shards each DiT through that DiT's own + ``_hsdp_shard_conditions``; an inner DiT's blocks are already covered by + its ancestor's conditions, so sharding it again would double-wrap them + (and the inner stack may not declare conditions at all). Keeping only the + outermost DiTs leaves the ancestor's conditions as the single source of + truth for every nested block. Order follows ``dit_names``. + """ + + def is_nested_in_another_dit(dit: nn.Module) -> bool: + for other in self.dits: + if other is dit: + continue + for submodule in other.modules(): + if submodule is dit: + return True + return False + + outer_names: list[str] = [] + outer_modules: list[nn.Module] = [] + for name, dit in zip(self.dit_names, self.dits): + if is_nested_in_another_dit(dit): + continue + outer_names.append(name) + outer_modules.append(dit) + return outer_names, outer_modules + class ModuleDiscovery: """Discovers pipeline components. diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 8fb675b697e..0b031be3a69 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -197,20 +197,37 @@ class ModelLevelOffloadBackend(OffloadBackend): def __init__(self, config: OffloadConfig, device: torch.device): super().__init__(config, device) self._offload_modules: list[nn.Module] = [] # Track modules with hooks + self._custom_pipeline: nn.Module | None = None def enable(self, pipeline: nn.Module) -> None: if self.enabled: logger.warning("ModelLevelOffloadBackend already enabled") return - modules = ModuleDiscovery.discover(pipeline) - if not modules.dits: - logger.warning("No DiT/transformer modules found, skipping model-level offloading") - return - if not modules.encoders: - logger.warning("No encoder modules found, skipping model-level offloading") + # Pipelines with a nested transformer (e.g. Cosmos3's reasoner/generator + # pathways) own their own mutual-exclusion swaps and cannot be offloaded + # with the generic encoder<->DiT hooks. Delegate to the custom hook. + # TODO: Cosmos3 is the only implementer today, so we duck-type the + # enable/disable_omni_model_cpu_offload pair via getattr. Once a second + # pipeline needs it, formalize this as a Protocol (like + # SupportsComponentDiscovery) instead of getattr detection. + custom_enable = getattr(pipeline, "enable_omni_model_cpu_offload", None) + if callable(custom_enable): + custom_enable( + device=self.device, + pin_memory=self.config.pin_cpu_memory, + use_hsdp=self.config.use_hsdp, + ) + self._custom_pipeline = pipeline + self.enabled = True + logger.info( + "Model-level offloading enabled through %s.enable_omni_model_cpu_offload", + pipeline.__class__.__name__, + ) return + modules = ModuleDiscovery.discover(pipeline) + # Move encoders to GPU for enc in modules.encoders: enc.to(self.device) @@ -229,6 +246,17 @@ def enable(self, pipeline: nn.Module) -> None: except Exception as exc: logger.warning("Failed to move resident module '%s' to GPU: %s", name, exc) + if not modules.dits: + logger.warning("No DiT/transformer modules found, skipping model-level offloading") + return + + if not modules.encoders: + # Nothing to swap against — move DiTs to GPU and skip hooks. + for dit in modules.dits: + dit.to(self.device) + logger.warning("No encoder modules found, skipping model-level offloading") + return + # Apply sequential offloading hooks apply_sequential_offload( dit_modules=modules.dits, @@ -254,6 +282,15 @@ def disable(self) -> None: if not self.enabled: return + if self._custom_pipeline is not None: + custom_disable = getattr(self._custom_pipeline, "disable_omni_model_cpu_offload", None) + if callable(custom_disable): + custom_disable() + self._custom_pipeline = None + self.enabled = False + logger.info("Model-level offloading disabled") + return + remove_sequential_offload(self._offload_modules) self._offload_modules.clear() diff --git a/vllm_omni/diffusion/output_formatter.py b/vllm_omni/diffusion/output_formatter.py new file mode 100644 index 00000000000..a36305fceed --- /dev/null +++ b/vllm_omni/diffusion/output_formatter.py @@ -0,0 +1,356 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from typing import TypeGuard + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.io_support import supports_audio_output +from vllm_omni.diffusion.registry import DiffusionModelRegistry +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniPromptType +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.output_metadata import ( + DiffusionMetadata, + DiffusionMultimodalOutput, + DiffusionOutputEnvelope, + DiffusionPayload, + DiffusionPayloadValue, + DiffusionPostprocessRawOutput, + DiffusionTrajectoryPayload, + strip_internal_metadata, + validate_diffusion_metadata, + validate_public_diffusion_metadata, +) + + +@dataclass(frozen=True) +class DiffusionStepTimings: + preprocess_time_s: float + exec_time_s: float + postprocess_time_s: float + total_time_ms: float + + +@dataclass(frozen=True) +class DiffusionPostprocessOutput: + outputs: DiffusionPayload + metadata: DiffusionMetadata = field(default_factory=dict) + primary_key: str | None = None + + +def _is_output_envelope(outputs: DiffusionPostprocessRawOutput) -> TypeGuard[DiffusionOutputEnvelope]: + return isinstance(outputs, dict) and isinstance(outputs.get("payload"), dict) + + +def normalize_diffusion_postprocess_output( + outputs: DiffusionPostprocessRawOutput, +) -> DiffusionPostprocessOutput: + """Normalize diffusion postprocess output into payload and metadata.""" + + if _is_output_envelope(outputs): + payload: DiffusionPayload = outputs.get("payload") or {} + metadata = outputs.get("metadata") or {} + if not isinstance(metadata, dict): + metadata = {} + validate_diffusion_metadata(metadata) + + public_metadata = strip_internal_metadata(metadata) + if "text" in payload and "text" not in public_metadata: + public_metadata = {**public_metadata, "text": {}} + validate_public_diffusion_metadata(public_metadata) + + return DiffusionPostprocessOutput( + outputs=payload, + metadata=public_metadata, + primary_key=_infer_primary_payload_key(payload), + ) + + if isinstance(outputs, dict): + payload = {key: value for key, value in outputs.items() if key not in {"audio_sample_rate", "fps"}} + metadata = _metadata_from_legacy_payload(outputs) + if "text" in payload and "text" not in metadata: + metadata["text"] = {} + if metadata: + validate_public_diffusion_metadata(metadata) + return DiffusionPostprocessOutput( + outputs=payload, + metadata=metadata, + primary_key=_infer_primary_payload_key(payload), + ) + + return DiffusionPostprocessOutput( + outputs={"output": outputs}, + primary_key="output", + ) + + +def _metadata_from_legacy_payload(payload: dict[str, DiffusionPayloadValue]) -> DiffusionMetadata: + metadata: DiffusionMetadata = {} + audio_sample_rate = payload.get("audio_sample_rate") + if audio_sample_rate is not None: + metadata["audio"] = {"sample_rate": audio_sample_rate} + fps = payload.get("fps") + if fps is not None: + metadata["video"] = {"fps": fps} + return metadata + + +def _infer_primary_payload_key(payload: DiffusionPayload) -> str | None: + for key in ("video", "image", "text", "audio", "output"): + if key in payload: + return key + if set(payload).issubset({"actions", "trajectory"}): + return None + if "actions" in payload: + return None + return next(iter(payload), None) + + +def format_empty_diffusion_outputs( + request: OmniDiffusionRequest, + *, + finished: bool = True, +) -> list[OmniRequestOutput]: + return [ + OmniRequestOutput.from_diffusion( + request_id=request.request_id, + images=[], + prompt=request.prompt, + metrics={}, + latents=None, + finished=finished, + ) + ] + + +def format_diffusion_outputs( + *, + request: OmniDiffusionRequest, + od_config: OmniDiffusionConfig, + diffusion_output: DiffusionOutput, + output_data: DiffusionPostprocessRawOutput, + postprocess_output: DiffusionPostprocessOutput, + timings: DiffusionStepTimings, +) -> list[OmniRequestOutput]: + """Convert a finished diffusion model output into API-facing outputs.""" + + primary_payload = _primary_payload(postprocess_output) + outputs = _ensure_list(primary_payload) + metrics = { + "preprocess_time_ms": timings.preprocess_time_s * 1000, + "diffusion_engine_exec_time_ms": timings.exec_time_s * 1000, + "diffusion_engine_total_time_ms": timings.total_time_ms, + "image_num": int(request.sampling_params.num_outputs_per_prompt), + "resolution": int(request.sampling_params.resolution), + "postprocess_time_ms": timings.postprocess_time_s * 1000, + } + + # Detect text output: when the pipeline returns a string (e.g., + # SenseNova-U1 / BAGEL single-stage img2text / text2text), wrap it + # as a text-type response instead of an image. + is_text_output = postprocess_output.primary_key == "text" or "text" in postprocess_output.metadata + + is_audio_output = supports_audio_output(od_config.model_class_name) + audio_sample_rate = _metadata_audio_sample_rate(postprocess_output.metadata) + if is_audio_output and audio_sample_rate is None: + model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name) + audio_sample_rate = getattr(model_cls, "audio_sample_rate", None) + + return _format_single_prompt_output( + request=request, + prompt=request.prompt, + diffusion_output=diffusion_output, + outputs=outputs, + metrics=metrics, + postprocess_output=postprocess_output, + is_text_output=is_text_output, + is_audio_output=is_audio_output, + audio_sample_rate=audio_sample_rate, + finished=diffusion_output.finished, + ) + + +def _ensure_list(outputs: DiffusionPayloadValue) -> list[DiffusionPayloadValue]: + if isinstance(outputs, list): + return outputs + return [outputs] if outputs is not None else [] + + +def _primary_payload(postprocess_output: DiffusionPostprocessOutput) -> DiffusionPayloadValue | None: + if postprocess_output.primary_key is None: + return [] + return postprocess_output.outputs.get(postprocess_output.primary_key) + + +def _metadata_audio_sample_rate(metadata: DiffusionMetadata) -> int | None: + audio_metadata = metadata.get("audio") + if not isinstance(audio_metadata, dict): + return None + sample_rate = audio_metadata.get("sample_rate") + return sample_rate if isinstance(sample_rate, int) else None + + +def _metadata_video_fps(metadata: DiffusionMetadata) -> float | None: + video_metadata = metadata.get("video") + if not isinstance(video_metadata, dict): + return None + fps = video_metadata.get("fps") + return fps if isinstance(fps, (int, float)) and not isinstance(fps, bool) else None + + +def _format_audio_multimodal_output( + payload: DiffusionPayloadValue, + audio_sample_rate: int | None, + metadata: DiffusionMetadata | None = None, +) -> DiffusionMultimodalOutput: + mm_output: DiffusionMultimodalOutput = {"audio": payload} + if metadata: + mm_output["metadata"] = metadata + if audio_sample_rate is not None: + mm_output["audio_sample_rate"] = audio_sample_rate + return mm_output + + +def _has_non_audio_postprocess_payload(postprocess_output: DiffusionPostprocessOutput) -> bool: + return ( + "video" in postprocess_output.outputs + or "image" in postprocess_output.outputs + or "text" in postprocess_output.metadata + or "actions" in postprocess_output.outputs + or "trajectory" in postprocess_output.outputs + or _metadata_video_fps(postprocess_output.metadata) is not None + ) + + +def _build_multimodal_output( + postprocess_output: DiffusionPostprocessOutput, + audio_sample_rate: int | None, +) -> DiffusionMultimodalOutput: + mm_output: DiffusionMultimodalOutput = {} + if postprocess_output.metadata: + mm_output["metadata"] = postprocess_output.metadata + for key, value in postprocess_output.outputs.items(): + if key in {"audio", "actions", "trajectory"}: + mm_output[key] = value + if audio_sample_rate is not None: + mm_output["audio_sample_rate"] = audio_sample_rate + if (fps := _metadata_video_fps(postprocess_output.metadata)) is not None: + mm_output["fps"] = fps + return mm_output + + +def _format_single_prompt_output( + *, + request: OmniDiffusionRequest, + prompt: OmniPromptType, + diffusion_output: DiffusionOutput, + outputs: list[DiffusionPayloadValue], + metrics: dict[str, object], + postprocess_output: DiffusionPostprocessOutput, + is_text_output: bool, + is_audio_output: bool, + audio_sample_rate: int | None, + finished: bool = True, +) -> list[OmniRequestOutput]: + request_id = request.request_id + mm_output = _build_multimodal_output(postprocess_output, audio_sample_rate) + if is_text_output: + mm_output["text"] = outputs[0] if len(outputs) == 1 else outputs + trajectory_payload = _trajectory_payload(postprocess_output, diffusion_output) + trajectory_latents = trajectory_payload.get("latents") + trajectory_timesteps = trajectory_payload.get("timesteps") + trajectory_log_probs = trajectory_payload.get("log_probs") + trajectory_decoded = trajectory_payload.get("decoded") + + if is_text_output: + return [ + OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + prompt=prompt, + metrics=metrics, + multimodal_output=mm_output, + final_output_type="text", + stage_durations=diffusion_output.stage_durations, + peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, + ), + ] + + if is_audio_output and not _has_non_audio_postprocess_payload(postprocess_output): + request_audio_payload = postprocess_output.outputs.get("audio") + if request_audio_payload is None: + request_audio_payload = outputs[0] if len(outputs) == 1 else outputs + return [ + OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + prompt=prompt, + metrics=metrics, + latents=trajectory_latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, + multimodal_output=_format_audio_multimodal_output( + request_audio_payload, + audio_sample_rate, + postprocess_output.metadata, + ), + final_output_type="audio", + stage_durations=diffusion_output.stage_durations, + peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, + ), + ] + + return [ + OmniRequestOutput.from_diffusion( + request_id=request_id, + images=outputs, + prompt=prompt, + metrics=metrics, + latents=trajectory_latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, + multimodal_output=mm_output, + stage_durations=diffusion_output.stage_durations, + peak_memory_mb=diffusion_output.peak_memory_mb, + finished=finished, + ), + ] + + +def _trajectory_payload( + postprocess_output: DiffusionPostprocessOutput, + diffusion_output: DiffusionOutput, +) -> DiffusionTrajectoryPayload: + trajectory: DiffusionTrajectoryPayload = {} + payload = postprocess_output.outputs.get("trajectory") + if isinstance(payload, Mapping): + for source_key, target_key in ( + ("latents", "latents"), + ("timesteps", "timesteps"), + ("log_probs", "log_probs"), + ("decoded", "decoded"), + ): + value = payload.get(source_key) + if value is not None: + trajectory[target_key] = value + + fallback_fields = ( + ("latents", diffusion_output.trajectory_latents), + ("timesteps", diffusion_output.trajectory_timesteps), + ("log_probs", diffusion_output.trajectory_log_probs), + ("decoded", diffusion_output.trajectory_decoded), + ) + for key, value in fallback_fields: + if key not in trajectory and value is not None: + trajectory[key] = value + return trajectory diff --git a/vllm_omni/diffusion/output_metadata.py b/vllm_omni/diffusion/output_metadata.py new file mode 100644 index 00000000000..abea8759eca --- /dev/null +++ b/vllm_omni/diffusion/output_metadata.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility imports for diffusion output metadata. + +The shared output payload and metadata schema lives in +``vllm_omni.outputs.output_metadata``. Keep this module as a stable import +path for existing diffusion callers. +""" + +from vllm_omni.outputs.output_metadata import * # noqa: F403 diff --git a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py index f67e9f692fe..5a293b9053c 100644 --- a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py +++ b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py @@ -14,6 +14,7 @@ def profiler(name: str, func: Callable, instance: Any) -> Callable: """Timing a function execution.""" + metric_name = f"{name.rsplit('.', 1)[0]}.diffuse" if name.endswith(".denoise_step") else name @functools.wraps(func) def wrapper(*args, **kwargs) -> Any: @@ -28,10 +29,10 @@ def wrapper(*args, **kwargs) -> Any: if current_omni_platform.is_available(): current_omni_platform.synchronize() duration = time.perf_counter() - start_time - logger.info(f"[DiffusionPipelineProfiler] {name} took {duration:.6f}s") + logger.info(f"[DiffusionPipelineProfiler] {metric_name} took {duration:.6f}s") # record the profiling data: duration of stages with instance._profiler_lock: - instance._stage_durations[name] = instance._stage_durations.get(name, 0.0) + duration + instance._stage_durations[metric_name] = instance._stage_durations.get(metric_name, 0.0) + duration return wrapper @@ -81,7 +82,9 @@ class DiffusionPipelineProfilerMixin: _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"] def setup_diffusion_pipeline_profiler( - self, profiler_targets: list[str] | None = None, enable_diffusion_pipeline_profiler: bool = False + self, + profiler_targets: list[str] | None = None, + enable_diffusion_pipeline_profiler: bool = False, ) -> None: self.enable_diffusion_pipeline_profiler = enable_diffusion_pipeline_profiler if not enable_diffusion_pipeline_profiler: @@ -91,7 +94,11 @@ def setup_diffusion_pipeline_profiler( self._stage_durations: dict[str, float] = {} targets = profiler_targets if profiler_targets is not None else self._PROFILER_TARGETS if not targets: - return + targets = [] + else: + targets = list(targets) + if not profiler_targets and hasattr(self, "denoise_step"): + targets.append("denoise_step") targets = ["forward"] + [ t for t in targets if t != "forward" diff --git a/vllm_omni/diffusion/quantization/__init__.py b/vllm_omni/diffusion/quantization/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/diffusion/quantization/__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/diffusion/quantization/hsdp_fp8.py b/vllm_omni/diffusion/quantization/hsdp_fp8.py new file mode 100644 index 00000000000..7c4666e29c6 --- /dev/null +++ b/vllm_omni/diffusion/quantization/hsdp_fp8.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HSDP/FSDP2 compatibility for online FP8 quantization. + +vllm's ``Fp8LinearMethod.process_weights_after_loading`` ends with +``layer.weight = qweight.t()`` so that the Cutlass FP8 GEMM kernel sees its +B operand as column-major ``[K, N]`` (the TN layout required by Hopper FP8 +``wgmma`` instructions). The resulting tensor is a non-contiguous transpose +view of a ``(out_features, in_features)`` row-major storage. + +FSDP2 ``fully_shard`` rejects non-contiguous parameters because dim-0 sharding +on a column-major view cannot be a contiguous memcpy. The two requirements are +fundamentally at odds with the same physical buffer. + +This module reconciles them by separating the views: keep the parameter as +the underlying ``(out, in)`` row-major contiguous storage (FSDP-friendly), +and inject the equivalent ``.t()`` at the GEMM call site so the Cutlass +kernel still receives a column-major B (zero-copy stride flip). + +The transpose is injected inside ``ScaledMMLinearKernel._get_layer_params``, +which is the single place where ``apply_weights`` reads the weight tensor. +The rest of ``apply_weights`` -- including its ``output_shape = w.shape[1]`` +computation and the eventual ``apply_scaled_mm(B=w, ...)`` call -- then +operates on a tensor whose shape ``(K, N)`` and stride ``(1, K)`` match +what the upstream FP8 kernel expects. No other upstream code is overridden, +which keeps us robust against future changes inside ``apply_weights``. +""" + +from __future__ import annotations + +import types + +from torch import nn +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod + +logger = init_logger(__name__) + + +def _build_transposed_get_layer_params(original_bound_method): + """Return a ``_get_layer_params`` replacement that transposes ``w``. + + ``original_bound_method`` is the kernel instance's existing + ``_get_layer_params`` (already bound to ``self``); we close over it and + apply ``.t()`` to the returned weight, leaving scales untouched. + """ + + def _get_layer_params(self, layer): + w, w_s, x_s, x_s_ub = original_bound_method(layer) + return w.t(), w_s, x_s, x_s_ub + + return _get_layer_params + + +def prepare_fp8_layers_for_fsdp(model: nn.Module) -> int: + """Make online-FP8 linear layers in ``model`` FSDP2-compatible. + + For every layer whose ``quant_method`` is an :class:`Fp8LinearMethod` + and whose weight is currently a non-contiguous transpose view, this + function: + + 1. Replaces ``layer.weight`` with the underlying ``(out, in)`` row-major + contiguous storage so FSDP2 ``fully_shard`` accepts it. + 2. Patches the per-layer GEMM kernel's bound ``_get_layer_params`` method + to return a ``.t()`` view of the weight, so ``apply_weights`` and the + downstream ``apply_scaled_mm`` continue to see a column-major + ``(in, out)`` B with zero copies. + + Layers whose weight is already contiguous (e.g. Marlin FP8, offline- + quantized checkpoints) or that use a different quant method are skipped. + + Returns: + Number of layers rewritten. + """ + n_patched = 0 + patched_kernel_ids: set[int] = set() + for module in model.modules(): + qm = getattr(module, "quant_method", None) + if not isinstance(qm, Fp8LinearMethod): + continue + + weight = getattr(module, "weight", None) + if weight is None or weight.is_contiguous(): + continue + + # ``weight`` here is ``qweight.t()`` (a non-contiguous (in, out) view + # of an (out, in) row-major qweight storage). ``weight.t()`` recovers + # that storage; ``.contiguous()`` is a zero-copy alias when the source + # is already row-major contiguous, but defensively materializes if + # upstream ever stacks another view. + contig = weight.data.t().contiguous() + new_param = nn.Parameter(contig, requires_grad=False) + module.weight = new_param + + kernel = qm.fp8_linear + # Each linear layer is expected to have its own kernel instance, but + # guard against shared instances to avoid stacking multiple ``.t()`` + # patches on the same object (which would compose into identity). + if id(kernel) not in patched_kernel_ids: + original_get = kernel._get_layer_params + kernel._get_layer_params = types.MethodType(_build_transposed_get_layer_params(original_get), kernel) + patched_kernel_ids.add(id(kernel)) + + n_patched += 1 + + if n_patched: + logger.info( + "Rewrote %d FP8 linear layer(s) into FSDP2-compatible storage; " + "transpose to column-major B is now applied at GEMM time.", + n_patched, + ) + return n_patched diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 9cc65486a7a..22c36f32b3f 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -73,17 +73,17 @@ ), "LTX2ImageToVideoPipeline": ( "ltx2", - "pipeline_ltx2_image2video", + "pipeline_ltx2", "LTX2ImageToVideoPipeline", ), "LTX2TwoStagesPipeline": ( "ltx2", - "pipeline_ltx2", + "pipeline_ltx2_two_stage", "LTX2TwoStagesPipeline", ), "LTX2ImageToVideoTwoStagesPipeline": ( "ltx2", - "pipeline_ltx2_image2video", + "pipeline_ltx2_two_stage", "LTX2ImageToVideoTwoStagesPipeline", ), "LTX2T2VDMD2Pipeline": ( @@ -93,17 +93,17 @@ ), "LTX2I2VDMD2Pipeline": ( "ltx2", - "pipeline_ltx2_image2video", + "pipeline_ltx2", "LTX2I2VDMD2Pipeline", ), "LTX23Pipeline": ( "ltx2", - "pipeline_ltx2_3", + "pipeline_ltx2", "LTX23Pipeline", ), "LTX23ImageToVideoPipeline": ( "ltx2", - "pipeline_ltx2_3_image2video", + "pipeline_ltx2", "LTX23ImageToVideoPipeline", ), "StableAudioPipeline": ( @@ -156,6 +156,11 @@ "pipeline_internvla_a1", "InternVLAA1Pipeline", ), + "Gr00tN1d7Pipeline": ( + "gr00t", + "pipeline_gr00t", + "Gr00tN1d7Pipeline", + ), "LongCatImageEditPipeline": ( "longcat_image", "pipeline_longcat_image_edit", @@ -271,6 +276,21 @@ "pipeline_cosmos3", "Cosmos3OmniDiffusersPipeline", ), + "Cosmos3OmniPipeline": ( + "cosmos3", + "pipeline_cosmos3", + "Cosmos3OmniDiffusersPipeline", + ), + "SoulXSingerPipeline": ( + "soulx_singer", + "pipeline_soulx_singer_svs", + "PipelineSoulXSingerSVS", + ), + "SoulXSingerSVCPipeline": ( + "soulx_singer", + "pipeline_soulx_singer_svc", + "PipelineSoulXSingerSVC", + ), "DiffusersAdapterPipeline": ( "diffusers_adapter", "pipeline_diffusers_adapter", @@ -286,6 +306,16 @@ "pipeline_dreamzero", "DreamZeroPipeline", ), + "StableDiffusionXLPipeline": ( + "sdxl", + "pipeline_sdxl", + "StableDiffusionXLPipeline", + ), + "Krea2Pipeline": ( + "krea2", + "pipeline_krea2", + "Krea2Pipeline", + ), } @@ -370,7 +400,7 @@ def initialize_model( model.vae.use_tiling = od_config.vae_use_tiling if is_distributed_vae: - model.vae.set_parallel_size(vae_pp_size) + model.vae.set_parallel_size(vae_pp_size, mode=od_config.parallel_config.vae_parallel_mode) # Apply sequence parallelism if enabled # This follows diffusers' pattern where enable_parallelism() is called @@ -479,6 +509,8 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "LTX23Pipeline": "get_ltx2_post_process_func", "LTX23ImageToVideoPipeline": "get_ltx2_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", + "SoulXSingerPipeline": "get_soulxsinger_post_process_func", + "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func", "AudioXPipeline": "get_audiox_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanS2VPipeline": "get_wan22_s2v_post_process_func", @@ -509,14 +541,10 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", "SenseNovaU1Pipeline": "get_sensenova_u1_post_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_post_process_func", + "Cosmos3OmniPipeline": "get_cosmos3_post_process_func", "HiDreamImagePipeline": "get_hidream_image_post_process_func", -} - -_DIFFUSION_ACTION_POST_PROCESS_FUNCS = { - # arch: action_post_process_func - # `action_post_process_func` function must be placed in {mod_folder}/{mod_relname}.py, - # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key. - "Cosmos3OmniDiffusersPipeline": "get_cosmos3_action_post_process_func", + "StableDiffusionXLPipeline": "get_sdxl_image_post_process_func", + "Krea2Pipeline": "get_krea2_post_process_func", } _DIFFUSION_IR_OP_PRIORITY_FUNCS = { @@ -524,6 +552,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - # `ir_op_priority_func` function must be placed in {mod_folder}/{mod_relname}.py, # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key. "Cosmos3OmniDiffusersPipeline": "get_cosmos3_ir_op_priority_func", + "Cosmos3OmniPipeline": "get_cosmos3_ir_op_priority_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { @@ -548,6 +577,9 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_pre_process_func", "MagiHumanPipeline": "get_magi_human_pre_process_func", "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func", + "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func", + "SoulXSingerPipeline": "get_soulxsinger_pre_process_func", + "SoulXSingerSVCPipeline": "get_soulxsinger_svc_pre_process_func", } @@ -557,8 +589,8 @@ def register_diffusion_model( class_name: str, pre_process_func_name: str | None = None, post_process_func_name: str | None = None, - action_post_process_func_name: str | None = None, ir_op_priority_func_name: str | None = None, + action_post_process_func_name: str | None = None, ) -> None: """Register a diffusion model pipeline from an out-of-tree plugin. @@ -577,13 +609,23 @@ def register_diffusion_model( post_process_func_name: Optional name of the post-process function located in *module_name*. Pass ``None`` to keep the existing entry when replacing a built-in model. - action_post_process_func_name: Optional name of the action post-process - function located in *module_name*. Pass ``None`` to keep the - existing entry when replacing a built-in model. ir_op_priority_func_name: Optional name of the IR op priority merge function located in *module_name*. Pass ``None`` to keep the existing entry when replacing a built-in model. + action_post_process_func_name: Deprecated compatibility-only keyword + for out-of-tree plugins. Action postprocess hooks are no longer + registered separately; move action handling into + ``post_process_func_name`` and return a payload/metadata envelope. """ + if action_post_process_func_name is not None: + logger.warning( + "Ignoring deprecated action_post_process_func_name=%r for diffusion " + "model %s. Move action postprocess logic into post_process_func_name " + "and return payload/metadata output.", + action_post_process_func_name, + model_arch, + ) + # Register model class in DiffusionModelRegistry DiffusionModelRegistry.register_model( model_arch, @@ -600,8 +642,6 @@ def register_diffusion_model( _DIFFUSION_PRE_PROCESS_FUNCS[model_arch] = pre_process_func_name if post_process_func_name is not None: _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name - if action_post_process_func_name is not None: - _DIFFUSION_ACTION_POST_PROCESS_FUNCS[model_arch] = action_post_process_func_name if ir_op_priority_func_name is not None: _DIFFUSION_IR_OP_PRIORITY_FUNCS[model_arch] = ir_op_priority_func_name @@ -634,13 +674,6 @@ def get_diffusion_post_process_func(od_config: OmniDiffusionConfig): return _load_process_func(od_config, func_name) -def get_diffusion_action_post_process_func(od_config: OmniDiffusionConfig): - if od_config.model_class_name not in _DIFFUSION_ACTION_POST_PROCESS_FUNCS: - return None - func_name = _DIFFUSION_ACTION_POST_PROCESS_FUNCS[od_config.model_class_name] - return _load_process_func(od_config, func_name) - - def get_diffusion_ir_op_priority_func(od_config: OmniDiffusionConfig): if od_config.model_class_name not in _DIFFUSION_IR_OP_PRIORITY_FUNCS: return None diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index 79198c3ec44..42181419c6f 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -13,10 +13,11 @@ @dataclass class OmniDiffusionRequest: """ - Complete state passed through the pipeline execution. + Input payload for a single diffusion request. - This dataclass contains the prompts and sampling parameters for the diffusion pipeline + This dataclass contains the prompt and sampling parameters for the diffusion pipeline execution. It also contains a request_id for other components to trace this request and its outputs. + The runner wraps one or more requests into a DiffusionRequestBatch before pipeline execution. """ # TODO(will): double check that args are separate from server_args @@ -24,7 +25,7 @@ class OmniDiffusionRequest: # specific arguments. # data_type: DataType - prompts: list[OmniPromptType] # Actually supporting str-based prompts + prompt: OmniPromptType sampling_params: OmniDiffusionSamplingParams request_id: str kv_sender_info: dict | None = None @@ -50,8 +51,8 @@ def __post_init__(self): self.sampling_params.guidance_scale = 1.0 # Set do_classifier_free_guidance based on guidance scale and negative prompt - if self.sampling_params.guidance_scale > 1.0 and any( - (not isinstance(p, str) and p.get("negative_prompt")) for p in self.prompts + if self.sampling_params.guidance_scale > 1.0 and ( + not isinstance(self.prompt, str) and self.prompt.get("negative_prompt") ): self.sampling_params.do_classifier_free_guidance = True diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 812873a1cdd..9e4b6af5b6b 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -16,6 +16,7 @@ DiffusionRequestStatus, DiffusionSchedulerOutput, NewRequestData, + RequestBatchSamplingParamsKey, SamplingParamsKey, SchedulerInterface, ) @@ -24,22 +25,31 @@ # LoRA identity is derived from `sampling.lora_request`, not a same-named field # on sampling params, so it must be resolved separately from the bulk lookup. -_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) - {"lora_int_id"} +_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(f.name for f in fields(SamplingParamsKey)) - {"lora_int_id"} +_REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES = frozenset(f.name for f in fields(RequestBatchSamplingParamsKey)) - { + "lora_int_id" +} -def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey | None: +def get_sampling_params_key(request: OmniDiffusionRequest) -> SamplingParamsKey: """Build a batch-compatibility key from the request's sampling params.""" - if len(request.prompts) != 1: - return None - sampling = request.sampling_params lora_request = getattr(sampling, "lora_request", None) return SamplingParamsKey( lora_int_id=lora_request.lora_int_id if lora_request is not None else None, - **{name: getattr(sampling, name) for name in _KEY_FIELD_NAMES}, + **{name: getattr(sampling, name) for name in _SAMPLING_PARAMS_KEY_FIELD_NAMES}, ) +def get_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey: + """Build a request-batch compatibility key from the request's sampling params.""" + sampling = request.sampling_params + lora_request = getattr(sampling, "lora_request", None) + key_kwargs = {name: getattr(sampling, name) for name in _REQUEST_BATCH_SAMPLING_PARAMS_KEY_FIELD_NAMES} + key_kwargs["lora_int_id"] = lora_request.lora_int_id if lora_request is not None else None + return RequestBatchSamplingParamsKey(**key_kwargs) + + class _BaseScheduler(SchedulerInterface): """Shared queue/state bookkeeping for diffusion schedulers.""" @@ -49,9 +59,10 @@ def __init__(self) -> None: self._step_id: int = 0 self._waiting: deque[str] = deque() self._running: list[str] = [] - self._running_sampling_params_key: SamplingParamsKey | None = None + self._running_sampling_params_key: SamplingParamsKey | RequestBatchSamplingParamsKey | None = None self._finished_req_ids: set[str] = set() self.max_num_running_reqs: int = 1 + self._prefetch_enabled: bool = False def initialize(self, od_config: OmniDiffusionConfig) -> None: self.od_config = od_config @@ -66,6 +77,8 @@ def initialize(self, od_config: OmniDiffusionConfig) -> None: self.max_num_running_reqs = max(1, int(max_num_seqs)) except (TypeError, ValueError): self.max_num_running_reqs = 1 + omni_kv = getattr(od_config, "omni_kv_config", None) or {} + self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False)) self._reset_scheduler_state() def add_request(self, request: OmniDiffusionRequest) -> str: @@ -111,6 +124,22 @@ def schedule(self) -> DiffusionSchedulerOutput: else: scheduled_cached_request_ids.append(request_id) + # Expose the next waiting request (serial mode) so the runner can + # prefetch its KV during this forward. Skip a request without + # kv_sender_info (would target the wrong sender under multi-replica) or + # one already finished/aborted (would consume its sender buffer for + # nothing). + kv_prefetch_jobs: dict | None = None + if self._prefetch_enabled and self._waiting: + nxt = self._request_states.get(self._waiting[0]) + if nxt is not None and not nxt.is_finished(): + sender_info = getattr(nxt.req, "kv_sender_info", None) + if sender_info: + kv_prefetch_jobs = { + "request_id": nxt.request_id, + "kv_sender_info": sender_info, + } + scheduler_output = DiffusionSchedulerOutput( step_id=self._step_id, scheduled_new_reqs=scheduled_new_reqs, @@ -118,6 +147,7 @@ def schedule(self) -> DiffusionSchedulerOutput: finished_req_ids=set(self._finished_req_ids), num_running_reqs=len(self._running), num_waiting_reqs=len(self._waiting), + kv_prefetch_jobs=kv_prefetch_jobs, ) # update after schedule @@ -128,6 +158,12 @@ def schedule(self) -> DiffusionSchedulerOutput: def has_requests(self) -> bool: return bool(self._waiting or self._running) + def num_waiting_requests(self) -> int: + return len(self._waiting) + + def num_running_requests(self) -> int: + return len(self._running) + def get_request_state(self, request_id: str) -> DiffusionRequestState | None: return self._request_states.get(request_id) @@ -230,7 +266,7 @@ def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> return DiffusionRequestState( request_id=request_id, req=request, - sampling_params_key=get_sampling_params_key(request), + sampling_params_key=self._build_sampling_params_key(request), ) def _can_schedule_waiting(self, state: DiffusionRequestState) -> bool: @@ -240,9 +276,14 @@ def _can_schedule_waiting(self, state: DiffusionRequestState) -> bool: current_key = self._current_sampling_params_key() return current_key is not None and current_key == state.sampling_params_key - def _current_sampling_params_key(self) -> SamplingParamsKey | None: + def _current_sampling_params_key(self) -> SamplingParamsKey | RequestBatchSamplingParamsKey | None: if self._running_sampling_params_key is not None or not self._running: return self._running_sampling_params_key state = self._request_states.get(self._running[0]) self._running_sampling_params_key = None if state is None else state.sampling_params_key return self._running_sampling_params_key + + def _build_sampling_params_key( + self, request: OmniDiffusionRequest + ) -> SamplingParamsKey | RequestBatchSamplingParamsKey | None: + return get_sampling_params_key(request) diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index b50e3555afb..695641b5d94 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -35,7 +35,7 @@ def is_finished(status: DiffusionRequestStatus) -> bool: @dataclass(frozen=True, eq=True) class SamplingParamsKey: - """Batch-compatibility key derived from ``OmniDiffusionSamplingParams``. + """Denoise step level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``. Only requests with the same key can be batched together. Fields not included here are treated as request-local and do not @@ -60,19 +60,71 @@ class SamplingParamsKey: true_cfg_scale: float | None = None cfg_normalize: bool = False + # Output count. Requests with different num_outputs_per_prompt produce + # differently shaped outputs and cannot share a batch. + num_outputs_per_prompt: int = 1 + # LoRA identity. Requests with different adapters or scales must run in # separate batches so the worker can activate exactly one adapter per step. lora_int_id: int | None = None lora_scale: float = 1.0 +@dataclass(frozen=True, eq=True) +class RequestBatchSamplingParamsKey: + """Request level Batch-compatibility key derived from ``OmniDiffusionSamplingParams``. + + Only request-batch-wide fields belong here. Request-local values such as + seeds, generators, latent tensors, timesteps, and pipeline-specific + ``extra_args`` are read per request from + ``DiffusionRequestBatch.sampling_params_list``. + """ + + # Spatial / temporal shape. + height: object = None + width: object = None + num_frames: int = 1 + resolution: object = 640 + fps: object = None + frame_rate: object = None + boundary_ratio: object = None + + # CFG / guidance. + do_classifier_free_guidance: bool = False + guidance_scale: float = 0.0 + guidance_scale_provided: bool = False + guidance_scale_2: object = None + guidance_rescale: float = 0.0 + true_cfg_scale: object = None + cfg_normalize: bool = False + strength: object = None + + # Scheduling / output shape. + num_inference_steps: object = None + sigmas: object = None + max_sequence_length: object = None + num_outputs_per_prompt: int = 1 + eta: float = 0.0 + decode_timestep: object = None + decode_noise_scale: object = None + output_type: object = None + + # Model-specific batch defaults used by request-mode pipelines. + layers: int = 4 + use_en_prompt: bool = False + + # LoRA identity. + lora_int_id: int | None = None + lora_scale: float = 1.0 + + @dataclass class DiffusionRequestState: """Scheduler-owned state for one queued OmniDiffusionRequest.""" request_id: str req: OmniDiffusionRequest - sampling_params_key: SamplingParamsKey | None = None + sampling_params_key: SamplingParamsKey | RequestBatchSamplingParamsKey | None = None status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING error: str | None = None @@ -82,7 +134,12 @@ def is_finished(self) -> bool: @dataclass class NewRequestData: - """Full request payload for a newly scheduled diffusion request.""" + """Payload for a newly scheduled diffusion request. + + Carries the already-initialized request object so executors and workers do + not re-run ``OmniDiffusionRequest.__post_init__`` and mutate sentinel-based + fields like ``guidance_scale_provided``. + """ request_id: str req: OmniDiffusionRequest @@ -113,6 +170,8 @@ class DiffusionSchedulerOutput: finished_req_ids: set[str] num_running_reqs: int num_waiting_reqs: int + # next request to background-prefetch KV + kv_prefetch_jobs: dict | None = None @cached_property def scheduled_request_ids(self) -> list[str]: @@ -160,6 +219,14 @@ def get_request_state(self, request_id: str) -> DiffusionRequestState | None: def has_requests(self) -> bool: """Return whether the scheduler still owns runnable requests.""" + @abstractmethod + def num_waiting_requests(self) -> int: + """Return the number of requests waiting to be scheduled.""" + + @abstractmethod + def num_running_requests(self) -> int: + """Return the number of requests currently running.""" + @abstractmethod def pop_request_state(self, request_id: str) -> DiffusionRequestState | None: """Remove and return request state if present.""" diff --git a/vllm_omni/diffusion/sched/request_scheduler.py b/vllm_omni/diffusion/sched/request_scheduler.py index ba9cda61615..9a02eb8f3c0 100644 --- a/vllm_omni/diffusion/sched/request_scheduler.py +++ b/vllm_omni/diffusion/sched/request_scheduler.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler +from vllm_omni.diffusion.sched.base_scheduler import _BaseScheduler, get_request_batch_sampling_params_key from vllm_omni.diffusion.sched.interface import ( DiffusionRequestStatus, DiffusionSchedulerOutput, @@ -19,6 +19,9 @@ class RequestScheduler(_BaseScheduler): """Diffusion scheduler with vLLM-style waiting/running queues.""" + def _build_sampling_params_key(self, request: OmniDiffusionRequest): + return get_request_batch_sampling_params_key(request) + def add_request(self, request: OmniDiffusionRequest) -> str: return super().add_request(request) @@ -41,6 +44,9 @@ def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: Run if result is None: terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR terminal_errors[request_id] = "No output result" + elif result.aborted: + terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ABORTED + terminal_errors[request_id] = None elif result.error: terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR terminal_errors[request_id] = result.error diff --git a/vllm_omni/diffusion/sched/step_scheduler.py b/vllm_omni/diffusion/sched/step_scheduler.py index 0bc2d4a3cda..6cb6a9a28f1 100644 --- a/vllm_omni/diffusion/sched/step_scheduler.py +++ b/vllm_omni/diffusion/sched/step_scheduler.py @@ -88,6 +88,10 @@ def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: Run continue req_result = req_output.result + if req_result is not None and req_result.aborted: + terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ABORTED + terminal_errors[request_id] = None + continue output_error = req_result.error if req_result is not None else None if output_error is not None: terminal_statuses[request_id] = DiffusionRequestStatus.FINISHED_ERROR diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 613fb64c71d..2b13ddc9d66 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -1,8 +1,7 @@ """Stage Diffusion Client for vLLM-Omni multi-stage runtime. -Spawns StageDiffusionProc in a subprocess and communicates via ZMQ -(PUSH/PULL) to expose the same interface the Orchestrator expects -from any stage client. +Owns the frontend-side ZMQ sockets for StageDiffusionProc and exposes the +interface the Orchestrator expects from a stage client. """ from __future__ import annotations @@ -23,15 +22,14 @@ from vllm_omni.diffusion.stage_diffusion_proc import ( StageDiffusionProc, - complete_diffusion_handshake, - spawn_diffusion_proc, + StageDiffusionProcManager, ) from vllm_omni.distributed.omni_connectors.utils.serialization import ( OmniMsgpackDecoder, OmniMsgpackEncoder, ) from vllm_omni.engine.stage_client import StageClientBase -from vllm_omni.engine.stage_init_utils import StageMetadata, terminate_alive_proc +from vllm_omni.engine.stage_init_utils import StageMetadata from vllm_omni.outputs import OmniRequestOutput if TYPE_CHECKING: @@ -39,6 +37,7 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType logger = init_logger(__name__) +_MISSING_RPC_RESULT = object() def create_diffusion_client( @@ -54,8 +53,17 @@ def create_diffusion_client( from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size) - return StageDiffusionClient( - model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size + proc_manager = StageDiffusionProcManager( + model=model, + od_config=od_config, + stage_init_timeout=stage_init_timeout, + ) + return StageDiffusionClient.from_addresses( + metadata, + request_address=proc_manager.addresses.inputs[0], + response_address=proc_manager.addresses.outputs[0], + proc_manager=proc_manager, + batch_size=batch_size, ) @@ -74,16 +82,20 @@ class StageDiffusionClient(StageClientBase): def __init__( self, - model: str, - od_config: OmniDiffusionConfig, metadata: StageMetadata, - stage_init_timeout: int, + request_address: str, + response_address: str, + *, + proc_manager: StageDiffusionProcManager | None = None, batch_size: int = 1, ) -> None: - # Spawn StageDiffusionProc subprocess and wait for READY. - proc, handshake_address, request_address, response_address = spawn_diffusion_proc(model, od_config) - complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) - self._initialize_client(metadata, request_address, response_address, proc=proc, batch_size=batch_size) + self._initialize_client( + metadata, + request_address, + response_address, + proc_manager=proc_manager, + batch_size=batch_size, + ) @classmethod def from_addresses( @@ -92,19 +104,17 @@ def from_addresses( request_address: str, response_address: str, *, - proc: Any = None, + proc_manager: StageDiffusionProcManager | None = None, batch_size: int = 1, ) -> StageDiffusionClient: """Create a client for an already-running diffusion subprocess.""" - client = cls.__new__(cls) - client._initialize_client( + return cls( metadata, request_address, response_address, - proc=proc, + proc_manager=proc_manager, batch_size=batch_size, ) - return client def _initialize_client( self, @@ -112,20 +122,44 @@ def _initialize_client( request_address: str, response_address: str, *, - proc: Any, + proc_manager: StageDiffusionProcManager | None = None, batch_size: int, ) -> None: + self._set_stage_metadata(metadata) + self._proc_manager = proc_manager + self._connect_transport(request_address, response_address) + + self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() + self._rpc_results: dict[str, Any] = {} + self._pending_rpcs: set[str] = set() + self._tasks: dict[str, asyncio.Task] = {} + self._shutting_down = False + self._engine_dead: bool = False + + if self._proc_manager is not None: + self._start_proc_monitor() + + logger.info( + "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)", + self.stage_id, + self.replica_id, + self._proc_manager is not None, + batch_size, + ) + + def _set_stage_metadata(self, metadata: StageMetadata) -> None: self.stage_id = metadata.stage_id self.replica_id = metadata.replica_id self.final_output = metadata.final_output self.final_output_type = metadata.final_output_type - self.model_stage = getattr(metadata, "model_stage", None) + self.model_stage = metadata.model_stage self.default_sampling_params = metadata.default_sampling_params + self.prompt_expand_func = metadata.prompt_expand_func self.requires_multimodal_data = getattr(metadata, "requires_multimodal_data", False) self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None) self.engine_input_source = getattr(metadata, "engine_input_source", []) - self._proc = proc - self._owns_process = proc is not None + + def _connect_transport(self, request_address: str, response_address: str) -> None: # Expose the ZMQ addresses on the instance so callers (e.g. # ``StagePool._client_input_addr``) can identify the diffusion # replica by its bound address. @@ -134,9 +168,9 @@ def _initialize_client( self._zmq_ctx = zmq.Context() self._request_socket = self._zmq_ctx.socket(zmq.PUSH) - self._request_socket.connect(request_address) + self._request_socket.bind(request_address) self._response_socket = self._zmq_ctx.socket(zmq.PULL) - self._response_socket.connect(response_address) + self._response_socket.bind(response_address) self._response_poller = zmq.asyncio.Poller() self._response_poller.register(self._response_socket, zmq.POLLIN) @@ -144,26 +178,6 @@ def _initialize_client( self._encoder = OmniMsgpackEncoder() self._decoder = OmniMsgpackDecoder() - self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() - self._rpc_results: dict[str, Any] = {} - self._pending_rpcs: set[str] = set() - self._tasks: dict[str, asyncio.Task] = {} - self._shutting_down = False - self._engine_dead: bool = False - - # Background thread to detect silent process death (SIGKILL, segfault) - # where the subprocess cannot send the ZMQ death sentinel. - # Mirrors MPClient.start_engine_core_monitor() in vLLM. - self._start_proc_monitor() - - logger.info( - "[StageDiffusionClient] stage-%s [rep-%s] initialized (owns_process=%s, batch_size=%d)", - self.stage_id, - self.replica_id, - self._owns_process, - batch_size, - ) - # ------------------------------------------------------------------ # Process monitor (mirrors vLLM's MPClient.start_engine_core_monitor) # ------------------------------------------------------------------ @@ -175,7 +189,10 @@ def _start_proc_monitor(self) -> None: (e.g. SIGKILL, segfault), this thread sets ``_engine_dead`` so subsequent calls raise ``EngineDeadError``. """ - proc = self._proc + # Background thread to detect silent process death (SIGKILL, segfault) + # where the subprocess cannot send the ZMQ death sentinel. + # Mirrors MPClient.start_engine_core_monitor() in vLLM. + proc = self._proc_manager.proc self_ref = weakref.ref(self) def _monitor() -> None: @@ -338,72 +355,6 @@ async def add_request_async( ) ) - # TODO(Long): Temporary solution to boost performance of diffusion stages. - # Remove this after scheduling algorithm is implemented - async def add_batch_request_async( - self, - request_id: str, - prompts: list[OmniPromptType], - sampling_params: OmniDiffusionSamplingParams, - kv_sender_info: dict[int, dict[str, Any]] | None = None, - ) -> None: - """Submit a list of prompts as a single batched engine call. - - All prompts are processed in one ``DiffusionEngine.step()`` call - and the combined result is placed on the output queue with a single - *request_id*. - """ - if self._engine_dead: - raise EngineDeadError() - logger.info( - "[StageDiffusionClient] stage-%s [rep-%s] add batch request: %s (%d prompts)", - self.stage_id, - self.replica_id, - request_id, - len(prompts), - ) - task = asyncio.create_task( - self._run_batch( - request_id, - prompts, - sampling_params, - kv_sender_info, - ), - name=f"diffusion-batch-{request_id}", - ) - self._tasks[request_id] = task - - async def _run_batch( - self, - request_id: str, - prompts: list[OmniPromptType], - sampling_params: OmniDiffusionSamplingParams, - kv_sender_info: dict[int, dict[str, Any]] | None = None, - ) -> None: - try: - self._request_socket.send( - self._encoder.encode( - { - "type": "add_batch_request", - "request_id": request_id, - "prompts": prompts, - "sampling_params": self._sampling_params_to_dict(sampling_params), - "kv_sender_info": kv_sender_info, - } - ) - ) - except Exception as e: - logger.exception( - "[StageDiffusionClient] stage-%s [rep-%s] batch req=%s failed: %s", - self.stage_id, - self.replica_id, - request_id, - e, - ) - await self._output_queue.put(OmniRequestOutput.from_error(request_id, str(e))) - finally: - self._tasks.pop(request_id, None) - def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: self._drain_responses() try: @@ -413,9 +364,12 @@ def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: if self._shutting_down: return None raise EngineDeadError() - if not self._shutting_down and self._owns_process and self._proc is not None and not self._proc.is_alive(): + if self._proc_manager is None: + return None + proc = self._proc_manager.proc + if not self._shutting_down and not proc.is_alive(): self._engine_dead = True - exitcode = self._proc.exitcode + exitcode = proc.exitcode # One final drain – the last ZMQ frame may have arrived # between the first drain and the is_alive() check. self._drain_responses() @@ -487,14 +441,15 @@ async def collective_rpc_async( try: while True: self._drain_responses() - result = self._rpc_results.pop(rpc_id, None) - if result is not None: + result = self._rpc_results.pop(rpc_id, _MISSING_RPC_RESULT) + if result is not _MISSING_RPC_RESULT: return result - if self._engine_dead or (self._owns_process and self._proc is not None and not self._proc.is_alive()): + proc = self._proc_manager.proc + if self._engine_dead or not proc.is_alive(): self._engine_dead = True raise EngineDeadError( f"StageDiffusionProc died while waiting for " - f"collective_rpc '{method}' (exit code {self._proc.exitcode})" + f"collective_rpc '{method}' (exit code {proc.exitcode})" ) if deadline is not None and time.monotonic() > deadline: raise TimeoutError(f"collective_rpc_async '{method}' timed out after {timeout}s") @@ -518,10 +473,13 @@ def check_health(self) -> None: """ if self._engine_dead: raise EngineDeadError(f"Stage-{self.stage_id} diffusion subprocess is dead") - if self._proc is not None and not self._proc.is_alive(): + if self._proc_manager is None: + return + proc = self._proc_manager.proc + if not proc.is_alive(): self._engine_dead = True raise EngineDeadError( - f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {self._proc.exitcode})." + f"Stage-{self.stage_id} diffusion subprocess is not alive (exit code: {proc.exitcode})." ) def shutdown(self) -> None: @@ -531,9 +489,8 @@ def shutdown(self) -> None: except Exception: pass - if self._owns_process and self._proc is not None and self._proc.is_alive(): - self._proc.join(timeout=10) - terminate_alive_proc(self._proc) + if self._proc_manager is not None and self._proc_manager.proc.is_alive(): + self._proc_manager.shutdown(timeout=10) self._request_socket.close(linger=0) self._response_socket.close(linger=0) diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 27c6f16a066..e74a34b597b 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -8,20 +8,21 @@ import asyncio import contextlib +import multiprocessing.connection import signal -import time +from collections.abc import AsyncGenerator from concurrent.futures import ThreadPoolExecutor -from multiprocessing.process import BaseProcess +from types import SimpleNamespace from typing import TYPE_CHECKING, Any import msgspec -import torch import zmq import zmq.asyncio -from PIL import Image from vllm.logger import init_logger from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx from vllm.utils.system_utils import get_mp_context +from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.engine.utils import CoreEngine, EngineZmqAddresses, wait_for_engine_startup from vllm.v1.utils import shutdown from vllm_omni.diffusion.data import DiffusionRequestAbortedError @@ -158,7 +159,7 @@ async def _process_request( sampling_params = self._reconstruct_sampling_params(sampling_params_dict) request = OmniDiffusionRequest( - prompts=[prompt], + prompt=prompt, sampling_params=sampling_params, request_id=request_id, kv_sender_info=kv_sender_info, @@ -170,81 +171,28 @@ async def _process_request( result.request_id = request_id return result - async def _process_batch_request( + async def _process_streaming_request( self, request_id: str, - prompts: list[Any], + prompt: Any, sampling_params_dict: dict, kv_sender_info: dict[str, Any] | None = None, - ) -> OmniRequestOutput: - """Build a batched diffusion request and run DiffusionEngine.step(). - - All prompts are processed in a single step() call. The per-prompt - results are merged into one :class:`OmniRequestOutput` whose - ``images`` list contains every generated image, matching the - contract expected by the orchestrator and tests. - """ + ) -> AsyncGenerator[OmniRequestOutput, None]: + """Process a streaming diffusion request and yield the results from DiffusionEngine.step_streaming().""" sampling_params = self._reconstruct_sampling_params(sampling_params_dict) request = OmniDiffusionRequest( - prompts=prompts, + prompt=prompt, sampling_params=sampling_params, request_id=request_id, kv_sender_info=kv_sender_info, ) - results = await self._engine.step(request) - - # Merge per-prompt results into a single combined output. - all_images: list = [] - merged_mm: dict[str, Any] = {} - merged_metrics: dict[str, Any] = {} - merged_durations: dict[str, float] = {} - merged_custom: dict[str, Any] = {} - peak_mem = 0.0 - latents = None - trajectory_latents: list[torch.Tensor] | None = None - trajectory_timesteps: list[torch.Tensor] | None = None - trajectory_log_probs: torch.Tensor | None = None - trajectory_decoded: list[Image.Image] | None = None - final_output_type = "image" - - for r in results: - all_images.extend(r.images) - merged_mm.update(r._multimodal_output) - merged_metrics.update(r.metrics) - merged_durations.update(r.stage_durations) - merged_custom.update(r._custom_output) - peak_mem = max(peak_mem, r.peak_memory_mb) - if latents is None and r.latents is not None: - latents = r.latents - if trajectory_latents is None: - trajectory_latents = r.trajectory_latents - if trajectory_timesteps is None: - trajectory_timesteps = r.trajectory_timesteps - if trajectory_log_probs is None: - trajectory_log_probs = r.trajectory_log_probs - if trajectory_decoded is None: - trajectory_decoded = r.trajectory_decoded - if r.final_output_type != "image": - final_output_type = r.final_output_type - - return OmniRequestOutput.from_diffusion( - request_id=request_id, - images=all_images, - prompt=prompts[0] if len(prompts) == 1 else None, - metrics=merged_metrics, - latents=latents, - trajectory_latents=trajectory_latents, - trajectory_timesteps=trajectory_timesteps, - trajectory_log_probs=trajectory_log_probs, - trajectory_decoded=trajectory_decoded, - custom_output=merged_custom or None, - multimodal_output=merged_mm or None, - final_output_type=final_output_type, - stage_durations=merged_durations, - peak_memory_mb=peak_mem, - ) + async for results in self._engine.step_streaming(request): # pyright: ignore[reportOptionalMemberAccess] + result = results[0] + if not result.request_id: + result.request_id = request_id + yield result # ------------------------------------------------------------------ # Collective RPC dispatch @@ -360,10 +308,10 @@ async def run_loop( ctx = zmq.asyncio.Context() request_socket = ctx.socket(zmq.PULL) - request_socket.bind(request_address) + request_socket.connect(request_address) response_socket = ctx.socket(zmq.PUSH) - response_socket.bind(response_address) + response_socket.connect(response_address) encoder = OmniMsgpackEncoder() decoder = OmniMsgpackDecoder() @@ -386,13 +334,22 @@ async def _dispatch_request( ) -> None: """Process a single diffusion request and send the response.""" try: - result = await self._process_request( - request_id, - prompt, - sampling_params_dict, - kv_sender_info=kv_sender_info, - ) - await response_socket.send(encoder.encode({"type": "result", "output": result})) + if not self._od_config.streaming_output: + result = await self._process_request( + request_id, + prompt, + sampling_params_dict, + kv_sender_info=kv_sender_info, + ) + await response_socket.send(encoder.encode({"type": "result", "output": result})) + else: + async for result in self._process_streaming_request( + request_id, + prompt, + sampling_params_dict, + kv_sender_info=kv_sender_info, + ): + await response_socket.send(encoder.encode({"type": "result", "output": result})) except DiffusionRequestAbortedError as e: logger.info( "request_id: %s aborted: %s", @@ -462,61 +419,6 @@ async def _dispatch_request( ) tasks[request_id] = task - elif msg_type == "add_batch_request": - request_id = msg["request_id"] - - async def _dispatch_batch( - rid: str, - prompts: list, - sp_dict: dict, - kv_sender_info: dict[str, Any] | None = None, - ) -> None: - try: - result = await self._process_batch_request( - rid, - prompts, - sp_dict, - kv_sender_info=kv_sender_info, - ) - await response_socket.send(encoder.encode({"type": "result", "output": result})) - except DiffusionRequestAbortedError as e: - logger.info( - "request_id: %s aborted: %s", - rid, - str(e), - ) - 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, - } - ) - ) - # Same rationale as the single-request path: a - # closed executor turns every subsequent batch - # into a 500, so escalate now. - if self._is_executor_dead(): - self._signal_fatal_engine_failure(f"add_batch_request {rid}: {e!s}") - finally: - tasks.pop(rid, None) - - task = asyncio.create_task( - _dispatch_batch( - request_id, - msg["prompts"], - msg["sampling_params"], - msg.get("kv_sender_info"), - ) - ) - tasks[request_id] = task - elif msg_type == "abort": for rid in msg.get("request_ids", []): task = tasks.pop(rid, None) @@ -612,15 +514,51 @@ def close(self) -> None: # Subprocess entry point # ------------------------------------------------------------------ + @staticmethod + def _open_startup_handshake( + handshake_address: str, + *, + local_client: bool, + headless: bool, + ) -> tuple[zmq.Context, zmq.Socket, EngineZmqAddresses]: + ctx = zmq.Context() + socket = ctx.socket(zmq.DEALER) + socket.setsockopt(zmq.IDENTITY, (0).to_bytes(2, "little")) + socket.connect(handshake_address) + addresses = EngineCoreProc.startup_handshake( + socket, + local_client=local_client, + headless=headless, + parallel_config=None, + ) + return ctx, socket, addresses + + @staticmethod + def _send_startup_ready( + handshake_socket: zmq.Socket, + *, + local_client: bool, + headless: bool, + ) -> None: + handshake_socket.send( + msgspec.msgpack.encode( + { + "status": "READY", + "local": local_client, + "headless": headless, + } + ) + ) + @classmethod def run_diffusion_proc( cls, model: str, od_config: OmniDiffusionConfig, handshake_address: str, - request_address: str, - response_address: str, *, + local_client: bool, + headless: bool, omni_coordinator_address: str | None = None, omni_stage_id: int | None = None, omni_replica_id: int = 0, @@ -651,21 +589,32 @@ def signal_handler(signum: int, frame: Any) -> None: proc = cls(model, od_config) coord_client: OmniCoordClientForStage | None = None + handshake_ctx: zmq.Context | None = None + handshake_socket: zmq.Socket | None = None try: + handshake_ctx, handshake_socket, addresses = cls._open_startup_handshake( + handshake_address, + local_client=local_client, + headless=headless, + ) + request_address = addresses.inputs[0] + response_address = addresses.outputs[0] + proc.initialize() - # Send READY via handshake socket - handshake_ctx = zmq.Context() - handshake_socket = handshake_ctx.socket(zmq.DEALER) - handshake_socket.connect(handshake_address) - handshake_socket.send(msgspec.msgpack.encode({"status": "READY"})) + cls._send_startup_ready( + handshake_socket, + local_client=local_client, + headless=headless, + ) handshake_socket.close() handshake_ctx.term() + handshake_socket = None + handshake_ctx = None - # Wire OmniCoordClientForStage *after* READY so that the head - # has bound its head-side request/response sockets — the - # address pair we report is the same pair this proc binds to - # (request/response addresses passed in). + # Wire OmniCoordClientForStage *after* READY. The address pair is + # owned by the frontend client; this proc connects to it as the + # backend runtime. if omni_coordinator_address is not None: if omni_stage_id is None: raise ValueError("omni_stage_id must be provided when omni_coordinator_address is set") @@ -688,7 +637,6 @@ def _refresh_queue_length() -> None: omni_coordinator_address, ) - # Run async event loop asyncio.run(proc.run_loop(request_address, response_address)) except SystemExit: @@ -698,105 +646,138 @@ def _refresh_queue_length() -> None: logger.exception("StageDiffusionProc encountered a fatal error.") raise finally: + if handshake_socket is not None: + handshake_socket.close(linger=0) + if handshake_ctx is not None: + handshake_ctx.term() if coord_client is not None: with contextlib.suppress(RuntimeError): coord_client.close() proc.close() -# -- Free functions for backward compatibility with StageDiffusionClient ------ +class StageDiffusionProcManager: + """Owns a StageDiffusionProc subprocess. + Mirrors the small process-lifecycle surface used by vLLM's + CoreEngineProcManager while keeping diffusion's custom wire protocol. + """ + + def __init__( + self, + *, + model: str, + od_config: OmniDiffusionConfig, + stage_init_timeout: int, + handshake_address: str | None = None, + addresses: EngineZmqAddresses | None = None, + omni_coordinator_address: str | None = None, + omni_stage_id: int | None = None, + omni_replica_id: int = 0, + ) -> None: + handshake_address = handshake_address or get_open_zmq_ipc_path() + addresses = addresses or EngineZmqAddresses( + inputs=[get_open_zmq_ipc_path()], + outputs=[get_open_zmq_ipc_path()], + ) -def spawn_diffusion_proc( - model: str, - od_config: OmniDiffusionConfig, - handshake_address: str | None = None, - request_address: str | None = None, - response_address: str | None = None, - *, - omni_coordinator_address: str | None = None, - omni_stage_id: int | None = None, - omni_replica_id: int = 0, -) -> tuple[BaseProcess, str, str, str]: - """Spawn a StageDiffusionProc subprocess. + ctx = get_mp_context() + proc = ctx.Process( + target=StageDiffusionProc.run_diffusion_proc, + name="StageDiffusionProc", + kwargs={ + "model": model, + "od_config": od_config, + "handshake_address": handshake_address, + "local_client": True, + "headless": False, + "omni_coordinator_address": omni_coordinator_address, + "omni_stage_id": omni_stage_id, + "omni_replica_id": omni_replica_id, + }, + ) + proc.start() + self.proc = proc + self.addresses = addresses + self.manager_stopped = False + self.failed_proc_name: str | None = None - Returns ``(proc, handshake_address, request_address, response_address)``. + self._wait_until_started(handshake_address, stage_init_timeout) - Pass ``omni_coordinator_address`` / ``omni_stage_id`` / ``omni_replica_id`` - to have the subprocess publish heartbeats to an OmniCoordinator. - """ - handshake_address = handshake_address or get_open_zmq_ipc_path() - request_address = request_address or get_open_zmq_ipc_path() - response_address = response_address or get_open_zmq_ipc_path() - - ctx = get_mp_context() - proc = ctx.Process( - target=StageDiffusionProc.run_diffusion_proc, - name="StageDiffusionProc", - kwargs={ - "model": model, - "od_config": od_config, - "handshake_address": handshake_address, - "request_address": request_address, - "response_address": response_address, - "omni_coordinator_address": omni_coordinator_address, - "omni_stage_id": omni_stage_id, - "omni_replica_id": omni_replica_id, - }, - ) - proc.start() - # Wait for the process to become alive before returning. - deadline = time.monotonic() + 10 - while not proc.is_alive(): - if proc.exitcode is not None: - raise RuntimeError(f"StageDiffusionProc failed to start (exit code {proc.exitcode})") - if time.monotonic() > deadline: - raise TimeoutError("StageDiffusionProc did not become alive within 10s") - time.sleep(0.01) - return proc, handshake_address, request_address, response_address - - -def complete_diffusion_handshake( - proc: BaseProcess, - handshake_address: str, - handshake_timeout: int, -) -> None: - """Wait for the diffusion subprocess to signal READY. - - On failure the process is terminated before re-raising. - """ - try: - _perform_diffusion_handshake(proc, handshake_address, handshake_timeout) - except Exception: - shutdown([proc]) - raise - - -def _perform_diffusion_handshake( - proc: BaseProcess, - handshake_address: str, - handshake_timeout: int, -) -> None: - """Run the handshake with the diffusion subprocess.""" - with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: - poller = zmq.Poller() - poller.register(handshake_socket, zmq.POLLIN) - poller.register(proc.sentinel, zmq.POLLIN) - - timeout_ms = handshake_timeout * 1000 - while True: - events = dict(poller.poll(timeout=timeout_ms)) - if not events: - raise TimeoutError( - f"Timed out waiting for READY from StageDiffusionProc after {handshake_timeout}s. " - f"This typically indicates model loading or warmup is taking too long. " - f"Consider increasing `stage_init_timeout` for large models." + @classmethod + def launch_headless( + cls, + *, + model: str, + od_config: OmniDiffusionConfig, + handshake_address: str, + addresses: EngineZmqAddresses, + omni_coordinator_address: str | None, + omni_stage_id: int, + omni_replica_id: int, + ) -> StageDiffusionProcManager: + """Launch a headless diffusion backend that connects to head-owned sockets.""" + self = cls.__new__(cls) + ctx = get_mp_context() + proc = ctx.Process( + target=StageDiffusionProc.run_diffusion_proc, + name="StageDiffusionProc", + kwargs={ + "model": model, + "od_config": od_config, + "handshake_address": handshake_address, + "local_client": False, + "headless": True, + "omni_coordinator_address": omni_coordinator_address, + "omni_stage_id": omni_stage_id, + "omni_replica_id": omni_replica_id, + }, + ) + proc.start() + self.proc = proc + self.addresses = addresses + self.manager_stopped = False + self.failed_proc_name = None + return self + + def _wait_until_started(self, handshake_address: str, stage_init_timeout: int) -> None: + try: + with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: + wait_for_engine_startup( + handshake_socket, + self.addresses, + [CoreEngine(index=0, local=True)], + SimpleNamespace( + data_parallel_size_local=1, + data_parallel_hybrid_lb=False, + data_parallel_external_lb=False, + ), + False, + None, + self, + None, ) - if handshake_socket in events: - identity, raw = handshake_socket.recv_multipart() - msg = msgspec.msgpack.decode(raw) - if msg.get("status") == "READY": - return - raise RuntimeError(f"Expected READY, got: {msg}") - if proc.exitcode is not None: - raise RuntimeError(f"StageDiffusionProc died during handshake (exit code {proc.exitcode})") + except Exception: + shutdown([self.proc]) + raise + + def shutdown(self, timeout: float | None = None) -> None: + self.manager_stopped = True + shutdown([self.proc], timeout=timeout) + + def sentinels(self) -> list[int]: + return [self.proc.sentinel] + + def finished_procs(self) -> dict[str, int]: + if self.proc.exitcode is None: + return {} + return {self.proc.name: self.proc.exitcode} + + def monitor_engine_liveness(self) -> None: + try: + multiprocessing.connection.wait([self.proc.sentinel]) + except Exception: + return + if self.proc.exitcode not in (None, 0) and not self.manager_stopped: + self.failed_proc_name = self.proc.name + self.shutdown() diff --git a/vllm_omni/diffusion/utils/flow_matching.py b/vllm_omni/diffusion/utils/flow_matching.py new file mode 100644 index 00000000000..e81d405ec23 --- /dev/null +++ b/vllm_omni/diffusion/utils/flow_matching.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import torch +from torch import nn + + +def swish(x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(x) + + +class SinusoidalPositionalEncoding(nn.Module): + def __init__(self, embedding_dim: int) -> None: + super().__init__() + self.embedding_dim = embedding_dim + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + timesteps = timesteps.float() + B, T = timesteps.shape + device = timesteps.device + + half_dim = self.embedding_dim // 2 + exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * (math.log(10000.0) / half_dim) + freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) + enc = torch.cat([torch.sin(freqs), torch.cos(freqs)], dim=-1) # (B, T, embedding_dim) + return enc diff --git a/vllm_omni/diffusion/utils/kv_utils.py b/vllm_omni/diffusion/utils/kv_utils.py new file mode 100644 index 00000000000..9ba231386fa --- /dev/null +++ b/vllm_omni/diffusion/utils/kv_utils.py @@ -0,0 +1,39 @@ +"""Utilities for batching variable-length tensors in diffusion attention.""" + +import torch + + +def left_pad_stack( + tensors: list[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Left-pad and stack variable-length tensors. Only dim 0 may vary, + and it's assumed that all tensors are the same dtype & use the same + device. + + Returns (stacked, mask) where mask is a 2D boolean mask, and both + tensors are on the device of the provided tensors. If all tensors are + the same length, None is returned for mask. + """ + trailing_dims = {ts.shape[1:] for ts in tensors} + if len(trailing_dims) != 1: + raise ValueError("Tensors must be non-empty and can only vary in dim 0") + trailing = trailing_dims.pop() + + seq_lens = [ts.shape[0] for ts in tensors] + max_len = max(seq_lens) + + # If everything is the same length, we don't need a mask / varlen + if all(sl == max_len for sl in seq_lens): + return torch.stack(tensors), None + + device = tensors[0].device + dtype = tensors[0].dtype + stacked = torch.zeros(len(tensors), max_len, *trailing, dtype=dtype, device=device) + # Create the boolean mask for the input sequences + mask = torch.zeros(len(tensors), max_len, dtype=torch.bool, device=device) + for idx, (ts, sl) in enumerate(zip(tensors, seq_lens)): + pad = max_len - sl + stacked[idx, pad:] = ts + mask[idx, pad:] = True + + return stacked, mask diff --git a/vllm_omni/diffusion/utils/media_utils.py b/vllm_omni/diffusion/utils/media_utils.py index a09cd459539..314a9d2bfba 100644 --- a/vllm_omni/diffusion/utils/media_utils.py +++ b/vllm_omni/diffusion/utils/media_utils.py @@ -6,11 +6,124 @@ import io from fractions import Fraction +from typing import Any, cast import av import numpy as np +class FragmentedMP4Muxer: + """Incrementally mux video frames into one fragmented MP4 byte stream.""" + + def __init__( + self, + *, + width: int, + height: int, + fps: float = 25.0, + video_codec: str = "h264", + crf: str = "18", + video_codec_options: dict[str, str] | None = None, + ) -> None: + self._buf = io.BytesIO() + self._closed = False + self._container = av.open( + self._buf, + mode="w", + format="mp4", + options={"movflags": "+frag_every_frame+empty_moov+default_base_moof"}, + ) + + self._stream: av.VideoStream = cast( + av.VideoStream, + self._container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)), + ) + self._stream.width = width + self._stream.height = height + self._stream.pix_fmt = "yuv420p" + + options: dict[str, object] = {"crf": str(crf)} + if video_codec_options: + options.update(video_codec_options) + self._stream.options = options + + try: + self._stream.codec_context.max_b_frames = 0 + except AttributeError: + pass + + def mux_video_frames(self, video_frames: np.ndarray) -> bytes: + """Mux a batch of ``uint8`` RGB frames and return newly written MP4 bytes.""" + if self._closed: + raise RuntimeError("Cannot mux frames after FragmentedMP4Muxer.close().") + if video_frames.ndim != 4 or video_frames.shape[-1] != 3: + raise ValueError("video_frames must have shape (T, H, W, 3).") + if video_frames.dtype != np.uint8: + raise ValueError("video_frames must be uint8.") + if video_frames.shape[1] != self._stream.height or video_frames.shape[2] != self._stream.width: + raise ValueError("All fragmented MP4 chunks in a session must use the same frame size.") + + for frame_data in video_frames: + frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") + for packet in self._stream.encode(frame): + self._container.mux(packet) + return self._read_new_bytes() + + def close(self) -> bytes: + """Flush delayed encoder packets, close the container, and return final bytes.""" + if self._closed: + return b"" + for packet in self._stream.encode(): + self._container.mux(packet) + self._container.close() + self._closed = True + return self._read_new_bytes() + + def _read_new_bytes(self) -> bytes: + """Return newly muxed bytes in the current video container, + then clear the buffer to prepare for the next chunk.""" + chunk = self._buf.getvalue() + self._buf.seek(0) + self._buf.truncate() + return chunk + + +def finalize_streaming_video_bytes( + video_bytes: bytes, + *, + input_format: str, + fps: float = 25.0, + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Convert streamed video bytes into a progressive MP4 for local playback.""" + if not video_bytes: + return video_bytes + + normalized_format = input_format.lower() + if normalized_format == "m4s": + demux_format = "mp4" + else: + raise ValueError(f"Unsupported streaming video format: {input_format}") + + try: + with cast(Any, av.open(io.BytesIO(video_bytes), format=demux_format)) as container: + stream = container.streams.video[0] + frame_arrays = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)] + except Exception: + return video_bytes + + if not frame_arrays: + return video_bytes + + frames_u8 = np.ascontiguousarray(np.stack(frame_arrays, axis=0), dtype=np.uint8) + return mux_video_audio_bytes( + frames_u8, + None, + fps=float(fps), + video_codec_options=video_codec_options, + ) + + def mux_video_audio_bytes( video_frames: np.ndarray, audio_waveform: np.ndarray | None = None, @@ -39,17 +152,19 @@ def mux_video_audio_bytes( buf = io.BytesIO() container = av.open(buf, mode="w", format="mp4") - v_stream = container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000)) + v_stream = cast(av.VideoStream, container.add_stream(video_codec, rate=Fraction(fps).limit_denominator(10000))) v_stream.width = video_frames.shape[2] v_stream.height = video_frames.shape[1] v_stream.pix_fmt = "yuv420p" - options = {"crf": str(crf)} + options: dict[str, object] = {"crf": str(crf)} if video_codec_options: options.update(video_codec_options) v_stream.options = options - a_stream = None + a_stream: av.AudioStream | None = None + samples: np.ndarray | None = None + layout: str | None = None if audio_waveform is not None: samples = audio_waveform.astype(np.float32) if samples.ndim == 1: @@ -58,7 +173,7 @@ def mux_video_audio_bytes( samples = np.ascontiguousarray(samples.T) num_channels = samples.shape[0] layout = "stereo" if num_channels >= 2 else "mono" - a_stream = container.add_stream(audio_codec, rate=audio_sample_rate) + a_stream = cast(av.AudioStream, container.add_stream(audio_codec, rate=audio_sample_rate)) a_stream.layout = layout for frame_data in video_frames: @@ -69,6 +184,8 @@ def mux_video_audio_bytes( container.mux(packet) if a_stream is not None and audio_waveform is not None: + if samples is None or layout is None: + raise ValueError("Audio samples were not prepared for muxing.") audio_frame = av.AudioFrame.from_ndarray(samples, format="fltp", layout=layout) audio_frame.sample_rate = audio_sample_rate for packet in a_stream.encode(audio_frame): diff --git a/vllm_omni/diffusion/utils/param_utils.py b/vllm_omni/diffusion/utils/param_utils.py new file mode 100644 index 00000000000..e56fcd072a5 --- /dev/null +++ b/vllm_omni/diffusion/utils/param_utils.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + +def apply_declared_extra_args( + sampling_params: OmniDiffusionSamplingParams, + declared_params: frozenset[str], + user_kwargs: dict[str, object], +) -> None: + """Route pipeline-declared request params into ``sampling_params.extra_args``. + + Both online serving and offline examples call this so that model-specific + keys (e.g. ``cfg_text_scale`` for BAGEL) end up in ``extra_args`` instead + of being silently dropped. + + This is a no-op when no declared params are present in ``user_kwargs``, so + it is safe to call on non-diffusion (e.g. AR) sampling params whose + ``extra_args`` defaults to ``None``. + """ + declared = {key: user_kwargs[key] for key in declared_params if user_kwargs.get(key) is not None} + if not declared: + return + sampling_params.extra_args = {**(sampling_params.extra_args or {}), **declared} diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 278f1df40f5..cb16bdeac6b 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -38,7 +38,16 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents -from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput +from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.diffusion.worker.utils import ( + BatchRunnerOutput, + DiffusionRequestState, + RunnerOutput, + attach_stage_durations, + clear_pipeline_stage_durations, + consume_pipeline_stage_durations, + merge_stage_durations, +) from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin @@ -46,6 +55,43 @@ logger = init_logger(__name__) +def _normalize_pipeline_outputs( + outputs: object, + *, + expected_count: int, + allow_single_output: bool, + pipeline_name: str, +) -> list[DiffusionOutput]: + if isinstance(outputs, DiffusionOutput): + if allow_single_output and expected_count == 1: + return [outputs] + raise RuntimeError( + f"{pipeline_name}.forward returned a single DiffusionOutput; " + "request-batch forward must return list[DiffusionOutput]." + ) + + if not isinstance(outputs, list): + raise RuntimeError( + f"{pipeline_name}.forward returned {type(outputs).__name__}; " + "expected DiffusionOutput or list[DiffusionOutput]." + ) + + if len(outputs) != expected_count: + raise RuntimeError( + f"{pipeline_name}.forward returned {len(outputs)} outputs for {expected_count} requests; " + "expected exactly one DiffusionOutput per request." + ) + + bad_index = next((idx for idx, output in enumerate(outputs) if not isinstance(output, DiffusionOutput)), None) + if bad_index is not None: + raise RuntimeError( + f"{pipeline_name}.forward returned list item {bad_index} with type " + f"{type(outputs[bad_index]).__name__}; expected DiffusionOutput." + ) + + return outputs + + class DiffusionModelRunner(OmniConnectorModelRunnerMixin): """ Model runner that handles model loading and execution for diffusion models. @@ -80,14 +126,29 @@ def __init__( # Cache for per-request stepwise state. self.state_cache: dict[str, DiffusionRequestState] = {} - # Initialize KV cache manager for connector management + # Initialize KV cache manager for connector management. self.kv_transfer_manager = OmniKVTransferManager.from_od_config(od_config) + # Prefetch covers TP / SP / CFG-Parallel / HSDP. Disabled when a CFG + # companion KV collector is set (that KV is not backgrounded). + has_cfg_companion_kv = getattr(od_config, "cfg_kv_collect_func", None) is not None + + self._kv_prefetch_enabled = ( + bool(self.kv_transfer_manager.config.enable_kv_async_prefetch) + and not has_cfg_companion_kv + and self.kv_transfer_manager.config.need_recv_cache + ) + + @property + def target_device(self) -> torch.device | None: + return getattr(self.pipeline, "device", None) + def _compile_transformer(self, attr_name: str) -> None: """Compile a transformer attribute on the pipeline with torch.compile.""" model = getattr(self.pipeline, attr_name, None) if model is None: return + try: setattr(self.pipeline, attr_name, regionally_compile(model, dynamic=True)) logger.info("Model runner: %s compiled with torch.compile.", attr_name) @@ -152,12 +213,21 @@ def get_memory_context(): ) logger.info("Model runner: Model loaded successfully.") + if self.od_config.streaming_output and not getattr(self.od_config, "step_execution", False): + logger.warning("streaming_output=True requires step_execution=True; enabling step execution.") + self.od_config.step_execution = True + if getattr(self.od_config, "step_execution", False) and not self.supports_step_mode(): raise ValueError( "step_execution=True requires a pipeline implementing " "prepare_encode(), denoise_step(), step_scheduler(), and post_decode(); " f"{self.od_config.model_class_name} does not support that contract." ) + if self.od_config.streaming_output and not self.supports_step_mode(): + raise ValueError( + "streaming_output=True requires step execution support; " + f"{self.od_config.model_class_name} does not support that contract." + ) # Apply CPU offloading self.offload_backend = get_offload_backend(self.od_config, device=self.device) @@ -168,8 +238,17 @@ def get_memory_context(): # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: if current_omni_platform.supports_torch_inductor(): - self._compile_transformer("transformer") - self._compile_transformer("transformer_2") + if hasattr(self.pipeline, "setup_compile"): + try: + self.pipeline.setup_compile() + except Exception as exc: + logger.warning( + "Model runner: setup_compile() failed (%s); running without compile.", + exc, + ) + else: + self._compile_transformer("transformer") + self._compile_transformer("transformer_2") else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", @@ -223,12 +302,13 @@ def get_prompt_embed_cache_stats(self) -> dict | None: return None return self.prompt_embed_cache.stats() - def _record_peak_memory(self, output: DiffusionOutput) -> None: - """Record peak GPU memory for the current forward pass into output. + def _sample_peak_memory_mb(self) -> float: + """Return peak GPU memory for the current forward pass in MB. - Must be called immediately after pipeline.forward(), with + Must be called immediately after the measured forward/step work, with reset_peak_memory_stats() called just before it, so the measurement - reflects this request only and not the global historical maximum. + reflects the current execution slice and not the global historical + maximum. Uses max_memory_reserved (CUDA memory pool high-water mark) rather than max_memory_allocated so that allocator fragmentation is also visible. @@ -237,7 +317,7 @@ def _record_peak_memory(self, output: DiffusionOutput) -> None: peak_reserved_bytes = current_omni_platform.max_memory_reserved() peak_allocated_bytes = current_omni_platform.max_memory_allocated() - output.peak_memory_mb = peak_reserved_bytes / (1024**2) + peak_memory_mb = peak_reserved_bytes / (1024**2) peak_reserved_gb = peak_reserved_bytes / (1024**3) peak_allocated_gb = peak_allocated_bytes / (1024**3) pool_overhead_gb = peak_reserved_gb - peak_allocated_gb @@ -249,100 +329,232 @@ def _record_peak_memory(self, output: DiffusionOutput) -> None: pool_overhead_gb, pool_overhead_gb / peak_reserved_gb * 100 if peak_reserved_gb > 0 else 0.0, ) + return peak_memory_mb - def execute_model(self, req: OmniDiffusionRequest) -> DiffusionOutput: - """ - Execute a forward pass for the given requests. + def _prepare_request_for_forward( + self, + req: OmniDiffusionRequest, + *, + od_config: OmniDiffusionConfig, + kv_prefetch_jobs: dict | None = None, + use_prefetch: bool = False, + ) -> None: + # Receive AR KV. Single-request execution can use the prefetch path: + # consume prior-forward payload, sync-fallback on miss; request-batch + # execution keeps the synchronous per-request receive path. + kv_recv_t0 = time.perf_counter() + if use_prefetch and self._kv_prefetch_enabled: + self.kv_transfer_manager.consume_and_distribute_kv_cache( + req, + target_device=self.target_device, + ) + else: + self.kv_transfer_manager.receive_multi_kv_cache_distributed( + req, + cfg_kv_collect_func=getattr(od_config, "cfg_kv_collect_func", None), + target_device=self.target_device if use_prefetch else getattr(self.pipeline, "device", None), + ) + kv_recv_ms = (time.perf_counter() - kv_recv_t0) * 1000 + logger.debug("KV recv for %s %.1fms", req.request_id, kv_recv_ms) + + # Kick off the next request's prefetch (+ H2D) to overlap this forward. + if use_prefetch and self._kv_prefetch_enabled and kv_prefetch_jobs is not None: + self.kv_transfer_manager.start_prefetch(kv_prefetch_jobs, self.target_device) + + if req.sampling_params.generator is None and req.sampling_params.seed is not None: + if req.sampling_params.generator_device is not None: + gen_device = req.sampling_params.generator_device + elif self.device.type == "cpu": + gen_device = "cpu" + else: + gen_device = self.device + req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed) - Args: - req: A diffusion request containing a list of prompts to process. + def _refresh_cache_for_requests( + self, + reqs: list[OmniDiffusionRequest], + *, + od_config: OmniDiffusionConfig, + ) -> None: + first_req = reqs[0] + if ( + getattr(first_req, "skip_cache_refresh", False) + or self.cache_backend is None + or not self.cache_backend.is_enabled() + ): + return - Returns: - DiffusionOutput with generated results. + # Refresh cache context if needed. Batch admission groups requests by + # SamplingParamsKey, so the first request's num_inference_steps applies + # to the whole runner batch. + num_inference_steps = first_req.sampling_params.num_inference_steps + if num_inference_steps is None and od_config.cache_backend in ( + "tea_cache", + "step_cache", + ): + # When num_inference_steps is None, some pipelines defer to their + # own defaults. TeaCache refresh ignores this value; step_cache + # refresh is a no-op because per-chunk state resets in the denoise + # loop. Use the pipeline default when available to keep refresh + # behavior aligned with single-request execution. + num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0 + + if num_inference_steps is not None: + self.cache_backend.refresh(self.pipeline, num_inference_steps) + else: + logger.warning( + "Failed to refresh the diffusion transformer cache; backend %s " + "currently requires num_inference_steps to be passed explicitly", + od_config.cache_backend, + ) - Note: - We use torch.no_grad() for HSDP because HSDP2's fully_shard requires access - to tensor version counters in pre_forward hooks, which inference tensors do - not track. For non-HSDP inference, we use torch.inference_mode() for better - performance. - """ - assert self.pipeline is not None, "Model not loaded. Call load_model() first." - if len(req.prompts) == 0: - raise ValueError("Cannot execute model with empty request list") + def _runner_output_from_outputs( + self, + reqs: list[OmniDiffusionRequest], + outputs: list[DiffusionOutput], + ) -> BatchRunnerOutput: + return BatchRunnerOutput.from_list( + [ + RunnerOutput( + request_id=reqs[i].request_id, + step_index=None, + finished=True, + result=outputs[i], + ) + for i in range(len(reqs)) + ] + ) - # Use no_grad() for HSDP compatibility, inference_mode() otherwise for better perf - use_hsdp = self.od_config.parallel_config.use_hsdp + def _execute_request_list( + self, + reqs: list[OmniDiffusionRequest], + *, + od_config: OmniDiffusionConfig, + allow_single_output: bool, + require_request_batch_support: bool, + kv_prefetch_jobs: dict | None = None, + record_name: str, + ) -> BatchRunnerOutput: + assert self.pipeline is not None, "Model not loaded. Call load_model() first." + if not reqs: + return BatchRunnerOutput.from_list([]) + for req in reqs: + if req.prompt is None: + raise ValueError("Cannot execute model with empty prompt") + if require_request_batch_support and not getattr(self.pipeline, "supports_request_batch", False): + raise RuntimeError(f"{type(self.pipeline).__name__} does not support request-batch forward.") + + # Use no_grad() for HSDP compatibility, inference_mode() otherwise for + # better perf. HSDP2's fully_shard pre-forward hooks need tensor version + # counters, which inference tensors do not track. + use_hsdp = od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: - # The manager handles the check for need_recv_cache internally - self.kv_transfer_manager.receive_multi_kv_cache_distributed( - req, - cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), - target_device=getattr(self.pipeline, "device", None), - ) - - if req.sampling_params.generator is None and req.sampling_params.seed is not None: - if req.sampling_params.generator_device is not None: - gen_device = req.sampling_params.generator_device - elif self.device.type == "cpu": - gen_device = "cpu" - else: - gen_device = self.device - req.sampling_params.generator = torch.Generator(device=gen_device).manual_seed(req.sampling_params.seed) + for req in reqs: + self._prepare_request_for_forward( + req, + od_config=od_config, + kv_prefetch_jobs=kv_prefetch_jobs, + use_prefetch=allow_single_output, + ) - # Refresh cache context if needed - if ( - not getattr(req, "skip_cache_refresh", False) - and self.cache_backend is not None - and self.cache_backend.is_enabled() - ): - # FIXME (Alex): When num_inference_steps is None, we defer to - # pipelines for default, but don't refresh the cache; the right - # way to do this is to merge the sampling params first. - # - # For now, if num_inference_steps is not set, we pass 0 to allow - # TeaCache to refresh to align with the param signature. This is - # okay to force refresh TeaCache because the refresh does not use - # num_inference_steps at all (i.e., just resets state and clears - # stale residuals). - num_inference_steps = req.sampling_params.num_inference_steps - if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None: - num_inference_steps = 0 - - if num_inference_steps is not None: - self.cache_backend.refresh(self.pipeline, num_inference_steps) - else: - logger.warning( - "Failed to refresh the diffusion transformer cache; backend %s " - "currently requires num_inference_steps to be passed explicitly", - self.od_config.cache_backend, - ) + self._refresh_cache_for_requests(reqs, od_config=od_config) + batch = DiffusionRequestBatch(requests=reqs) is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 if is_primary: current_omni_platform.reset_peak_memory_stats() - with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=self.od_config): - with record_function("pipeline_forward"): - output = self.pipeline.forward(req) + with set_forward_context(vllm_config=self.vllm_config, omni_diffusion_config=od_config): + with record_function(record_name): + raw_outputs = self.pipeline.forward(batch) + outputs = _normalize_pipeline_outputs( + raw_outputs, + expected_count=len(reqs), + allow_single_output=allow_single_output, + pipeline_name=type(self.pipeline).__name__, + ) - if is_primary: - self._record_peak_memory(output) + if is_primary and outputs: + batch_peak_memory_mb = self._sample_peak_memory_mb() + for output in outputs: + output.peak_memory_mb = max(output.peak_memory_mb, batch_peak_memory_mb) - # Log prompt-embed cache activity (hits/misses accumulate across requests). - if is_primary and self.prompt_embed_cache is not None: - logger.debug("prompt-embed cache: %s", self.prompt_embed_cache.stats()) + # Log prompt-embed cache activity; hits/misses accumulate across requests. + prompt_embed_cache = getattr(self, "prompt_embed_cache", None) + if is_primary and prompt_embed_cache is not None: + logger.debug("prompt-embed cache: %s", prompt_embed_cache.stats()) - # NOTE: if ( self.cache_backend is not None and self.cache_backend.is_enabled() - and self.od_config.cache_backend == "cache_dit" - and self.od_config.enable_cache_dit_summary + and od_config.cache_backend == "cache_dit" + and od_config.enable_cache_dit_summary ): cache_summary(self.pipeline, details=True) - return output + return self._runner_output_from_outputs(reqs, outputs) + + def _attach_stepwise_metrics( + self, + state: DiffusionRequestState, + output: DiffusionOutput, + *, + is_primary: bool, + ) -> None: + merge_stage_durations( + state, + consume_pipeline_stage_durations(self.pipeline), + ) + attach_stage_durations(state, output) + + def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput: + """ + Execute a forward pass for the given requests. + + Args: + req: A diffusion request containing a list of prompts to process. + + Returns: + DiffusionOutput with generated results. + + Note: + We use torch.no_grad() for HSDP because HSDP2's fully_shard requires access + to tensor version counters in pre_forward hooks, which inference tensors do + not track. For non-HSDP inference, we use torch.inference_mode() for better + performance. + """ + runner_output = self._execute_request_list( + [req], + od_config=self.od_config, + allow_single_output=True, + require_request_batch_support=False, + kv_prefetch_jobs=kv_prefetch_jobs, + record_name="pipeline_forward", + ) + output = runner_output.runner_outputs[0].result + assert output is not None + return output + + def execute_model_batch( + self, + scheduler_output: DiffusionSchedulerOutput, + od_config: OmniDiffusionConfig, + ) -> BatchRunnerOutput: + """Execute scheduled request-mode requests through the batch forward path. + + Builds a ``DiffusionRequestBatch`` from scheduled new requests, runs + per-request setup, and calls ``pipeline.forward(batch)``. The pipeline + must declare ``supports_request_batch = True``. + """ + reqs = [nr.req for nr in scheduler_output.scheduled_new_reqs] + return self._execute_request_list( + reqs, + od_config=od_config, + allow_single_output=False, + require_request_batch_support=True, + record_name="pipeline_forward_batch", + ) # ------------------------------------------------------------------ # Step-wise execution @@ -365,21 +577,21 @@ def _update_states( # process new requests for sched_new_req in scheduler_output.scheduled_new_reqs: request_id = sched_new_req.request_id - req = sched_new_req.req new_request_ids.append(request_id) if request_id in self.state_cache: raise ValueError(f"Received duplicate new-request payload for cached request {request_id}.") new_state = DiffusionRequestState( request_id=request_id, - sampling=copy.deepcopy(req.sampling_params), - prompts=req.prompts, + sampling=copy.deepcopy(sched_new_req.req.sampling_params), + prompt=sched_new_req.req.prompt, + kv_sender_info=sched_new_req.req.kv_sender_info, ) - state_req = copy.copy(req) + state_req = copy.copy(sched_new_req.req) state_req.sampling_params = new_state.sampling self.kv_transfer_manager.receive_multi_kv_cache_distributed( state_req, cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), - target_device=getattr(self.pipeline, "device", None), + target_device=self.target_device, ) self.state_cache[request_id] = new_state resolved.append(new_state) @@ -410,8 +622,13 @@ def _prepare_batch_inputs(self, states: list[DiffusionRequestState], new_request else: gen_device = self.device state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed) + clear_pipeline_stage_durations(self.pipeline) # encode self.pipeline.prepare_encode(state) + merge_stage_durations( + state, + consume_pipeline_stage_durations(self.pipeline), + ) input_batch = InputBatch.make_batch( states, @@ -441,7 +658,7 @@ def _update_states_after( scatter_latents(states, input_batch) for state in states: - if interrupted or state.denoise_completed: + if interrupted or state.request_denoise_completed: self.state_cache.pop(state.request_id, None) def _prepare_attn_metadata(self, input_batch: InputBatch) -> Any: @@ -467,7 +684,11 @@ def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchR use_hsdp = self.od_config.parallel_config.use_hsdp grad_context = torch.no_grad() if use_hsdp else torch.inference_mode() with grad_context: + had_active_states = bool(self.state_cache) states, new_request_ids = self._update_states(scheduler_output) + is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available(): + current_omni_platform.reset_peak_memory_stats() input_batch = self._prepare_batch_inputs(states, new_request_ids) attn_metadata = self._prepare_attn_metadata(input_batch) @@ -476,7 +697,14 @@ def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchR omni_diffusion_config=self.od_config, attn_metadata=attn_metadata, ): - noise_pred = self.pipeline.denoise_step(input_batch) + clear_pipeline_stage_durations(self.pipeline) + noise_pred = self.pipeline.denoise_step(input_batch, states=states) + denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline) + for state in states: + merge_stage_durations( + state, + denoise_stage_durations, + ) runner_output_list = [] pipeline_interrupted = getattr(self.pipeline, "interrupt", False) @@ -495,22 +723,57 @@ def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchR offset = 0 for req in states: row_num = req.latents.shape[0] - self.pipeline.step_scheduler( - req, noise_pred[offset : offset + row_num] if noise_pred is not None else None - ) - offset = offset + row_num - if req.denoise_completed: - result = self.pipeline.post_decode(req) - else: - result = None - runner_output_list.append( - RunnerOutput( - request_id=req.request_id, - step_index=req.step_index, - finished=req.denoise_completed, - result=result, + try: + self.pipeline.step_scheduler( + req, noise_pred[offset : offset + row_num] if noise_pred is not None else None + ) + if self.od_config.streaming_output: + should_decode = req.chunk_denoise_completed + else: + should_decode = req.denoise_completed + + if should_decode: + clear_pipeline_stage_durations(self.pipeline) + result = self.pipeline.post_decode(req) + if result is not None: + self._attach_stepwise_metrics( + req, + result, + is_primary=is_primary, + ) + else: + result = None + # finished should be computed after post_decode() advanced chunk_index + finished = ( + req.request_denoise_completed + if self.od_config.streaming_output + else req.denoise_completed + ) + runner_output_list.append( + RunnerOutput( + request_id=req.request_id, + step_index=req.step_index, + finished=finished, + result=result, + ) + ) + offset = offset + row_num + except Exception as per_req_exc: + offset = offset + row_num + logger.error( + "Stepwise per-request error for %s: %s", + req.request_id, + per_req_exc, + exc_info=True, + ) + runner_output_list.append( + RunnerOutput( + request_id=req.request_id, + step_index=req.step_index, + finished=True, + result=DiffusionOutput(error=str(per_req_exc)), + ) ) - ) if noise_pred is not None and offset != noise_pred.shape[0]: raise ValueError( @@ -518,6 +781,22 @@ def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchR f"but batched noise_pred has {noise_pred.shape[0]} rows." ) + if is_primary: + batch_peak_memory_mb = self._sample_peak_memory_mb() + states_by_id = {state.request_id: state for state in states} + for state in states: + state.peak_memory_mb = max(state.peak_memory_mb, batch_peak_memory_mb) + for runner_output in runner_output_list: + if runner_output.result is None: + continue + state = states_by_id.get(runner_output.request_id) + if state is None: + continue + runner_output.result.peak_memory_mb = max( + runner_output.result.peak_memory_mb, + state.peak_memory_mb, + ) + self._update_states_after(states, input_batch, pipeline_interrupted) return BatchRunnerOutput.from_list(runner_output_list) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 0b4e877180a..006bbae485d 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -48,7 +48,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner -from vllm_omni.diffusion.worker.utils import BaseRunnerOutput +from vllm_omni.diffusion.worker.utils import BaseRunnerOutput, BatchRunnerOutput from vllm_omni.engine.stage_init_utils import set_death_signal from vllm_omni.lora.request import LoRARequest from vllm_omni.platforms import current_omni_platform @@ -81,6 +81,10 @@ def is_model_moe(self) -> bool: def is_nvfp4_quantized(self) -> bool: return self.quantization == "modelopt_fp4" + @property + def is_diffusion(self) -> bool: + return False + def _make_diffusion_vllm_model_config(od_config: OmniDiffusionConfig) -> _DiffusionVllmModelConfig: quant_config = getattr(od_config, "quantization_config", None) @@ -161,6 +165,12 @@ def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDi return vllm_config +def _get_cumem_allocator_class() -> type: + from vllm.device_allocator.cumem import CuMemAllocator + + return CuMemAllocator + + def _resolve_ir_op_priority(od_config: OmniDiffusionConfig, vllm_config: VllmConfig) -> Any: ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config) ir_op_priority_func = get_diffusion_ir_op_priority_func(od_config) @@ -203,8 +213,32 @@ def __init__( self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {} self.stage_id = getattr(od_config, "stage_id", 0) self.init_device() - # Create model runner using the platform-specified class - model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls() + # Create model runner — one decision chain, in precedence order: + # 1. explicit od_config.diffusion_model_runner_cls (user override), + # 2. the runner declared by the engine class that engine_backend + # selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner), + # 3. the platform default. + # Routing policy therefore lives on the engine class / config surface; + # engines never mutate od_config. Overrides must be import-path + # strings — guard with isinstance so a non-string (e.g. a Mock + # od_config in tests) doesn't shadow the platform hook. + runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None) + engine_runner = None + if not (isinstance(runner_override, str) and runner_override): + try: + from vllm_omni.diffusion.diffusion_engine import DiffusionEngine + + engine_cls = DiffusionEngine.resolve_engine_class(self.od_config) + engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None) + except Exception: + logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank) + engine_runner = None + if isinstance(runner_override, str) and runner_override: + model_runner_cls_path = runner_override + elif isinstance(engine_runner, str) and engine_runner: + model_runner_cls_path = engine_runner + else: + model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls() model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path) self.model_runner = model_runner_cls( vllm_config=self.vllm_config, @@ -236,19 +270,35 @@ def init_device(self) -> None: # Create vllm_config for parallel configuration. Pass explicit device_config # so DeviceConfig does not rely on current_platform in worker subprocesses. vllm_config = _create_diffusion_worker_vllm_config(self.device, self.od_config) - vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size - vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size - vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel + parallel_config = self.od_config.parallel_config + vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size + vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size + if parallel_config.enable_expert_parallel and self.od_config.is_moe: + # Diffusion uses its own DP/CFG/SP groups normally. vLLM groups are + # only remapped for expert-parallel runtimes that consume vLLM's + # FusedMoE/EP semantics. + vllm_config.parallel_config.data_parallel_size = ( + parallel_config.data_parallel_size * parallel_config.cfg_parallel_size + ) + vllm_config.parallel_config.prefill_context_parallel_size = parallel_config.sequence_parallel_size + vllm_config.parallel_config.enable_expert_parallel = parallel_config.enable_expert_parallel vllm_config.profiler_config = self.od_config.profiler_config vllm_config.model_config = _make_diffusion_vllm_model_config(self.od_config) # type: ignore[assignment] vllm_config.quant_config = self.od_config.quantization_config # Since vLLM v0.20.0, IR wraps GPU ops. Set IR op priority preference to enforce GPU op fusion during wrapping. # Also need to log, because vLLM internally logs another line in VllmConfig.__post_init__. Avoid confusion. vllm_config.kernel_config.ir_op_priority = _resolve_ir_op_priority(self.od_config, vllm_config) + if self.od_config.moe_backend != "auto": + logger.warning( + "Overriding MoE backend from default 'auto' to '%s' per deploy config.", + self.od_config.moe_backend, + ) + vllm_config.kernel_config.moe_backend = self.od_config.moe_backend logger.info( "Final IR op priority after setting vLLM-Omni overrides: %s", vllm_config.kernel_config.ir_op_priority ) self.vllm_config = vllm_config + current_omni_platform.init_diffusion_worker_vllm_config(vllm_config) # Initialize distributed environment with ( @@ -362,7 +412,12 @@ def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> N else: profiler.stop() - def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfig) -> DiffusionOutput: + def execute_model( + self, + req: OmniDiffusionRequest, + od_config: OmniDiffusionConfig, + kv_prefetch_jobs: dict | None = None, + ) -> DiffusionOutput: """Execute a forward pass by delegating to the model runner.""" assert self.model_runner is not None, "Model runner not initialized" if self.lora_manager is not None: @@ -375,7 +430,29 @@ def execute_model(self, req: OmniDiffusionRequest, od_config: OmniDiffusionConfi profiler = self._get_profiler() ctx = profiler.annotate_context_manager("diffusion_forward") if profiler else nullcontext() with ctx: - output = self.model_runner.execute_model(req) + output = self.model_runner.execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs) + if profiler: + profiler.step() + return output + + def execute_model_batch( + self, scheduler_output: DiffusionSchedulerOutput, od_config: OmniDiffusionConfig + ) -> BatchRunnerOutput: + """Batch forward: LoRA activate once, delegate to model runner.""" + assert self.model_runner is not None, "Model runner not initialized" + # LoRA: same adapter/scale within batch guaranteed by SamplingParamsKey + if self.lora_manager is not None and scheduler_output.scheduled_new_reqs: + sp = scheduler_output.scheduled_new_reqs[0].req.sampling_params + try: + self.lora_manager.set_active_adapter(sp.lora_request, sp.lora_scale) + except Exception as exc: + if sp.lora_request is not None: + raise + logger.warning("LoRA activation skipped: %s", exc) + profiler = self._get_profiler() + ctx = profiler.annotate_context_manager("diffusion_forward_batch") if profiler else nullcontext() + with ctx: + output = self.model_runner.execute_model_batch(scheduler_output, od_config) if profiler: profiler.step() return output @@ -455,8 +532,7 @@ def sleep(self, level: int = 1) -> bool: Args: level: Sleep level. Level 1 offloads weights, level 2 also saves buffers. """ - from vllm.device_allocator.cumem import CuMemAllocator - + CuMemAllocator = _get_cumem_allocator_class() allocator = CuMemAllocator.get_instance() usage_before = allocator.get_current_usage() @@ -508,8 +584,7 @@ def wake_up(self, tags: list[str] | None = None) -> bool: tags: List of memory pool tags to re-activate (e.g., ["weights"] to match Level 1 sleep). If None, all pools are re-activated. """ - from vllm.device_allocator.cumem import CuMemAllocator - + CuMemAllocator = _get_cumem_allocator_class() allocator = CuMemAllocator.get_instance() allocator.wake_up(tags) current_omni_platform.synchronize() @@ -631,8 +706,7 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: if is_sleep_enabled: current_omni_platform.synchronize() gc.collect() - from vllm.device_allocator.cumem import CuMemAllocator - + CuMemAllocator = _get_cumem_allocator_class() allocator = CuMemAllocator.get_instance() if tag == "weights": assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process." @@ -642,6 +716,10 @@ def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager: def shutdown(self) -> None: """Shutdown the worker and cleanup distributed environment.""" + if self.model_runner is not None: + mgr = getattr(self.model_runner, "kv_transfer_manager", None) + if mgr is not None: + mgr.shutdown_prefetch() destroy_distributed_env() @@ -874,7 +952,7 @@ def worker_busy_loop(self) -> None: continue else: - # Handle generation request + # Handle direct generation requests. try: output = self.worker.execute_model(msg, self.od_config) except Exception as e: @@ -1073,18 +1151,24 @@ def generate(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput: """ return self.worker.generate(requests) - def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput: + def execute_model( + self, + reqs: list[OmniDiffusionRequest], + od_config: OmniDiffusionConfig, + kv_prefetch_jobs: dict | None = None, + ) -> DiffusionOutput: """ Execute a forward pass. Args: reqs: List of diffusion requests od_config: OmniDiffusionConfig configuration + kv_prefetch_jobs: Optional next-request KV prefetch descriptor. Returns: DiffusionOutput with generated results """ - return self.worker.execute_model(reqs, od_config) + return self.worker.execute_model(reqs, od_config, kv_prefetch_jobs=kv_prefetch_jobs) def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput: """Execute one diffusion step.""" diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py index ef502d0d433..511f2dd6963 100644 --- a/vllm_omni/diffusion/worker/input_batch.py +++ b/vllm_omni/diffusion/worker/input_batch.py @@ -253,8 +253,15 @@ def _prepare_request_prompt_field( f"{embeds_attr} for request {state.request_id} requires seq_len " f"{max_actual_seq_len}, got target {target_seq_len}." ) - return embeds[:, :target_seq_len], None if mask is None else mask[:, :target_seq_len] + embeds = embeds[:, :target_seq_len] + mask = None if mask is None else mask[:, :target_seq_len] + setattr(state, embeds_attr, embeds) + if mask is not None: + setattr(state, mask_attr, mask) + setattr(state, seq_lens_attr, [int(embeds.shape[1])] * int(embeds.shape[0])) + return embeds, mask + setattr(state, seq_lens_attr, [int(embeds.shape[1])] * int(embeds.shape[0])) return embeds, mask @@ -747,4 +754,7 @@ def scatter_latents( ) -DiffusionInputBatch = InputBatch +# Alias: InputBatch is the step/tensor-level batch. +# DiffusionRequestBatch (in request_batch.py) is the request-level batch. +StepInputBatch = InputBatch +DiffusionInputBatch = StepInputBatch diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py new file mode 100644 index 00000000000..1f7c2acaabd --- /dev/null +++ b/vllm_omni/diffusion/worker/request_batch.py @@ -0,0 +1,266 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Request-level batch abstraction for diffusion runner.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType + + +def _slice_request_output(value: Any, start: int, stop: int) -> Any: + if isinstance(value, tuple): + return tuple(_slice_request_output(item, start, stop) for item in value) + if isinstance(value, list): + return value[start:stop] + if isinstance(value, torch.Tensor): + return value[start:stop] + return value + + +def split_diffusion_output_by_request( + result: DiffusionOutput, + req: DiffusionRequestBatch, + *, + num_outputs_per_prompt: int, +) -> list[DiffusionOutput]: + """Split a batched DiffusionOutput into one output per request.""" + if num_outputs_per_prompt <= 0: + raise ValueError(f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}.") + + return [ + DiffusionOutput( + output=_slice_request_output( + result.output, + idx * num_outputs_per_prompt, + (idx + 1) * num_outputs_per_prompt, + ), + error=result.error, + finished=result.finished, + stage_durations=result.stage_durations, + peak_memory_mb=result.peak_memory_mb, + chunk_index=result.chunk_index, + total_chunks=result.total_chunks, + ) + for idx in range(req.num_reqs) + ] + + +@dataclass +class DiffusionRequestBatch: + """Request-level batch wrapping original diffusion requests. + + Each :class:`~vllm_omni.diffusion.request.OmniDiffusionRequest` represents + one logical diffusion request with one prompt. The scheduler and runner use + this wrapper to present a compatible request batch to pipeline + ``forward()`` methods without reintroducing list-shaped request payloads. + + This is distinct from ``InputBatch`` (aliased as ``StepInputBatch``), + which manages step/tensor-level data for stepwise execution. + + Args: + requests: Independent diffusion requests scheduled together for + request-mode execution. + + Attributes: + requests: Original request objects in scheduler order. + num_reqs: Number of requests in the batch. + request_ids: Request IDs in the same order as ``requests``. + prompts: Prompt list assembled from each request's single ``prompt``. + sampling_params_list: Per-request sampling parameters in scheduler + order. Request-batch pipelines read request-local values here. + sampling_params: Sampling parameters for single-request legacy paths. + request_id: First request ID, kept as a compatibility convenience for + code paths that handle a single-request batch. + kv_sender_info: KV-transfer metadata from the first request. + """ + + requests: list[OmniDiffusionRequest] + + @property + def num_reqs(self) -> int: + return len(self.requests) + + @property + def request_ids(self) -> list[str]: + return [req.request_id for req in self.requests] + + @property + def prompts(self) -> list[OmniPromptType]: + return [req.prompt for req in self.requests] + + @property + def sampling_params_list(self) -> list[OmniDiffusionSamplingParams]: + return [req.sampling_params for req in self.requests] + + @property + def sampling_params(self) -> OmniDiffusionSamplingParams: + # Legacy pipelines do not accept RequestBatch, so they are invoked with + # one request at a time. In that path, this batch is expected to contain + # a single request, and we expose its sampling params for compatibility. + assert len(self.requests) == 1, "RequestBatch with multiple requests does not have a single sampling_params" + return self.requests[0].sampling_params + + @property + def request_id(self) -> str: + return self.requests[0].request_id + + @property + def kv_sender_info(self) -> dict | None: + return self.requests[0].kv_sender_info + + def is_dummy_run(self) -> bool: + return self.requests[0].is_dummy_run_request_id(self.request_id) + + def get(self, request_id: str) -> OmniDiffusionRequest | None: + for req in self.requests: + if req.request_id == request_id: + return req + return None + + def collate_request_generators( + self, + num_outputs_per_prompt: int, + default_generator: torch.Generator | list[torch.Generator] | None, + ) -> torch.Generator | list[torch.Generator] | None: + return self.collate_sampling_param_generators( + self.sampling_params_list, + num_outputs_per_prompt, + default_generator, + ) + + def collate_request_tensors( + self, + attr: str, + default_tensor: torch.Tensor | None, + ) -> torch.Tensor | None: + return self.collate_tensors( + [getattr(sampling, attr) for sampling in self.sampling_params_list], + attr, + default_tensor, + ) + + @staticmethod + def collate_tensors( + tensors: list[torch.Tensor | None], + name: str, + default_tensor: torch.Tensor | None, + ) -> torch.Tensor | None: + validated_tensors = DiffusionRequestBatch._validate_tensor_sequence(tensors, name) + if validated_tensors is None: + return default_tensor + return torch.cat(validated_tensors, dim=0) + + @staticmethod + def collate_prompt_tensors( + tensors: list[torch.Tensor | None], + name: str, + default_tensor: torch.Tensor | None, + ) -> torch.Tensor | None: + validated_tensors = DiffusionRequestBatch._validate_tensor_sequence(tensors, name) + if validated_tensors is None: + return default_tensor + return torch.stack(validated_tensors, dim=0) + + @staticmethod + def get_prompt_field(prompt: OmniPromptType, name: str) -> Any: + if isinstance(prompt, str): + return None + value = prompt.get(name) + if value is None: + additional = prompt.get("additional_information") + if isinstance(additional, dict): + value = additional.get(name) + if isinstance(value, list): + return value[0] if value else None + return value + + @staticmethod + def collate_prompt_fields( + prompts: list[OmniPromptType], + name: str, + default_tensor: torch.Tensor | None, + ) -> torch.Tensor | None: + return DiffusionRequestBatch.collate_prompt_tensors( + [DiffusionRequestBatch.get_prompt_field(prompt, name) for prompt in prompts], + name, + default_tensor, + ) + + @staticmethod + def get_prompt_field_with_aliases(prompt: OmniPromptType, names: Sequence[str]) -> Any: + for name in names: + value = DiffusionRequestBatch.get_prompt_field(prompt, name) + if value is not None: + return value + return None + + @staticmethod + def collate_prompt_field_map( + prompts: list[OmniPromptType], + field_defaults: Mapping[str, torch.Tensor | None], + field_aliases: Mapping[str, Sequence[str]] | None = None, + ) -> dict[str, torch.Tensor | None]: + collated_fields: dict[str, torch.Tensor | None] = {} + for name, default_tensor in field_defaults.items(): + aliases = field_aliases.get(name, (name,)) if field_aliases is not None else (name,) + collated_fields[name] = DiffusionRequestBatch.collate_prompt_tensors( + [DiffusionRequestBatch.get_prompt_field_with_aliases(prompt, aliases) for prompt in prompts], + name, + default_tensor, + ) + return collated_fields + + @staticmethod + def _validate_tensor_sequence( + tensors: list[torch.Tensor | None], + name: str, + ) -> list[torch.Tensor] | None: + if not any(tensor is not None for tensor in tensors): + return None + if not all(isinstance(tensor, torch.Tensor) for tensor in tensors): + raise ValueError(f"Cannot batch requests with a mix of provided and missing {name}.") + + first = tensors[0] + assert isinstance(first, torch.Tensor) + for tensor in tensors[1:]: + assert isinstance(tensor, torch.Tensor) + if tensor.shape != first.shape or tensor.dtype != first.dtype or tensor.device != first.device: + raise ValueError( + f"Batched request {name} must have matching shape, dtype, and device; " + f"got {tensor.shape}/{tensor.dtype}/{tensor.device} and " + f"{first.shape}/{first.dtype}/{first.device}." + ) + return tensors + + @staticmethod + def collate_sampling_param_generators( + sampling_params_list: list[Any], + num_outputs_per_prompt: int, + default_generator: torch.Generator | list[torch.Generator] | None, + ) -> torch.Generator | list[torch.Generator] | None: + request_generators = [sampling.generator for sampling in sampling_params_list] + if not any(generator is not None for generator in request_generators): + return default_generator + if not all(generator is not None for generator in request_generators): + raise ValueError("Cannot batch requests with a mix of provided and missing generators.") + + generators: list[torch.Generator] = [] + for generator in request_generators: + if isinstance(generator, list): + if len(generator) != num_outputs_per_prompt: + raise ValueError( + "Per-request generator lists must match num_outputs_per_prompt, " + f"got {len(generator)} and {num_outputs_per_prompt}." + ) + generators.extend(generator) + else: + generators.extend([generator] * num_outputs_per_prompt) + return generators diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index a0bfcea9b42..ba6165a9d9f 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -15,6 +15,41 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType +def clear_pipeline_stage_durations(pipeline: Any) -> None: + clear_records = getattr(pipeline, "clear_profiler_records", None) + if getattr(pipeline, "enable_diffusion_pipeline_profiler", False) and callable(clear_records): + clear_records() + + +def consume_pipeline_stage_durations(pipeline: Any) -> dict[str, float]: + if not getattr(pipeline, "enable_diffusion_pipeline_profiler", False): + return {} + stage_durations = getattr(pipeline, "stage_durations", None) + if not isinstance(stage_durations, dict): + return {} + result = {stage: float(duration) for stage, duration in stage_durations.items()} + clear_pipeline_stage_durations(pipeline) + return result + + +def merge_stage_durations( + state: DiffusionRequestState, + stage_durations: dict[str, float], +) -> None: + if not stage_durations: + return + for stage, duration in stage_durations.items(): + state.stage_durations[stage] = float(state.stage_durations.get(stage, 0.0)) + float(duration) + + +def attach_stage_durations( + state: DiffusionRequestState, + output: DiffusionOutput, +) -> None: + if state.stage_durations: + output.stage_durations = dict(state.stage_durations) + + @dataclass class DiffusionRequestState: """Per-request mutable state across all pipeline stages. @@ -40,7 +75,8 @@ class DiffusionRequestState: # ── Identity / request-level inputs ── request_id: str sampling: OmniDiffusionSamplingParams - prompts: list[OmniPromptType] | None = None + prompt: OmniPromptType | None = None + kv_sender_info: dict | None = None # ── Encoded prompts (set once by prepare_encode) ── prompt_embeds: torch.Tensor | None = None @@ -55,6 +91,12 @@ class DiffusionRequestState: timesteps: torch.Tensor | list[torch.Tensor] | None = None step_index: int = 0 + # ── Optional chunked streaming progress ── + chunk_index: int = 0 + step_in_chunk: int = 0 + total_chunks: int = 1 + chunk_num_steps: int | None = None + # ── Per-request scheduler instance (set once by prepare_encode) ── scheduler: Any | None = None @@ -72,6 +114,12 @@ class DiffusionRequestState: # For example: Wan condition tensors / masks, or Bagel KV contexts. extra: dict[str, Any] = field(default_factory=dict) + # ── Runner-owned profiling metadata ── + stage_durations: dict[str, float] = field(default_factory=dict) + + # Peak device memory observed while this request is active in step mode. + peak_memory_mb: float = 0.0 + # ── Properties ── @property @@ -103,6 +151,18 @@ def denoise_completed(self) -> bool: return False return self.step_index >= total_steps + @property + def chunk_denoise_completed(self) -> bool: + if self.chunk_num_steps is None: + return False + return self.step_in_chunk >= self.chunk_num_steps + + @property + def request_denoise_completed(self) -> bool: + if self.chunk_num_steps is None: + return self.denoise_completed + return self.chunk_index >= self.total_chunks + @property def new_request(self) -> bool: # TODO: this is only an approximation for current stepwise mode. diff --git a/vllm_omni/distributed/kv_transfer/__init__.py b/vllm_omni/distributed/kv_transfer/__init__.py deleted file mode 100644 index f8914a428ba..00000000000 --- a/vllm_omni/distributed/kv_transfer/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Patched KV transfer connectors for PD disaggregation. - -This package provides monkey-patched versions of vLLM's native KV transfer -connectors (e.g. MooncakeConnector) that fix the request-ID mismatch problem -in prefill-decode disaggregation. - -vLLM's ``InputProcessor.assign_request_id()`` appends a random 8-char suffix -to each request ID internally. The prefill engine stores KV under its own -suffix, but the decode engine generates a *different* suffix — so it can never -find the KV data. The patched connector threads the prefill engine's internal -``remote_request_id`` through ``kv_transfer_params`` so the decode side can -reference the correct KV entry. -""" diff --git a/vllm_omni/distributed/kv_transfer/monkey_patch.py b/vllm_omni/distributed/kv_transfer/monkey_patch.py deleted file mode 100644 index 5c3b97755b8..00000000000 --- a/vllm_omni/distributed/kv_transfer/monkey_patch.py +++ /dev/null @@ -1,193 +0,0 @@ -"""Monkey-patch vLLM's MooncakeConnector to fix request-ID mismatch in PD disaggregation. - -vLLM's InputProcessor appends a random suffix to each request ID. The prefill -engine stores KV under its suffix, but the decode engine generates a different -suffix. This patch threads ``remote_request_id`` through ``kv_transfer_params`` -so the decode side references the correct KV entry. -""" - -from __future__ import annotations - -import importlib -import logging -import sys -from dataclasses import dataclass -from typing import Any - -logger = logging.getLogger(__name__) - -_patched: bool = False - - -@dataclass -class PatchedRecvReqMeta: - """Receive-request metadata carrying the prefill engine's request ID.""" - - request_id: str - remote_request_id: str - local_block_ids: list[int] - kv_transfer_params: dict[str, Any] - - -def _import_mooncake_module(): - """Import MooncakeConnector module, supporting both vLLM >=0.16 and older.""" - for mod_path in ( - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector", - "vllm.distributed.kv_transfer.kv_connector.v1.mooncake_connector", - ): - try: - return importlib.import_module(mod_path) - except (ImportError, ModuleNotFoundError): - continue - return None - - -def _create_patched_mooncake_connector(): - """Return a subclass of MooncakeConnector with remote_request_id support.""" - _mc_mod = _import_mooncake_module() - if _mc_mod is None: - raise ImportError("Cannot import MooncakeConnector from upstream vLLM") - _OriginalMooncakeConnector = _mc_mod.MooncakeConnector - - class PatchedMooncakeConnector(_OriginalMooncakeConnector): - """Fixes request-ID mismatch in PD disaggregation by injecting - remote_request_id on the prefill side and using it for KV lookup - on the decode side. - """ - - def __init__(self, *args: Any, **kwargs: Any) -> None: - super().__init__(*args, **kwargs) - self.remote_to_local_req: dict[str, str] = {} - logger.info("[PatchedMooncakeConnector] Initialized") - - def request_finished( - self, - request: Any, - block_ids: list[int], - ) -> tuple[bool, dict[str, Any] | None]: - result = super().request_finished(request, block_ids) - - if isinstance(result, tuple) and len(result) == 2: - delay_free, kv_params = result - else: - delay_free, kv_params = False, result - - # Normalise _reqs_need_send values - req_id = getattr(request, "request_id", None) - if req_id and hasattr(self, "_reqs_need_send"): - entry = self._reqs_need_send.get(req_id) - if isinstance(entry, tuple) and len(entry) == 2: - self._reqs_need_send[req_id] = entry[1] - - # Inject remote_request_id into kv_transfer_params - if kv_params is not None and isinstance(kv_params, dict): - kv_params["remote_request_id"] = req_id or "NOT_SET" - if hasattr(self, "side_channel_host"): - kv_params.setdefault("remote_host", self.side_channel_host) - if hasattr(self, "side_channel_port"): - kv_params.setdefault("remote_port", self.side_channel_port) - - return delay_free, kv_params - - def add_new_req( - self, - request_id: str, - local_block_ids: list[int], - kv_transfer_params: dict[str, Any] | None = None, - **kwargs: Any, - ) -> None: - super().add_new_req(request_id, local_block_ids, kv_transfer_params, **kwargs) - - kv_transfer_params = kv_transfer_params or {} - load_remote_cache = kv_transfer_params.get( - "do_remote_prefill", - kv_transfer_params.get("load_remote_cache", False), - ) - - if load_remote_cache: - remote_request_id = kv_transfer_params.get("remote_request_id", request_id) - meta = PatchedRecvReqMeta( - request_id=request_id, - remote_request_id=remote_request_id, - local_block_ids=local_block_ids, - kv_transfer_params=kv_transfer_params, - ) - if not hasattr(self, "_reqs_need_recv"): - self._reqs_need_recv = {} - self._reqs_need_recv[request_id] = meta - - def group_kv_pull(self, metadata: Any | None = None) -> None: - """Use remote_request_id as ZMQ lookup key via save-patch-restore.""" - if not hasattr(self, "_reqs_need_recv") or not self._reqs_need_recv: - return - - original_recv = self._reqs_need_recv.copy() - patched_recv: dict[str, Any] = {} - - for local_id, meta in original_recv.items(): - if isinstance(meta, PatchedRecvReqMeta): - remote_id = meta.remote_request_id - self.remote_to_local_req[remote_id] = local_id - patched_meta = type(meta)( - request_id=remote_id, - remote_request_id=remote_id, - local_block_ids=meta.local_block_ids, - kv_transfer_params=meta.kv_transfer_params, - ) - patched_recv[remote_id] = patched_meta - else: - patched_recv[local_id] = meta - - self._reqs_need_recv = patched_recv - super().group_kv_pull(metadata) - - # Restore unconsumed entries to original local keys - for remote_id, local_id in list(self.remote_to_local_req.items()): - if remote_id in self._reqs_need_recv: - entry = self._reqs_need_recv.pop(remote_id) - self._reqs_need_recv[local_id] = original_recv.get(local_id, entry) - - def receive_kv(self, path: Any = None, req_blocks: Any = None) -> Any: - result = super().receive_kv(path, req_blocks) - - if self.remote_to_local_req: - completed = [ - rid - for rid, lid in self.remote_to_local_req.items() - if not hasattr(self, "_reqs_need_recv") or lid not in self._reqs_need_recv - ] - for remote_id in completed: - self.remote_to_local_req.pop(remote_id, None) - - return result - - PatchedMooncakeConnector.__qualname__ = _OriginalMooncakeConnector.__qualname__ - - return PatchedMooncakeConnector - - -def apply_mooncake_connector_patch() -> bool: - """Replace vLLM's MooncakeConnector with the patched version.""" - global _patched - if _patched: - return True - - _mc_module = _import_mooncake_module() - if _mc_module is None: - logger.warning("[monkey_patch] Cannot import MooncakeConnector — patch NOT applied.") - return False - - _OriginalClass = _mc_module.MooncakeConnector - - PatchedClass = _create_patched_mooncake_connector() - - _mc_module.MooncakeConnector = PatchedClass - # Snapshot sys.modules: hasattr() may trigger lazy submodule imports that - # mutate sys.modules during iteration. - for _, module in list(sys.modules.items()): - if hasattr(module, "MooncakeConnector") and module.MooncakeConnector is _OriginalClass: - module.MooncakeConnector = PatchedClass - - _patched = True - logger.info("[monkey_patch] MooncakeConnector patch applied") - return True diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index 297d4ef59ce..34d19062a0a 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -692,12 +692,20 @@ def get( return None # 1. Allocate Destination Buffer from Pool + # For non-fast-path (serialized objects), allocate a few extra bytes + # as padding. Mooncake's TCP/memcpy transport may prepend a leading + # protocol byte (0x01) before the payload. The extra padding allows + # the receiver to detect and skip this byte without truncating the + # actual payload. The padding is harmless for fast-path (ManagedBuffer) + # because the caller reads exact-size slices. + _MOONCAKE_TCP_PADDING = 4 if not is_fast_path else 0 + alloc_size = data_size + _MOONCAKE_TCP_PADDING try: - offset = self.allocator.alloc(data_size) - recv_buffer = ManagedBuffer(self.allocator, offset, data_size, self.pool) + offset = self.allocator.alloc(alloc_size) + recv_buffer = ManagedBuffer(self.allocator, offset, alloc_size, self.pool) dst_ptr = self.base_ptr + offset except MemoryError: - logger.error(f"Failed to allocate {data_size} bytes in receive pool") + logger.error(f"Failed to allocate {alloc_size} bytes in receive pool") return None _t2 = _time_mod.perf_counter() @@ -775,8 +783,19 @@ def get( _t_copy_end = _time_mod.perf_counter() _copy_ms = (_t_copy_end - _t_copy_start) * 1000 + # Try deserializing the first data_size bytes. Mooncake's + # TCP/memcpy transport may prepend a leading protocol byte + # (0x01) before the payload — handle this by retrying from + # offset 1 if the first attempt raises DecodeError. _t_deser_start = _time_mod.perf_counter() - val = OmniSerializer.deserialize(raw_bytes) + payload = raw_bytes[:data_size] + try: + val = OmniSerializer.deserialize(payload) + except msgspec.DecodeError: + # Likely a leading Mooncake TCP protocol byte — retry + # from offset 1, using the extra padding bytes if needed. + payload = raw_bytes[1 : data_size + 1] + val = OmniSerializer.deserialize(payload) _t_deser_end = _time_mod.perf_counter() _deser_ms = (_t_deser_end - _t_deser_start) * 1000 diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 63466d11613..5de0d54db5d 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -116,6 +116,14 @@ def _get_by_key(self, get_key: str) -> tuple[Any, int] | None: return result except FileNotFoundError: return None + except ValueError as e: + # A receiver can observe a newly-created POSIX SHM object before + # the writer has finished sizing it. Treat that as "not ready yet" + # so async polling can retry without a traceback. + if "empty file" in str(e): + return None + logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True) + return None except Exception: logger.debug("_get_by_key: unexpected error reading SHM segment %s", get_key, exc_info=True) return None diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 3b234ad1cc1..51e604336f8 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unified OmniConnector and KV cache transfer management.""" +import enum import json import struct import time from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from dataclasses import asdict, dataclass from typing import Any @@ -13,6 +15,7 @@ from vllm.logger import init_logger from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.platforms import current_omni_platform from .factory import OmniConnectorFactory from .utils.config import TRANSFER_ENGINE_CONNECTOR_NAMES, ConnectorSpec @@ -32,11 +35,64 @@ slice_layer_blocks, slice_received_rank_shard, ) +from .utils.serialization import OmniSerializer logger = init_logger(__name__) LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] + +class KVPrefetchConsumeError(RuntimeError): + """Payload consumed from connector but post-get processing failed — sync retry impossible.""" + + +class ReceiveRole(enum.Enum): + """How a rank obtains KV: LOCAL pulls its own shard, LEADER pulls then + distributes, FOLLOWER receives via collective.""" + + LOCAL = "local" + LEADER = "leader" + FOLLOWER = "follower" + + +@dataclass(frozen=True) +class _TransferTopoConfig: + """Snapshot of the current rank's parallel topology for KV receive routing.""" + + role: ReceiveRole + tp_active: bool + cfg_size: int + cfg_rank: int + cfg_group: Any + sp_size: int + sp_rank: int + sp_group: Any + world: Any + + @property + def is_leader(self) -> bool: + return self.role == ReceiveRole.LEADER + + @property + def is_follower(self) -> bool: + return self.role == ReceiveRole.FOLLOWER + + @property + def is_local(self) -> bool: + return self.role == ReceiveRole.LOCAL + + @property + def cfg_active(self) -> bool: + return self.cfg_size > 1 + + @property + def sp_active(self) -> bool: + return self.sp_size > 1 + + +# Placeholder for the heavy primary KV in the side-payload dict; receiver swaps in the rebuilt object from the blob. +_KV_PLACEHOLDER = "__kv_placeholder__" + _SAFE_TORCH_DTYPES = { name: dtype for name in ( @@ -75,6 +131,8 @@ class OmniKVCacheConfig: recv_timeout: float = 30.0 from_tp: int = 1 to_tp: int = 1 + enable_kv_async_prefetch: bool = False + kv_prefetch_min_free_mem_ratio: float = 0.0 @dataclass @@ -290,7 +348,7 @@ class OmniKVTransferManager: - KV cache receiving with timeout """ - def __init__(self, config: OmniKVCacheConfig): + def __init__(self, config: OmniKVCacheConfig, *, async_prefetch: bool = False): self.config = config self._connector = None @@ -336,6 +394,19 @@ def __init__(self, config: OmniKVCacheConfig): self._sender_base_host: str | None = None self._sender_base_zmq_port: int | None = None + # Prefetch + self._async_prefetch = async_prefetch + self._prefetch_min_free_mem_ratio: float = max(0.0, config.kv_prefetch_min_free_mem_ratio) + # Single-worker: serial mode ensures at most one outstanding prefetch + # and avoids stream-creation races on _bg_copy_stream. + self._prefetch_executor: ThreadPoolExecutor | None = ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="kv-prefetch") if async_prefetch else None + ) + self._prefetch_futures: dict[str, Any] = {} + self._bg_copy_stream: current_omni_platform.Stream | None = None + + self._topo_config: _TransferTopoConfig | None = None + if config.need_send_cache and config.connector_config: try: _ = self.connector @@ -343,15 +414,22 @@ def __init__(self, config: OmniKVCacheConfig): except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) + @property + def topo_config(self) -> _TransferTopoConfig: + """Lazily build the parallel topology snapshot on first access.""" + if self._topo_config is None: + self._topo_config = self._build_topo_config() + return self._topo_config + # ------------------------------------------------------------------ # # Factory helpers # ------------------------------------------------------------------ # @classmethod - def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": + def _create(cls, cfg: dict | None, *, async_prefetch: bool = False) -> "OmniKVTransferManager": """Create manager from raw config dict.""" if not cfg or not isinstance(cfg, dict): - return cls(OmniKVCacheConfig()) + return cls(OmniKVCacheConfig(), async_prefetch=async_prefetch) rank_mapping = cfg.get("rank_mapping", {}) if not isinstance(rank_mapping, dict): @@ -369,38 +447,66 @@ def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": recv_timeout=cfg.get("recv_timeout", 30.0), from_tp=int(rank_mapping.get("from_tp", 1)), to_tp=int(rank_mapping.get("to_tp", 1)), - ) + enable_kv_async_prefetch=async_prefetch, + kv_prefetch_min_free_mem_ratio=cfg.get("kv_prefetch_min_free_mem_ratio", 0.0), + ), + async_prefetch=async_prefetch, ) + @staticmethod + def _receiver_pool_on_device(omni_kv: Any) -> bool: + """True when the receiver connector pool is on a device (non-CPU).""" + if not isinstance(omni_kv, dict): + return False + cc = omni_kv.get("connector_config") + if not isinstance(cc, dict): + return False + return str(cc.get("memory_pool_device", "cpu")).lower() != "cpu" + + @classmethod + def _resolve_async_prefetch(cls, omni_kv: Any, *, has_companion: bool) -> bool: + """Effective prefetch flag: requested by config and not force-disabled.""" + if not isinstance(omni_kv, dict) or not bool(omni_kv.get("enable_kv_async_prefetch", False)): + return False + return not (has_companion or cls._receiver_pool_on_device(omni_kv)) + @classmethod def from_od_config(cls, config: Any) -> "OmniKVTransferManager": """Create from model or OmniDiffusion config.""" - return cls._create(getattr(config, "omni_kv_config", None)) + omni_kv = getattr(config, "omni_kv_config", None) + has_companion = getattr(config, "cfg_kv_collect_func", None) is not None + async_prefetch = cls._resolve_async_prefetch(omni_kv, has_companion=has_companion) + return cls._create(omni_kv, async_prefetch=async_prefetch) from_model_config = from_od_config @classmethod def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager": """Create from vllm config with fallback to kv_transfer_config.""" - # Primary: omni_kv_config from model_config - omni_kv = getattr(model_config, "omni_kv_config", None) - if isinstance(omni_kv, dict): - return cls._create(omni_kv) - - # Fallback: check kv_transfer_config - kv_cfg = getattr(vllm_config, "kv_transfer_config", None) - if kv_cfg: - direct = getattr(kv_cfg, "omni_connector_config", None) - if isinstance(direct, dict) and direct: - return cls._create({"connector_config": direct}) - extra = getattr(kv_cfg, "kv_connector_extra_config", None) - if isinstance(extra, dict): - omni = extra.get("omni_connector_config") - if isinstance(omni, dict) and omni: - return cls._create({"connector_config": omni}) + # Primary: omni_kv_config on model_config — same shape as from_od_config. + if isinstance(getattr(model_config, "omni_kv_config", None), dict): + return cls.from_od_config(model_config) + connector_cfg = cls._connector_cfg_from_kv_transfer(getattr(vllm_config, "kv_transfer_config", None)) + if connector_cfg is not None: + return cls._create({"connector_config": connector_cfg}) return cls(OmniKVCacheConfig()) + @staticmethod + def _connector_cfg_from_kv_transfer(kv_cfg: Any) -> dict | None: + """Extract the omni connector config dict nested in a kv_transfer_config.""" + if not kv_cfg: + return None + direct = getattr(kv_cfg, "omni_connector_config", None) + if isinstance(direct, dict) and direct: + return direct + extra = getattr(kv_cfg, "kv_connector_extra_config", None) + if isinstance(extra, dict): + omni = extra.get("omni_connector_config") + if isinstance(omni, dict) and omni: + return omni + return None + @property def connector(self): """Lazy initialization of connector.""" @@ -459,6 +565,88 @@ def connector(self): get_connector = property(lambda self: self.connector) + def _build_topo_config(self) -> _TransferTopoConfig: + """Build the parallel topology snapshot (called once at init).""" + from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_sequence_parallel_rank, + get_sequence_parallel_world_size, + get_sp_group, + get_world_group, + ) + + try: + world = get_world_group() + world_size = world.world_size + world_rank = world.rank_in_group + except Exception: + logger.exception("World group unavailable; defaulting to LOCAL") + return _TransferTopoConfig( + role=ReceiveRole.LOCAL, + tp_active=False, + cfg_size=1, + cfg_rank=0, + cfg_group=None, + sp_size=1, + sp_rank=0, + sp_group=None, + world=None, + ) + + if world_size <= 1: + return _TransferTopoConfig( + role=ReceiveRole.LOCAL, + tp_active=False, + cfg_size=1, + cfg_rank=0, + cfg_group=None, + sp_size=1, + sp_rank=0, + sp_group=None, + world=world, + ) + + tp_active = self._tp_topo.source_tp_size > 1 or self._tp_topo.target_tp_size > 1 + + # CFG/SP groups may legitimately be uninitialized on pure-TP or + # world-only distributed paths; treat that as "not active" (size 1) + # rather than raising, so those paths can still take the local/world + # receive route. + try: + cfg_size = get_classifier_free_guidance_world_size() + cfg_rank = get_classifier_free_guidance_rank() + cfg_group = get_cfg_group() + except Exception: + cfg_size, cfg_rank, cfg_group = 1, 0, None # CFG-parallel not enabled + + try: + sp_size = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + sp_group = get_sp_group() + except Exception: + sp_size, sp_rank, sp_group = 1, 0, None # SP not enabled + + if tp_active and cfg_size <= 1 and sp_size <= 1: + role = ReceiveRole.LOCAL + elif tp_active and (cfg_size > 1 or sp_size > 1): + role = ReceiveRole.LEADER if (cfg_rank == 0 and sp_rank == 0) else ReceiveRole.FOLLOWER + else: + role = ReceiveRole.LEADER if world_rank == 0 else ReceiveRole.FOLLOWER + + return _TransferTopoConfig( + role=role, + tp_active=tp_active, + cfg_size=cfg_size, + cfg_rank=cfg_rank, + cfg_group=cfg_group, + sp_size=sp_size, + sp_rank=sp_rank, + sp_group=sp_group, + world=world, + ) + def _resolve_sender_info( self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None ) -> dict[str, Any] | None: @@ -502,19 +690,6 @@ def _resolve_sender_info( ) return None - @staticmethod - def _clone_received_payload_tensors(data: dict[str, Any]) -> dict[str, Any]: - if not isinstance(data, dict) or "layer_blocks" not in data: - return data - - layer_blocks = data["layer_blocks"] - for cache_name in ("key_cache", "value_cache"): - cache_list = layer_blocks.get(cache_name, []) - for idx, tensor in enumerate(cache_list): - if isinstance(tensor, torch.Tensor): - cache_list[idx] = tensor.clone() - return data - def _slice_transfer_data_for_target(self, kv_data: KVCacheTransferData, target_rank: int) -> KVCacheTransferData: """Pre-slice sender payload for one target rank when sender TP < receiver TP.""" topo = self._tp_topo @@ -681,8 +856,20 @@ def _build_cfg_rank_local_payloads(cls, req: Any, cfg_size: int) -> list[dict[st payloads.append(local_payload) + padding = None + if not branch_roles and cfg_size == 2 and main_payload: + # AR didn't split KV by branch; cfg follower needs the same + # positive KV for shared-prefix reuse. Only safe for cfg_size == 2: + # for cfg_size > 2 every rank must get its own KV, which we can't + # assign here, so padding stays None and the runner raises. + # Shallow dict copy: shares PKV tensors with main_payload. Safe + # only because the send path (cfg_scatter -> pack -> to_gpu_tensor) + # reads tensors to serialize, never mutates in place. Any async/ + # KV-pool path that mutates req.past_key_values mid-scatter must + # clone tensors here instead. + padding = dict(main_payload) while len(payloads) < cfg_size: - payloads.append(None) + payloads.append(padding) return payloads[:cfg_size] @@ -984,24 +1171,183 @@ def _transfer_with_retry( return False, 0, None + def _resolve_sender_base(self, sender_info: dict[str, Any] | None) -> tuple[str | None, int | None]: + """Resolve (host, base_zmq_port) from sender_info.""" + actual = self._resolve_sender_info(sender_info, sender_stage_id=self.recv_stages[0]) + if not actual or "host" not in actual: + return None, None + port = actual.get("zmq_port") + return actual.get("host"), (int(port) if port is not None else None) + + def _has_enough_prefetch_device_mem(self, device: torch.device | None) -> bool: + """False when free-device-memory fraction is below the prefetch threshold. + + CUDA/NPU only; other platforms report unconstrained (throttle skipped). + """ + ratio = self._prefetch_min_free_mem_ratio + if ratio <= 0.0 or device is None: + return True + try: + free, total = current_omni_platform.get_device_memory(device) + except Exception: + return True + return total <= 0 or (free / total) >= ratio + + @staticmethod + def _release_pool_buffers(buffers: list[Any]) -> None: + """Release all pool buffers in *buffers* and clear the list.""" + if not buffers: + return + for buf in buffers: + try: + buf.release() + except Exception: + logger.exception("Failed to release KV pool buffer") + buffers.clear() + + def start_prefetch( + self, kv_prefetch_jobs: dict[str, Any] | None, target_device: torch.device | None = None + ) -> None: + """Kick off a background KV load (non-blocking). No-op unless prefetch enabled.""" + if not (self._async_prefetch and self.config.need_recv_cache) or not kv_prefetch_jobs: + return + # Followers receive via collective distribute; bg pull would consume the owner's payload. + if self.topo_config.is_follower: + return + rid = kv_prefetch_jobs.get("request_id") + if not rid: + return + # Memory pressure → skip prefetch; sync receive handles it. + if not self._has_enough_prefetch_device_mem(target_device): + logger.warning( + "Skip KV prefetch for %s: device free mem below %.2f", rid, self._prefetch_min_free_mem_ratio + ) + return + # Serial mode: at most one outstanding prefetch; drop any leftover request's future. + for stale_rid in [k for k in self._prefetch_futures if k != rid]: + self._discard_future(stale_rid) + if rid in self._prefetch_futures: + return + sender_info = kv_prefetch_jobs.get("kv_sender_info") + if not sender_info: + # No explicit endpoint → bg receive would target wrong sender under multi-replica. + logger.debug("Skip KV prefetch for %s: stub has no kv_sender_info", rid) + return + try: + self._prefetch_futures[rid] = self._prefetch_executor.submit( + self._prefetch_payload, rid, sender_info, target_device + ) + except Exception: + logger.exception("Failed to submit KV prefetch for %s", rid) + + def _prefetch_payload( + self, + request_id: str, + sender_info: dict[str, Any] | None, + target_device: torch.device | None, + ) -> tuple[dict[str, Any] | None, int]: + """Bg-thread body: get + deserialize + H2D on the dedicated ``_bg_copy_stream``. + + Raises on failure (payload may be consumed → no sync retry). + """ + try: + on_device = target_device is not None and target_device.type != "cpu" + if on_device: + # bg thread doesn't inherit the main thread's current device. + torch.accelerator.set_device_index(target_device.index) + if self._bg_copy_stream is None: + assert self._prefetch_executor._max_workers == 1 + self._bg_copy_stream = current_omni_platform.Stream() + with current_omni_platform.stream(self._bg_copy_stream): + data, size = self.receive_kv_cache_for_request( + request_id, target_device=target_device, sender_info=sender_info + ) + else: + data, size = self.receive_kv_cache_for_request( + request_id, target_device=target_device, sender_info=sender_info + ) + except Exception: + logger.exception("KV prefetch payload failed for %s (payload may be lost)", request_id) + raise + + if data is None: + return None, 0 + return data, size + + def consume_prefetched_kv(self, req: Any) -> tuple[dict[str, Any] | None, int]: + """Consume a prefetched KV payload; (None, 0) on miss/recoverable failure. + + Raises ``KVPrefetchConsumeError`` when consumed but post-get failed (no fallback). + """ + request_id = self._resolve_request_id(req) + if not request_id: + return None, 0 + + fut = self._prefetch_futures.pop(request_id, None) + # Serial mode: any other request still in the table is an orphan — drop it. + for stale_rid in list(self._prefetch_futures): + self._discard_future(stale_rid) + + if fut is None: + return None, 0 + + try: + return fut.result() + except KVPrefetchConsumeError: + logger.exception("KV load failed for %s (payload consumed, cannot retry)", request_id) + raise + except Exception: + logger.exception("KV load failed for %s; falling back to sync receive", request_id) + return None, 0 + + def _discard_future(self, request_id: str) -> None: + """Cancel an unstarted prefetch or attach a callback to drop a running one.""" + fut = self._prefetch_futures.pop(request_id, None) + if fut is None: + return + if not fut.cancel(): + fut.add_done_callback(_drop_prefetch_result) + + def shutdown_prefetch(self) -> None: + """Cancel pending prefetches and stop the executor (call on teardown).""" + for rid in list(self._prefetch_futures): + self._discard_future(rid) + if self._prefetch_executor is not None: + try: + self._prefetch_executor.shutdown(wait=True, cancel_futures=True) + except Exception: + logger.exception("Failed to shut down KV prefetch executor") + self._prefetch_executor = None + + @staticmethod + def _record_stream_for_prefetched(data: dict[str, Any]) -> None: + """``record_stream(current_stream)`` on GPU tensors in *data*. + + Protects prefetch-thread tensors from allocator reuse on LEADER collective paths. + """ + if not isinstance(data, dict) or "layer_blocks" not in data: + return + if current_omni_platform.get_device_count() < 1: + return + current_stream = current_omni_platform.current_stream() + layer_blocks = data["layer_blocks"] + for cache_list in (layer_blocks.get("key_cache", []), layer_blocks.get("value_cache", [])): + for tensor in cache_list: + if isinstance(tensor, torch.Tensor) and tensor.device.type != "cpu": + tensor.record_stream(current_stream) + @torch.inference_mode() def receive_kv_cache_for_request( self, request_id: str, target_device: torch.device | None = None, + *, + sender_info: dict[str, Any] | None = None, ) -> tuple[dict[str, Any] | None, int]: - """Receive KV cache for a specific request. - - This implements the receiving logic from gpu_diffusion_model_runner.py. + """Receive KV cache for *request_id*; returns (data dict, size) or (None, 0). - Args: - request_id: The request ID to receive KV cache for - target_device: Optional device to move tensors to - - Returns: - Tuple of (data dict, size) if successful, (None, 0) otherwise + ``sender_info`` overrides instance ``_sender_base_*`` (per-call endpoint). """ - # Check if we should receive KV cache based on config if not self.config.need_recv_cache: logger.debug("Skip receiving KV cache for %s (need_recv_cache=False)", request_id) return None, 0 @@ -1031,6 +1377,13 @@ def receive_kv_cache_for_request( ) pending_pairs = list(recv_key_pairs) received_payloads: dict[str, tuple[dict[str, Any], int]] = {} + deferred_memory: list[Any] = [] + + # Per-call sender_info takes precedence over instance fields. + if sender_info is not None: + base_host, base_port = self._resolve_sender_base(sender_info) + else: + base_host, base_port = self._sender_base_host, self._sender_base_zmq_port logger.info( "Wait for KV cache for request %s from stage %s to %s via %s key(s)...", @@ -1044,16 +1397,21 @@ def receive_kv_cache_for_request( while True: link_start = time.perf_counter() for get_key, from_rank in list(pending_pairs): - # Construct per-rank metadata so the connector queries - # the correct sender endpoint (heterogeneous TP path). - # When from_rank is None (TP<=1), metadata stays None - # and the connector falls back to its default sender. rank_metadata: dict[str, Any] | None = None - if from_rank is not None and self._sender_base_host and self._sender_base_zmq_port is not None: + build_meta = from_rank is not None or sender_info is not None + if build_meta and base_host and base_port is not None: rank_metadata = { - "source_host": self._sender_base_host, - "source_port": self._sender_base_zmq_port + from_rank * KV_RANK_PORT_STRIDE, + "source_host": base_host, + "source_port": base_port + (from_rank or 0) * KV_RANK_PORT_STRIDE, } + elif sender_info is not None: + logger.error( + "KV receive for %s: unresolved per-call sender_info (host=%s, port=%s); aborting", + request_id, + base_host, + base_port, + ) + return None, 0 result = self.connector.get( from_stage=from_stage, @@ -1065,25 +1423,19 @@ def receive_kv_cache_for_request( continue raw_data, size = result - managed_buffer = None if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"): - managed_buffer = raw_data try: buf_tensor = raw_data.tensor if getattr(buf_tensor.device, "type", "cpu") != "cpu": data = KVCacheTransferData.from_bytes_device(buf_tensor) raw_data.release() - managed_buffer = None else: data = KVCacheTransferData.from_bytes(memoryview(buf_tensor.numpy())) - data = self._clone_received_payload_tensors(data) - raw_data.release() - managed_buffer = None + deferred_memory.append(raw_data) except Exception as e: logger.error("Failed to deserialize KV cache from ManagedBuffer: %s", e) - if managed_buffer is not None: - raw_data.release() + raw_data.release() return None, 0 elif isinstance(raw_data, (bytes, bytearray)): data = KVCacheTransferData.from_bytes(raw_data) @@ -1110,20 +1462,27 @@ def receive_kv_cache_for_request( data = merge_received_rank_shards(ordered_payloads, merger=self.kv_payload_merger) data = slice_received_rank_shard(data, topo, slicer=self.kv_payload_slicer) + needs_clone = bool(deferred_memory) try: if isinstance(data, dict) and "layer_blocks" in data: layer_blocks = data["layer_blocks"] - for cache_list in [ + cache_lists = [ layer_blocks.get("key_cache", []), layer_blocks.get("value_cache", []), - ]: + ] + for cache_list in cache_lists: for i, tensor in enumerate(cache_list): if not isinstance(tensor, torch.Tensor): continue if target_device is not None and tensor.device != target_device: cache_list[i] = tensor.to(target_device).contiguous() - except Exception: - logger.exception("Failed to move KV cache tensors to target device") + elif needs_clone: + cache_list[i] = tensor.clone() + except Exception as exc: + logger.exception("Failed to detach/move KV cache tensors for %s", request_id) + raise KVPrefetchConsumeError( + f"Post-get processing failed for {request_id} (payload already consumed)" + ) from exc logger.info( "Successfully received KV cache for %s, %s bytes across %s key(s), wait=%.3fs, link=%.1fms", @@ -1142,9 +1501,13 @@ def receive_kv_cache_for_request( time.sleep(poll_interval) poll_interval = min(poll_interval * 2, max_poll_interval) + except KVPrefetchConsumeError: + raise except Exception: logger.exception("Error receiving KV cache for %s", request_id) return None, 0 + finally: + self._release_pool_buffers(deferred_memory) def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None: """Apply received KV cache data to a request object. @@ -1246,130 +1609,253 @@ def receive_multi_kv_cache( return primary_ok - def receive_multi_kv_cache_distributed( - self, - req: Any, - cfg_kv_collect_func: Callable | None = None, - target_device: torch.device | None = None, - ) -> bool: - """Distributed wrapper around :meth:`receive_multi_kv_cache`. - - TP-aware path selection: - - world size 1: direct receive - - TP active, cfg size 1: each rank independently receives - - TP active, cfg size > 1: cfg-rank 0 receives, then broadcasts to - peers that share the same TP rank - - TP inactive: legacy rank-0 receive then world broadcast - """ - from vllm_omni.diffusion.distributed.parallel_state import ( - get_cfg_group, - get_classifier_free_guidance_rank, - get_classifier_free_guidance_world_size, - get_sequence_parallel_rank, - get_sequence_parallel_world_size, - get_sp_group, - get_world_group, - ) - - world = get_world_group() + def _extract_primary_kv_obj(self, kv_payload: dict[str, Any]) -> Any | None: + """Return the primary KV object (``SimpleNamespace(key_cache, value_cache)``).""" + for key in ("past_key_values", "sp.past_key_values"): + obj = kv_payload.get(key) + if obj is not None and hasattr(obj, "key_cache"): + return obj + return None - if world.world_size <= 1: - return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) + def _pack_kv_payload(self, kv_payload: dict[str, Any] | None, device: torch.device) -> torch.Tensor | None: + """Pack a KV payload into one uint8 tensor on device + Layout: ``[4B side_len][msgpack side_bytes][gpu blob]``. + Returns None on failure.""" + if not kv_payload: + return None + primary = self._extract_primary_kv_obj(kv_payload) + if primary is None: + return None - topo = self._tp_topo - tp_active = topo.source_tp_size > 1 or topo.target_tp_size > 1 - cfg_size = 1 - cfg_rank = 0 - cfg_group = None - sp_size = 1 - sp_rank = 0 - sp_group = None + key_cache = _move_to_device(getattr(primary, "key_cache", []) or [], device) + value_cache = _move_to_device(getattr(primary, "value_cache", []) or [], device) + if not any(isinstance(t, torch.Tensor) for t in key_cache): + return None + td = KVCacheTransferData( + request_id="", + layer_blocks={"key_cache": key_cache, "value_cache": value_cache}, + block_ids=[], + metadata={}, + ) try: - cfg_size = get_classifier_free_guidance_world_size() - cfg_rank = get_classifier_free_guidance_rank() - cfg_group = get_cfg_group() + blob = td.to_gpu_tensor() except Exception: - cfg_size = 1 - cfg_rank = 0 - cfg_group = None + logger.exception("D2D: failed to pack primary KV; falling back to object transport") + return None + + # Side-payload: replace the heavy primary with a sentinel, msgpack the rest. + side = dict(kv_payload) + for key in ("past_key_values", "sp.past_key_values"): + if side.get(key) is primary: + side[key] = _KV_PLACEHOLDER try: - sp_size = get_sequence_parallel_world_size() - sp_rank = get_sequence_parallel_rank() - sp_group = get_sp_group() + side_bytes = OmniSerializer.serialize(side) except Exception: - sp_size = 1 - sp_rank = 0 - sp_group = None + logger.exception("D2D: failed to serialize side-payload; falling back to object transport") + return None + side_len = len(side_bytes) + side_tensor = torch.frombuffer(bytearray(side_bytes), dtype=torch.uint8).to(device) - if tp_active and (cfg_size <= 1 and sp_size <= 1): - logger.info( - "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)", - topo.local_rank, - topo.source_tp_size, - topo.target_tp_size, + # [4B side_len][side_bytes][blob] + header = torch.tensor([side_len], dtype=torch.int32, device=device).view(torch.uint8) + combined = torch.cat([header, side_tensor, blob]) + return combined + + def _unpack_kv_payload(self, combined: torch.Tensor) -> dict[str, Any] | None: + """Inverse of :meth:`_pack_kv_payload`.""" + if combined is None or combined.numel() == 0: + return None + + offset = 4 + side_len = int(combined[:offset].cpu().view(torch.int32).item()) + if side_len <= 0 and combined.numel() > offset: + # Corrupt/zeroed blob (sender-side failure) → treat as miss; raising + # would crash only the receiver and desync the group. + logger.warning( + "D2D: corrupt KV payload (side_len=%d, total_bytes=%d, likely sender-side failure); treating as miss", + side_len, + combined.numel(), ) - return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) + return None + side_bytes = combined[offset : offset + side_len].cpu().numpy().tobytes() + offset += side_len + side: dict[str, Any] = OmniSerializer.deserialize(side_bytes) + + # Rebuild primary KV from the trailing blob and restore it in place. + primary: Any = None + blob = combined[offset:] + if blob.numel() > 0: + from types import SimpleNamespace - if tp_active and (cfg_size > 1 or sp_size > 1): - kv_payload: dict[str, object] | None = None - is_owner = cfg_rank == 0 and sp_rank == 0 - - # step1: only owner need to receive - if is_owner: - received = self.receive_multi_kv_cache( - req, - cfg_kv_collect_func, - torch.device("cpu"), - ) - if received: - if cfg_size > 1: # build cfg-local payloads - cfg_rank_payloads = self._build_cfg_rank_local_payloads( - req, - cfg_size, - ) - kv_payload = cfg_rank_payloads[0] - # send to other cfg - for dst_cfg_rank in range(1, cfg_size): - cfg_group.send_object( - cfg_rank_payloads[dst_cfg_rank], - dst_cfg_rank, - ) - elif sp_size > 1: - kv_payload = self._collect_request_kv_payload(req) - else: - # Owner receive failed: send None sentinel to CFG followers to avoid deadlock - if cfg_size > 1: - for dst_cfg_rank in range(1, cfg_size): - cfg_group.send_object(None, dst_cfg_rank) - elif sp_rank == 0 and cfg_size > 1: - kv_payload = cfg_group.recv_object(0) - # sp broadcast - if sp_size > 1 and sp_group is not None: - kv_payload = sp_group.broadcast_object( - kv_payload, - src=0, - ) + data = KVCacheTransferData.from_bytes_device(blob) + lb = data["layer_blocks"] + primary = SimpleNamespace(key_cache=lb.get("key_cache"), value_cache=lb.get("value_cache")) + + for key in ("past_key_values", "sp.past_key_values"): + if side.get(key) == _KV_PLACEHOLDER: + side[key] = primary + return side + + def _broadcast_kv_payload( + self, group: Any, kv_payload: dict[str, Any] | None, device: torch.device, src: int = 0 + ) -> dict[str, Any] | None: + """Broadcast KV payload via packed tensor; falls back to ``broadcast_object`` on packing failure. - if not kv_payload: - return False + The packed-tensor path uses a two-phase collective (size then data). + Both phases must always be called together — a rank entering only + one phase will deadlock the group. + """ + is_src = group.rank_in_group == src + if is_src: + combined = self._pack_kv_payload(kv_payload, device) + if combined is None: + group.broadcast_object(kv_payload, src=src) # packing failed -> object + return kv_payload + group.broadcast_object(int(combined.numel()), src=src) + group.broadcast(combined, src=src) + return kv_payload + + size_obj = group.broadcast_object(None, src=src) + if size_obj is None or not isinstance(size_obj, int): + return size_obj # src fell back to object + combined = torch.empty(size_obj, dtype=torch.uint8, device=device) + group.broadcast(combined, src=src) + return self._unpack_kv_payload(combined) + + def _send_kv_payload(self, group: Any, kv_payload: dict[str, Any] | None, dst: int, device: torch.device) -> None: + """Point-to-point send; falls back to ``send_object`` on packing failure.""" + combined = self._pack_kv_payload(kv_payload, device) + if combined is None: + group.send_object(kv_payload, dst) + return + size = int(combined.numel()) + group.send_object(size, dst) + try: + group.send(combined, dst) + except Exception: + logger.exception("D2D: send(combined) to rank %d failed; sending zero sentinel", dst) + combined.zero_() + group.send(combined, dst) + raise + + def _recv_kv_payload(self, group: Any, src: int, device: torch.device) -> dict[str, Any] | None: + """Inverse of :meth:`_send_kv_payload`.""" + size_or_obj = group.recv_object(src) + if not isinstance(size_or_obj, int): + return size_or_obj + combined = group.recv(torch.Size([size_or_obj]), torch.uint8, src) + combined = combined.to(device) + return self._unpack_kv_payload(combined) + def receive_multi_kv_cache_distributed( + self, + req: Any, + cfg_kv_collect_func: Callable | None = None, + target_device: torch.device | None = None, + ) -> bool: + """Synchronous receive: fetch → distribute → apply (FOLLOWER skips fetch).""" + received = False + if not self.topo_config.is_follower: + received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) + kv_payload = self.distribute_kv_cache(req, target_device, received=received) + if kv_payload is not None: self._apply_request_kv_payload(req, kv_payload, target_device) - return True + return received or kv_payload is not None + + def consume_and_distribute_kv_cache(self, req: Any, target_device: torch.device | None = None) -> bool: + """Consume prefetched KV → apply → distribute; sync-receive on miss (no retry on consume-error).""" + received = False + payload_consumed = False + if self._async_prefetch and not self.topo_config.is_follower: + try: + data, _ = self.consume_prefetched_kv(req) + if data is not None: + self._record_stream_for_prefetched(data) + self.apply_kv_cache_to_request(req, data) + received = True + except KVPrefetchConsumeError: + logger.error( + "KV prefetch consumed payload for %s but post-get failed; request cannot recover", + self._resolve_request_id(req), + ) + payload_consumed = True + # Only fall back to sync receive on miss, not when payload was + # already consumed from the connector (sync receive would block + # until timeout as the data is gone). + if not received and not self.topo_config.is_follower and not payload_consumed: + logger.debug("KV prefetch miss for %s; falling back to sync receive", self._resolve_request_id(req)) + received = self.receive_multi_kv_cache(req, None, target_device) + kv_payload = self.distribute_kv_cache(req, target_device, received=received) + if kv_payload is not None: + self._apply_request_kv_payload(req, kv_payload, target_device) + return received or kv_payload is not None - kv_payload: dict[str, object] | None = None - if world.rank_in_group == 0: - received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) - if received: - kv_payload = self._collect_request_kv_payload(req) + def _cfg_scatter( + self, + group: Any, + payloads: list[Any], + cfg_size: int, + device: torch.device | None, + request_id: str | None = None, + ) -> None: + """Scatter one payload per CFG follower.""" + for dst in range(1, cfg_size): + try: + self._send_kv_payload(group, payloads[dst], dst, device) + except Exception: + logger.warning("CFG scatter to rank %d failed for %s; that rank runs without KV", dst, request_id) + continue - kv_payload = world.broadcast_object(kv_payload, src=0) + def distribute_kv_cache( + self, req: Any, target_device: torch.device | None = None, *, received: bool = False + ) -> dict[str, Any] | None: + """Distribute KV LEADER→FOLLOWER. No-op for LOCAL/pure TP.""" + pt = self.topo_config - if not kv_payload: - return False + # LOCAL: no distribution. + if pt.is_local: + return None - self._apply_request_kv_payload(req, kv_payload, target_device) - return True + device = target_device if target_device is not None else torch.device("cpu") + # TP + CFG/SP + if pt.tp_active and (pt.cfg_active or pt.sp_active): + kv_payload: dict[str, object] | None = None + if pt.is_leader: + if pt.cfg_active: + payloads = ( + self._build_cfg_rank_local_payloads(req, pt.cfg_size) if received else [None] * pt.cfg_size + ) + kv_payload = payloads[0] + self._cfg_scatter(pt.cfg_group, payloads, pt.cfg_size, device, self._resolve_request_id(req)) + elif pt.sp_active: + kv_payload = self._collect_request_kv_payload(req) if received else None + if pt.sp_active and pt.sp_group is not None: + kv_payload = self._broadcast_kv_payload(pt.sp_group, kv_payload, device, src=0) + elif pt.is_follower: + if pt.cfg_active and pt.sp_rank == 0: + kv_payload = self._recv_kv_payload(pt.cfg_group, 0, device) + if pt.sp_active and pt.sp_group is not None: + kv_payload = self._broadcast_kv_payload(pt.sp_group, kv_payload, device, src=0) + else: + logger.error("distribute_kv_cache: unexpected role %s in TP+CFG/SP path", pt.role) + return None + return kv_payload or None + + # TP inactive, world broadcast + kv_payload = None + if pt.is_leader and received: + kv_payload = self._collect_request_kv_payload(req) + kv_payload = self._broadcast_kv_payload(pt.world, kv_payload, device, src=0) + return kv_payload or None + + +def _drop_prefetch_result(fut: Any) -> None: + try: + if not fut.cancelled(): + fut.result() + except Exception: + pass def _move_to_device(obj: object, device: torch.device) -> object: diff --git a/vllm_omni/distributed/omni_coordinator/__init__.py b/vllm_omni/distributed/omni_coordinator/__init__.py index 326fb9b9a16..b37c2bfb21d 100644 --- a/vllm_omni/distributed/omni_coordinator/__init__.py +++ b/vllm_omni/distributed/omni_coordinator/__init__.py @@ -8,11 +8,13 @@ RandomBalancer, RoundRobinBalancer, Task, - build_load_balancer_factory, ) from .messages import ReplicaEvent, ReplicaInfo, ReplicaList, ReplicaStatus from .omni_coord_client_for_hub import OmniCoordClientForHub -from .omni_coord_client_for_stage import OmniCoordClientForStage +from .omni_coord_client_for_stage import ( + OmniCoordClientForStage, + create_stage_coord_client, +) from .omni_coordinator import OmniCoordinator from .runtime import OmniCoordinatorRuntime @@ -24,6 +26,7 @@ "ReplicaInfo", "ReplicaList", "OmniCoordClientForStage", + "create_stage_coord_client", "OmniCoordClientForHub", "Task", "LoadBalancer", @@ -31,5 +34,4 @@ "RandomBalancer", "RoundRobinBalancer", "LeastQueueLengthBalancer", - "build_load_balancer_factory", ] diff --git a/vllm_omni/distributed/omni_coordinator/load_balancer.py b/vllm_omni/distributed/omni_coordinator/load_balancer.py index bcd9822ec69..9f21d0b9773 100644 --- a/vllm_omni/distributed/omni_coordinator/load_balancer.py +++ b/vllm_omni/distributed/omni_coordinator/load_balancer.py @@ -6,7 +6,6 @@ import random import threading from abc import ABC, abstractmethod -from collections.abc import Callable from enum import Enum from typing import Any, TypedDict @@ -122,22 +121,6 @@ def select(self, task: Task, replicas: list[ReplicaInfo]) -> int: # noqa: ARG00 return random.choice(candidates) -def build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: - """Translate ``--omni-lb-policy`` (string) into a per-pool LB factory.""" - try: - normalized = LoadBalancingPolicy(policy) - except ValueError as exc: - valid = ", ".join(p.value for p in LoadBalancingPolicy) - raise ValueError(f"unknown --omni-lb-policy {policy!r} (valid: {valid})") from exc - if normalized is LoadBalancingPolicy.RANDOM: - return RandomBalancer - if normalized is LoadBalancingPolicy.ROUND_ROBIN: - return RoundRobinBalancer - if normalized is LoadBalancingPolicy.LEAST_QUEUE_LENGTH: - return LeastQueueLengthBalancer - raise ValueError(f"unhandled load balancing policy {normalized!r}") - - __all__ = [ "Task", "LoadBalancingPolicy", @@ -145,5 +128,4 @@ def build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: "RandomBalancer", "RoundRobinBalancer", "LeastQueueLengthBalancer", - "build_load_balancer_factory", ] diff --git a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py index aee25e584a6..fb9a9cc7cfd 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coord_client_for_stage.py @@ -230,3 +230,30 @@ def close(self) -> None: except zmq.ZMQError: pass self._closed = True + + +def create_stage_coord_client( + *, + coord_zmq_addr: str, + input_addr: str, + output_addr: str, + stage_id: int, + queue_length_getter: Callable[[], int] | None = None, +) -> OmniCoordClientForStage: + """Create a stage coordinator client with an optional heartbeat hook.""" + client = OmniCoordClientForStage( + coord_zmq_addr=coord_zmq_addr, + input_addr=input_addr, + output_addr=output_addr, + stage_id=stage_id, + ) + if queue_length_getter is not None: + + def _refresh_queue_length() -> None: + try: + client._queue_length = max(int(queue_length_getter()), 0) + except Exception: + pass + + client._on_heartbeat = _refresh_queue_length + return client diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py index bbedcfeaa60..224204f2055 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -190,6 +190,15 @@ def close(self) -> None: except zmq.ZMQError: pass + def wait_for_shutdown(self) -> None: + """Block until the coordinator is asked to stop. + + ``OmniCoordinatorRuntime`` runs the coordinator in a child process. + The child needs a small, explicit blocking lifecycle API so it does + not return immediately after construction. + """ + self._stop_event.wait() + def _parse_replica_event(self, data: dict[str, Any]) -> ReplicaEvent | None: """Parse wire payload dict into ReplicaEvent. Returns None if invalid.""" try: diff --git a/vllm_omni/distributed/omni_coordinator/runtime.py b/vllm_omni/distributed/omni_coordinator/runtime.py index 71df7c0521c..2cc0d2d272c 100644 --- a/vllm_omni/distributed/omni_coordinator/runtime.py +++ b/vllm_omni/distributed/omni_coordinator/runtime.py @@ -3,35 +3,91 @@ """Lifecycle wrapper around :class:`OmniCoordinator`. -``OmniCoordinatorRuntime`` is the single-purpose owner of the head-side -coordinator process artifacts: it picks two free TCP ports, constructs an -:class:`OmniCoordinator` bound to them, exposes the resulting addresses, and -provides a single ``close()`` method to tear everything down. +``OmniCoordinatorRuntime`` spawns the coordinator as an independent process +(matching vLLM's DPCoordinator pattern). Physical isolation prevents GIL +contention and makes direct-object coupling impossible. The ROUTER address is later handed to :class:`OmniMasterServer` so it can be published to registering replicas; the PUB address is handed to the -``Orchestrator``, which constructs its :class:`OmniCoordClientForHub` against -it. +``MembershipController``, which constructs its :class:`OmniCoordClientForHub` +against it. """ from __future__ import annotations import logging +import multiprocessing +import multiprocessing.connection +import os +import signal +import weakref +from typing import Any from vllm.utils.network_utils import get_open_ports_list -from .omni_coordinator import OmniCoordinator +from vllm_omni.distributed.omni_coordinator.omni_coordinator import OmniCoordinator logger = logging.getLogger(__name__) +def run_omni_coordinator_proc( + router_zmq_addr: str, + pub_zmq_addr: str, + heartbeat_timeout: float, + ready_pipe: Any, +) -> None: + """Main loop running inside the coordinator child process.""" + signal.signal(signal.SIGINT, signal.SIG_IGN) + + coordinator = OmniCoordinator( + router_zmq_addr=router_zmq_addr, + pub_zmq_addr=pub_zmq_addr, + heartbeat_timeout=heartbeat_timeout, + ) + + ready_pipe.send("ready") + ready_pipe.close() + + coordinator.wait_for_shutdown() + + +def _get_coordinator_mp_context() -> multiprocessing.context.BaseContext: + """Return the multiprocessing context used for OmniCoordinator. + + For the current vllm-omni startup path, ``spawn`` is too expensive: the + child re-imports a heavy CLI / model stack before it can acknowledge the + ready pipe, which can exceed the coordinator startup timeout. Prefer + ``fork`` on platforms that support it. + + ``fork`` must happen before the parent initializes CUDA or owns long-lived + ZMQ sockets; otherwise the child inherits unsafe state. If coordinator + startup moves later, switch this to ``spawn``/``forkserver``. + + TODO: make the coordinator child entry cheap and spawn-safe, then revisit + whether ``spawn`` is still needed here. + """ + if os.name != "nt" and "fork" in multiprocessing.get_all_start_methods(): + return multiprocessing.get_context("fork") + return multiprocessing.get_context("spawn") + + +def _shutdown_proc(proc: multiprocessing.Process) -> None: + """Best-effort process termination for weakref finalizer.""" + if proc.is_alive(): + proc.terminate() + proc.join(timeout=5) + if proc.is_alive(): + proc.kill() + proc.join(timeout=2) + + class OmniCoordinatorRuntime: - """Own one :class:`OmniCoordinator` and the two ports it binds. + """Own one :class:`OmniCoordinator` running in a child process. - Constructor binds; :meth:`close` tears down. The class deliberately does - not expose the coordinator instance — callers should consume the - coordinator only via its wire protocol through - :class:`OmniCoordClientForStage` and :class:`OmniCoordClientForHub`. + Constructor spawns the process; :meth:`close` tears it down. + The class deliberately does not expose the coordinator instance — + callers consume it only via ZMQ through :class:`OmniCoordClientForStage` + and :class:`OmniCoordClientForHub`. """ def __init__( @@ -50,25 +106,54 @@ def __init__( self.pub_address: str = f"tcp://{host}:{pub_port}" self._closed = False - self._coordinator = OmniCoordinator( - router_zmq_addr=self.router_address, - pub_zmq_addr=self.pub_address, - heartbeat_timeout=heartbeat_timeout, + + ctx = _get_coordinator_mp_context() + parent_conn, child_conn = ctx.Pipe(duplex=False) + + self._proc: multiprocessing.Process = ctx.Process( + target=run_omni_coordinator_proc, + kwargs={ + "router_zmq_addr": self.router_address, + "pub_zmq_addr": self.pub_address, + "heartbeat_timeout": heartbeat_timeout, + "ready_pipe": child_conn, + }, + daemon=True, + name="OmniCoordinator", ) + self._proc.start() + child_conn.close() + + ready = multiprocessing.connection.wait([parent_conn, self._proc.sentinel], timeout=30) + if not ready: + self._proc.terminate() + self._proc.join(timeout=5) + raise RuntimeError("OmniCoordinator process failed to start within 30s") + + try: + status = parent_conn.recv() + except EOFError: + raise RuntimeError("OmniCoordinator process died during startup") from None + finally: + parent_conn.close() + + if status != "ready": + raise RuntimeError(f"OmniCoordinator unexpected status: {status}") + + self._finalizer = weakref.finalize(self, _shutdown_proc, self._proc) logger.info( - "[OmniCoordinatorRuntime] Started (router=%s pub=%s heartbeat_timeout=%.1fs)", + "[OmniCoordinatorRuntime] Started (pid=%d router=%s pub=%s heartbeat_timeout=%.1fs)", + self._proc.pid, self.router_address, self.pub_address, heartbeat_timeout, ) def close(self) -> None: - """Tear down the underlying coordinator. Idempotent.""" + """Tear down the coordinator process. Idempotent.""" if self._closed: return self._closed = True - try: - self._coordinator.close() - except Exception: - logger.exception("[OmniCoordinatorRuntime] coordinator close failed") + _shutdown_proc(self._proc) + self._finalizer.detach() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 148412012cf..bdf43711cd5 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -9,7 +9,7 @@ from vllm.logger import init_logger from vllm_omni.config import OmniModelConfig -from vllm_omni.engine.output_modality import OutputModality +from vllm_omni.outputs.output_modality import OutputModality from vllm_omni.platforms import current_omni_platform from vllm_omni.plugins import load_omni_general_plugins @@ -20,6 +20,8 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = { "CosyVoice3Model": "cosyvoice3", "GLMTTSForConditionalGeneration": "glm_tts", + "IndexTTS2S2MelDecoder": "indextts2", + "IndexTTS2TalkerForConditionalGeneration": "indextts2", "OmniVoiceModel": "omnivoice", "VoxCPM2TalkerForConditionalGeneration": "voxcpm2", } @@ -35,7 +37,17 @@ 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.indextts2.configuration_indextts2 import ( + IndexTTS2Config, + ) + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + MingDenseConfig, + MingMoeConfig, + ) + from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( + MossTTSLocalConfig, + MossTTSRealtimeConfig, + ) from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import ( Qwen3TTSConfig, ) @@ -57,6 +69,10 @@ def _register_omni_hf_configs() -> None: for model_type, config_cls in [ ("dense", MingDenseConfig), + ("bailingmm", MingMoeConfig), + ("indextts2", IndexTTS2Config), + ("moss_tts_local", MossTTSLocalConfig), + ("moss_tts_realtime", MossTTSRealtimeConfig), ("qwen3_tts", Qwen3TTSConfig), ("cosyvoice3", CosyVoice3Config), ("glm_tts", GLMTTSConfig), @@ -156,6 +172,8 @@ class OmniEngineArgs(EngineArgs): # in __post_init__ based on worker_type (ar/generation), so None is safe here. enable_sleep_mode: bool = False omni: bool = False + # Diffusion request-mode batch admission (forwarded to OmniDiffusionConfig). + request_batch_max_wait_ms: float = 0.0 @classmethod def _add_omni_specific_args(cls, parser: argparse.ArgumentParser) -> argparse.ArgumentParser: @@ -416,12 +434,16 @@ class OrchestratorArgs: stage_configs_path: str | None = None deploy_config: str | None = None stage_overrides: str | None = None # raw JSON string; parsed downstream + # Optional composable-parallel strategy.yaml; orchestrator reads it, overlays + # derived sizing onto merged stages, then drops it before per-stage engine args. + strategy_config: str | None = None # === Mode Switches (orchestrator reads, DeployConfig redistributes) === async_chunk: bool | None = None # === Observability === log_stats: bool = False + enable_orch_monitor: bool = False # === Headless Mode (also forwarded to engine — see SHARED_FIELDS) === stage_id: int | None = None @@ -461,6 +483,7 @@ class OrchestratorArgs: diffusion_kv_cache_skip_layers: str | None = None cfg_parallel_size: int = 1 vae_patch_parallel_size: int = 1 + vae_parallel_mode: str = "tile" default_sampling_params: str | None = None max_generated_image_size: int | None = None tts_max_instructions_length: int | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index fc3d9542707..fd3d63d7e22 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -9,18 +9,15 @@ import asyncio import concurrent.futures -import copy import dataclasses import json -import os import queue import threading import time import uuid import weakref -from collections.abc import Awaitable, Callable, Mapping, Sequence -from contextlib import ExitStack -from dataclasses import asdict, dataclass +from collections.abc import Mapping, Sequence +from dataclasses import asdict from typing import Any, Literal, cast import janus @@ -33,22 +30,10 @@ from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.input_processor import InputProcessor -from vllm_omni.config.stage_config import StageConfigFactory, strip_parent_engine_args +from vllm_omni.config.config_factory import StageConfigFactory +from vllm_omni.config.stage_config import strip_parent_engine_args from vllm_omni.diffusion.data import DiffusionParallelConfig, parse_attention_config from vllm_omni.diffusion.diffusion_engine import supports_audio_output -from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient -from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient -from vllm_omni.diffusion.stage_diffusion_proc import ( - complete_diffusion_handshake, - spawn_diffusion_proc, -) -from vllm_omni.distributed.omni_connectors.utils.initialization import ( - resolve_omni_kv_config_for_stage, -) -from vllm_omni.distributed.omni_coordinator import ( - LoadBalancer, - build_load_balancer_factory, -) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -57,73 +42,33 @@ CollectiveRPCResultMessage, EngineQueueMessage, ErrorMessage, - RegisterRemoteReplicaMessage, ShutdownRequestMessage, StageSubmissionMessage, ) from vllm_omni.engine.orchestrator import Orchestrator -from vllm_omni.engine.output_modality import FinalOutputModalityType from vllm_omni.engine.serialization import ( deserialize_additional_information, serialize_additional_information, ) from vllm_omni.engine.stage_client import StageClient -from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase -from vllm_omni.engine.stage_engine_core_proc import ( - complete_stage_handshake, - spawn_stage_core, -) -from vllm_omni.engine.stage_engine_startup import ( - OmniMasterServer, - connect_remote_engine_cores, - launch_omni_core_engines, - register_stage_with_omni_master, +from vllm_omni.engine.stage_init_utils import build_stage0_input_processor +from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.engine.stage_runtime import ( + StageRuntimeInfo, + create_stage_runtime, ) -from vllm_omni.engine.stage_init_utils import ( - LogicalStageInitPlan, - ReplicaInitPlan, - StageRemoteFactoryContext, - _inject_inferred_kv_tp_topology, - acquire_device_locks, - acquire_diffusion_device_locks, - build_diffusion_config, - build_engine_args_dict, - build_llm_stage_output_processor, - build_stage0_input_processor, - build_vllm_config, - capture_stage_factory_contexts, - compute_replica_layout, - extract_stage_metadata, - get_stage_connector_spec, - initialize_diffusion_stage, - inject_kv_stage_info, - load_omni_transfer_config_for_model, - prepare_engine_environment, - release_device_locks, - stage_runtime_setup, - terminate_alive_proc, -) -from vllm_omni.engine.stage_pool import StagePool, StagePoolClient from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import ( - inject_omni_kv_config, load_and_resolve_stage_configs, + parse_stage_overrides, ) from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.metrics.prometheus import OmniRequestCounter -from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) _STARTUP_POLL_INTERVAL_S = 1.0 - - -@dataclass(frozen=True, slots=True) -class StageRuntimeInfo: - final_output: bool - final_output_type: FinalOutputModalityType | None - stage_type: str - model_stage: str | None = None +_REQUEST_QUEUE_MAXSIZE = 256 # ============================================================================ @@ -231,7 +176,7 @@ def _weak_shutdown_async_omni_engine( try: if orchestrator_thread is not None and orchestrator_thread.is_alive(): - orchestrator_thread.join(timeout=10) + orchestrator_thread.join() except Exception: pass @@ -263,6 +208,7 @@ class AsyncOmniEngine: _log_stats: bool = False _coordinator_runtime: Any = None _transfer_emitter: Any = None + _enable_orch_monitor: bool = False def __init__( self, @@ -280,6 +226,8 @@ def __init__( self.model = model self.tokenizer = tokenizer self.diffusion_batch_size = diffusion_batch_size + # Cached by get_diffusion_od_config(). + self._diffusion_od_config_view: Any = None startup_timeout = int(init_timeout) # Forwarded into Orchestrator so its _forward_to_next_stage path can # emit per-edge transfer_tx_s / transfer_size_bytes histograms. @@ -291,6 +239,7 @@ def __init__( # replica) vllm:* wrap stays registered but reads zero. Respects the # --log-stats CLI flag set by the user via OmniBase. self._log_stats = log_stats + self._enable_orch_monitor = bool(kwargs.pop("enable_orch_monitor", False)) logger.info(f"[AsyncOmniEngine] Initializing with model {model}") @@ -309,7 +258,6 @@ def __init__( ) self._omni_master_address: str | None = kwargs.get("omni_master_address") self._omni_master_port: int | None = kwargs.get("omni_master_port") - self._omni_master_server: OmniMasterServer | None = None # New omni-coordinator flags. Consumed only in single_stage_mode. # ``omni_dp_size_local`` is process-local: each invocation (head and @@ -321,12 +269,6 @@ def __init__( self._omni_heartbeat_timeout: float = float(kwargs.get("omni_heartbeat_timeout") or 30.0) if self._omni_heartbeat_timeout <= 0: raise ValueError(f"--omni-heartbeat-timeout must be > 0, got {self._omni_heartbeat_timeout}") - # Coordinator runtime (head-distributed only). - self._coordinator_runtime: Any | None = None - # Per-stage construction context, captured after _initialize_stages - # and used by ``_build_remote_replica`` (the RemoteReplicaFactory - # passed to Orchestrator) when a headless replica registers. - self._stage_remote_factory_contexts: dict[int, StageRemoteFactoryContext] = {} if single_stage_mode: logger.info( @@ -347,9 +289,11 @@ def __init__( ) kwargs["trust_remote_code"] = trust_remote_code - - self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs) - self._validate_single_stage_mode_replica_constraints() + self.config_path, self.stage_configs = self._resolve_stage_configs( + model, + kwargs, + trust_remote_code=trust_remote_code, + ) self.num_stages = len(self.stage_configs) stage0_args = getattr(self.stage_configs[0], "engine_args", None) if self.num_stages > 0 else None @@ -357,6 +301,7 @@ def __init__( self.stage_pools: list[StagePool] = [] self.stage_clients: list[StageClient] = [] # logical-stage view for external readers self.input_processor: InputProcessor | None = None + self.prompt_expand_func: Any | None = None self.supported_tasks: tuple[str, ...] = ("generate",) self.default_sampling_params_list: list[OmniSamplingParams] = [] self.stage_metadata: list[StageRuntimeInfo] = [] @@ -365,11 +310,10 @@ def __init__( # sees a non-None ``self.request_queue`` when on_register fires. # ``async_q`` lazily binds to whatever event loop first awaits on # it (the orchestrator loop), so cross-thread use stays correct. - self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue() + self.request_queue: janus.Queue[EngineQueueMessage] = janus.Queue(maxsize=_REQUEST_QUEUE_MAXSIZE) self.output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self.rpc_output_queue: janus.Queue[EngineQueueMessage] = janus.Queue() self._shutdown_called = False - self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None self._weak_finalizer: weakref.finalize | None = None self._rpc_lock = threading.Lock() self._running_counter = OmniRequestCounter() @@ -403,977 +347,75 @@ def __init__( logger.info(f"[AsyncOmniEngine] Orchestrator ready with {self.num_stages} stages") - @staticmethod - def _cleanup_launched_llm_resources( - *, - stage_id: int, - proc: Any = None, - engine_manager: Any = None, - coordinator: Any = None, - ) -> None: - """Release launch-only LLM resources when client creation never completed.""" - - if proc is not None: - try: - terminate_alive_proc(proc) - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to terminate process for stage %s: %s", - stage_id, - cleanup_error, - ) - - for resource, resource_name in ( - (engine_manager, "engine manager"), - (coordinator, "coordinator"), - ): - if resource is None: - continue - shutdown = getattr(resource, "shutdown", None) - close = getattr(resource, "close", None) - try: - if callable(shutdown): - shutdown() - elif callable(close): - close() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to cleanup launched %s for stage %s: %s", - resource_name, - stage_id, - cleanup_error, - ) - - @staticmethod - def _collect_initialized_clients_for_cleanup( - stage_pools: Sequence[StagePool], - initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]], - ) -> list[StageClient]: - """Collect initialized clients exactly once for failure cleanup.""" - - collected: list[StageClient] = [] - seen: set[int] = set() - - def _add_client(client: StageClient | None) -> None: - if client is None: - return - client_id = id(client) - if client_id in seen: - return - seen.add(client_id) - collected.append(client) - - for pool in stage_pools: - for client in getattr(pool, "clients", ()): - _add_client(client) - - for clients in initialized_clients_by_stage.values(): - for client in clients: - _add_client(client) - - return collected - - @staticmethod - def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None: - """Best-effort shutdown for attached clients after init failure.""" - - for client in reversed(list(clients)): - if client is None: - continue - try: - client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to shutdown initialized client after init failure: %s", - cleanup_error, - ) - - def _validate_single_stage_mode_replica_constraints(self) -> None: - """Apply --omni-dp-size-local to the local stage's runtime.num_replicas. - - In the previous revision this method rejected LLM stages with - ``num_replicas > 1``. The whole point of ``--omni-dp-size-local`` is - to lift that restriction for the *local* stage, so the rejection is - gone. We now use this hook to write ``--omni-dp-size-local`` onto - the self-stage's runtime config so downstream code - (``compute_replica_layout`` → ``_build_logical_stage_init_plans``) - sees a consistent view. - """ - if not self.single_stage_mode: - return - target_stage_id = self._single_stage_id_filter - if target_stage_id is None: - return - - for idx, stage_cfg in enumerate(self.stage_configs): - stage_id = int(getattr(stage_cfg, "stage_id", idx)) - runtime_cfg = getattr(stage_cfg, "runtime", None) - if runtime_cfg is None: - continue - if stage_id == target_stage_id: - # Self stage: take --omni-dp-size-local from this process. - try: - runtime_cfg.num_replicas = self._omni_dp_size_local - except Exception: - if hasattr(runtime_cfg, "__setitem__"): - runtime_cfg["num_replicas"] = self._omni_dp_size_local - # Other stages keep their config-declared num_replicas; in - # head-distributed mode they will be launched as ``launch_mode - # == "remote"`` with the configured count. - - def _build_logical_stage_init_plans( - self, - omni_transfer_config: Any, - replicas_per_stage: Sequence[int], - replica_devices_map: Mapping[int, Sequence[str]], - ) -> tuple[list[LogicalStageInitPlan], Any]: - """Build startup plans for every logical stage and replica.""" - - prompt_expand_func = None - stage_plans: list[LogicalStageInitPlan] = [] - - for stage_idx, stage_cfg in enumerate(self.stage_configs): - base_metadata = extract_stage_metadata(stage_cfg) - configured_stage_id = base_metadata.stage_id - if base_metadata.prompt_expand_func is not None: - prompt_expand_func = base_metadata.prompt_expand_func - - stage_connector_spec = get_stage_connector_spec( - omni_transfer_config=omni_transfer_config, - stage_id=configured_stage_id, - async_chunk=self.async_chunk, - ) - omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, configured_stage_id) - num_replicas = replicas_per_stage[stage_idx] - launch_mode = "local" - if ( - self.single_stage_mode - and self._single_stage_id_filter is not None - and configured_stage_id != self._single_stage_id_filter - ): - launch_mode = "remote" - - replicas: list[ReplicaInitPlan] = [] - stage_vllm_config = None - executor_class = None - if base_metadata.stage_type != "diffusion": - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - cli_tokenizer=getattr(self, "tokenizer", None), - ) - omni_conn_cfg, omni_from, omni_to = omni_kv_connector - if omni_conn_cfg: - omni_kv = engine_args_dict.get("omni_kv_config") or {} - if not isinstance(omni_kv, dict): - omni_kv = dict(omni_kv) - omni_kv["connector_config"] = omni_conn_cfg - omni_kv["omni_from_stage"] = omni_from - omni_kv["omni_to_stage"] = omni_to - omni_kv.setdefault("stage_id", configured_stage_id) - engine_args_dict["omni_kv_config"] = omni_kv - if self.stage_configs: - _inject_inferred_kv_tp_topology( - engine_args_dict.get("omni_kv_config"), - configured_stage_id, - self.stage_configs, - ) - stage_vllm_config, executor_class = build_vllm_config( - stage_cfg, - self.model, - stage_connector_spec=stage_connector_spec, - engine_args_dict=engine_args_dict, - ) - - for replica_id in range(num_replicas): - replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg - if stage_idx in replica_devices_map: - replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] - - replica_metadata = extract_stage_metadata(replica_cfg) - replica_metadata.replica_id = replica_id - # In single_stage_mode the head only owns its self stage's - # replicas; the remote-stage metadata exists only so the - # orchestrator can route requests through StagePool. Wiping - # ``runtime_cfg`` there makes sense (we don't manage those - # devices). For the *self stage* we MUST keep the - # per-replica runtime so ``setup_stage_devices`` can apply - # the device split from ``compute_replica_layout`` — - # otherwise every replica inherits the parent's full - # CUDA_VISIBLE_DEVICES and stacks on cuda:0 (OOM with any - # model whose footprint exceeds ~1/(2N) of the card). - if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": - replica_metadata.runtime_cfg = None - - replicas.append( - ReplicaInitPlan( - replica_id=replica_id, - num_replicas=num_replicas, - launch_mode=launch_mode, - stage_cfg=replica_cfg, - metadata=replica_metadata, - stage_connector_spec=stage_connector_spec, - omni_kv_connector=omni_kv_connector, - stage_vllm_config=stage_vllm_config, - executor_class=executor_class, - ) - ) - - stage_plans.append( - LogicalStageInitPlan( - stage_idx=stage_idx, - configured_stage_id=configured_stage_id, - replicas=replicas, - ) - ) - - return stage_plans, prompt_expand_func - - def _start_omni_master_server(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: - """Start OmniMasterServer for single-stage mode.""" - - if not self._omni_master_address or not self._omni_master_port: - raise ValueError( - "AsyncOmniEngine single_stage_mode requires both omni_master_address and omni_master_port to be set." - ) - - all_stage_ids: list[int] = [] - stage_replica_counts: dict[int, int] = {} - # Slots that the head itself will register (launch_mode == "local") - # — auto-assign on the master must not hand these out to remote - # headless registrations even within the race window before the - # head's own register_stage_with_omni_master call completes. - head_local_replicas: dict[int, list[int]] = {} - seen_stage_ids: set[int] = set() - for plan in stage_plans: - stage_id = plan.configured_stage_id - if stage_id in seen_stage_ids: - raise ValueError( - f"Duplicate stage_id {stage_id!r} detected among configured stages; stage_ids must be unique." - ) - seen_stage_ids.add(stage_id) - all_stage_ids.append(stage_id) - stage_replica_counts[stage_id] = len(plan.replicas) - local_rids = [rep.replica_id for rep in plan.replicas if rep.launch_mode == "local"] - if local_rids: - head_local_replicas[stage_id] = local_rids - - # Start the OmniCoordinator runtime first so its router address is - # available to publish in every registration reply. - from vllm_omni.distributed.omni_coordinator import OmniCoordinatorRuntime - - self._coordinator_runtime = OmniCoordinatorRuntime( - host=self._omni_master_address, - heartbeat_timeout=self._omni_heartbeat_timeout, - ) - - self._omni_master_server = OmniMasterServer( - master_address=self._omni_master_address, - master_port=self._omni_master_port, - stage_ids=all_stage_ids, - stage_replica_counts=stage_replica_counts, - coordinator_router_address=self._coordinator_runtime.router_address, - on_register=self._dispatch_master_register, - head_local_replicas=head_local_replicas, - ) - self._omni_master_server.start() - logger.info( - "[AsyncOmniEngine] OmniMasterServer started for stages %s", - all_stage_ids, - ) - - # ------------------------------------------------------------------ - # Remote replica factory (head-side client construction) - # ------------------------------------------------------------------ + def get_diffusion_od_config(self) -> Any: + """Expose the diffusion ``model_class_name`` to client-side model-extras. - async def _build_remote_replica(self, stage_id: int, replica_id: int) -> Any: - """Construct a head-side stage client for a newly-registered remote replica. - - Used by :class:`Orchestrator` as its ``remote_replica_factory``. - The orchestrator awaits this from its own asyncio loop, so the - client is created in the same loop that owns ZMQ sockets — no - cross-thread setup is required. - - Raises if the stage is not known or if its construction context - was not captured (e.g. the stage was empty at bring-up time). + The worker holds the full config; here we just resolve the pipeline class + name from the model config (cached). ``model_class_name`` may be ``None``. """ - ctx = self._stage_remote_factory_contexts.get(stage_id) - if ctx is None: - raise RuntimeError( - f"no factory context captured for stage {stage_id}; " - f"known stages: {sorted(self._stage_remote_factory_contexts.keys())}" - ) - if self._omni_master_server is None: - raise RuntimeError("OmniMasterServer is not running; cannot build remote replica") + if self._diffusion_od_config_view is None: + from types import SimpleNamespace - alloc = self._omni_master_server.get_allocation(stage_id, replica_id=replica_id) + from vllm_omni.diffusion.data import resolve_model_class_name - # Build a per-replica copy of the base metadata so ``replica_id`` - # is correct (StageMetadata is a plain dataclass-like). - metadata = copy.copy(ctx.base_metadata) - try: - metadata.replica_id = replica_id - except Exception: - # Best-effort: if the metadata object is frozen / unusual, the - # downstream client will fall back to ``replica_id = 0``. - pass - - if ctx.stage_type == "diffusion": - client = StageDiffusionClient.from_addresses( - metadata, - request_address=alloc.input_bind_address, - response_address=alloc.output_bind_address, - batch_size=ctx.diffusion_batch_size, - ) - logger.info( - "[AsyncOmniEngine] Built remote diffusion client for stage=%d replica=%d (req=%s resp=%s)", - stage_id, - replica_id, - alloc.input_bind_address, - alloc.output_bind_address, - ) - return client - - # LLM path - if ctx.vllm_config is None or ctx.executor_class is None: - raise RuntimeError(f"stage {stage_id} factory context is missing vllm_config / executor_class") - - # The headless's StageEngineCoreProc subprocess calls - # vllm.v1.engine.core.startup_handshake at boot and blocks until the - # head's handshake ROUTER answers — without this step it hits the - # built-in 5-minute timeout and exits. The bootstrap (pre-allocated) - # path runs `connect_remote_engine_cores` to perform that - # handshake; dynamic attach must do the same, otherwise every - # replica that comes in via `on_register` (auto-assigned or - # explicit-id-beyond-pre-alloc) deadlocks. Run the blocking - # handshake in a thread so the orchestrator loop stays responsive, - # then build the async client on this loop where `make_async_mp_client` - # expects to be invoked. - master_server = self._omni_master_server - ctx_vllm_config = ctx.vllm_config - ctx_executor_class = ctx.executor_class - - def _run_handshake() -> Any: - with connect_remote_engine_cores( - vllm_config=ctx_vllm_config, - omni_master_server=master_server, - stage_id=stage_id, - replica_id=replica_id, - ) as remote_resources: - _engine_manager, _coordinator, addresses, _ = remote_resources - return _engine_manager, _coordinator, addresses - - engine_manager, coordinator, addresses = await asyncio.to_thread(_run_handshake) - - client_addresses: dict[str, str] = { - "input_address": addresses.inputs[0], - "output_address": addresses.outputs[0], - } - if addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address - - client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=ctx_vllm_config, - executor_class=ctx_executor_class, - metadata=metadata, - client_addresses=client_addresses, - proc=None, - engine_manager=engine_manager, - coordinator=coordinator, - ) - logger.info( - "[AsyncOmniEngine] Built remote LLM client for stage=%d replica=%d (input=%s)", - stage_id, - replica_id, - client_addresses["input_address"], - ) - return client - - # ------------------------------------------------------------------ - # OmniCoordinator on_register proxy - # ------------------------------------------------------------------ - - def _dispatch_master_register(self, stage_id: int, replica_id: int, alloc: Any) -> None: - """Forward a master-server registration to the orchestrator queue. - - Called on :class:`OmniMasterServer`'s ROUTER thread (not the - orchestrator loop). Must return promptly. ``self.request_queue`` - is initialized in :meth:`_initialize_janus_queues` (running inside - the orchestrator thread) before the master server starts, so it - is non-None here; ``janus.Queue.sync_q`` is thread-safe by - construction. - """ - if self.request_queue is None: - logger.warning( - "[AsyncOmniEngine] request_queue not initialized; dropping register_remote_replica stage=%d replica=%d", - stage_id, - replica_id, - ) - return - msg = RegisterRemoteReplicaMessage( - stage_id=int(stage_id), - replica_id=int(replica_id), - ) - try: - self.request_queue.sync_q.put_nowait(msg) - except Exception: - logger.exception( - "[AsyncOmniEngine] Failed to enqueue register_remote_replica for stage=%d replica=%d", - stage_id, - replica_id, - ) - - def _initialize_llm_replica( - self, - plan: ReplicaInitPlan, - stage_init_timeout: int, - llm_stage_launch_lock: threading.Lock, - ) -> StageEngineCoreClientBase: - """Initialize one LLM replica end-to-end.""" - - proc = None - engine_manager = None - coordinator = None - stage_client = None - lock_fds: list[int] = [] - device_control_env = current_omni_platform.device_control_env_var - stage_cfg = plan.stage_cfg - - try: - if plan.launch_mode == "remote": - assert self._omni_master_server is not None - raw_stage_cfg = self._omni_master_server.get_stage_config( - plan.metadata.stage_id, - timeout_s=stage_init_timeout, - replica_id=plan.replica_id, - ) - if raw_stage_cfg is None: - raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") - vllm_config = plan.stage_vllm_config - executor_class = plan.executor_class - assert vllm_config is not None - assert executor_class is not None - vllm_config.parallel_config.data_parallel_size_local = 0 - launch_cm = connect_remote_engine_cores( - vllm_config=vllm_config, - omni_master_server=self._omni_master_server, - stage_id=plan.metadata.stage_id, - replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s remote engine handshake started", - plan.metadata.stage_id, - ) - with launch_cm as remote_resources: - engine_manager, coordinator, addresses, _tensor_queue = remote_resources - - logger.info( - "[AsyncOmniEngine] Stage %s remote engine startup completed", - plan.metadata.stage_id, - ) - client_addresses: dict[str, str] = { - "input_address": addresses.inputs[0], - "output_address": addresses.outputs[0], - } - if addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address - stage_client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=vllm_config, - executor_class=executor_class, - metadata=plan.metadata, - client_addresses=client_addresses, - engine_manager=engine_manager, - coordinator=coordinator, - ) - else: - handshake_address = None - with ExitStack() as launch_stack: - with llm_stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - with stage_runtime_setup( - plan.metadata.stage_id, - plan.metadata.runtime_cfg, - ): - vllm_config = plan.stage_vllm_config - executor_class = plan.executor_class - assert vllm_config is not None - assert executor_class is not None - engine_args_dict = build_engine_args_dict( - stage_cfg, - self.model, - stage_connector_spec=plan.stage_connector_spec, - cli_tokenizer=getattr(self, "tokenizer", None), - ) - lock_fds = acquire_device_locks( - plan.metadata.stage_id, - engine_args_dict, - stage_init_timeout, - ) - if self.single_stage_mode and self._omni_master_server is not None: - coord_router_addr: str | None = ( - self._coordinator_runtime.router_address - if self._coordinator_runtime is not None - else None - ) - engine_manager, coordinator, addresses = launch_stack.enter_context( - launch_omni_core_engines( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - omni_master_server=self._omni_master_server, - stage_id=plan.metadata.stage_id, - stage_config=stage_cfg, - replica_id=plan.replica_id, - omni_coordinator_address=coord_router_addr, - ) - ) - else: - addresses, proc, handshake_address = spawn_stage_core( - vllm_config=vllm_config, - executor_class=executor_class, - log_stats=self._log_stats, - ) - logger.info( - "[AsyncOmniEngine] Stage %s engine launch started", - plan.metadata.stage_id, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - - if self.single_stage_mode and self._omni_master_server is not None: - launch_stack.close() - else: - assert proc is not None - assert handshake_address is not None - complete_stage_handshake(proc, handshake_address, addresses, vllm_config, stage_init_timeout) - logger.info( - "[AsyncOmniEngine] Stage %s engine startup completed", - plan.metadata.stage_id, - ) - - client_addresses: dict[str, str] = { - "input_address": addresses.inputs[0], - "output_address": addresses.outputs[0], - } - if addresses.frontend_stats_publish_address is not None: - client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address - stage_client = StageEngineCoreClientBase.make_async_mp_client( - vllm_config=vllm_config, - executor_class=executor_class, - metadata=plan.metadata, - client_addresses=client_addresses, - proc=proc, - engine_manager=engine_manager, - coordinator=coordinator, - ) - - logger.info("[AsyncOmniEngine] Stage %s initialized", plan.metadata.stage_id) - return stage_client - except Exception: - if stage_client is not None: - try: - stage_client.shutdown() - except Exception as cleanup_error: - logger.warning( - "[AsyncOmniEngine] Failed to cleanup stage %s after attach failure: %s", - plan.metadata.stage_id, - cleanup_error, - ) - else: - self._cleanup_launched_llm_resources( - stage_id=plan.metadata.stage_id, - proc=proc, - engine_manager=engine_manager, - coordinator=coordinator, - ) - raise - finally: - if lock_fds: - release_device_locks(lock_fds) - - def _initialize_diffusion_replica( - self, - plan: ReplicaInitPlan, - stage_init_timeout: int, - stage_launch_lock: threading.Lock, - ) -> StageDiffusionClient | InlineStageDiffusionClient: - """Initialize one diffusion replica end-to-end.""" - - client = None - proc = None - lock_fds: list[int] = [] - try: - if plan.launch_mode == "remote": - assert self._omni_master_server is not None - remote_stage_cfg = OmegaConf.create( - self._omni_master_server.get_stage_config( - plan.metadata.stage_id, - timeout_s=stage_init_timeout, - replica_id=plan.replica_id, - ) - ) - remote_metadata = extract_stage_metadata(remote_stage_cfg) - addresses = self._omni_master_server.get_zmq_addresses( - plan.metadata.stage_id, - replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s remote diffusion startup completed", - plan.metadata.stage_id, - ) - client = StageDiffusionClient.from_addresses( - remote_metadata, - request_address=addresses.inputs[0], - response_address=addresses.outputs[0], - batch_size=self.diffusion_batch_size, - ) - else: - device_control_env = current_omni_platform.device_control_env_var - with stage_launch_lock: - previous_visible_devices = os.environ.get(device_control_env) - try: - with stage_runtime_setup( - plan.metadata.stage_id, - plan.metadata.runtime_cfg, - ): - omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector - if omni_conn_cfg: - inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) - inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self.stage_configs) - if self.single_stage_mode: - assert self._omni_master_server is not None - od_config = build_diffusion_config(self.model, plan.stage_cfg, plan.metadata) - lock_fds = acquire_diffusion_device_locks( - plan.metadata.stage_id, - od_config, - stage_init_timeout, - ) - handshake_address, request_address, response_address = register_stage_with_omni_master( - omni_master_address=self._omni_master_server.address, - omni_master_port=self._omni_master_server.port, - omni_stage_id=plan.metadata.stage_id, - omni_stage_config=plan.stage_cfg, - return_addresses=True, - replica_id=plan.replica_id, - ) - logger.info( - "[AsyncOmniEngine] Stage %s diffusion registration completed", - plan.metadata.stage_id, - ) - coord_router_addr: str | None = ( - self._coordinator_runtime.router_address - if self._coordinator_runtime is not None - else None - ) - proc, _, _, _ = spawn_diffusion_proc( - self.model, - od_config, - handshake_address=handshake_address, - request_address=request_address, - response_address=response_address, - omni_coordinator_address=coord_router_addr, - omni_stage_id=plan.metadata.stage_id, - omni_replica_id=plan.replica_id, - ) - complete_diffusion_handshake(proc, handshake_address, stage_init_timeout) - logger.info( - "[AsyncOmniEngine] Stage %s diffusion startup completed", - plan.metadata.stage_id, - ) - client = StageDiffusionClient.from_addresses( - plan.metadata, - request_address=request_address, - response_address=response_address, - proc=proc, - batch_size=self.diffusion_batch_size, - ) - else: - client = initialize_diffusion_stage( - plan.metadata.stage_id, - self.model, - plan.stage_cfg, - plan.metadata, - stage_init_timeout=stage_init_timeout, - batch_size=self.diffusion_batch_size, - use_inline=self.num_stages == 1 and plan.num_replicas == 1, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - - logger.info( - "[AsyncOmniEngine] Stage %s replica %s initialized (diffusion, batch_size=%d, devices=%s)", - plan.metadata.stage_id, - plan.replica_id, - self.diffusion_batch_size, - getattr(getattr(plan.stage_cfg, "runtime", None), "devices", "default"), - ) - return client - except Exception: - if proc is not None: - terminate_alive_proc(proc) - raise - finally: - if lock_fds: - release_device_locks(lock_fds) - - def _initialize_replica( - self, - plan: ReplicaInitPlan, - stage_init_timeout: int, - stage_launch_lock: threading.Lock, - ) -> StagePoolClient: - """Initialize one replica, regardless of backend type.""" - - if plan.metadata.stage_type == "diffusion": - return self._initialize_diffusion_replica(plan, stage_init_timeout, stage_launch_lock) - return self._initialize_llm_replica(plan, stage_init_timeout, stage_launch_lock) - - def _initialize_stage_replicas( - self, - stage_plans: Sequence[LogicalStageInitPlan], - stage_init_timeout: int, - ) -> dict[int, list[StagePoolClient | None]]: - """Initialize all stage replicas with parallel startup. - - Diffusion replicas are launched **inline on the orchestrator thread** - (the long-lived daemon thread created in ``__init__``). Their - ``mp.Process`` workers are therefore parented by a thread whose - lifetime equals the engine's lifetime. Submitting diffusion init to a - scoped ``ThreadPoolExecutor`` causes the clone-parent Python thread to - be destroyed at the end of init, which under Ray's actor subreaper - leads the spawned ``DiffusionWorker`` processes to be silently - ``SIGKILL``ed (exitcode -9). - - LLM replicas are submitted to a **long-lived** engine-level executor - (``_stage_init_executor``) whose threads outlive the init phase, - preventing ``prctl(PR_SET_PDEATHSIG, SIGTERM)`` from firing when the - spawning thread exits. - - To achieve parallel startup, LLM replicas are submitted to the - executor first so they begin loading in background, then diffusion - replicas run inline on the orchestrator thread concurrently, and - finally LLM futures are awaited. - """ - - stage_launch_lock = threading.Lock() - initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { - plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans - } - primary_exc: Exception | None = None - - # Partition replicas: diffusion runs inline on the orchestrator thread; - # LLM replicas are submitted to the engine-level executor. - diffusion_replicas: list[tuple[int, ReplicaInitPlan]] = [] - llm_replicas: list[tuple[int, ReplicaInitPlan]] = [] - for plan in stage_plans: - for replica in plan.replicas: - if replica.metadata.stage_type == "diffusion": - diffusion_replicas.append((plan.stage_idx, replica)) - else: - llm_replicas.append((plan.stage_idx, replica)) - - # Lazily create the engine-level executor (once per engine lifetime). - if llm_replicas and getattr(self, "_stage_init_executor", None) is None: - self._stage_init_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=max(1, len(llm_replicas)), - thread_name_prefix="stage-init", - ) - - # --- 1) Submit LLM replicas to the long-lived executor (starts - # loading in background). --- - future_to_replica: dict[concurrent.futures.Future[StagePoolClient], tuple[int, int]] = {} - for stage_idx, replica in llm_replicas: - future = self._stage_init_executor.submit( - self._initialize_replica, - replica, - stage_init_timeout, - stage_launch_lock, - ) - future_to_replica[future] = (stage_idx, replica.replica_id) - - # --- 2) Diffusion replicas: inline on the orchestrator thread. - # LLM is already loading in background — both run concurrently. --- - if primary_exc is None: - for stage_idx, replica in diffusion_replicas: - try: - initialized_clients_by_stage[stage_idx][replica.replica_id] = self._initialize_replica( - replica, - stage_init_timeout, - stage_launch_lock, - ) - except Exception as exc: - primary_exc = exc - break - - # --- 3) Await LLM futures. --- - if primary_exc is None and future_to_replica: - for future in concurrent.futures.as_completed(future_to_replica): - stage_idx, replica_id = future_to_replica[future] - try: - initialized_clients_by_stage[stage_idx][replica_id] = future.result() - except concurrent.futures.CancelledError: - continue - except Exception as exc: - if primary_exc is None: - primary_exc = exc - for other_future in future_to_replica: - if other_future is future: - continue - other_future.cancel() - - if primary_exc is not None: - setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage) - raise primary_exc - - return initialized_clients_by_stage - - def _assemble_stage_pools( - self, - stage_plans: Sequence[LogicalStageInitPlan], - initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]], - ) -> list[StagePool]: - """Assemble logical stage pools and update top-level stage metadata.""" - - stage_pools: list[StagePool] = [] - default_sampling_params_list: list[OmniSamplingParams] = [] - stage_metadata_list: list[StageRuntimeInfo] = [] - - for plan in stage_plans: - replica_clients = initialized_clients_by_stage[plan.stage_idx] - first_client = replica_clients[0] if replica_clients else None - if first_client is None: - raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client") - - clients: list[StagePoolClient] = [client for client in replica_clients if client is not None] - stage_vllm_config = None - output_processor = None - if plan.replicas[0].metadata.stage_type != "diffusion": - stage_vllm_config = plan.replicas[0].stage_vllm_config - assert stage_vllm_config is not None - output_processor = build_llm_stage_output_processor(plan, stage_vllm_config, log_stats=self._log_stats) - - stage_pools.append( - StagePool( - plan.stage_idx, - clients, - output_processor=output_processor, - stage_vllm_config=stage_vllm_config, - ) - ) - default_sampling_params_list.append(first_client.default_sampling_params) - stage_metadata_list.append( - StageRuntimeInfo( - final_output=first_client.final_output, - final_output_type=first_client.final_output_type, - stage_type=first_client.stage_type, - model_stage=getattr(first_client, "model_stage", None), - ) - ) - - self.default_sampling_params_list = list(default_sampling_params_list) - self.stage_metadata = list(stage_metadata_list) - return stage_pools + self._diffusion_od_config_view = SimpleNamespace(model_class_name=resolve_model_class_name(self.model)) + return self._diffusion_od_config_view def _initialize_stages(self, stage_init_timeout: int) -> None: - """Initialize stage clients/processors in orchestrator thread and assign to self. - - Phases: - 1. Compute replica layout (counts + device splits). - 2. Build per-stage/per-replica startup plans. - 3. Initialize all replicas in parallel via backend-specific launchers. - 4. Build logical StagePools and finalize runtime metadata. - - TODO(stage-pool): move per-stage launch + attach logic into a - StagePool.build_from_config() classmethod so this method only - iterates stage_configs, collects pools, and finalizes metadata. - """ - num_stages = len(self.stage_configs) - self.num_stages = num_stages - self._validate_single_stage_mode_replica_constraints() - - replicas_per_stage, replica_devices_map = compute_replica_layout(self.stage_configs) - - prepare_engine_environment() - omni_transfer_config = load_omni_transfer_config_for_model(self.model, self.config_path) - stage_plans, prompt_expand_func = self._build_logical_stage_init_plans( - omni_transfer_config, - replicas_per_stage, - replica_devices_map, + """Initialize stage clients/processors via StageRuntime and assign to self.""" + self._runtime = create_stage_runtime( + stage_configs=self.stage_configs, + model=self.model, + config_path=self.config_path, + single_stage_mode=self.single_stage_mode, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=self.diffusion_batch_size, + async_chunk=self.async_chunk, + tokenizer=self.tokenizer, + single_stage_id_filter=self._single_stage_id_filter, + omni_master_address=self._omni_master_address, + omni_master_port=self._omni_master_port, + omni_dp_size_local=self._omni_dp_size_local, + omni_heartbeat_timeout=self._omni_heartbeat_timeout, + omni_lb_policy=self._omni_lb_policy, + request_queue=self.request_queue, ) - # Capture per-stage context now (before _start_omni_master_server) - # so the on_register proxy can build head-side clients for - # registrations that arrive immediately after the server starts. - if self.single_stage_mode: - self._stage_remote_factory_contexts = capture_stage_factory_contexts( - stage_plans, diffusion_batch_size=self.diffusion_batch_size - ) - self._start_omni_master_server(stage_plans) - - stage_pools: list[StagePool] = [] - input_processor: InputProcessor | None = None - initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { - plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans - } + self._runtime.initialize() - try: - initialized_clients_by_stage = self._initialize_stage_replicas(stage_plans, stage_init_timeout) - if stage_plans and stage_plans[0].replicas[0].metadata.stage_type != "diffusion": - stage0_vllm_config = stage_plans[0].replicas[0].stage_vllm_config - assert stage0_vllm_config is not None - input_processor = build_stage0_input_processor(stage0_vllm_config) - stage_pools = self._assemble_stage_pools(stage_plans, initialized_clients_by_stage) - except Exception as exc: - initialized_clients_by_stage = getattr( - exc, - "_initialized_clients_by_stage", - initialized_clients_by_stage, - ) - cleanup_clients = self._collect_initialized_clients_for_cleanup( - stage_pools, - initialized_clients_by_stage, - ) - logger.exception( - "[AsyncOmniEngine] Stage initialization failed; shutting down %s initialized client(s)", - len(cleanup_clients), - ) - self._shutdown_initialized_clients(cleanup_clients) - if self._omni_master_server is not None: - try: - self._omni_master_server.stop() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during stage-init cleanup") - if self._coordinator_runtime is not None: - try: - self._coordinator_runtime.close() - except Exception: - logger.exception( - "[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during stage-init cleanup" - ) - self._coordinator_runtime = None - raise - - self.stage_pools = stage_pools - self.input_processor = input_processor - self.prompt_expand_func = prompt_expand_func - - # Derive logical-stage views for external readers (entrypoints/async_omni.py). - self.stage_clients = [cast(StageClient, pool.stage_client) for pool in self.stage_pools] + self.num_stages = len(self.stage_configs) + self.stage_pools = self._runtime.stage_pools + self.stage_clients = [ + cast(StageClient, pool.stage_client) for pool in self.stage_pools if pool.stage_client is not None + ] self.stage_vllm_configs = [pool.stage_vllm_config for pool in self.stage_pools] self.output_processors = [pool.output_processor for pool in self.stage_pools] - - # TODO(Peiqi): Hack here + self.input_processor = ( + build_stage0_input_processor(self.stage_vllm_configs[0]) + if self.stage_vllm_configs and self.stage_vllm_configs[0] is not None + else None + ) + self.prompt_expand_func = next( + ( + getattr(client, "prompt_expand_func", None) + for client in self.stage_clients + if getattr(client, "prompt_expand_func", None) is not None + ), + None, + ) + self.default_sampling_params_list = [client.default_sampling_params for client in self.stage_clients] + self.stage_metadata = [ + StageRuntimeInfo( + final_output=client.final_output, + final_output_type=client.final_output_type, + stage_type=client.stage_type, + model_stage=getattr(client, "model_stage", None), + ) + for client in self.stage_clients + ] supported_tasks: set[str] = set() - if any(getattr(pool.stage_client, "is_comprehension", False) for pool in self.stage_pools): + if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): supported_tasks.add("generate") - if any(m.final_output_type == "audio" for m in self.stage_metadata): + if any(meta.final_output_type == "audio" for meta in self.stage_metadata): supported_tasks.add("speech") self.supported_tasks = tuple(supported_tasks) if supported_tasks else ("generate",) @@ -1391,13 +433,9 @@ 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 - remote_replica_factory: Callable[[int, int], Awaitable[Any]] | None = None - if self._coordinator_runtime is not None: - coordinator_pub_address = self._coordinator_runtime.pub_address - load_balancer_factory = build_load_balancer_factory(self._omni_lb_policy) - remote_replica_factory = self._build_remote_replica + + membership_controller = self._runtime.create_membership_controller() + orchestrator = Orchestrator( request_async_queue=self.request_queue.async_q, output_async_queue=self.output_queue.async_q, @@ -1405,12 +443,11 @@ async def _run_orchestrator() -> None: stage_pools=self.stage_pools, async_chunk=self.async_chunk, pd_config=pd_config, + membership_controller=membership_controller, running_counter=self._running_counter, - coordinator_pub_address=coordinator_pub_address, - load_balancer_factory=load_balancer_factory, - remote_replica_factory=remote_replica_factory, transfer_emitter=self._transfer_emitter, log_stats=self._log_stats, + enable_orch_monitor=self._enable_orch_monitor, ) if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) @@ -1506,7 +543,7 @@ def _ensure_stage_replica_mm_uuids( return mm_data = prompt.get("multi_modal_data") - if not isinstance(mm_data, Mapping) or not mm_data: + if not isinstance(mm_data, dict) or not mm_data: return from vllm.multimodal.hasher import MultiModalHasher @@ -1787,7 +824,7 @@ def _enqueue_cfg_companions( # Registration of this companion on stage-0's output processor is # deferred to Orchestrator._handle_add_companion, which routes # admission through StagePool.submit_initial(..., affinity_request_id=...). - self.request_queue.sync_q.put_nowait( + self.request_queue.sync_q.put( AddCompanionRequestMessage( companion_id=cid, parent_id=parent_id, @@ -1828,6 +865,14 @@ def _get_default_cache_config(cache_backend: str | None) -> dict[str, Any] | Non "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1, } + if cache_backend in ("step_cache"): + return { + "step_cache_dit_enabled": True, + "velocity_sim_thresholds": [0.95, 0.93], + "velocity_skip_countdowns": [4, 2], + "step_cache_dit_min_history": 2, + "step_cache_dit_max_history": 2, + } return None @staticmethod @@ -1931,10 +976,11 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 + vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile" enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False - use_hsdp = normalized_kwargs.get("use_hsdp") or False - hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size") or -1 - hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size") or 1 + use_hsdp = normalized_kwargs.get("use_hsdp", False) + hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) + hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) if sequence_parallel_size is None: sequence_parallel_size = ulysses_degree * ring_degree @@ -1949,6 +995,7 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: ulysses_mode=ulysses_mode, cfg_parallel_size=cfg_parallel_size, vae_patch_parallel_size=vae_patch_parallel_size, + vae_parallel_mode=vae_parallel_mode, use_hsdp=use_hsdp, hsdp_shard_size=hsdp_shard_size, hsdp_replicate_size=hsdp_replicate_size, @@ -1976,6 +1023,7 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: "model_config": kwargs.get("model_config", None), "additional_config": kwargs.get("additional_config", None), "step_execution": kwargs.get("step_execution", False), + "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), "vae_use_slicing": kwargs.get("vae_use_slicing", False), "vae_use_tiling": kwargs.get("vae_use_tiling", False), "cache_backend": cache_backend, @@ -1999,12 +1047,14 @@ def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), "quantization": kwargs.get("quantization", None), + "quantization_config": kwargs.get("quantization_config", None), "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), + "streaming_output": kwargs.get("diffusion_streaming_output", False), "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), "extras": { "auxiliary_text_encoder": kwargs.get("auxiliary_text_encoder", None), @@ -2081,11 +1131,47 @@ def _strip_single_engine_args(kwargs: dict[str, Any]) -> dict[str, Any]: return result - def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]: + def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: + """Apply a strategy-derived ``omni_lb_policy`` to the engine. + + Precedence: an explicit ``--omni-lb-policy`` always wins. ``"random"`` is + the engine default and is treated as "unset" (indistinguishable from no + flag), so a strategy value overrides it. If the user explicitly passed a + non-default policy that conflicts with the strategy-derived one, raise so + the mismatch is not silently ignored. + """ + if not derived: + return + explicit = kwargs.get("omni_lb_policy") + user_set = explicit is not None and str(explicit) != "random" + if user_set: + if str(explicit) != str(derived): + raise ValueError( + f"Conflicting load-balancer policy: --omni-lb-policy={explicit!r} was given " + f"but the composable-parallel strategy derived omni_lb_policy={derived!r}. " + "Drop --omni-lb-policy to use the strategy value, or make them match." + ) + return + if self._omni_lb_policy != str(derived): + logger.info( + "[composable_parallel] applying strategy-derived omni_lb_policy=%r (was %r).", + derived, + self._omni_lb_policy, + ) + self._omni_lb_policy = str(derived) + + def _resolve_stage_configs( + self, + model: str, + kwargs: dict[str, Any], + *, + trust_remote_code: bool, + ) -> tuple[str, list[Any]]: """Resolve stage configs and inject defaults shared by orchestrator/headless.""" stage_configs_path = kwargs.get("stage_configs_path", None) deploy_config_path = kwargs.pop("deploy_config", None) + strategy_config_path = kwargs.pop("strategy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) explicit_stage_configs = kwargs.pop("stage_configs", None) if explicit_stage_configs is not None: @@ -2100,27 +1186,24 @@ def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[st base_kwargs = kwargs # Parse --stage-overrides JSON string if provided - stage_overrides = None - if stage_overrides_json: - if isinstance(stage_overrides_json, str): - try: - stage_overrides = json.loads(stage_overrides_json) - except json.JSONDecodeError as exc: - raise ValueError( - f"--stage-overrides is not valid JSON: {exc}. Got: {stage_overrides_json!r}" - ) from exc - else: - stage_overrides = stage_overrides_json - - config_path, stage_configs = load_and_resolve_stage_configs( + stage_overrides = parse_stage_overrides(stage_overrides_json) + + config_path, stage_configs, strategy_lb_policy = load_and_resolve_stage_configs( model, stage_configs_path, base_kwargs, + trust_remote_code=trust_remote_code, default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, ) + # A strategy.yaml may derive a pipeline-wide load-balancer policy. It is + # an orchestrator-level knob (read once at construction), so apply it here + # rather than as a per-stage config field. + self._apply_strategy_lb_policy(strategy_lb_policy, kwargs) + # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config. for cfg in stage_configs: try: @@ -2162,7 +1245,7 @@ def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[st kwargs.get("diffusion_attention_config"), attention_backend=kwargs.get("diffusion_attention_backend"), ) - quantization_config = kwargs.get("diffusion_quantization_config") + quantization_config = kwargs.get("diffusion_quantization_config") or kwargs.get("quantization_config") if quantization_config is not None: if ( not hasattr(cfg.engine_args, "quantization_config") @@ -2253,7 +1336,7 @@ def add_request( reasoning_ended=reasoning_ended, resumable=resumable, ) - self.request_queue.sync_q.put_nowait(msg) + self.request_queue.sync_q.put(msg) # CFG companion expansion: create and enqueue companion requests # so the AR stage also generates their KV caches. @@ -2324,7 +1407,7 @@ def add_streaming_update( resumable=resumable, message_type="streaming_update", ) - self.request_queue.sync_q.put_nowait(msg) + self.request_queue.sync_q.put(msg) async def add_streaming_update_async( self, @@ -2376,7 +1459,7 @@ def abort(self, request_ids: list[str]) -> None: """Send abort message to the Orchestrator.""" if self.request_queue is None: raise RuntimeError("request_queue is not initialized") - self.request_queue.sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids)) + self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) async def abort_async(self, request_ids: list[str]) -> None: """Async abort API.""" @@ -2406,7 +1489,7 @@ def collective_rpc( ) with self._rpc_lock: - self.request_queue.sync_q.put_nowait(msg) + self.request_queue.sync_q.put(msg) deadline = None if timeout is None else time.monotonic() + timeout while True: @@ -2471,20 +1554,10 @@ def shutdown(self) -> None: finalizer.detach() logger.info("[AsyncOmniEngine] Shutting down Orchestrator") - try: - if self.request_queue is not None: - self.request_queue.sync_q.put_nowait(ShutdownRequestMessage()) - except Exception: - pass + if self.request_queue is not None: + self.request_queue.sync_q.put_nowait(ShutdownRequestMessage()) if self.is_alive(): - self.orchestrator_thread.join(timeout=10) - if self.orchestrator_thread.is_alive(): - logger.warning("[AsyncOmniEngine] Orchestrator thread did not exit in time") - - stage_init_executor = getattr(self, "_stage_init_executor", None) - if stage_init_executor is not None: - stage_init_executor.shutdown(wait=False, cancel_futures=True) - self._stage_init_executor = None + self.orchestrator_thread.join() for q in (self.request_queue, self.output_queue, self.rpc_output_queue): try: @@ -2492,19 +1565,11 @@ def shutdown(self) -> None: except Exception: pass - if self._omni_master_server is not None: - try: - self._omni_master_server.stop() - except Exception: - logger.exception("[AsyncOmniEngine] Failed to stop OmniMasterServer during shutdown") - self._omni_master_server = None - - if self._coordinator_runtime is not None: + if hasattr(self, "_runtime") and self._runtime is not None: try: - self._coordinator_runtime.close() + self._runtime.shutdown() except Exception: - logger.exception("[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during shutdown") - self._coordinator_runtime = None + logger.exception("[AsyncOmniEngine] Failed to shutdown StageRuntime") # ── Release CuMem allocator memory pool ────────────────────────────── # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory @@ -2514,16 +1579,18 @@ def shutdown(self) -> None: # and can cause CUDA OOM for subsequent engine instances (especially # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). # - # Discard mode (level=2) is correct at shutdown: there is no benefit to - # keeping a CPU backup when the engine is being torn down. + # CuMemAllocator.sleep() is NOT idempotent — calling it on already- + # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator + # cuMemRelease (double-free of the memory handle). Use release_pools() + # instead, which is the designed cleanup path: it drops MemPool refs + # and lets the destructor/free path handle asleep entries correctly + # (returns a null handle so the C extension skips unmap/release). try: from vllm.device_allocator.cumem import CuMemAllocator, cumem_available if cumem_available: allocator = CuMemAllocator.get_instance() - # Sleep at level 2 discards all pool memory from the GPU - # without creating CPU backups — cheapest and fastest. - allocator.sleep() + allocator.release_pools() logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") except Exception: pass diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py new file mode 100644 index 00000000000..48b6fe55959 --- /dev/null +++ b/vllm_omni/engine/membership_controller.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""MembershipController: distributed replica lifecycle management. + +Extracted from Orchestrator to keep request-flow code free of distributed +concerns. Owns the OmniCoordClientForHub, watches for replica +disappearances, and handles register/unregister by building head-side +clients via an injected factory and mutating StagePool membership. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.distributed.omni_coordinator import ( + LoadBalancer, +) +from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus +from vllm_omni.distributed.omni_coordinator.omni_coord_client_for_hub import OmniCoordClientForHub +from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage +from vllm_omni.engine.stage_pool import StagePool + +logger = init_logger(__name__) + +RemoteReplicaFactory = Callable[[int, int], Any] + + +class MembershipController: + """Manages dynamic replica attach/detach for distributed mode. + + Constructed by DistStageRuntime and passed to Orchestrator. + The Orchestrator delegates register/unregister messages here. + """ + + WATCH_INTERVAL_S: float = 0.5 + + def __init__( + self, + stage_pools: list[StagePool], + coordinator_pub_address: str, + load_balancer_factory: Callable[[], LoadBalancer], + remote_replica_factory: RemoteReplicaFactory, + ) -> None: + self._stage_pools = stage_pools + self._remote_replica_factory = remote_replica_factory + self._membership_tasks: set[asyncio.Task[None]] = set() + self._shutdown_event = asyncio.Event() + self._watcher_task: asyncio.Task[None] | None = None + self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None + self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None + + self._hub = OmniCoordClientForHub(coordinator_pub_address) + factory = load_balancer_factory + for pool in self._stage_pools: + pool.attach_hub(self._hub) + pool.attach_load_balancer(factory()) + + def start(self) -> asyncio.Task[None]: + """Start the replica watcher as a background task. Returns the task.""" + self._watcher_task = asyncio.create_task(self._watch_replica_list(), name="membership-watcher") + return self._watcher_task + + async def handle_register(self, stage_id: int, replica_id: int) -> None: + """Handle a register_remote_replica message (fire-and-forget).""" + self._spawn_task( + self._do_register(stage_id, replica_id), + label=f"register-s{stage_id}-r{replica_id}", + ) + + async def handle_unregister( + self, + stage_id: int, + input_addr: str, + output_queue: asyncio.Queue[EngineQueueMessage] | None = None, + cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None, + ) -> None: + """Handle an unregister_remote_replica message.""" + pool = self._pool_for_stage_id(stage_id) + if pool is None: + return + effective_output_queue = output_queue if output_queue is not None else self._output_queue + effective_cleanup_callback = cleanup_callback if cleanup_callback is not None else self._cleanup_callback + affected = pool.invalidate_addr(input_addr) + self._detach_replica(stage_id, input_addr) + if affected and effective_cleanup_callback is not None: + await effective_cleanup_callback(affected) + if affected and effective_output_queue is not None: + for req_id in affected: + await effective_output_queue.put( + ErrorMessage(error="stage replica disappeared", request_id=req_id, stage_id=stage_id) + ) + + def shutdown(self) -> None: + """Signal stop and close the hub.""" + self._shutdown_event.set() + if self._hub is not None: + self._hub.close() + self._hub = None + if self._watcher_task is not None and not self._watcher_task.done(): + self._watcher_task.cancel() + + async def drain_tasks(self, timeout: float = 10.0) -> None: + """Wait for in-flight membership tasks to complete.""" + if self._membership_tasks: + await asyncio.wait(self._membership_tasks, timeout=timeout) + + # ---- Internal ---- + + def _spawn_task(self, coro: Awaitable[None], *, label: str) -> None: + task = asyncio.create_task(coro, name=f"membership-{label}") + self._membership_tasks.add(task) + + def _on_done(t: asyncio.Task[None]) -> None: + self._membership_tasks.discard(t) + if t.cancelled(): + return + exc = t.exception() + if exc is not None: + logger.error("[MembershipController] %s task crashed", label, exc_info=exc) + + task.add_done_callback(_on_done) + + async def _watch_replica_list(self) -> None: + """Convert hub replica disappearances into unregister actions.""" + last_up: set[tuple[int, str]] = set() + while not self._shutdown_event.is_set(): + try: + snap = self._hub.get_replica_list() + current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} + for stage_id, addr in last_up - current: + self._spawn_task( + self.handle_unregister(stage_id, addr), + label=f"unregister-s{stage_id}", + ) + last_up = current + except asyncio.CancelledError: + raise + except Exception: + logger.exception("[MembershipController] _watch_replica_list iteration failed") + + try: + await asyncio.sleep(self.WATCH_INTERVAL_S) + except asyncio.CancelledError: + raise + + async def _do_register(self, stage_id: int, replica_id: int) -> None: + pool = self._pool_for_stage_id(stage_id) + if pool is None: + logger.warning("[MembershipController] register: stage_id %d out of range", stage_id) + return + client = await asyncio.to_thread(self._remote_replica_factory, stage_id, replica_id) + input_addr = StagePool._client_input_addr(client) + if input_addr is None: + raise RuntimeError(f"remote replica factory for stage {stage_id} produced a client without input address") + pool.add_client(input_addr, client) + logger.info( + "[MembershipController] attached remote replica stage=%d replica=%d addr=%s", + stage_id, + replica_id, + input_addr, + ) + + def _detach_replica(self, stage_id: int, input_addr: str) -> None: + pool = self._pool_for_stage_id(stage_id) + if pool is None: + return + client = pool.remove_client(input_addr) + if client is None: + return + try: + client.shutdown() + except Exception: + logger.exception("[MembershipController] failed to shutdown client stage=%d addr=%s", stage_id, input_addr) + logger.info("[MembershipController] detached replica stage=%d addr=%s", stage_id, input_addr) + + def install_unregister_handlers( + self, + *, + output_queue: asyncio.Queue[EngineQueueMessage], + cleanup_callback: Callable[[list[str]], Awaitable[None]], + ) -> None: + """Install shared cleanup sinks for watcher-driven unregister events.""" + self._output_queue = output_queue + self._cleanup_callback = cleanup_callback + + def _pool_for_stage_id(self, stage_id: int) -> StagePool | None: + if not (0 <= stage_id < len(self._stage_pools)): + return None + return self._stage_pools[stage_id] diff --git a/vllm_omni/engine/mm_outputs.py b/vllm_omni/engine/mm_outputs.py index 3c9da2496fa..336bf7d9a9b 100644 --- a/vllm_omni/engine/mm_outputs.py +++ b/vllm_omni/engine/mm_outputs.py @@ -1,127 +1,9 @@ -"""Multimodal output data structures for vLLM-Omni. +import warnings -This module defines structured types for multimodal outputs. +from vllm_omni.outputs.mm_outputs import * # noqa: F401,F403 -""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping -from dataclasses import dataclass, field -from typing import Any - -import torch -from vllm.outputs import CompletionOutput - - -@dataclass(eq=False) -class MultimodalPayload(Mapping): - """Structured multimodal output payload. - - Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)`` - style checks in downstream code can be replaced with duck-typing, and - ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)`` - all work seamlessly for both tensors and metadata. - - Attributes: - tensors: Dictionary mapping modality/key names to their tensors. - metadata: Optional dictionary for non-tensor metadata - (e.g., sample rate for audio, image dimensions). - """ - - tensors: dict[str, torch.Tensor] = field(default_factory=dict) - metadata: dict[str, Any] = field(default_factory=dict) - - @property - def primary_tensor(self) -> torch.Tensor | None: - """Return the first tensor in the payload, or None if empty.""" - if self.tensors: - return next(iter(self.tensors.values())) - return None - - @property - def is_empty(self) -> bool: - """Return True if the payload has no tensors and no metadata.""" - return len(self.tensors) == 0 and len(self.metadata) == 0 - - def get(self, key: str, default: Any = None) -> Any: - """Get a value by key, searching tensors first then metadata.""" - if key in self.tensors: - return self.tensors[key] - return self.metadata.get(key, default) - - def __contains__(self, key: str) -> bool: - return key in self.tensors or key in self.metadata - - def __getitem__(self, key: str) -> Any: - """Dict-like indexing: search tensors first, then metadata.""" - if key in self.tensors: - return self.tensors[key] - if key in self.metadata: - return self.metadata[key] - raise KeyError(key) - - def __len__(self) -> int: - return len(self.tensors) + len(self.metadata) - - def __iter__(self) -> Iterator[str]: - yield from self.tensors - yield from self.metadata - - def __bool__(self) -> bool: - return bool(self.tensors) or bool(self.metadata) - - def __eq__(self, other: object) -> bool: - """Support equality with plain dicts and other Mappings.""" - if isinstance(other, MultimodalPayload): - return self.tensors == other.tensors and self.metadata == other.metadata - if isinstance(other, Mapping): - return dict(self) == dict(other) - return NotImplemented - - def to_dict(self) -> dict[str, Any]: - """Convert back to a plain dict (tensors + metadata merged).""" - result: dict[str, Any] = dict(self.tensors) - result.update(self.metadata) - return result - - @classmethod - def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None: - """Create a MultimodalPayload from a raw dictionary. - - Separates torch.Tensor values into tensors and everything - else into metadata. - """ - if not data: - return None - tensors: dict[str, torch.Tensor] = {} - metadata: dict[str, Any] = {} - for k, v in data.items(): - if isinstance(v, torch.Tensor): - tensors[k] = v - else: - metadata[k] = v - if not tensors and not metadata: - return None - return cls(tensors=tensors, metadata=metadata) - - -@dataclass -class MultimodalCompletionOutput(CompletionOutput): - """CompletionOutput with multimodal support. - - Inherits all CompletionOutput fields and adds multimodal_output. - As a CompletionOutput subclass, compatible with all existing vLLM consumers. - """ - - def __init__( - self, - multimodal_output: MultimodalPayload | None = None, - **kwargs: Any, - ): - super().__init__(**kwargs) - self.multimodal_output = multimodal_output - - def __repr__(self) -> str: - base = super().__repr__() - return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})" +warnings.warn( + "Importing from 'vllm_omni.engine.mm_outputs' is deprecated. Use 'vllm_omni.outputs.mm_outputs' instead.", + DeprecationWarning, + stacklevel=2, +) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index f8ff84de8fd..bdd678f2763 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -5,19 +5,14 @@ Owns logical request progression across stage pools and handles stage-to-stage transfer logic. -In distributed mode (``coordinator_pub_address`` provided), it also -owns the single :class:`OmniCoordClientForHub`, runs a -:meth:`_watch_replica_list` task that converts replica disappearances -into ``unregister_remote_replica`` control messages, and handles the -``register_remote_replica`` / ``unregister_remote_replica`` flow that -attaches / detaches head-side stage clients for headless replicas. +Distributed membership (replica attach/detach, hub monitoring) is +handled by :class:`MembershipController`, which is injected optionally. """ from __future__ import annotations import asyncio import time as _time -from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any @@ -25,20 +20,16 @@ import torch from vllm.config import ModelConfig from vllm.logger import init_logger +from vllm.outputs import CompletionOutput, RequestOutput from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreOutputs from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats -from vllm_omni.distributed.omni_coordinator import ( - LoadBalancer, - OmniCoordClientForHub, - RandomBalancer, - ReplicaStatus, -) from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker +from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, AddCompanionRequestMessage, @@ -53,25 +44,78 @@ StageSubmissionMessage, UnregisterRemoteReplicaMessage, ) +from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.engine.stage_pool import StagePool from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.outputs import OmniRequestOutput -# Factory signature for building a head-side stage client for a -# *dynamically attached* (auto-assigned) remote replica. -# -# Receives ``(stage_id, replica_id)`` and returns an awaitable yielding the -# constructed client (any type — it must satisfy the shape expected by the -# matching :class:`StagePool`, i.e. expose ``client_addresses["input_address"]`` -# or ``request_address``, plus the usual ``add_request_async`` / -# ``get_output_async`` / ``shutdown`` surface). -RemoteReplicaFactory = Callable[[int, int], Awaitable[Any]] - logger = init_logger(__name__) +def _build_terminal_empty_output( + request_id: str, + *, + final_output_type: str | None, + audio_sample_rate: int = 24000, +) -> RequestOutput: + """Build a terminal empty output when no downstream stage input exists.""" + completion = CompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + if final_output_type == "audio": + completion.multimodal_output = { + "audio": torch.zeros((0,), dtype=torch.float32), + "sr": audio_sample_rate, + } + return RequestOutput( + request_id=request_id, + prompt=None, + prompt_token_ids=[], + prompt_logprobs=None, + outputs=[completion], + finished=True, + ) + + +def _coerce_int_scalar(value: Any) -> int: + if value is None: + return 0 + if isinstance(value, (list, tuple)): + for item in value: + coerced = _coerce_int_scalar(item) + if coerced > 0: + return coerced + return 0 + if hasattr(value, "item"): + value = value.item() + try: + coerced = int(value) + except (TypeError, ValueError): + return 0 + return coerced if coerced > 0 else 0 + + +def _infer_stage_audio_sample_rate(stage_pool: StagePool, default: int = 24000) -> int: + """Infer the final audio stage sample rate from stage metadata when possible.""" + sample_rate_attrs = ("audio_sample_rate", "sample_rate", "sampling_rate", "output_sample_rate", "sr") + stage_client = getattr(stage_pool, "stage_client", None) + stage_config = getattr(stage_pool, "_stage_vllm_config", None) + for source in (stage_client, stage_config): + for attr in sample_rate_attrs: + sample_rate = _coerce_int_scalar(getattr(source, attr, None)) + if sample_rate > 0: + return sample_rate + return default + + def build_engine_core_request_from_tokens( request_id: str, prompt: dict[str, Any], @@ -159,10 +203,6 @@ class StreamingInputState: class Orchestrator: """Runs inside a background thread's asyncio event loop.""" - # Cadence at which the replica-list watcher polls for disappearances. - _WATCH_REPLICA_INTERVAL_S: float = 0.5 - _WATCH_REPLICA_IDLE_INTERVAL_S: float = 1.0 - # Class-level defaults so tests that bypass __init__ via object.__new__ # don't AttributeError when transfer / counter emit paths access them. _running_counter: OmniRequestCounter | None = None @@ -178,12 +218,11 @@ def __init__( *, async_chunk: bool = False, pd_config: dict[str, Any] | None = None, + membership_controller: MembershipController | None = None, running_counter: OmniRequestCounter | None = None, - coordinator_pub_address: str | None = None, - load_balancer_factory: Callable[[], LoadBalancer] | None = None, - remote_replica_factory: RemoteReplicaFactory | None = None, transfer_emitter: Any = None, log_stats: bool = False, + enable_orch_monitor: bool = False, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -192,6 +231,13 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + self._orch_monitor = create_orch_monitor( + enabled=enable_orch_monitor, + replica_sampler=self._sample_replica_metrics, + ) + for stage_id, pool in enumerate(self.stage_pools): + for replica_id in pool.live_replica_ids(): + self._orch_monitor.register_replica(stage_id, replica_id) # PD disaggregation state self._pd_pair: tuple[int, int] | None = None @@ -206,33 +252,15 @@ def __init__( self._init_metrics_state(stage_pools, running_counter, transfer_emitter, log_stats=log_stats) self._cfg_tracker = CfgCompanionTracker() + self._stage_input_processors: dict[int, Any] = {} self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._fatal_error: str | None = None self._fatal_error_stage_id: int | None = None - # Background tasks for fire-and-forget message handlers (currently - # only ``register_remote_replica`` and ``unregister_remote_replica``). - # Held as a set so each task's reference survives the loop and the - # task can self-deregister on completion. - self._membership_tasks: set[asyncio.Task[None]] = set() - - # Distributed-mode wiring. The hub is constructed on the - # orchestrator's asyncio loop because it spawns a SUB background - # thread; building it from another thread would race the - # ``_init_done`` event. - self._hub: OmniCoordClientForHub | None = ( - OmniCoordClientForHub(coordinator_pub_address) if coordinator_pub_address is not None else None - ) - self._remote_replica_factory = remote_replica_factory - # Inject hub + per-pool LB into each StagePool so they can run - # distributed dispatch via ``StagePool.pick``. - if self._hub is not None: - factory = load_balancer_factory or RandomBalancer - for pool in self.stage_pools: - pool.attach_hub(self._hub) - pool.attach_load_balancer(factory()) + # Distributed membership (optional, injected by DistStageRuntime) + self._membership = membership_controller def _init_metrics_state( self, @@ -308,16 +336,22 @@ async def run(self) -> None: self._orchestration_output_handler(), name="orchestrator-stage-output-handler", ) - # The replica watcher only runs in distributed mode. It's still - # created in both cases so ``run()`` has a uniform task graph; - # ``_watch_replica_list`` is a no-op poll when ``self._hub`` is None. - watch_task = asyncio.create_task( - self._watch_replica_list(), - name="orchestrator-replica-watcher", - ) + + # Start membership watcher if distributed mode is active. + membership_watcher: asyncio.Task[None] | None = None + if self._membership is not None: + self._membership.install_unregister_handlers( + output_queue=self.output_async_queue, + cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), + ) + membership_watcher = self._membership.start() + + tasks = [request_task, output_task] + if membership_watcher is not None: + tasks.append(membership_watcher) try: - await asyncio.gather(request_task, output_task, watch_task) + await asyncio.gather(*tasks) except asyncio.CancelledError: raise except EngineDeadError as e: @@ -332,46 +366,24 @@ async def run(self) -> None: raise finally: self._shutdown_event.set() - for task in (request_task, output_task, watch_task): + for task in tasks: if not task.done(): task.cancel() try: - await asyncio.gather(request_task, output_task, watch_task, return_exceptions=True) + await asyncio.gather(*tasks, return_exceptions=True) except Exception: pass - # If a fatal error caused the shutdown, drain any pending - # add_request messages that were never processed and broadcast - # fatal error responses so callers are not left hanging. if self._fatal_error is not None: await self._drain_pending_requests_on_fatal() - # Wait briefly for any in-flight membership handlers (register / - # unregister remote replica) to finish so they don't leave the - # head-side pool in a half-attached state. Cancel anything that - # hasn't completed in time; the generic pending-task sweep below - # will collect the cancellations. - if self._membership_tasks: - try: - await asyncio.wait_for( - asyncio.gather(*self._membership_tasks, return_exceptions=True), - timeout=10.0, - ) - except (asyncio.TimeoutError, Exception): - for t in self._membership_tasks: - if not t.done(): - t.cancel() + if self._membership is not None: + await self._membership.drain_tasks(timeout=10.0) + self._membership.shutdown() + self._orch_monitor.flush() self._shutdown_stages() - # Close the hub last so any in-flight dispatch still has access. - if self._hub is not None: - try: - self._hub.close() - except RuntimeError: - pass - self._hub = None - loop = asyncio.get_running_loop() pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()] for task in pending: @@ -379,27 +391,6 @@ async def run(self) -> None: if pending: await asyncio.gather(*pending, return_exceptions=True) - # ---- Background task helpers ---- - - def _spawn_membership_task(self, coro: Awaitable[None], *, label: str) -> None: - """Run a fire-and-forget membership-change coroutine. - - Holds a strong reference until completion (asyncio would otherwise - garbage-collect a bare task), and logs any uncaught exception. - """ - task = asyncio.create_task(coro, name=f"orchestrator-{label}") - self._membership_tasks.add(task) - - def _on_done(t: asyncio.Task[None]) -> None: - self._membership_tasks.discard(t) - if t.cancelled(): - return - exc = t.exception() - if exc is not None: - logger.error("[Orchestrator] %s task crashed", label, exc_info=exc) - - task.add_done_callback(_on_done) - # ---- Request handling ---- async def _request_handler(self) -> None: @@ -419,18 +410,12 @@ async def _request_handler(self) -> None: elif msg_type == "collective_rpc": await self._handle_collective_rpc(msg) elif isinstance(msg, RegisterRemoteReplicaMessage): - # Dynamic-attach involves a ~5s blocking handshake (run in a - # thread by ``_build_remote_replica``); ``await`` here would - # block the queue and stall the next ``add_request`` until - # the attach finishes. Dispatch as a background task so the - # main message loop keeps draining. - self._spawn_membership_task(self._handle_register_remote_replica(msg), label="register_remote_replica") + if self._membership is not None: + await self._membership.handle_register(msg.stage_id, msg.replica_id) + self._orch_monitor.register_replica(msg.stage_id, msg.replica_id) elif isinstance(msg, UnregisterRemoteReplicaMessage): - # Symmetric with register: keep the main queue flowing. - self._spawn_membership_task( - self._handle_unregister_remote_replica(msg), - label="unregister_remote_replica", - ) + if self._membership is not None: + await self._membership.handle_unregister(msg.stage_id, msg.input_addr) elif isinstance(msg, ShutdownRequestMessage): logger.info("[Orchestrator] Received shutdown signal") self._shutdown_event.set() @@ -441,7 +426,8 @@ async def _request_handler(self) -> None: for client in pool.clients: if hasattr(client, "_shutting_down"): client._shutting_down = True - self._shutdown_stages() + # Stage teardown runs once in run()'s finally after the + # orchestration loop observes _shutdown_event and exits. break else: logger.warning("[Orchestrator] Unknown message type: %s", msg_type) @@ -654,6 +640,14 @@ async def _handle_collective_rpc(self, msg: CollectiveRPCRequestMessage) -> None # ---- Orchestration loop ---- + def _sample_replica_metrics(self) -> dict[str, tuple[int, int]]: + samples: dict[str, tuple[int, int]] = {} + for stage_id, pool in enumerate(self.stage_pools): + for replica_id in pool.live_replica_ids(): + key = replica_key(stage_id, replica_id) + samples[key] = pool.replica_monitor_sample(replica_id) + return samples + async def _orchestration_output_handler(self) -> None: """Poll all stages, handle transfers, send final outputs to main.""" try: @@ -762,6 +756,7 @@ async def _orchestration_loop(self) -> None: await self._handle_processed_outputs(stage_id, replica_id, raw_output) idle = False + self._orch_monitor.note_loop(idle=idle) if idle: await asyncio.sleep(0.001) else: @@ -973,6 +968,77 @@ async def _route_output( def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: return (stage_id + 1) in req_state.stage_submit_ts + def _get_stage_input_processor(self, stage_id: int) -> Any: + processor = self._stage_input_processors.get(stage_id) + if processor is None: + from vllm_omni.engine.stage_init_utils import build_stage0_input_processor + + processor = build_stage0_input_processor(self.stage_pools[stage_id].stage_vllm_config) + self._stage_input_processors[stage_id] = processor + return processor + + def _upgrade_processed_stage_request(self, request: Any, raw_prompt: Any) -> Any: + prompt_embeds = getattr(request, "prompt_embeds", None) + additional_information = None + + if isinstance(raw_prompt, dict): + if prompt_embeds is None: + raw_prompt_embeds = raw_prompt.get("prompt_embeds") + if isinstance(raw_prompt_embeds, torch.Tensor): + prompt_embeds = raw_prompt_embeds + additional_information = serialize_additional_information( + raw_prompt.get("additional_information"), + log_prefix="Orchestrator stage input", + ) + + if prompt_embeds is None and additional_information is None: + return request + + return OmniEngineCoreRequest.from_request( + request, + prompt_embeds=prompt_embeds, + additional_information=additional_information, + ) + + def _next_stage_input_is_tokens(self, next_input: Any) -> bool: + return isinstance(next_input, dict) and "prompt_token_ids" in next_input + + def _build_next_stage_request( + self, + req_id: str, + next_stage_id: int, + next_input: Any, + params: SamplingParams | PoolingParams, + *, + mm_features: list | None = None, + resumable: bool = False, + ) -> Any: + next_pool = self.stage_pools[next_stage_id] + if self._next_stage_input_is_tokens(next_input): + request = build_engine_core_request_from_tokens( + request_id=req_id, + prompt=next_input, + params=params, + model_config=next_pool.stage_vllm_config.model_config, + mm_features=mm_features, + resumable=resumable, + ) + request.external_req_id = request.request_id + return request + + processor = self._get_stage_input_processor(next_stage_id) + request = processor.process_inputs( + request_id=req_id, + prompt=next_input, + params=params, + supported_tasks=("generate",), + arrival_time=_time.time(), + resumable=resumable, + ) + request = self._upgrade_processed_stage_request(request, next_input) + request.external_req_id = req_id + return request + async def _handle_cfg_companion_ready(self, req_id: str) -> None: """Mark a CFG companion as done; if all companions are done, flush deferred parent.""" parent_id = self._cfg_tracker.on_companion_completed(req_id) @@ -1170,6 +1236,31 @@ async def _forward_to_next_stage( src_stage_id, next_logical, ) + if diffusion_prompt is None: + error_output = OmniRequestOutput.from_error( + req_id, + f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", + ) + logger.warning( + "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " + "routing terminal error output", + req_id, + src_stage_id, + next_logical, + ) + await self.output_async_queue.put( + OutputMessage( + request_id=req_id, + stage_id=next_logical, + engine_outputs=error_output, + metrics=None, + finished=True, + ) + ) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + ) + return if isinstance(diffusion_prompt, list): if not diffusion_prompt: error_output = OmniRequestOutput.from_error( @@ -1196,7 +1287,7 @@ async def _forward_to_next_stage( [req_id, *self._cfg_tracker.cleanup_parent(req_id)], ) return - if already_submitted and len(diffusion_prompt) == 1: + if len(diffusion_prompt) == 1: diffusion_prompt = diffusion_prompt[0] else: diffusion_prompt = req_state.prompt @@ -1296,22 +1387,64 @@ async def _forward_to_next_stage( ) raise + if not next_inputs: + if not getattr(output, "finished", False): + logger.debug( + "[Orchestrator] req=%s stage-%s produced no inputs for stage-%s; waiting for more outputs", + req_id, + src_stage_id, + next_logical, + ) + return + + final_stage_id = req_state.final_stage_id + final_pool = self.stage_pools[final_stage_id] + final_output_type = getattr(final_pool.stage_client, "final_output_type", None) + terminal_output = _build_terminal_empty_output( + req_id, + final_output_type=final_output_type, + audio_sample_rate=_infer_stage_audio_sample_rate(final_pool), + ) + submit_ts = _time.time() + req_state.stage_submit_ts[final_stage_id] = submit_ts + logger.info( + "[Orchestrator] req=%s stage-%s produced no terminal inputs for stage-%s; " + "returning empty %s output from final stage-%s", + req_id, + src_stage_id, + next_logical, + final_output_type or "text", + final_stage_id, + ) + await self.output_async_queue.put( + OutputMessage( + request_id=req_id, + stage_id=final_stage_id, + replica_id=0, + engine_outputs=terminal_output, + metrics=None, + finished=True, + stage_submit_ts=submit_ts, + ) + ) + await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) + return + # Build and submit requests for each input for next_input in next_inputs: # Only AR thinker stages consume encoder mm_features; downstream # (talker/code2wav/…) must not see them (avoids encoder-cache misses). - model_stage = getattr(next_client, "model_stage", None) + model_stage = getattr(getattr(next_pool.stage_vllm_config, "model_config", None), "model_stage", None) mm_features = req_state.mm_features if model_stage == "thinker" else None - request = build_engine_core_request_from_tokens( - request_id=req_id, - prompt=next_input, + request = self._build_next_stage_request( + req_id, + next_logical, + next_input, params=params, - model_config=next_pool.stage_vllm_config.model_config, mm_features=mm_features, resumable=next_stage_resumable, ) - request.external_req_id = request.request_id if already_submitted: await next_pool.submit_update(req_id, req_state, request) else: @@ -1495,129 +1628,6 @@ async def _drain_pending_requests_on_fatal(self) -> None: ) self.request_states.pop(req_id, None) - # ---- Distributed-mode replica attach / detach ---- - - async def _watch_replica_list(self) -> None: - """Convert hub replica disappearances into unregister control messages.""" - last_up: set[tuple[int, str]] = set() - while not self._shutdown_event.is_set(): - if self._hub is None: - # No coordinator wired up; sleep coarsely and re-check shutdown. - try: - await asyncio.sleep(self._WATCH_REPLICA_IDLE_INTERVAL_S) - except asyncio.CancelledError: - raise - continue - - try: - snap = self._hub.get_replica_list() - current = {(rep.stage_id, rep.input_addr) for rep in snap.replicas if rep.status == ReplicaStatus.UP} - for stage_id, addr in last_up - current: - await self.request_async_queue.put( - UnregisterRemoteReplicaMessage( - stage_id=stage_id, - input_addr=addr, - ) - ) - last_up = current - except asyncio.CancelledError: - raise - except Exception: - logger.exception("[Orchestrator] _watch_replica_list iteration failed") - - try: - await asyncio.sleep(self._WATCH_REPLICA_INTERVAL_S) - except asyncio.CancelledError: - raise - - async def _handle_register_remote_replica(self, msg: RegisterRemoteReplicaMessage) -> None: - """Bind a head-side client for a newly registered remote replica.""" - stage_id = int(msg.stage_id) - replica_id = int(msg.replica_id) - if not (0 <= stage_id < self.num_stages): - logger.warning( - "[Orchestrator] register_remote_replica: stage_id %d out of range (num_stages=%d)", - stage_id, - self.num_stages, - ) - return - if self._remote_replica_factory is None: - logger.warning( - "[Orchestrator] register_remote_replica received for stage=%d replica=%d but no factory installed", - stage_id, - replica_id, - ) - return - - try: - await self._attach_remote_replica(stage_id, replica_id) - except Exception: - logger.exception( - "[Orchestrator] failed to attach remote replica stage=%d replica=%d", - stage_id, - replica_id, - ) - - async def _handle_unregister_remote_replica(self, msg: UnregisterRemoteReplicaMessage) -> None: - """Tear down the head-side client for a vanished remote replica.""" - stage_id = int(msg.stage_id) - input_addr = str(msg.input_addr) - if not (0 <= stage_id < self.num_stages): - return - pool = self.stage_pools[stage_id] - affected = pool.invalidate_addr(input_addr) - self._detach_remote_replica(stage_id, input_addr) - if affected: - await self._cleanup_request_ids(affected, abort=True) - for req_id in affected: - await self.output_async_queue.put( - ErrorMessage( - error="stage replica disappeared", - request_id=req_id, - stage_id=stage_id, - ) - ) - - async def _attach_remote_replica(self, stage_id: int, replica_id: int) -> None: - """Build a head-side stage client via the injected factory and register it.""" - factory = self._remote_replica_factory - if factory is None: - return - pool = self.stage_pools[stage_id] - client = await factory(stage_id, replica_id) - input_addr = StagePool._client_input_addr(client) - if input_addr is None: - raise RuntimeError( - f"remote replica factory for stage {stage_id} produced a client without a discoverable input address" - ) - pool.add_client(input_addr, client) - logger.info( - "[Orchestrator] attached remote replica stage=%d replica=%d addr=%s", - stage_id, - replica_id, - input_addr, - ) - - def _detach_remote_replica(self, stage_id: int, input_addr: str) -> None: - """Shut down + remove the head-side client at ``input_addr``.""" - pool = self.stage_pools[stage_id] - client = pool.remove_client(input_addr) - if client is None: - return - try: - client.shutdown() - except Exception: - logger.exception( - "[Orchestrator] failed to shutdown client for stage=%d addr=%s", - stage_id, - input_addr, - ) - logger.info( - "[Orchestrator] detached remote replica stage=%d addr=%s", - stage_id, - input_addr, - ) - def _shutdown_stages(self) -> None: """Shutdown all stage pools.""" if self._stages_shutdown: diff --git a/vllm_omni/engine/orchestrator_monitor.py b/vllm_omni/engine/orchestrator_monitor.py new file mode 100644 index 00000000000..8be903a88dd --- /dev/null +++ b/vllm_omni/engine/orchestrator_monitor.py @@ -0,0 +1,193 @@ +"""Optional orchestrator window monitor (diagnostic only). + +Enable with: + --enable-orch-monitor + +Optional output path override: + export VLLM_OMNI_ORCH_MONITOR_PATH=/path/to/monitor.json + +Each 1s window records: + - duration_s, loop_idle, loop_active (orchestrator poll-loop counts) + - per-replica outputs_queue_size (MP client outputs_queue backlog) + - per-replica inflight (requests currently bound/routed to the replica) + +See docs/contributing/profiling.md (Orchestrator Monitor). +""" + +from __future__ import annotations + +import json +import os +import time +from collections.abc import Callable +from pathlib import Path +from typing import Protocol, TypedDict + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_DEFAULT_OUTPUT_PREFIX = "vllm_omni_orch_monitor" +_WINDOW_S = 1.0 + +ReplicaSample = tuple[int, int] +ReplicaSampler = Callable[[], dict[str, ReplicaSample]] + + +class _ReplicaSeries(TypedDict): + outputs_queue_size: list[int] + inflight: list[int] + + +def replica_key(stage_id: int, replica_id: int) -> str: + return f"stage={stage_id},replica={replica_id}" + + +def _resolve_output_path() -> str: + env_path = os.environ.get("VLLM_OMNI_ORCH_MONITOR_PATH") + if env_path: + return env_path + timestamp = time.strftime("%m%d%H%M", time.localtime()) + return str(Path.cwd() / f"{_DEFAULT_OUTPUT_PREFIX}_{timestamp}.json") + + +class OrchestratorMonitorBase(Protocol): + def register_replica(self, stage_id: int, replica_id: int) -> None: ... + + def note_loop(self, *, idle: bool) -> None: ... + + def flush(self) -> None: ... + + +class _NullOrchestratorMonitor: + def register_replica(self, stage_id: int, replica_id: int) -> None: + return + + def note_loop(self, *, idle: bool) -> None: + return + + def flush(self) -> None: + return + + +_NULL_ORCH_MONITOR = _NullOrchestratorMonitor() + + +def create_orch_monitor( + *, + enabled: bool, + replica_sampler: ReplicaSampler, +) -> OrchestratorMonitorBase: + if not enabled: + return _NULL_ORCH_MONITOR + return OrchestratorMonitor(replica_sampler=replica_sampler) + + +class OrchestratorMonitor: + def __init__(self, *, replica_sampler: ReplicaSampler) -> None: + self._output_path = _resolve_output_path() + self._flushed = False + self._replica_sampler = replica_sampler + self._window_start_mono = time.monotonic() + self._loop_idle = 0 + self._loop_active = 0 + self._duration_s: list[float] = [] + self._loop_idle_buf: list[int] = [] + self._loop_active_buf: list[int] = [] + self._replicas: dict[str, _ReplicaSeries] = {} + + def register_replica(self, stage_id: int, replica_id: int) -> None: + self._ensure_replica(replica_key(stage_id, replica_id)) + + def note_loop(self, *, idle: bool) -> None: + now_mono = time.monotonic() + self._roll_window_if_needed(now_mono) + if idle: + self._loop_idle += 1 + else: + self._loop_active += 1 + + def flush(self) -> None: + if self._flushed: + return + self._roll_window(time.monotonic()) + self._flushed = True + payload = { + "configured_window_s": _WINDOW_S, + "windows": { + "duration_s": self._duration_s, + "loop_idle": self._loop_idle_buf, + "loop_active": self._loop_active_buf, + }, + "replicas": dict(sorted(self._replicas.items())), + } + try: + output_path = Path(self._output_path).expanduser() + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as f: + json.dump(payload, f, sort_keys=True) + f.write("\n") + self._log_summary(str(output_path)) + except OSError: + logger.exception("[OrchestratorMonitor] Failed to write log to %s", self._output_path) + + def _ensure_replica(self, key: str) -> None: + if key in self._replicas: + return + # Backfill zeros for windows recorded before this replica appeared. + n = len(self._duration_s) + self._replicas[key] = { + "outputs_queue_size": [0] * n, + "inflight": [0] * n, + } + + def _roll_window_if_needed(self, now_mono: float) -> None: + if now_mono - self._window_start_mono >= _WINDOW_S: + self._roll_window(now_mono) + + def _roll_window(self, now_mono: float) -> None: + try: + sampled = self._replica_sampler() + except Exception: + logger.exception("[OrchestratorMonitor] replica sampling failed") + sampled = {} + self._duration_s.append(max(now_mono - self._window_start_mono, 1e-9)) + self._loop_idle_buf.append(self._loop_idle) + self._loop_active_buf.append(self._loop_active) + for key, series in self._replicas.items(): + outputs_qsize, inflight = sampled.get(key, (0, 0)) + series["outputs_queue_size"].append(outputs_qsize) + series["inflight"].append(inflight) + + self._loop_idle = 0 + self._loop_active = 0 + self._window_start_mono = now_mono + + def _log_summary(self, output_path: str) -> None: + loop_idle = sum(self._loop_idle_buf) + loop_active = sum(self._loop_active_buf) + loop_total = loop_idle + loop_active + loop_active_pct = (float(loop_active) / float(loop_total) * 100.0) if loop_total else 0.0 + logger.info( + "[OrchestratorMonitor] wrote %s windows=%d duration=%.3fs loop_active=%.1f%%", + output_path, + len(self._duration_s), + sum(self._duration_s), + loop_active_pct, + ) + for key, series in sorted(self._replicas.items()): + queue_values = series["outputs_queue_size"] + inflight_values = series["inflight"] + queue_avg = sum(queue_values) / float(len(queue_values)) if queue_values else 0.0 + queue_max = max(queue_values) if queue_values else 0 + inflight_avg = sum(inflight_values) / float(len(inflight_values)) if inflight_values else 0.0 + inflight_max = max(inflight_values) if inflight_values else 0 + logger.info( + "[OrchestratorMonitor] %s outputs_queue_avg=%.2f outputs_queue_max=%d " + "inflight_avg=%.2f inflight_max=%d", + key, + queue_avg, + queue_max, + inflight_avg, + inflight_max, + ) diff --git a/vllm_omni/engine/output_modality.py b/vllm_omni/engine/output_modality.py index ac72910dd77..364519bedd9 100644 --- a/vllm_omni/engine/output_modality.py +++ b/vllm_omni/engine/output_modality.py @@ -1,123 +1,9 @@ -"""Output modality types for vLLM-Omni. +import warnings -This module defines the OutputModality enum and TensorAccumulationStrategy -for type-safe multimodal output routing and tensor merging. +from vllm_omni.outputs.output_modality import * # noqa: F401,F403 -""" - -from __future__ import annotations - -import re -from enum import Enum, Flag, StrEnum, auto -from typing import Literal, TypeAlias - -FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"] - -_MODALITY_ALIASES: dict[str, str] = { - "speech": "audio", - "images": "image", - "latents": "latent", - "wav": "audio", - "waveform": "audio", - "pixel_values": "image", - "pixels": "image", - "token_ids": "text", - "tokens": "text", -} - - -class OutputModalityNames(StrEnum): - """Keys for output modalities. - - TODO: (Alex) Integrate this with the big-flag enum below + throughout the code - for better type safety (currently only used for output processor). - """ - - TEXT = "text" - IMAGE = "image" - AUDIO = "audio" - LATENT = "latent" - - -# Specify which output modalities may be drained when handling delta messages. -# For some types, e.g., latents, we need to be careful to ensure the full context -# is passed as the stream yields due to assumptions in the I/O processing and model -# when async chunk isn't enabled. -NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT} -DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES} - - -class OutputModality(Flag): - """Bit-flag enum for output modalities. - - Compose freely with ``|`` — no need to enumerate every combination. - - Single: ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ... - Compound: ``OutputModality.TEXT | OutputModality.IMAGE`` (text+image) - - Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's - native path (pooling_output → PoolingRequestOutput), handled entirely - by the base OutputProcessor. vLLM-Omni's layer does not participate. - """ - - TEXT = auto() - IMAGE = auto() - AUDIO = auto() - LATENT = auto() - - @classmethod - def from_string(cls, s: str | None) -> OutputModality: - """Parse a free-text modality string into an OutputModality flag. - - Handles common aliases and compound strings separated by + or ,. - - Examples:: - - OutputModality.from_string("text+image") - # → OutputModality.TEXT | OutputModality.IMAGE - """ - if not s or not s.strip(): - return cls.TEXT - - parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())] - result = cls(0) - for p in parts: - p = _MODALITY_ALIASES.get(p, p) - try: - result |= cls[p.upper()] - except KeyError: - raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}") - return result - - @property - def has_text(self) -> bool: - return OutputModality.TEXT in self - - @property - def has_multimodal(self) -> bool: - return bool(self & ~OutputModality.TEXT) - - -class TensorAccumulationStrategy(Enum): - """Strategy for merging incremental multimodal tensors.""" - - CONCAT_DIM0 = "concat_dim0" - """Concatenate along dimension 0. Used for image/latent tensors.""" - - CONCAT_LAST = "concat_last" - """Concatenate along the last dimension. Used for audio waveforms.""" - - APPEND_LIST = "append_list" - """Append to a list (no tensor concatenation).""" - - REPLACE = "replace" - """Replace previous tensor entirely with the latest one.""" - - -def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy: - """Determine tensor merge strategy from the multimodal flags.""" - if OutputModality.AUDIO in modality: - return TensorAccumulationStrategy.CONCAT_LAST - if OutputModality.IMAGE in modality or OutputModality.LATENT in modality: - return TensorAccumulationStrategy.CONCAT_DIM0 - return TensorAccumulationStrategy.CONCAT_DIM0 # default +warnings.warn( + "Importing from 'vllm_omni.engine.output_modality' is deprecated. Use 'vllm_omni.outputs.output_modality' instead.", + DeprecationWarning, + stacklevel=2, +) diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index b84a9eb3ae8..0d87583a8d4 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,739 +1,10 @@ -from dataclasses import fields as dataclass_fields -from typing import Any +import warnings -import torch -from vllm.logger import init_logger -from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput -from vllm.sampling_params import RequestOutputKind -from vllm.tokenizers import TokenizerLike -from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason -from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor -from vllm.v1.engine.output_processor import ( - OutputProcessorOutput, - RequestOutputCollector, - RequestState, -) -from vllm.v1.engine.parallel_sampling import ParentRequest -from vllm.v1.metrics.stats import IterationStats, RequestStateStats +from vllm_omni.outputs.output_processor import * # noqa: F401,F403 -from vllm_omni.data_entry_keys import unflatten_payload -from vllm_omni.engine.mm_outputs import MultimodalCompletionOutput, MultimodalPayload -from vllm_omni.engine.output_modality import ( - DRAINABLE_MODALITIES, - OutputModality, - TensorAccumulationStrategy, - get_accumulation_strategy, +warnings.warn( + "Importing from 'vllm_omni.engine.output_processor' is deprecated. " + "Use 'vllm_omni.outputs.output_processor' instead.", + DeprecationWarning, + stacklevel=2, ) -from vllm_omni.outputs import OmniRequestOutput - -logger = init_logger(__name__) - - -# --------------------------------------------------------------------------- -# Module-level helpers -# --------------------------------------------------------------------------- - - -def _to_cpu(x: Any) -> Any: - """Move a tensor to CPU; pass through non-tensors unchanged.""" - if isinstance(x, torch.Tensor): - try: - return x.detach().to("cpu", non_blocking=True).contiguous() - except (RuntimeError, AttributeError): - return x - if isinstance(x, dict): - return {str(sk): _to_cpu(sv) for sk, sv in x.items()} - return x - - -def _merge_payload(target: MultimodalPayload, incoming: MultimodalPayload) -> None: - """Merge *incoming* into *target* using list-based deferred concatenation.""" - for k, v in incoming.tensors.items(): - if k not in target.tensors: - target.tensors[k] = v - else: - existing = target.tensors[k] - if isinstance(existing, list): - existing.append(v) - elif isinstance(existing, torch.Tensor) and isinstance(v, torch.Tensor): - target.tensors[k] = [existing, v] - else: - target.tensors[k] = v - - for k, v in incoming.metadata.items(): - if k not in target.metadata: - target.metadata[k] = v - else: - existing = target.metadata[k] - if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor): - target.metadata[k] = [existing, v] - elif isinstance(v, torch.Tensor) and isinstance(existing, list): - existing.append(v) - else: - # Metadata: replace with latest value - target.metadata[k] = v - - -def _cat_tensors( - tensors: list[torch.Tensor], - strategy: TensorAccumulationStrategy, -) -> torch.Tensor: - """Concatenate a list of tensors according to *strategy*.""" - if strategy == TensorAccumulationStrategy.CONCAT_LAST: - return torch.cat(tensors, dim=-1) - if strategy == TensorAccumulationStrategy.REPLACE: - return tensors[-1] - # CONCAT_DIM0 / APPEND_LIST / default - return torch.cat(tensors, dim=0) - - -def _modality_to_type_string(modality: OutputModality) -> str: - """Convert an OutputModality flag to a lowercase type string.""" - try: - if OutputModality.AUDIO in modality: - return "audio" - if OutputModality.IMAGE in modality: - return "image" - if OutputModality.LATENT in modality: - return "latent" - except TypeError: - # Flag identity mismatch (e.g. after module reload in tests). - name = getattr(modality, "name", "") or "" - lowered = name.lower() - if "audio" in lowered: - return "audio" - if "image" in lowered: - return "image" - if "latent" in lowered: - return "latent" - return "text" - - -class OmniRequestState(RequestState): - """Request state for omni models, tracking multimodal outputs. - - Extends the base RequestState with support for accumulating - multimodal tensor outputs (e.g., images, audio, latents) that - are produced incrementally during generation. - """ - - def __init__( - self, - *args, - **kwargs, - ): - # arrival_time is always passed as a keyword argument by the sole caller - # (from_new_request). The *args fallback via positional index was removed - # because upstream RequestState.__init__ reorders parameters between - # releases, making args[12] fragile. - arrival_time = kwargs.get("arrival_time") - super().__init__(*args, **kwargs) - self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) - # Omni-specific: multimodal output accumulation - # TODO: mm_type is per-request, not per-key. If a model ever produces - # both audio and latent outputs, the modality type would flip on each - # add_multimodal_tensor() call. Consider tracking per-key modality - # types (e.g. dict[str, str]) for future multi-output models. - self.mm_type: str | None = None - self.mm_accumulated: MultimodalPayload = MultimodalPayload() - - def apply_streaming_update(self, update) -> None: - super().apply_streaming_update(update) - self.native_text_stats.arrival_time = update.arrival_time - - @staticmethod - def _to_cpu(x): - """Try to convert to CPU tensor if needed.""" - if isinstance(x, torch.Tensor): - try: - return x.detach().to("cpu", non_blocking=True).contiguous() - except Exception: - return x - return x - - def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: - """Accumulate a multimodal tensor payload into the request state. - - Normalizes incoming payloads (dict or raw tensor) into a - MultimodalPayload and merges with any previously accumulated data. - Uses list-based deferred concatenation to avoid O(n²) repeated - torch.cat calls. - """ - if payload is None: - return - try: - if mm_type: - self.mm_type = (mm_type or "").lower() - - target_key = self.mm_type or "hidden" - - # Build incoming MultimodalPayload - if isinstance(payload, dict): - remapped: dict[str, Any] = {} - for k, v in payload.items(): - # Normalize producer keys to the modality name. - # AR runners produce {"hidden": ...} and generation - # runners produce {"model_outputs": ...}; remap both - # to the semantic modality key (e.g. "audio", "latent"). - if k == "model_outputs": - k = target_key - elif k == "hidden" and target_key != "hidden": - k = target_key - remapped[k] = _to_cpu(v) - incoming = MultimodalPayload.from_dict(remapped) - elif isinstance(payload, MultimodalPayload): - incoming = payload - else: - incoming = MultimodalPayload.from_dict({target_key: _to_cpu(payload)}) - - if incoming is None: - return - - if self.mm_accumulated.is_empty: - self.mm_accumulated = incoming - else: - _merge_payload(self.mm_accumulated, incoming) - except (ValueError, TypeError, RuntimeError): - logger.exception("Error accumulating multimodal tensor") - - def _consolidate_multimodal_tensors(self) -> None: - """Consolidate accumulated tensor lists into single tensors. - - Uses TensorAccumulationStrategy derived from the output modality - to determine concatenation behavior. Metadata values always use - REPLACE (keep latest). - """ - if self.mm_accumulated.is_empty: - return - - try: - modality = OutputModality.from_string(self.mm_type) - except (ValueError, KeyError): - modality = OutputModality.TEXT - strategy = get_accumulation_strategy(modality) - - try: - for k, v in list(self.mm_accumulated.tensors.items()): - if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): - try: - self.mm_accumulated.tensors[k] = _cat_tensors(v, strategy) - except RuntimeError: - if k == "audio": - try: - self.mm_accumulated.tensors[k] = torch.cat(v, dim=-1) - except RuntimeError: - self.mm_accumulated.tensors[k] = torch.cat([t.reshape(-1) for t in v], dim=0) - else: - logger.warning( - "Error concatenating tensor for key %s; keeping last tensor", - k, - ) - self.mm_accumulated.tensors[k] = v[-1] - - # Metadata: consolidate any deferred tensor lists (REPLACE semantics) - for k, v in list(self.mm_accumulated.metadata.items()): - if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): - self.mm_accumulated.metadata[k] = v[-1] - elif isinstance(v, dict): - for sk, sv in list(v.items()): - if isinstance(sv, list) and sv and isinstance(sv[0], torch.Tensor): - v[sk] = sv[-1] - except (RuntimeError, TypeError, KeyError): - logger.exception("Error consolidating multimodal tensors") - - # Override: do not route to pooling-only path; always create completion - # outputs, and attach pooling_result into the CompletionOutput. - def make_request_output( - self, - new_token_ids: list[int], - pooling_output: torch.Tensor | None, - finish_reason: FinishReason | None, - stop_reason: int | str | None, - kv_transfer_params: dict[str, Any] | None = None, - ) -> OmniRequestOutput | PoolingRequestOutput | None: - """Create a request output from generation results. - - Creates a RequestOutput or PoolingRequestOutput from the generated - tokens and accumulated multimodal outputs. Attaches multimodal - tensors to the completion output if available. - - Args: - new_token_ids: List of newly generated token IDs - pooling_output: Optional pooling output tensor - finish_reason: Optional finish reason indicating why generation stopped - stop_reason: Optional stop reason (token ID or stop string) - kv_transfer_params: Optional KV cache transfer parameters - - Returns: - OmniRequestOutput or PoolingRequestOutput if output should be - emitted (based on finish status and output kind), None otherwise - """ - # Pooling-only requests should follow base behavior. - if self.detokenizer is None and pooling_output is not None: - return super().make_request_output( - new_token_ids, - pooling_output, - finish_reason, - stop_reason, - kv_transfer_params, - ) - - finished = finish_reason is not None - final_only = self.output_kind == RequestOutputKind.FINAL_ONLY - is_delta = self.output_kind == RequestOutputKind.DELTA - - if not finished and final_only: - return None - - # Consolidate accumulated tensors when finishing, or for every - # CUMULATIVE step so the consumer always sees a single tensor. - if finished or not is_delta: - self._consolidate_multimodal_tensors() - - if self.stream_interval > 1 and self.detokenizer is not None: - # Send output request only when - # 1. It has finished, or - # 2. It is the first token, or - # 3. It has reached the stream interval number of tokens - if not ( - finished - or self.sent_tokens_offset == 0 - or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval - ): - return None - - if self.output_kind == RequestOutputKind.DELTA: - # Send tokens from the offset in DELTA mode, otherwise all - # tokens are sent. - new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] - self.sent_tokens_offset = self.detokenizer.num_output_tokens() - - external_req_id = self.external_req_id - - output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) - - if self.parent_req is None: - outputs = [output] - else: - outputs, finished = self.parent_req.get_outputs(self.request_id, output) - if not outputs: - return None - external_req_id = self.parent_req.external_req_id - - return self._new_request_output( - external_req_id, - outputs, - finished, - kv_transfer_params, - ) - - def _new_completion_output( - self, - token_ids: list[int], - finish_reason: FinishReason | None, - stop_reason: int | str | None, - ) -> MultimodalCompletionOutput | CompletionOutput: - """Create a completion output with multimodal data attached. - - Returns a MultimodalCompletionOutput when multimodal data has been - accumulated, otherwise returns the base CompletionOutput. Snapshots - the accumulated payload before draining DELTA-mode modality keys so - that callers receive an immutable view. - """ - # When there is no detokenizer (generation stages), build a minimal - # CompletionOutput directly since upstream asserts detokenizer != None. - if self.detokenizer is None: - finished = finish_reason is not None - base_output = CompletionOutput( - index=self.request_index, - text="", - token_ids=token_ids, - logprobs=None, - cumulative_logprob=None, - finish_reason=str(finish_reason) if finished else None, - stop_reason=stop_reason if finished else None, - ) - else: - base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) - - # Always provide cumulative token IDs for inter-stage processors. - if self.detokenizer is not None: - base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids) - else: - base_output.cumulative_token_ids = list(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 - - try: - if self.mm_accumulated and not self.mm_accumulated.is_empty: - # Snapshot: copy current tensors/metadata so drain doesn't - # mutate the payload already handed to the caller. - # Unflatten dotted keys (e.g. "hidden_states.layer_0") back - # to nested dicts so downstream consumers (thinker2talker etc.) - # can access with .get("hidden_states", {}).get("layers", {}). - merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata} - nested = unflatten_payload(merged) - snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload() - kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)} - output = MultimodalCompletionOutput( - multimodal_output=snapshot, - **kwargs, - ) - output.cumulative_token_ids = base_output.cumulative_token_ids - if hasattr(base_output, "cumulative_text"): - output.cumulative_text = base_output.cumulative_text - - # DELTA mode: drain modality keys (e.g. audio) so the next - # step only sees freshly accumulated data for those keys. - if self.output_kind == RequestOutputKind.DELTA: - for modality_key in DRAINABLE_MODALITIES: - self.mm_accumulated.tensors.pop(modality_key, None) - - return output - except (RuntimeError, TypeError, AttributeError): - logger.exception("Error creating MultimodalCompletionOutput") - return base_output - - def _new_request_output( - self, - external_req_id: str, - outputs: list, - finished: bool, - kv_transfer_params: dict[str, Any] | None = None, - ) -> RequestOutput | PoolingRequestOutput: - """Create request output, handling no-detokenizer generation stages. - - Upstream asserts ``self.logprobs_processor is not None`` which fails - for generation stages (talker, code2wav) that have no tokenizer. - When the logprobs processor is absent we build the RequestOutput - directly with ``prompt_logprobs=None``. - """ - if self.logprobs_processor is not None: - return super()._new_request_output( - external_req_id, - outputs, - finished, - kv_transfer_params, - ) - - # No-detokenizer path: build RequestOutput directly. - prompt_token_ids = self.prompt_token_ids - if prompt_token_ids is None and self.prompt_embeds is not None: - prompt_token_ids = [0] * len(self.prompt_embeds) - if prompt_token_ids is None: - prompt_token_ids = [] - - return RequestOutput( - request_id=external_req_id, - lora_request=self.lora_request, - prompt=self.prompt, - prompt_token_ids=prompt_token_ids, - prompt_logprobs=None, - outputs=outputs, - finished=finished, - kv_transfer_params=kv_transfer_params, - num_cached_tokens=self.num_cached_tokens, - metrics=self.stats, - ) - - -class MultimodalOutputProcessor(VLLMOutputProcessor): - """Handles multimodal output processing. - - Captures multimodal outputs from OmniEngineCoreOutput and accumulates - them as MultimodalPayload in OmniRequestState, before delegating to - the base vLLM OutputProcessor for text handling. - - The data flow is: - 1. For each EngineCoreOutput with multimodal_output: - - Capture into OmniRequestState.add_multimodal_tensor() - 2. Base vLLM OutputProcessor handles text detokenization - 3. On finish, _consolidate_multimodal_tensors() concatenates accumulated - tensors using strategy-based dispatch - 4. _new_completion_output() returns MultimodalCompletionOutput - """ - - def __init__( - self, - tokenizer: TokenizerLike | None, - *, - log_stats: bool, - stream_interval: int = 1, - tracing_enabled: bool = False, - engine_core_output_type: str | None = None, - output_modality: OutputModality = OutputModality.TEXT, - ): - """Initialize the multimodal output processor. - - Args: - tokenizer: Tokenizer for detokenizing text outputs - log_stats: Whether to log statistics - stream_interval: Stream interval for output generation - engine_core_output_type: Optional output type string (e.g., - "image", "audio", "latent"). Converted to OutputModality - internally. Kept for backward compatibility with - stage_init_utils. - output_modality: Type-safe output modality flag. Used to tag - multimodal outputs with the correct modality key when - per-output type info is unavailable. - """ - super().__init__( - tokenizer=tokenizer, - log_stats=log_stats, - stream_interval=stream_interval, - tracing_enabled=tracing_enabled, - ) - # Convert string-based engine_core_output_type to OutputModality - if engine_core_output_type is not None: - self.output_modality = OutputModality.from_string(engine_core_output_type) - else: - self.output_modality = output_modality - self.engine_core_output_type = engine_core_output_type - self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {} - - def _native_text_metric_record(self, request_id: str) -> dict[str, Any]: - return self._native_text_metrics_by_request.setdefault( - request_id, - { - "vllm_ttft_ms": 0.0, - "vllm_tpot_ms": 0.0, - "vllm_itl_ms": 0.0, - "vllm_itls_ms": [], - }, - ) - - def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: - return self._native_text_metrics_by_request.pop(request_id, {}) - - def abort_requests(self, request_ids, internal: bool) -> list[str]: - request_ids = list(request_ids) - for request_id in request_ids: - if internal: - req_state = self.request_states.get(request_id) - if req_state is not None: - self._native_text_metrics_by_request.pop(req_state.external_req_id, None) - else: - self._native_text_metrics_by_request.pop(request_id, None) - return super().abort_requests(request_ids, internal) - - def add_request( - self, - request: EngineCoreRequest, - prompt: str | None, - parent_req: ParentRequest | None = None, - request_index: int = 0, - queue: RequestOutputCollector | None = None, - ) -> None: - """Add a new request to be processed. - - Creates an OmniRequestState for the request and registers it - for output processing. - - Args: - request: Engine core request to add - prompt: Optional prompt string for the request - parent_req: Optional parent request for parallel sampling - request_index: Index of the request in the batch - queue: Optional queue for collecting outputs - - Raises: - ValueError: If the request ID is already registered - """ - request_id = request.request_id - req_state = self.request_states.get(request_id) - if req_state is not None: - self._update_streaming_request_state(req_state, request, prompt) - return - - req_state = OmniRequestState.from_new_request( - tokenizer=self.tokenizer, - request=request, - prompt=prompt, - parent_req=parent_req, - request_index=request_index, - queue=queue, - log_stats=self.log_stats, - stream_interval=self.stream_interval, - ) - self.request_states[request_id] = req_state - if parent_req: - self.parent_requests[parent_req.request_id] = parent_req - self.external_req_ids[req_state.external_req_id].append(request_id) - - def remove_request(self, request_id: str) -> None: - """Rollback one previously registered request if it was never submitted.""" - req_state = self.request_states.pop(request_id, None) - if req_state is None: - return - - external_req_id = getattr(req_state, "external_req_id", None) - if external_req_id is not None: - self._native_text_metrics_by_request.pop(external_req_id, None) - request_ids = self.external_req_ids.get(external_req_id) - if request_ids is not None: - self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id] - if not self.external_req_ids[external_req_id]: - self.external_req_ids.pop(external_req_id, None) - - parent_req = getattr(req_state, "parent_req", None) - if parent_req is not None: - self.parent_requests.pop(parent_req.request_id, None) - - def process_outputs( - self, - engine_core_outputs: list[EngineCoreOutput], - engine_core_timestamp: float | None = None, - iteration_stats: IterationStats | None = None, - ) -> OutputProcessorOutput: - default_mm_type = _modality_to_type_string(self.output_modality) - - # Separate outputs that upstream can handle (has detokenizer or - # pooling) from multimodal-only outputs (no detokenizer, no pooling) - # that would trigger upstream's `assert detokenizer is not None`. - upstream_outputs: list[EngineCoreOutput] = [] - mm_only_outputs: list[EngineCoreOutput] = [] - - for eco in engine_core_outputs: - req_state = self.request_states.get(eco.request_id) - if req_state is None: - continue - - # Accumulate multimodal tensors regardless of path. - if isinstance(req_state, OmniRequestState): - mm_output = getattr(eco, "multimodal_output", None) - if mm_output is not None: - mm_type = getattr(eco, "output_type", None) or default_mm_type - req_state.add_multimodal_tensor(mm_output, mm_type) - - # Route: if no detokenizer and no pooling output, handle locally - # to avoid upstream's assert on detokenizer. - if req_state.detokenizer is None and eco.pooling_output is None: - mm_only_outputs.append(eco) - else: - upstream_outputs.append(eco) - - # Handle multimodal-only outputs (generation stages) locally. - self._process_mm_only_outputs(mm_only_outputs) - - # Delegate text/pooling outputs to upstream. - return super().process_outputs( - upstream_outputs, - engine_core_timestamp=engine_core_timestamp, - iteration_stats=iteration_stats, - ) - - def _process_mm_only_outputs( - self, - engine_core_outputs: list[EngineCoreOutput], - ) -> None: - """Handle outputs from generation stages that have no detokenizer. - - These cannot go through upstream process_outputs because it asserts - detokenizer is not None when pooling_output is None. - """ - for eco in engine_core_outputs: - req_state = self.request_states.get(eco.request_id) - if req_state is None or not isinstance(req_state, OmniRequestState): - continue - - new_token_ids = eco.new_token_ids - finish_reason = eco.finish_reason - stop_reason = eco.stop_reason - kv_transfer_params = eco.kv_transfer_params - routed_experts = eco.routed_experts - req_state.num_cached_tokens = eco.num_cached_tokens - req_state.is_prefilling = False - - if request_output := req_state.make_request_output( - new_token_ids, - None, # pooling_output - finish_reason, - stop_reason, - kv_transfer_params, - routed_experts, - ): - if req_state.queue is not None: - req_state.queue.put(request_output) - - if finish_reason is not None: - self._finish_request(req_state) - - def _update_stats_from_output( - self, - req_state: RequestState, - engine_core_output: EngineCoreOutput, - engine_core_timestamp: float | None, - iteration_stats: IterationStats | None, - ): - was_prefilling = req_state.is_prefilling - native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None - previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 - - # NOTE: We pass ``None`` for *iteration_stats* to the parent so that - # the upstream's ``_update_stats_from_output`` logs stats via its own - # ``RequestStateStats`` (req_state.stats) without touching the omni - # ``IterationStats`` object. Below we re-apply the same update using - # *native_stats* (the omni-specific ``RequestStateStats`` attached to - # ``OmniRequestState``), which correctly feeds the omni metrics path. - # If the upstream adds new logic in ``_update_stats_from_output`` (e.g. - # additional stat fields), this override must be kept in sync. - super()._update_stats_from_output( - req_state, - engine_core_output, - engine_core_timestamp, - None, - ) - - if iteration_stats is None or engine_core_timestamp is None or native_stats is None: - return - - iteration_stats.update_from_output( - engine_core_output, - engine_core_timestamp, - was_prefilling, - native_stats, - self.lora_states, - req_state.lora_name, - ) - record = self._native_text_metric_record(req_state.external_req_id) - if was_prefilling: - record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) - return - - if previous_last_token_ts > 0: - itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0) - itls_ms = record.setdefault("vllm_itls_ms", []) - itls_ms.append(itl_ms) - record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) - - def _update_stats_from_finished( - self, - req_state: RequestState, - finish_reason: FinishReason | None, - iteration_stats: IterationStats | None, - ): - super()._update_stats_from_finished(req_state, finish_reason, None) - - native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None - if iteration_stats is None or native_stats is None or finish_reason is None: - return - - iteration_stats.update_from_finished_request( - finish_reason=finish_reason, - request_id=req_state.external_req_id, - num_prompt_tokens=req_state.prompt_len, - max_tokens_param=req_state.max_tokens_param, - req_stats=native_stats, - num_cached_tokens=req_state.num_cached_tokens, - ) - if not iteration_stats.finished_requests: - return - - finished_request = iteration_stats.finished_requests[-1] - if finished_request.request_id != req_state.external_req_id: - return - self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = ( - float(finished_request.mean_time_per_output_token) * 1000.0 - ) diff --git a/vllm_omni/engine/stage_client.py b/vllm_omni/engine/stage_client.py index 1f10a219c6a..6db0909557e 100644 --- a/vllm_omni/engine/stage_client.py +++ b/vllm_omni/engine/stage_client.py @@ -11,8 +11,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniTokensPrompt from vllm_omni.outputs import OmniRequestOutput -from vllm_omni.engine.output_modality import FinalOutputModalityType from vllm_omni.inputs.data import OmniSamplingParams +from vllm_omni.outputs.output_metadata import FinalOutputModalityType class StageClient(Protocol): @@ -26,9 +26,11 @@ class StageClient(Protocol): stage_id: int replica_id: int stage_type: str + model_stage: str | None final_output: bool final_output_type: FinalOutputModalityType | None default_sampling_params: OmniSamplingParams + prompt_expand_func: Any | None requires_multimodal_data: bool custom_process_input_func: Any | None engine_input_source: Sequence[int] @@ -66,8 +68,6 @@ async def collective_rpc_async( class StagePoolLLMClient(StagePoolClient, Protocol): """Pool-facing API for LLM-style stages.""" - model_stage: str | None - async def add_request_async(self, request: EngineCoreRequest) -> None: ... async def get_output_async(self) -> EngineCoreOutputs: ... @@ -100,12 +100,4 @@ async def add_request_async( kv_sender_info: dict[int, dict[str, Any]] | None = None, ) -> None: ... - async def add_batch_request_async( - self, - request_id: str, - prompts: list[OmniPromptType], - sampling_params: OmniDiffusionSamplingParams, - kv_sender_info: dict[int, dict[str, Any]] | None = None, - ) -> None: ... - def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: ... diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 2c6c006659a..5fbbbc46da1 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -6,18 +6,12 @@ from __future__ import annotations -import contextlib import inspect -import multiprocessing.connection import os import socket -import threading -import time -import weakref from typing import TYPE_CHECKING, Any from urllib.parse import urlparse -import psutil import vllm.v1.engine as _vllm_engine_module import vllm.v1.engine.core_client as _vllm_core_client_module from vllm.logger import init_logger @@ -43,8 +37,6 @@ logger = init_logger(__name__) -SHUTDOWN_TIMEOUT_S = 5 - def _default_process_engine_inputs( source_outputs: list[Any], @@ -78,10 +70,8 @@ class StageEngineCoreClientBase(StageClientBase): - outputs_queue, output_queue_task - All utility methods (get_output_async, abort_requests_async, etc.) - The subprocess is spawned externally via ``spawn_stage_core`` / - ``complete_stage_handshake`` from *stage_engine_core_proc.py*. - In single-stage CLI mode, the client may instead attach to an - ``engine_manager`` / ``coordinator`` pair created elsewhere. + The stage engine subprocesses are owned by vLLM-style engine managers + (`resources.engine_manager`), matching `MPClient` ownership. """ replica_id: int = 0 @@ -93,7 +83,6 @@ def make_async_mp_client( log_stats: bool = False, metadata: StageMetadata | None = None, client_addresses: dict[str, str] | None = None, - proc: Any = None, engine_manager: Any = None, coordinator: Any = None, client_count: int = 1, @@ -107,7 +96,6 @@ def make_async_mp_client( log_stats=log_stats, metadata=metadata, client_addresses=client_addresses, - proc=proc, engine_manager=engine_manager, coordinator=coordinator, client_count=client_count, @@ -125,7 +113,6 @@ def __init__( executor_class: type, log_stats: bool = False, client_addresses: dict[str, str] | None = None, - proc: Any = None, client_count: int = 1, client_index: int = 0, *, @@ -140,10 +127,8 @@ def __init__( via helpers in stage_init_utils.py. This constructor just stores metadata and calls super().__init__(). - The subprocess is spawned externally via ``spawn_stage_core`` / - ``complete_stage_handshake`` (see *stage_engine_core_proc.py*). - The resulting ``proc`` handle is passed in so this client can - manage the process lifecycle on shutdown. + The subprocess lifecycle is owned by the engine manager attached to + vLLM's background resources, not by this client directly. """ # -------- Stage metadata (public fields used at runtime) -------- self.replica_id = 0 @@ -151,19 +136,17 @@ def __init__( self.stage_id = metadata.stage_id self.replica_id = getattr(metadata, "replica_id", 0) self.stage_type = metadata.stage_type - self.engine_output_type = metadata.engine_output_type + self.model_stage = metadata.model_stage self.is_comprehension = metadata.is_comprehension self.requires_multimodal_data = metadata.requires_multimodal_data self.engine_input_source = metadata.engine_input_source self.final_output = metadata.final_output self.final_output_type = metadata.final_output_type self.default_sampling_params = metadata.default_sampling_params + self.prompt_expand_func = metadata.prompt_expand_func self.custom_process_input_func = metadata.custom_process_input_func - self.model_stage = metadata.model_stage self.engine_outputs: Any = None - self._proc = proc - self._shutting_down = False self.client_addresses = dict(client_addresses or {}) self._omni_kv_config = getattr(getattr(vllm_config, "model_config", None), "omni_kv_config", None) self._kv_sender_host = self._resolve_contact_host() @@ -203,6 +186,7 @@ def __init__( ) if engine_manager is not None: self.resources.engine_manager = engine_manager + self.start_engine_core_monitor() if coordinator is not None: self.resources.coordinator = coordinator except Exception: @@ -226,9 +210,6 @@ def __init__( self._initialize_kv_sender_endpoint() - if self._proc is not None: - self._start_proc_monitor() - logger.info( "[%s] stage-%s [rep-%s] EngineCore running", client_name, @@ -236,64 +217,6 @@ def __init__( self.replica_id, ) - def _start_proc_monitor(self) -> None: - """Start a daemon thread that watches the subprocess sentinel. - - When the subprocess dies without sending the ZMQ ``ENGINE_CORE_DEAD`` - sentinel (e.g. SIGKILL, segfault, OOM-killer), this thread sets - ``resources.engine_dead`` so subsequent calls raise - ``EngineDeadError``. - """ - proc = self._proc - resources_ref = weakref.ref(self.resources) - stage_id = self.stage_id - replica_id = self.replica_id - - def _monitor() -> None: - try: - multiprocessing.connection.wait([proc.sentinel]) - except Exception: - return - - # Give multiprocessing a brief chance to publish a stable exitcode - # before emitting the monitor log. Without this, the sentinel can - # fire first and the log may misleadingly print ``None``. - with contextlib.suppress(Exception): - proc.join(timeout=0) - exitcode = proc.exitcode - if exitcode is None: - deadline = time.monotonic() + 0.2 - while exitcode is None and time.monotonic() < deadline: - time.sleep(0.01) - with contextlib.suppress(Exception): - proc.join(timeout=0) - exitcode = proc.exitcode - - resources = resources_ref() - if resources is None or resources.engine_dead or self._shutting_down: - return - if exitcode == 0: - logger.info( - "[StageEngineCoreClient] stage-%s [rep-%s] subprocess exited cleanly.", - stage_id, - replica_id, - ) - return - resources.engine_dead = True - logger.error( - "[StageEngineCoreClient] stage-%s [rep-%s] subprocess died unexpectedly (exit code %s).", - stage_id, - replica_id, - exitcode, - ) - - t = threading.Thread( - target=_monitor, - daemon=True, - name=f"StageCoreProcMonitor-{stage_id}", - ) - t.start() - def check_health(self) -> None: """Raise ``EngineDeadError`` if the stage subprocess is dead. @@ -302,15 +225,12 @@ def check_health(self) -> None: """ if self.resources.engine_dead: raise EngineDeadError(f"Stage-{self.stage_id} engine core is dead") - if self._proc is not None and not self._proc.is_alive(): - self.resources.engine_dead = True - raise EngineDeadError(f"Stage-{self.stage_id} subprocess is not alive (exit code {self._proc.exitcode})") # ==================== Overrides ==================== async def add_request_async(self, request: EngineCoreRequest) -> None: """Add request to the stage engine core.""" - logger.info( + logger.debug( "[%s] stage-%s [rep-%s] add request: %s", self.__class__.__name__, self.stage_id, @@ -335,7 +255,16 @@ def _detect_local_ip() -> str | None: return None def _resolve_contact_host(self) -> str | None: - """Resolve a routable host for this stage from its client addresses.""" + """Resolve a routable host for this stage from its client addresses. + + For remote LLM replicas the ZMQ sockets are bound on the head node, + so the ZMQ addresses contain the head's IP. The ``replica_host`` key + (injected by DistStageRuntime) carries the actual replica IP where + the KV connector binds — prefer it when available. + """ + replica_host = self.client_addresses.get("replica_host") + if replica_host: + return replica_host for key in ("input_address", "output_address", "stats_update_address"): address = self.client_addresses.get(key) if not address: @@ -503,41 +432,6 @@ async def collective_rpc_async( kwargs=kwargs, ) - def shutdown(self, timeout: float | None = None) -> None: - """Shutdown managed resources and any externally spawned subprocess.""" - self._shutting_down = True - child_procs: list[psutil.Process] = [] - if self._proc is not None and self._proc.pid is not None: - try: - child_procs = psutil.Process(self._proc.pid).children(recursive=True) - except psutil.Error: - child_procs = [] - - try: - super().shutdown(timeout=timeout) - finally: - if self._proc is not None and self._proc.is_alive(): - self._proc.terminate() - self._proc.join(timeout=SHUTDOWN_TIMEOUT_S) - if self._proc.is_alive(): - self._proc.kill() - self._proc.join(timeout=SHUTDOWN_TIMEOUT_S) - - alive_children = [proc for proc in child_procs if proc.is_running()] - for proc in alive_children: - try: - proc.terminate() - except psutil.Error: - pass - _, still_alive = psutil.wait_procs(alive_children, timeout=SHUTDOWN_TIMEOUT_S) - for proc in still_alive: - try: - proc.kill() - except psutil.Error: - pass - # The process handle is no longer reliable after best-effort cleanup. - self._proc = None - class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient): """Stage async client backed by vLLM's ``AsyncMPClient``.""" diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 62c5b2a8f6f..e3f75d9e33e 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -10,40 +10,28 @@ import contextlib import os import signal -from multiprocessing.process import BaseProcess -from typing import TYPE_CHECKING, Any +from typing import Any -import msgspec import vllm.v1.engine.core as _vllm_engine_core_module -import zmq from vllm.logger import init_logger from vllm.transformers_utils.config import ( maybe_register_config_serialize_by_value, ) -from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx from vllm.utils.system_utils import ( decorate_logs, - get_mp_context, set_process_title, ) from vllm.v1.engine import EngineCoreRequestType from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState from vllm.v1.engine.utils import ( - EngineHandshakeMetadata, EngineZmqAddresses, SignalCallback, - get_engine_zmq_addresses, ) -from vllm.v1.utils import shutdown -from vllm_omni.distributed.omni_coordinator import OmniCoordClientForStage +from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.stage_init_utils import set_death_signal -if TYPE_CHECKING: - from vllm.config import VllmConfig - from vllm.v1.executor import Executor - logger = init_logger(__name__) @@ -106,7 +94,7 @@ def run_stage_core( ) engine_core: StageEngineCoreProc | None = None - coord_client: OmniCoordClientForStage | None = None + coord_client = None try: # NOTE: previous revisions hardcoded data_parallel_size=1 here # (TODO referencing issue #984). The hardcoding has been removed @@ -154,29 +142,17 @@ def run_stage_core( "EngineCore handshake did not populate input/output addresses; " "cannot start OmniCoordClientForStage" ) - coord_client = OmniCoordClientForStage( + scheduler = getattr(engine_core, "scheduler", None) + if scheduler is None: + raise RuntimeError("EngineCore scheduler is not initialized") + coord_client = create_stage_coord_client( coord_zmq_addr=omni_coordinator_address, input_addr=addresses.inputs[0], output_addr=addresses.outputs[0], stage_id=int(omni_stage_id), + queue_length_getter=scheduler.get_num_unfinished_requests, ) - def _refresh_queue_length() -> None: - """Pre-heartbeat hook: refresh queue_length from scheduler.""" - scheduler = getattr(engine_core, "scheduler", None) - if scheduler is None: - return - try: - coord_client._queue_length = int( # type: ignore[union-attr] - scheduler.get_num_unfinished_requests() - ) - except Exception: - # Live scheduler stats are best-effort — heartbeats - # must not fail because of a stats lookup error. - pass - - coord_client._on_heartbeat = _refresh_queue_length - def wakeup_engine() -> None: engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) @@ -212,109 +188,3 @@ def signal_handler(signum: int, frame: Any) -> None: coord_client.close() if engine_core is not None: engine_core.shutdown() - - -def spawn_stage_core( - vllm_config: VllmConfig, - executor_class: type[Executor], - log_stats: bool = False, -) -> tuple[EngineZmqAddresses, BaseProcess, str]: - """Spawn a *StageEngineCoreProc* subprocess without performing the handshake. - - Must be called while the correct device env vars are set (e.g. under - the stage-launch lock). Call ``complete_stage_handshake`` afterwards. - - Returns ``(addresses, process, handshake_address)``. - """ - addresses = get_engine_zmq_addresses(vllm_config) - handshake_address = get_open_zmq_ipc_path() - - ctx = get_mp_context() - proc = ctx.Process( - target=StageEngineCoreProc.run_stage_core, - name="StageEngineCoreProc", - kwargs={ - "vllm_config": vllm_config, - "local_client": True, - "handshake_address": handshake_address, - "executor_class": executor_class, - "log_stats": log_stats, - "dp_rank": 0, - "local_dp_rank": 0, - }, - ) - proc.start() - return addresses, proc, handshake_address - - -def complete_stage_handshake( - proc: BaseProcess, - handshake_address: str, - addresses: EngineZmqAddresses, - vllm_config: VllmConfig, - handshake_timeout: int, -) -> None: - """Perform the HELLO/INIT/READY handshake with an already-spawned proc. - - On failure the process is terminated before re-raising. - """ - try: - _perform_handshake(proc, handshake_address, addresses, vllm_config, handshake_timeout) - except Exception: - shutdown([proc]) - raise - - -def _perform_handshake( - proc: BaseProcess, - handshake_address: str, - addresses: EngineZmqAddresses, - vllm_config: VllmConfig, - handshake_timeout: int, -) -> None: - """Run the HELLO / INIT / READY handshake with the subprocess.""" - with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: - poller = zmq.Poller() - poller.register(handshake_socket, zmq.POLLIN) - poller.register(proc.sentinel, zmq.POLLIN) - - identity, msg = _recv(poller, handshake_socket, proc, "HELLO", handshake_timeout) - if msg.get("status") != "HELLO": - raise RuntimeError(f"Expected HELLO, got: {msg}") - - init_payload = EngineHandshakeMetadata( - addresses=addresses, - parallel_config={}, - ) - handshake_socket.send_multipart([identity, msgspec.msgpack.encode(init_payload)]) - - identity, msg = _recv(poller, handshake_socket, proc, "READY", handshake_timeout) - if msg.get("status") != "READY": - raise RuntimeError(f"Expected READY, got: {msg}") - num_gpu_blocks = msg.get("num_gpu_blocks") - if num_gpu_blocks is not None: - vllm_config.cache_config.num_gpu_blocks = num_gpu_blocks - - -def _recv( - poller: zmq.Poller, - handshake_socket: zmq.Socket, - proc: BaseProcess, - expected: str, - timeout_s: int = 600, -) -> tuple[bytes, dict]: - """Wait for one handshake message; raise if the process dies first.""" - timeout_ms = timeout_s * 1000 - while True: - events = dict(poller.poll(timeout=timeout_ms)) - if not events: - raise TimeoutError( - f"Timed out waiting for {expected} from StageEngineCoreProc after {timeout_s}s. " - f"This typically indicates model loading or initialization is taking too long. " - f"Consider increasing `stage_init_timeout` for large models." - ) - if handshake_socket in events: - identity, raw = handshake_socket.recv_multipart() - return identity, msgspec.msgpack.decode(raw) - if proc.exitcode is not None: - raise RuntimeError(f"StageEngineCoreProc died during {expected} (exit code {proc.exitcode})") diff --git a/vllm_omni/engine/omni_core_engine_proc_manager.py b/vllm_omni/engine/stage_engine_core_proc_manager.py similarity index 92% rename from vllm_omni/engine/omni_core_engine_proc_manager.py rename to vllm_omni/engine/stage_engine_core_proc_manager.py index 8df5fd05326..6b613bf8e23 100644 --- a/vllm_omni/engine/omni_core_engine_proc_manager.py +++ b/vllm_omni/engine/stage_engine_core_proc_manager.py @@ -38,16 +38,14 @@ logger = init_logger(__name__) -try: - # ``set_device_control_env_var`` lives next to CoreEngineProcManager and - # is only required for non-CUDA DP, so we tolerate its absence on - # older / future vLLM revisions. - from vllm.v1.engine.utils import set_device_control_env_var # type: ignore -except ImportError: # pragma: no cover - depends on vLLM build - set_device_control_env_var = None # type: ignore[assignment] +# ``set_device_control_env_var`` was removed from upstream vllm. +# It was only required for non-CUDA DP; set to None so the existing +# guard at the call-site (vllm_config is not None and ... is not None) +# skips the call transparently. +set_device_control_env_var = None -class OmniCoreEngineProcManager(CoreEngineProcManager): +class StageEngineCoreProcManager(CoreEngineProcManager): """Spawn :class:`StageEngineCoreProc` subprocesses with omni kwargs. The body mirrors :class:`CoreEngineProcManager.__init__` because the @@ -102,7 +100,7 @@ def __init__( # Intra-replica vLLM DP mesh (i.e. ``data_parallel_size`` ranks sharing # one engine, one DPCoordinator, one set of weights). Distinct from # the omni-level notion of multiple independent replicas of a stage — - # those each spawn their own OmniCoreEngineProcManager and never join + # those each spawn their own StageEngineCoreProcManager and never join # a vLLM DP group across replicas. has_intra_replica_dp = vllm_config.parallel_config.data_parallel_size > 1 diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index a2744eea47b..7a55bb93f2b 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -4,30 +4,42 @@ import contextlib import dataclasses +import os import socket import threading from collections.abc import Callable, Iterator -from dataclasses import dataclass +from dataclasses import dataclass, field +from multiprocessing import connection +from types import SimpleNamespace from typing import Any import msgspec import zmq from omegaconf import OmegaConf -from vllm.config import CacheConfig, VllmConfig +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.network_utils import get_open_ports_list, zmq_socket_ctx from vllm.v1.engine.coordinator import DPCoordinator from vllm.v1.engine.utils import ( - STARTUP_POLL_PERIOD_MS, CoreEngine, CoreEngineProcManager, - CoreEngineState, - EngineHandshakeMetadata, EngineZmqAddresses, + get_engine_zmq_addresses, wait_for_engine_startup, ) from vllm.v1.executor import Executor +from vllm_omni.distributed.omni_connectors.utils import initialization +from vllm_omni.engine import stage_init_utils +from vllm_omni.engine.stage_init_utils import ( + acquire_device_locks, + build_diffusion_config, + initialize_diffusion_stage, + release_device_locks, +) +from vllm_omni.entrypoints.utils import inject_omni_kv_config +from vllm_omni.platforms import current_omni_platform + logger = init_logger(__name__) StageRoute = tuple[int, int] @@ -95,6 +107,11 @@ class StageAllocation: # Output channel: client binds PULL, engine connects PUSH output_bind_address: str output_connect_address: str + # The replica's routable IP (from registration). For LLM remote replicas + # the ZMQ sockets are bound on the head, but the KV connector runs on + # the replica host — this field preserves the replica's actual IP so the + # orchestrator can advertise the correct KV sender endpoint. + replica_host: str | None = None @dataclass(frozen=True) @@ -106,11 +123,36 @@ class StageCoordinatorAddresses: frontend_stats_publish_address: str | None = None +@dataclass +class StageReplicaResources: + """Resources created while launching one stage replica.""" + + manager: Any | None = None + coordinator: Any | None = None + addresses: EngineZmqAddresses | None = None + lock_fds: list[int] = field(default_factory=list) + + # --------------------------------------------------------------------------- # OmniMasterServer # --------------------------------------------------------------------------- +def _port_from_zmq_address(address: str | None) -> int | None: + """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``. + + Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return + ``None`` so callers can skip seeding them into the port-dedup set. + """ + if not address: + return None + tail = address.rsplit(":", 1)[-1] + try: + return int(tail) + except ValueError: + return None + + class OmniMasterServer: """Registration server for single-stage engine startup.""" @@ -131,10 +173,25 @@ def __init__( self._stage_configs: dict[StageRoute, Any] = {} self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {} self._stage_config_events: dict[StageRoute, threading.Event] = {} + # Ports already handed out by *this* server. ``get_open_ports_list`` only + # guarantees uniqueness *within* a single call; a per-route call cannot + # see ports drawn for other routes, so two stages/replicas can draw the + # same ephemeral port and the second engine to ``bind()`` it dies with + # ``zmq.error.ZMQError: Address already in use``. Multi-stage models + # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit + # this intermittently. We dedup every allocated port against this set, + # seeded with the registration port (and coordinator ROUTER port) that + # are already bound on the same host. + self._allocated_ports: set[int] = set() + if master_port: + self._allocated_ports.add(int(master_port)) # Coordinator ROUTER address echoed back in every registration reply # so OmniCoordClientForStage knows where to connect from inside the # engine subprocess. self._coordinator_router_address = coordinator_router_address + coord_port = _port_from_zmq_address(coordinator_router_address) + if coord_port is not None: + self._allocated_ports.add(coord_port) # Fires only for *newly assigned* (auto-assigned) replicas, not for # head-side pre-allocated slots that already have head-side clients. self._on_register = on_register @@ -142,12 +199,11 @@ def __init__( # registrations from multiple headless processes for the same stage # don't race on the routing table. self._alloc_lock = threading.Lock() - self._stage_ids_known: set[int] = set(int(sid) for sid in stage_ids) self._thread: threading.Thread | None = None self._stop_event = threading.Event() stage_replica_counts = dict(stage_replica_counts or {}) - # Slots the *head* itself will fill via ``launch_omni_core_engines`` + # Slots the *head* itself will fill via ``_launch_omni_core_engines`` # / its own ``register_stage_with_omni_master`` call. Auto-assigning # headless registrations must skip these even when they appear # ``_stage_configs``-unfilled — otherwise a fast headless on the same @@ -187,11 +243,6 @@ def port(self) -> int: """Return the registration port exposed to stage launchers.""" return self._port - @property - def coordinator_router_address(self) -> str | None: - """Return the OmniCoordinator ROUTER address echoed to replicas.""" - return self._coordinator_router_address - def get_allocation(self, stage_id: int, replica_id: int = 0) -> StageAllocation: """Return the full address allocation for *stage_id*.""" return self._stage_routes[(stage_id, replica_id)] @@ -214,7 +265,7 @@ def _allocate_route_locked(self, stage_id: int, replica_id: int) -> StageAllocat self._stage_config_events[route] = threading.Event() self._stage_coordinator_addresses[route] = StageCoordinatorAddresses() - hs_port, inp_port, out_port = get_open_ports_list(count=3) + hs_port, inp_port, out_port = self._alloc_unique_ports(3) alloc = StageAllocation( handshake_bind_address=f"tcp://{self._address}:{hs_port}", handshake_connect_address=f"tcp://{self._address}:{hs_port}", @@ -226,6 +277,32 @@ def _allocate_route_locked(self, stage_id: int, replica_id: int) -> StageAllocat self._stage_routes[route] = alloc return alloc + def _alloc_unique_ports(self, count: int) -> list[int]: + """Return ``count`` open ports unique across every route this server owns. + + ``get_open_ports_list`` dedups only within its own call, so we redraw any + port that collides with one already recorded in ``self._allocated_ports`` + (registration/coordinator ports plus every previously-allocated route + port) and register the winners before returning. Callers already hold the + relevant allocation context (``__init__`` is single-threaded; + registration holds ``self._alloc_lock``). + """ + picked: list[int] = [] + # Bounded retry budget so a pathological host (near port exhaustion) + # fails loudly instead of spinning forever. + for _ in range(64): + for port in get_open_ports_list(count=count - len(picked)): + if port in self._allocated_ports: + continue + self._allocated_ports.add(port) + picked.append(port) + if len(picked) == count: + return picked + raise RuntimeError( + f"[OmniMasterServer] Could not allocate {count} unique open ports " + f"(host={self._address}); {len(self._allocated_ports)} already in use." + ) + def _next_free_replica_id(self, stage_id: int) -> int: """Return the next replica id to assign for an auto-assign registration. @@ -237,7 +314,7 @@ def _next_free_replica_id(self, stage_id: int) -> int: this stage has been filled do we allocate a fresh id. Slots in ``_head_local_slots`` are reserved for the head's own - ``launch_omni_core_engines`` registration. Auto-assign must skip + ``_launch_omni_core_engines`` registration. Auto-assign must skip them even when ``_stage_configs`` shows them unfilled — otherwise a same-host headless that registers before the head's own ``register_stage_with_omni_master`` call would steal slot 0. @@ -314,14 +391,6 @@ def get_stage_coordinator_addresses( return self._stage_coordinator_addresses[key] - def get_client_addresses(self, stage_id: int, replica_id: int = 0) -> dict[str, str]: - """Return the addresses the client-side sockets should *bind* to.""" - alloc = self.get_allocation(stage_id, replica_id) - return { - "input_address": alloc.input_bind_address, - "output_address": alloc.output_bind_address, - } - def get_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> EngineZmqAddresses: """Return EngineZmqAddresses using the *bind* (client) side addresses.""" alloc = self.get_allocation(stage_id, replica_id) @@ -338,6 +407,13 @@ def get_engine_zmq_addresses(self, stage_id: int, replica_id: int = 0) -> Engine outputs=[alloc.output_connect_address], ) + def get_replica_host(self, stage_id: int, replica_id: int = 0) -> str | None: + """Return the replica's routable IP if it registered one.""" + alloc = self._stage_routes.get((stage_id, replica_id)) + if alloc is None: + return None + return alloc.replica_host + # ------------------------------------------------------------------ # Lifecycle # ------------------------------------------------------------------ @@ -473,59 +549,16 @@ def _handle_registration(self, reg_socket: zmq.Socket) -> StageRoute | None: # else: alloc = self._stage_routes[(stage_id, replica_id)] - # Cross-host override: when the registering replica advertised - # its own bind address + ports, rewrite the StageAllocation so - # each socket is rooted on the host that actually binds it - # (the master's pre-allocated ports are unreachable from a - # remote replica's host). - # - # Diffusion and LLM stages have different binder ownership: - # - # Diffusion remote replica (StageDiffusionProc): - # - handshake: replica binds -> rewrite to replica IP - # - input : replica binds -> rewrite to replica IP - # - output : replica binds -> rewrite to replica IP - # - # LLM remote replica (CoreClient on head): - # - handshake: head binds (``connect_remote_engine_cores``) - # -> keep on master IP, worker TCP-connects - # - input : head binds (``CoreClient`` ROUTER) - # -> keep on master IP, worker TCP-connects - # - output : head binds (``CoreClient`` PULL — default - # bind=True for PULL in ``make_zmq_socket``) - # -> keep on master IP, worker TCP-connects - # - # The registrant indicates which case via the boolean - # ``replica_binds_sockets`` payload flag. It defaults to - # True (the diffusion / single-host case) so older callers - # still get the previous full-rewrite semantics. For LLM - # remote replicas, the master keeps every address on its - # own host and the remote worker establishes 3 outbound - # TCP connections to the master. + # Cross-host remote replicas connect to head-owned sockets, but + # still advertise their routable host for KV connector endpoints. new_bind_address = msg.get("replica_bind_address") if new_bind_address: - replica_binds_sockets = bool(msg.get("replica_binds_sockets", True)) - if replica_binds_sockets: - hs_port = int(msg["replica_handshake_port"]) - inp_port = int(msg["replica_input_port"]) - out_port = int(msg["replica_output_port"]) - hs_bind_addr = f"tcp://{new_bind_address}:{hs_port}" - inp_bind_addr = f"tcp://{new_bind_address}:{inp_port}" - out_bind_addr = f"tcp://{new_bind_address}:{out_port}" - alloc = StageAllocation( - handshake_bind_address=hs_bind_addr, - handshake_connect_address=hs_bind_addr, - input_bind_address=inp_bind_addr, - input_connect_address=inp_bind_addr, - output_bind_address=out_bind_addr, - output_connect_address=out_bind_addr, - ) - self._stage_routes[(stage_id, replica_id)] = alloc + alloc.replica_host = new_bind_address logger.info( - "[OmniMasterServer] Stage %d replica %d cross-host bind (sockets bound on %s; replica_ip=%s)", + "[OmniMasterServer] Stage %d replica %d registered from host %s " + "(serving sockets remain head-owned)", stage_id, replica_id, - "replica" if replica_binds_sockets else "master", new_bind_address, ) @@ -599,11 +632,9 @@ def _detect_local_bind_address(master_address: str, master_port: int) -> str: Uses a connected UDP socket as a routing-table probe: ``connect()`` on SOCK_DGRAM sends no packets but forces a route lookup, after which ``getsockname()[0]`` exposes the source IP that an outbound packet to - ``(master_address, master_port)`` would carry. For a co-located master - this returns the loopback or eth0 IP (same effect as the legacy - ``self._address`` behaviour); for a remote master it returns the - NIC IP that's actually reachable from the master — which is exactly - the address the headless's per-stage ZMQ sockets must bind on. + ``(master_address, master_port)`` would carry. For a remote master this + returns the NIC IP that's reachable from the master, which is exactly the + address the headless's per-stage ZMQ sockets must bind on. """ s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) try: @@ -613,6 +644,21 @@ def _detect_local_bind_address(master_address: str, master_port: int) -> str: s.close() +def _engine_count(parallel_config: Any) -> int: + data_parallel_size_local = getattr(parallel_config, "data_parallel_size_local", None) + if data_parallel_size_local is not None and data_parallel_size_local > 0: + return int(data_parallel_size_local) + return max(1, int(parallel_config.data_parallel_size)) + + +def _single_diffusion_parallel_config(*, local_client: bool) -> Any: + return SimpleNamespace( + data_parallel_size_local=1 if local_client else 0, + data_parallel_hybrid_lb=False, + data_parallel_external_lb=False, + ) + + def register_stage_with_omni_master( *, omni_master_address: str, @@ -620,21 +666,11 @@ def register_stage_with_omni_master( omni_stage_id: int, omni_stage_config: Any = None, coordinator: DPCoordinator | None = None, - return_addresses: bool = False, replica_id: int | None = 0, - return_full_response: bool = False, replica_bind_address: str | None = None, - replica_binds_sockets: bool = True, -) -> str | tuple[str, str, str] | StageRegistrationResponse: +) -> StageRegistrationResponse: """Register a stage with the omni master server. - Returns the per-stage handshake address by default. When - ``return_addresses`` is true, also returns the stage input/output - addresses allocated by the master. When ``return_full_response`` is - true, returns the full :class:`StageRegistrationResponse` including the - assigned ``replica_id`` and the OmniCoordinator ROUTER address (if - published by the master). - Pass ``replica_id=None`` to request auto-assignment of a free replica id by the master (used by headless launchers). """ @@ -660,30 +696,11 @@ def register_stage_with_omni_master( payload["coordinator_output"] = coordinator_output payload["frontend_stats_publish_address"] = coordinator.get_stats_publish_address() - # Always advertise THIS host's local bind address + 3 locally - # free ports so the master can root the per-stage socket - # allocation on the replica's own interface. For a co-located - # replica the detected IP matches the master's address and - # the override is a no-op semantically; for a cross-host - # replica it's what makes the headless's ROUTER bind succeed - # (otherwise the master would hand back ``tcp://<master_ip>:port`` - # and ``zmq.bind`` would EADDRNOTAVAIL on the remote host). + # Advertise this host for KV connector routing. Serving/control + # sockets stay head-owned; remote workers only connect to them. if replica_bind_address is None: replica_bind_address = _detect_local_bind_address(omni_master_address, omni_master_port) - hs_port, inp_port, out_port = get_open_ports_list(count=3) payload["replica_bind_address"] = replica_bind_address - payload["replica_handshake_port"] = hs_port - payload["replica_input_port"] = inp_port - payload["replica_output_port"] = out_port - # ``False`` only for LLM headless replicas: the head's - # ``connect_remote_engine_cores`` is the binder for the - # handshake ROUTER, and ``CoreClient.__init__`` binds the - # input ROUTER and the output PULL (``make_zmq_socket`` - # defaults bind=True for PULL). The master must keep all - # three addresses on the master's host so the head can - # ``bind`` them; the remote LLM worker TCP-connects across - # hosts on all three. - payload["replica_binds_sockets"] = bool(replica_binds_sockets) reg_sock.send(msgspec.msgpack.encode(payload)) timeout_ms = _DEFAULT_STARTUP_TIMEOUT_S * 1_000 @@ -712,81 +729,13 @@ def register_stage_with_omni_master( finally: reg_ctx.term() - if return_full_response: - return StageRegistrationResponse( - handshake_address=handshake_address, - input_address=input_address, - output_address=output_address, - replica_id=assigned_replica_id, - coordinator_router_address=coord_router_addr, - ) - if return_addresses: - return handshake_address, input_address, output_address - return handshake_address - - -def _wait_for_omni_engine_startup( - handshake_socket: zmq.Socket, - engine_addresses: EngineZmqAddresses, - engines: list[CoreEngine], - cache_config: CacheConfig, -) -> None: - """Wait for omni-managed engines to finish the HELLO/READY handshake.""" - conn_pending = len(engines) - start_pending = 0 - - poller = zmq.Poller() - poller.register(handshake_socket, zmq.POLLIN) - - while conn_pending or start_pending: - events = poller.poll(STARTUP_POLL_PERIOD_MS) - if not events: - logger.debug( - "[omni] Waiting for %d engine(s) to connect, %d to start.", - conn_pending, - start_pending, - ) - continue - - eng_identity, msg_bytes = handshake_socket.recv_multipart() - eng_index = int.from_bytes(eng_identity, "little") - engine = next((e for e in engines if e.identity == eng_identity), None) - if engine is None: - raise RuntimeError(f"[omni] Handshake message from unexpected engine rank: {eng_index}") - - msg = msgspec.msgpack.decode(msg_bytes) - status: str = msg["status"] - - if status == "HELLO" and engine.state == CoreEngineState.NEW: - init_message = msgspec.msgpack.encode( - EngineHandshakeMetadata(addresses=engine_addresses, parallel_config={}) - ) - handshake_socket.send_multipart((eng_identity, init_message), copy=False) - conn_pending -= 1 - start_pending += 1 - engine.state = CoreEngineState.CONNECTED - logger.debug("[omni] HELLO from engine %d", eng_index) - - elif status == "READY" and engine.state == CoreEngineState.CONNECTED: - # Upstream vllm >=0.19 dropped `num_gpu_blocks` from the READY - # handshake payload (the field is now communicated out-of-band via - # stats/health topics); tolerate both legacy and new message - # shapes so the omni handshake keeps working across rebases. - reported_blocks = msg.get("num_gpu_blocks") - if reported_blocks is not None: - cache_config.num_gpu_blocks = (cache_config.num_gpu_blocks or 0) + int(reported_blocks) - if engine_addresses.frontend_stats_publish_address is None: - engine_addresses.frontend_stats_publish_address = msg.get("dp_stats_address") - start_pending -= 1 - engine.state = CoreEngineState.READY - logger.debug( - "[omni] READY from engine %d (num_gpu_blocks=%s)", - eng_index, - "unknown" if reported_blocks is None else reported_blocks, - ) - - else: - raise RuntimeError(f"[omni] Unexpected status '{status}' from engine {eng_index} in state {engine.state}.") + return StageRegistrationResponse( + handshake_address=handshake_address, + input_address=input_address, + output_address=output_address, + replica_id=assigned_replica_id, + coordinator_router_address=coord_router_addr, + ) @contextlib.contextmanager @@ -795,16 +744,11 @@ def connect_remote_engine_cores( omni_master_server: OmniMasterServer, stage_id: int, replica_id: int = 0, -) -> Iterator[tuple[None, DPCoordinator | None, EngineZmqAddresses, None]]: +) -> Iterator[StageReplicaResources]: """Wait for remote engine cores to connect through the omni handshake.""" addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config - # Mirror the engine-count logic from launch_omni_core_engines. - remote_engine_count = ( - parallel_config.data_parallel_size_local - if parallel_config.data_parallel_size_local is not None and parallel_config.data_parallel_size_local > 0 - else max(1, parallel_config.data_parallel_size) - ) + remote_engine_count = _engine_count(parallel_config) start_index = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 coordinator = None @@ -828,18 +772,75 @@ def connect_remote_engine_cores( handshake_bind_address = omni_master_server.get_allocation(stage_id, replica_id=replica_id).handshake_bind_address with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: - yield None, coordinator, addresses, None - - _wait_for_omni_engine_startup( + yield StageReplicaResources( + coordinator=coordinator, + addresses=addresses, + ) + wait_for_engine_startup( handshake_socket, addresses, engines_to_handshake, + vllm_config.parallel_config, + False, # coordinated_dp vllm_config.cache_config, + None, # proc_manager (remote — no local procs) + None, # coord_process + ) + + +@contextlib.contextmanager +def connect_remote_diffusion_proc( + omni_master_server: OmniMasterServer, + stage_id: int, + replica_id: int = 0, +) -> Iterator[StageReplicaResources]: + """Wait for a remote headless diffusion proc to connect to head-owned sockets.""" + addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) + handshake_bind_address = omni_master_server.get_allocation( + stage_id, + replica_id=replica_id, + ).handshake_bind_address + + logger.info("Waiting for remote diffusion proc for stage %d replica %d", stage_id, replica_id) + with zmq_socket_ctx(handshake_bind_address, zmq.ROUTER, bind=True) as handshake_socket: + yield StageReplicaResources(addresses=addresses) + wait_for_engine_startup( + handshake_socket, + addresses, + [CoreEngine(index=0, local=False)], + _single_diffusion_parallel_config(local_client=False), + False, + None, + None, + None, ) @contextlib.contextmanager -def launch_omni_core_engines( +def scoped_spawn_device_env( + stage_visible_devices: str | None, + spawn_device_lock: threading.Lock | None, +) -> Iterator[None]: + """Briefly scope device visibility while spawning a stage subprocess.""" + if stage_visible_devices is None or spawn_device_lock is None: + yield + return + + device_control_env = current_omni_platform.device_control_env_var + with spawn_device_lock: + previous_visible_devices = os.environ.get(device_control_env) + try: + current_omni_platform.set_device_control_env_var(stage_visible_devices) + yield + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + + +@contextlib.contextmanager +def _launch_omni_core_engines( vllm_config: VllmConfig, executor_class: type[Executor], log_stats: bool, @@ -849,22 +850,19 @@ def launch_omni_core_engines( replica_id: int = 0, *, omni_coordinator_address: str | None = None, + stage_visible_devices: str | None = None, + spawn_device_lock: threading.Lock | None = None, ) -> Iterator[tuple[CoreEngineProcManager, DPCoordinator | None, EngineZmqAddresses]]: """Launch local engine cores using the omni registration flow. When ``omni_coordinator_address`` is provided, the spawned engine - subprocesses use :class:`OmniCoreEngineProcManager` and each + subprocesses use :class:`StageEngineCoreProcManager` and each instantiates an :class:`OmniCoordClientForStage` after the handshake completes so the head's :class:`OmniCoordinator` knows about them. """ addresses = omni_master_server.get_zmq_addresses(stage_id, replica_id=replica_id) parallel_config = vllm_config.parallel_config - # Determine the number of local engines and their ranks. - local_engine_count = ( - parallel_config.data_parallel_size_local - if parallel_config.data_parallel_size_local is not None and parallel_config.data_parallel_size_local > 0 - else max(1, parallel_config.data_parallel_size) - ) + local_engine_count = _engine_count(parallel_config) dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 local_start_index = 0 start_index = dp_rank @@ -903,7 +901,7 @@ def launch_omni_core_engines( # Register the stage once and reuse the returned per-stage handshake # address for all local engine-core processes. - handshake_address = register_stage_with_omni_master( + registration = register_stage_with_omni_master( omni_master_address=omni_master_server.address, omni_master_port=omni_master_server.port, omni_stage_id=stage_id, @@ -911,6 +909,7 @@ def launch_omni_core_engines( coordinator=coordinator, replica_id=replica_id, ) + handshake_address = registration.handshake_address # One CoreEngine entry per local engine so wait_for_engine_startup can # track the HELLO/READY handshake for each of them. @@ -923,45 +922,557 @@ def launch_omni_core_engines( if omni_coordinator_address is not None: # Use the omni subclass so each spawned subprocess instantiates # an OmniCoordClientForStage and heartbeats to the coordinator. - from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager - - local_engine_manager: CoreEngineProcManager = OmniCoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - omni_stage_id=stage_id, - omni_coordinator_address=omni_coordinator_address, - omni_replica_base_id=replica_id, - ) + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager + + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + local_engine_manager: CoreEngineProcManager = StageEngineCoreProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=omni_coordinator_address, + omni_replica_base_id=replica_id, + ) else: - local_engine_manager = CoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=start_index, - local_start_index=local_start_index, - vllm_config=vllm_config, - local_client=True, - handshake_address=handshake_address, - executor_class=executor_class, - log_stats=log_stats, - ) + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + local_engine_manager = CoreEngineProcManager( + local_engine_count=local_engine_count, + start_index=start_index, + local_start_index=local_start_index, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + ) yield local_engine_manager, coordinator, addresses - - # Wait for all local engine-core processes to complete the - # standard HELLO/READY handshake — mirrors launch_core_engines. - coordinated_dp = parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe wait_for_engine_startup( handshake_socket, addresses, engines_to_handshake, parallel_config, - coordinated_dp, + parallel_config.data_parallel_size > 1 and vllm_config.model_config.is_moe, vllm_config.cache_config, local_engine_manager, coordinator.proc if coordinator else None, ) + + +@contextlib.contextmanager +def launch_stage_replica( + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + stage_id: int, + *, + replica_id: int = 0, + stage_config: Any = None, + omni_master_server: OmniMasterServer | None = None, + omni_coordinator_address: str | None = None, + stage_visible_devices: str | None = None, + spawn_device_lock: threading.Lock | None = None, +) -> Iterator[StageReplicaResources]: + """Launch a local LLM stage replica. + + This is the common entry point for colocated and distributed local LLM + replicas. Distributed launches delegate to ``_launch_omni_core_engines`` so + registration/address ownership stays centralized in ``OmniMasterServer``. + Colocated launches use an IPC handshake without master registration but + keep the same returned resource bundle. + """ + if omni_master_server is not None: + with _launch_omni_core_engines( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=log_stats, + omni_master_server=omni_master_server, + stage_id=stage_id, + stage_config=stage_config, + replica_id=replica_id, + omni_coordinator_address=omni_coordinator_address, + stage_visible_devices=stage_visible_devices, + spawn_device_lock=spawn_device_lock, + ) as resources: + engine_manager, coordinator, addresses = resources + yield StageReplicaResources( + manager=engine_manager, + coordinator=coordinator, + addresses=addresses, + ) + return + + from vllm.utils.network_utils import get_open_zmq_ipc_path + + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager + + addresses = get_engine_zmq_addresses(vllm_config) + handshake_address = get_open_zmq_ipc_path() + engines_to_handshake = [CoreEngine(index=0, local=True)] + with scoped_spawn_device_env(stage_visible_devices, spawn_device_lock): + engine_manager = StageEngineCoreProcManager( + local_engine_count=1, + start_index=0, + local_start_index=0, + vllm_config=vllm_config, + local_client=True, + handshake_address=handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=omni_coordinator_address, + omni_replica_base_id=replica_id, + ) + + with zmq_socket_ctx(handshake_address, zmq.ROUTER, bind=True) as handshake_socket: + yield StageReplicaResources( + manager=engine_manager, + addresses=addresses, + ) + wait_for_engine_startup( + handshake_socket, + addresses, + engines_to_handshake, + vllm_config.parallel_config, + False, # coordinated_dp + vllm_config.cache_config, + engine_manager, + None, # coordinator_proc + ) + + +def launch_headless_llm_replica( + *, + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + omni_master_address: str, + omni_master_port: int, + stage_id: int, + stage_config: Any, + coordinator: DPCoordinator | None = None, + replica_bind_address: str | None = None, +) -> CoreEngineProcManager: + """Register and launch one headless LLM replica. + + Headless LLM workers are connectors: the head binds the handshake/input/ + output sockets, while this process connects to the addresses allocated by + ``OmniMasterServer``. Keep that registration/manager wiring in the startup + layer so the CLI only handles argument/config normalization. + """ + from vllm_omni.engine.stage_engine_core_proc_manager import StageEngineCoreProcManager + + parallel_config = vllm_config.parallel_config + local_engine_count = parallel_config.data_parallel_size_local + if local_engine_count <= 0: + raise ValueError("data_parallel_size_local must be > 0 in headless mode") + dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 + + response = register_stage_with_omni_master( + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_stage_id=stage_id, + omni_stage_config=stage_config, + coordinator=coordinator, + replica_id=None, + replica_bind_address=replica_bind_address, + ) + + manager = StageEngineCoreProcManager( + local_engine_count=local_engine_count, + start_index=dp_rank, + local_start_index=0, + vllm_config=vllm_config, + local_client=False, + handshake_address=response.handshake_address, + executor_class=executor_class, + log_stats=log_stats, + omni_stage_id=stage_id, + omni_coordinator_address=response.coordinator_router_address, + omni_replica_base_id=response.replica_id, + ) + logger.info( + "[Headless] Stage %d replica id=%d up (coord=%s)", + stage_id, + response.replica_id, + response.coordinator_router_address, + ) + return manager + + +def launch_headless_llm_replicas( + *, + vllm_config: VllmConfig, + executor_class: type[Executor], + log_stats: bool, + omni_master_address: str, + omni_master_port: int, + stage_id: int, + stage_config: Any, + omni_dp_size_local: int, + per_replica_devices: list[str | None], + replica_bind_address: str | None = None, +) -> None: + """Launch, monitor, and clean up all local headless LLM replicas.""" + parallel_config = vllm_config.parallel_config + local_engine_count = parallel_config.data_parallel_size_local + if local_engine_count <= 0: + raise ValueError("data_parallel_size_local must be > 0 in headless mode") + + dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 + coordinator = None + if vllm_config.needs_dp_coordinator and dp_rank == 0: + coordinator = DPCoordinator( + parallel_config, + enable_wave_coordination=vllm_config.model_config.is_moe, + ) + logger.info( + "[Headless] Started DP Coordinator process for stage %d (PID: %d)", + stage_id, + coordinator.proc.pid, + ) + + logger.info( + "[Headless] Launching %d omni replica(s) (vLLM dp_size_local=%d each) for stage %d " + "via OmniMasterServer at %s:%d", + omni_dp_size_local, + local_engine_count, + stage_id, + omni_master_address, + omni_master_port, + ) + + try: + + def _launch_one(rep_idx: int) -> Any: + return launch_headless_llm_replica( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=log_stats, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + stage_id=stage_id, + stage_config=stage_config, + coordinator=coordinator, + replica_bind_address=replica_bind_address, + ) + + launch_headless_replica_group( + stage_id=stage_id, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + launch_one=_launch_one, + ) + finally: + if coordinator is not None: + coordinator.shutdown() + + +def launch_headless_diffusion_replica( + *, + model: str, + od_config: Any, + stage_config: Any, + stage_id: int, + omni_master_address: str, + omni_master_port: int, + replica_bind_address: str | None = None, +) -> Any: + """Register and launch one headless diffusion replica. + + Headless diffusion follows the LLM remote-attach model: the head binds + handshake/input/output sockets, while this backend process only connects. + """ + response = register_stage_with_omni_master( + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_stage_id=stage_id, + omni_stage_config=stage_config, + replica_id=None, + replica_bind_address=replica_bind_address, + ) + from vllm_omni.diffusion import stage_diffusion_proc + + manager = stage_diffusion_proc.StageDiffusionProcManager.launch_headless( + model=model, + od_config=od_config, + handshake_address=response.handshake_address, + addresses=EngineZmqAddresses( + inputs=[response.input_address], + outputs=[response.output_address], + ), + omni_coordinator_address=response.coordinator_router_address, + omni_stage_id=stage_id, + omni_replica_id=response.replica_id, + ) + logger.info( + "[Headless] Diffusion replica id=%d for stage %d is up (coord=%s)", + response.replica_id, + stage_id, + response.coordinator_router_address, + ) + return manager + + +@contextlib.contextmanager +def replica_device_env(stage_id: int, devices: str | None): + """Temporarily scope device visibility for one replica launch.""" + device_control_env = current_omni_platform.device_control_env_var + previous_visible_devices = os.environ.get(device_control_env) + try: + if devices is not None: + stage_init_utils.setup_stage_devices(stage_id, {"devices": devices}) + yield + finally: + if previous_visible_devices is None: + current_omni_platform.unset_device_control_env_var() + else: + current_omni_platform.set_device_control_env_var(previous_visible_devices) + + +def get_headless_replica_devices( + stage_cfg: Any, + stage_id: int, + omni_dp_size_local: int, +) -> list[str | None]: + """Return per-replica device slices for a headless stage.""" + runtime_cfg = getattr(stage_cfg, "runtime", None) + devices_str: str | None = None + if runtime_cfg is not None: + devices_str = ( + runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + ) + if not devices_str: + return [None] * omni_dp_size_local + + devices_per_replica = stage_init_utils.get_stage_devices_per_replica(stage_cfg) + per_replica_devices = stage_init_utils.split_devices_for_replicas( + devices_str, omni_dp_size_local, devices_per_replica, stage_id + ) + logger.info( + "[Headless] Stage %d: %d local replicas, devices_per_replica=%d, per-replica devices: %s", + stage_id, + omni_dp_size_local, + devices_per_replica, + per_replica_devices, + ) + return per_replica_devices + + +def wait_for_manager_liveness(engine_managers: list[Any]) -> None: + """Block until one or more engine managers report process exit.""" + if len(engine_managers) == 1: + engine_managers[0].monitor_engine_liveness() + return + + def _monitor_target(mgr: Any) -> None: + try: + mgr.monitor_engine_liveness() + except Exception: + logger.exception("[Headless] monitor_engine_liveness raised") + + monitor_threads: list[threading.Thread] = [] + for mgr in engine_managers: + t = threading.Thread( + target=_monitor_target, + args=(mgr,), + name=f"omni-replica-monitor-{id(mgr):x}", + ) + t.start() + monitor_threads.append(t) + for t in monitor_threads: + t.join() + + +def wait_for_diffusion_manager_liveness(managers: list[Any]) -> None: + """Block until one diffusion manager exits and surface non-zero exits.""" + sentinel_to_proc = {manager.proc.sentinel: manager.proc for manager in managers} + died = connection.wait(list(sentinel_to_proc.keys())) + first = sentinel_to_proc[died[0]] + logger.info( + "[Headless] Diffusion replica %s exited (code=%s).", + first.name, + first.exitcode, + ) + if first.exitcode not in (None, 0): + raise RuntimeError(f"Diffusion replica {first.name!r} exited with code {first.exitcode}") + + +def launch_headless_replica_group( + *, + stage_id: int, + omni_dp_size_local: int, + per_replica_devices: list[str | None], + launch_one: Callable[[int], Any], + wait_for_replicas: Callable[[list[Any]], None] = wait_for_manager_liveness, +) -> None: + """Launch, monitor, and clean up a group of local headless replicas.""" + managers: list[Any] = [] + try: + for rep_idx in range(omni_dp_size_local): + with replica_device_env(stage_id, per_replica_devices[rep_idx]): + managers.append(launch_one(rep_idx)) + wait_for_replicas(managers) + finally: + logger.info("[Headless] Shutting down stage %d (%d manager(s)).", stage_id, len(managers)) + for manager in managers: + try: + manager.shutdown() + except Exception: + logger.exception("[Headless] manager shutdown failed") + + +def launch_headless_diffusion_replicas( + *, + model: str, + stage_cfg: Any, + stage_configs: list[Any], + stage_id: int, + omni_master_address: str, + omni_master_port: int, + omni_dp_size_local: int, + per_replica_devices: list[str | None], + config_path: str, + replica_bind_address: str | None = None, +) -> None: + """Prepare diffusion config, launch replicas, monitor, and clean up.""" + omni_transfer_config = stage_init_utils.load_omni_transfer_config_for_model(model, config_path) + omni_conn_cfg, omni_from, omni_to = initialization.resolve_omni_kv_config_for_stage( + omni_transfer_config, + stage_id, + ) + + metadata = stage_init_utils.extract_stage_metadata(stage_cfg) + if omni_conn_cfg: + inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) + # Headless single-stage launch must still infer cross-stage TP topology + # from the loaded deploy config so heterogeneous KV routing keys match the + # head process (e.g. from_tp=2, to_tp=1). + stage_init_utils.inject_kv_stage_info(stage_cfg, stage_id, stage_configs) + od_config = stage_init_utils.build_diffusion_config(model, stage_cfg, metadata) + + logger.info( + "[Headless] Launching %d diffusion replica(s) for stage %d via OmniMasterServer at %s:%d", + omni_dp_size_local, + stage_id, + omni_master_address, + omni_master_port, + ) + + def _launch_one(rep_idx: int) -> Any: + # Keep torch.distributed ports away from the ZMQ ephemeral range that + # OmniMasterServer pre-allocates for sibling headless replicas. + if omni_dp_size_local > 1: + od_config.master_port = od_config.settle_port( + 61000 + rep_idx * 100, + port_inc=37, + ) + return launch_headless_diffusion_replica( + model=model, + od_config=od_config, + stage_config=stage_cfg, + stage_id=stage_id, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + replica_bind_address=replica_bind_address, + ) + + launch_headless_replica_group( + stage_id=stage_id, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + launch_one=_launch_one, + wait_for_replicas=wait_for_diffusion_manager_liveness, + ) + + +def launch_diffusion_stage_replica( + *, + model: str, + stage_config: Any, + metadata: Any, + stage_init_timeout: int, + batch_size: int, + use_inline: bool, + replica_id: int = 0, + omni_master_server: OmniMasterServer | None = None, + omni_coordinator_address: str | None = None, +) -> tuple[Any, StageReplicaResources]: + """Launch a local diffusion stage replica. + + Colocated mode delegates to ``initialize_diffusion_stage``. Distributed + local mode registers with ``OmniMasterServer`` and spawns a + ``StageDiffusionProc`` that heartbeats to ``OmniCoordinator``. + """ + if omni_master_server is None: + client = initialize_diffusion_stage( + metadata.stage_id, + model, + stage_config, + metadata, + stage_init_timeout=stage_init_timeout, + batch_size=batch_size, + use_inline=use_inline, + ) + return client, StageReplicaResources() + + from vllm_omni.diffusion import stage_diffusion_proc + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + od_config = build_diffusion_config(model, stage_config, metadata) + parallel_config = getattr(od_config, "parallel_config", None) + world_size = getattr(parallel_config, "world_size", 1) + try: + world_size = max(1, int(world_size)) + except (TypeError, ValueError): + world_size = 1 + lock_fds = acquire_device_locks( + metadata.stage_id, + {"tensor_parallel_size": world_size}, + stage_init_timeout, + ) + proc_manager = None + try: + registration = register_stage_with_omni_master( + omni_master_address=omni_master_server.address, + omni_master_port=omni_master_server.port, + omni_stage_id=metadata.stage_id, + omni_stage_config=stage_config, + replica_id=replica_id, + ) + proc_manager = stage_diffusion_proc.StageDiffusionProcManager( + model=model, + od_config=od_config, + stage_init_timeout=stage_init_timeout, + handshake_address=registration.handshake_address, + addresses=EngineZmqAddresses( + inputs=[registration.input_address], + outputs=[registration.output_address], + ), + omni_coordinator_address=omni_coordinator_address, + omni_stage_id=metadata.stage_id, + omni_replica_id=replica_id, + ) + client = StageDiffusionClient.from_addresses( + metadata, + request_address=proc_manager.addresses.inputs[0], + response_address=proc_manager.addresses.outputs[0], + proc_manager=proc_manager, + batch_size=batch_size, + ) + return client, StageReplicaResources( + manager=proc_manager, + addresses=proc_manager.addresses, + lock_fds=lock_fds, + ) + except Exception: + if proc_manager is not None: + proc_manager.shutdown() + if lock_fds: + release_device_locks(lock_fds) + raise diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 53267e23b9e..d4316089586 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -27,11 +27,11 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.engine.arg_utils import OmniEngineArgs -from vllm_omni.engine.output_processor import MultimodalOutputProcessor from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams from vllm_omni.inputs.preprocess import OmniInputPreprocessor +from vllm_omni.outputs.output_processor import MultimodalOutputProcessor from vllm_omni.platforms import current_omni_platform from vllm_omni.quantization.inc_config import OmniINCConfig @@ -51,6 +51,7 @@ class ReplicaInitPlan: omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None] stage_vllm_config: Any | None = None executor_class: type | None = None + engine_args_dict: dict[str, Any] | None = None @dataclass @@ -58,65 +59,8 @@ class LogicalStageInitPlan: """Startup plan for one logical stage.""" stage_idx: int - configured_stage_id: int - replicas: list[ReplicaInitPlan] - - -@dataclass -class StageRemoteFactoryContext: - """Per-stage context cached by AsyncOmniEngine for dynamic replica attach. - - Populated once during ``_bootstrap_orchestrator`` from the per-stage - init plans. ``_build_remote_replica`` consumes it to construct the - right head-side stage client when a headless replica registers. - """ - stage_id: int - stage_type: str - stage_cfg: Any - base_metadata: Any - # LLM-only fields: - vllm_config: Any | None = None - executor_class: type | None = None - # Diffusion-only fields: - diffusion_batch_size: int = 1 - - -def capture_stage_factory_contexts( - stage_plans: Sequence[LogicalStageInitPlan], - diffusion_batch_size: int, -) -> dict[int, StageRemoteFactoryContext]: - """Snapshot per-stage construction context for dynamic replica attach. - - Called once after ``_initialize_stages`` finishes. The captured - context holds everything ``_build_remote_replica`` needs to build a - fresh head-side client when a new headless replica registers - (vllm_config / executor_class for LLM, batch_size for diffusion, - plus the base stage metadata). - - Per-replica fields like ``replica_id`` are filled in at build time, - not at capture time. - """ - contexts: dict[int, StageRemoteFactoryContext] = {} - for plan in stage_plans: - if not plan.replicas: - # Stage was declared but has zero replicas locally; we still - # want to be able to attach incoming headless ones, so use - # the stage_cfg-derived context if any replica plan exists. - continue - template = plan.replicas[0] - stage_id = int(plan.configured_stage_id) - stage_type = template.metadata.stage_type or "llm" - contexts[stage_id] = StageRemoteFactoryContext( - stage_id=stage_id, - stage_type=stage_type, - stage_cfg=template.stage_cfg, - base_metadata=template.metadata, - vllm_config=template.stage_vllm_config, - executor_class=template.executor_class, - diffusion_batch_size=diffusion_batch_size, - ) - return contexts + replicas: list[ReplicaInitPlan] def _resolve_model_to_local_path(model: str) -> str: @@ -359,6 +303,26 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[ logger.debug("Failed to inject stage info into omni_kv_config: %s", e) +def inject_omni_kv_connector_config( + engine_args_dict: dict[str, Any], + omni_kv_connector: tuple[dict[str, Any] | None, str | None, str | None], + stage_id: int, +) -> None: + """Inject resolved connector config into a stage engine-args dict.""" + omni_conn_cfg, omni_from, omni_to = omni_kv_connector + if not omni_conn_cfg: + return + + omni_kv = engine_args_dict.get("omni_kv_config") or {} + if not isinstance(omni_kv, dict): + omni_kv = dict(omni_kv) + omni_kv["connector_config"] = omni_conn_cfg + omni_kv["omni_from_stage"] = omni_from + omni_kv["omni_to_stage"] = omni_to + omni_kv.setdefault("stage_id", stage_id) + engine_args_dict["omni_kv_config"] = omni_kv + + @dataclass class StageMetadata: """Lightweight stage attributes extracted from stage_config.""" @@ -566,23 +530,24 @@ def get_stage_tp_size(stage_cfg: Any) -> int: def get_stage_devices_per_replica(stage_cfg: Any) -> int: """Return the number of devices consumed by one replica of *stage_cfg*.""" - if getattr(stage_cfg, "stage_type", "llm") != "diffusion": - return get_stage_tp_size(stage_cfg) + engine_args = getattr(stage_cfg, "engine_args", {}) + if getattr(stage_cfg, "stage_type", "llm") == "diffusion": + parallel_config = _get_attr_or_item(engine_args, "parallel_config") + if parallel_config is None: + return 1 - parallel_config = _get_attr_or_item(getattr(stage_cfg, "engine_args", {}), "parallel_config") - if parallel_config is None: - return 1 + world_size = _get_attr_or_item(parallel_config, "world_size") + if world_size is not None: + return max(1, int(world_size)) - world_size = _get_attr_or_item(parallel_config, "world_size") - if world_size is not None: - return max(1, int(world_size)) + try: + from vllm_omni.diffusion.data import DiffusionParallelConfig - try: - from vllm_omni.diffusion.data import DiffusionParallelConfig + return max(1, int(DiffusionParallelConfig.from_dict(_to_dict(parallel_config)).world_size)) + except Exception: + return 1 - return max(1, int(DiffusionParallelConfig.from_dict(_to_dict(parallel_config)).world_size)) - except Exception: - return 1 + return get_stage_tp_size(stage_cfg) def compute_replica_layout( @@ -734,6 +699,7 @@ def build_engine_args_dict( stage_id = stage_config.stage_id engine_args_dict = _to_dict(engine_args) + model = engine_args_dict.pop("model", None) or model stage_defines_tokenizer = ( engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None ) @@ -1061,32 +1027,6 @@ def release_device_locks(lock_fds: list[int]) -> None: pass -def acquire_diffusion_device_locks( - stage_id: int, - od_config: Any, - stage_init_timeout: int, -) -> list[int]: - """Acquire init locks for the GPU set used by a diffusion stage. - - Diffusion stages express their device count via ``OmniDiffusionConfig``'s - ``parallel_config.world_size`` rather than the LLM-style - ``tensor_parallel_size`` knob, so adapt to the shape that - ``acquire_device_locks`` understands. - """ - parallel_config = getattr(od_config, "parallel_config", None) - world_size = getattr(parallel_config, "world_size", 1) - try: - world_size = max(1, int(world_size)) - except (TypeError, ValueError): - world_size = 1 - - return acquire_device_locks( - stage_id, - {"tensor_parallel_size": world_size}, - stage_init_timeout, - ) - - def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any: """Load omni transfer config from an explicit path or resolved model config. diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 4e1f3d5a55b..34af459ab19 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -4,6 +4,7 @@ import asyncio import time as _time +from collections.abc import Mapping from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast @@ -429,6 +430,37 @@ def get_bound_llm_client(self, request_id: str) -> StagePoolLLMClient | None: return None return cast(StagePoolLLMClient, client) + def replica_monitor_sample(self, replica_id: int) -> tuple[int, int]: + """Return (outputs_queue_size, inflight) for orchestrator load diagnostics.""" + if replica_id < 0 or replica_id >= len(self.clients): + return (0, 0) + client = self.clients[replica_id] + if client is None: + return (0, 0) + + outputs_qsize = 0 + outputs_queue = getattr(client, "outputs_queue", None) + if outputs_queue is not None: + try: + outputs_qsize = max(int(outputs_queue.qsize()), 0) + except (AttributeError, NotImplementedError, RuntimeError, ValueError): + pass + + if self._hub is not None: + input_addr = self._client_input_addr(client) + if input_addr is not None: + for replica in self._hub.get_replicas_for_stage(self.stage_id).replicas: + if replica.input_addr == input_addr: + return (outputs_qsize, max(int(replica.queue_length), 0)) + # Distributed dispatch (pick): request_id -> input_addr in _affinity. + inflight = sum( + 1 for input_addr in self._affinity.values() if self._addr_to_replica_id.get(input_addr) == replica_id + ) + else: + # Legacy local dispatch (select_replica_id): request_id -> replica_id. + inflight = sum(1 for bound_replica_id in self._request_bindings.values() if bound_replica_id == replica_id) + return (outputs_qsize, inflight) + def release_binding(self, request_id: str) -> None: """Drop the route binding for *request_id* in this stage.""" self._request_bindings.pop(request_id, None) @@ -566,7 +598,6 @@ def build_stage_metrics( batch_id = metrics.batch_seq metrics.agg_total_tokens += num_tokens_out metrics.agg_total_gen_time_ms += stage_gen_time_ms - audio_rtf = defs.compute_audio_rtf(stage_gen_time_ms / 1000.0, audio_duration_s) return StageRequestMetrics( num_tokens_in=num_tokens_in, @@ -585,7 +616,6 @@ def build_stage_metrics( audio_generated_frames=audio_generated_frames, audio_sample_rate=audio_sample_rate, audio_duration_s=audio_duration_s, - audio_rtf=audio_rtf, image_pixels=image_pixels, denoise_step_latency_ms=denoise_step_latency_ms, output_unit_type=output_unit_type, @@ -655,7 +685,7 @@ def _count_output_units( def _has_audio_output(self, request_outputs: list[Any]) -> bool: for ro in request_outputs: for mm_output in self._iter_multimodal_outputs(ro): - if isinstance(mm_output, dict) and mm_output.get("audio") is not None: + if isinstance(mm_output, Mapping) and mm_output.get("audio") is not None: return True return False @@ -664,7 +694,7 @@ def _collect_audio_metrics(self, request_outputs: list[Any]) -> tuple[int, int, sample_rate = 0 for ro in request_outputs: for mm_output in self._iter_multimodal_outputs(ro): - if not isinstance(mm_output, dict): + if not isinstance(mm_output, Mapping): continue if sample_rate <= 0: sample_rate = self._infer_audio_sample_rate(mm_output) @@ -748,7 +778,7 @@ def _has_latent_output(self, request_outputs: list[Any]) -> bool: def _count_images(self, request_output: Any) -> int: total_images = self._count_value_units(getattr(request_output, "images", None)) for mm_output in self._iter_multimodal_outputs(request_output): - if isinstance(mm_output, dict): + if isinstance(mm_output, Mapping): total_images += self._count_value_units(mm_output.get("image")) total_images += self._count_value_units(mm_output.get("images")) return total_images @@ -758,7 +788,7 @@ def _count_image_pixels(self, request_outputs: list[Any]) -> int: for ro in request_outputs: total_pixels += self._count_image_value_pixels(getattr(ro, "images", None)) for mm_output in self._iter_multimodal_outputs(ro): - if isinstance(mm_output, dict): + if isinstance(mm_output, Mapping): total_pixels += self._count_image_value_pixels(mm_output.get("image")) total_pixels += self._count_image_value_pixels(mm_output.get("images")) return total_pixels @@ -793,7 +823,7 @@ def _count_videos(self, request_output: Any) -> int: total_videos = self._count_video_units(getattr(request_output, "video", None)) total_videos += self._count_video_units(getattr(request_output, "videos", None)) for mm_output in self._iter_multimodal_outputs(request_output): - if isinstance(mm_output, dict): + if isinstance(mm_output, Mapping): total_videos += self._count_video_units(mm_output.get("video")) total_videos += self._count_video_units(mm_output.get("videos")) return total_videos @@ -818,7 +848,7 @@ def _has_non_empty_output(self, request_output: Any) -> bool: if isinstance(custom_output, dict) and custom_output: return True for mm_output in self._iter_multimodal_outputs(request_output): - if isinstance(mm_output, dict) and any(self._is_non_empty_value(value) for value in mm_output.values()): + if isinstance(mm_output, Mapping) and any(self._is_non_empty_value(value) for value in mm_output.values()): return True for output in getattr(request_output, "outputs", None) or []: if getattr(output, "text", None): @@ -886,14 +916,14 @@ def record_output_timestamps(self, request_outputs: list[Any], *, output_ts: flo if self._audio_sample_rate_by_request.get(rid, 0) <= 0 and audio_sample_rate > 0: self._audio_sample_rate_by_request[rid] = audio_sample_rate - def _iter_multimodal_outputs(self, request_output: Any) -> list[dict[str, Any]]: - multimodal_outputs: list[dict[str, Any]] = [] + def _iter_multimodal_outputs(self, request_output: object) -> list[Mapping[str, Any]]: + multimodal_outputs: list[Mapping[str, Any]] = [] outer_mm = getattr(request_output, "multimodal_output", None) - if isinstance(outer_mm, dict) and outer_mm: + if isinstance(outer_mm, Mapping) and outer_mm: multimodal_outputs.append(outer_mm) for output in getattr(request_output, "outputs", None) or []: inner_mm = getattr(output, "multimodal_output", None) - if isinstance(inner_mm, dict) and inner_mm: + if isinstance(inner_mm, Mapping) and inner_mm: multimodal_outputs.append(inner_mm) return multimodal_outputs @@ -918,15 +948,17 @@ async def submit_initial( params = OmniDiffusionSamplingParams() submit_kwargs = dict(submit_kwargs or {}) if self.stage_type == "diffusion": + if isinstance(request, list): + raise ValueError( + "Diffusion list-prompt batch requests are no longer supported. " + "Submit multiple independent requests to use scheduler batching." + ) replica_id = await self._pick_or_select( request_id, affinity_request_id=affinity_request_id, ) client = self._diffusion_client(replica_id) - if isinstance(request, list): - await client.add_batch_request_async(request_id, request, params, **submit_kwargs) - else: - await client.add_request_async(request_id, request, params, **submit_kwargs) + await client.add_request_async(request_id, request, params, **submit_kwargs) return replica_id replica_id = await self._pick_or_select( @@ -987,6 +1019,11 @@ async def submit_update( raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") if self.stage_type == "diffusion": + if isinstance(request, list): + raise ValueError( + "Diffusion list-prompt batch requests are no longer supported. " + "Submit multiple independent requests to use scheduler batching." + ) await self._diffusion_client(replica_id).add_request_async(request_id, request, params) else: # Refresh the shared output-processor state before yielding to the diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py new file mode 100644 index 00000000000..8897043ca81 --- /dev/null +++ b/vllm_omni/engine/stage_runtime.py @@ -0,0 +1,1116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Stage Runtime implementations for single-node and distributed omni stages.""" + +from __future__ import annotations + +import concurrent.futures +import copy +import os +import threading +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +import janus +from omegaconf import OmegaConf +from vllm.logger import init_logger + +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + resolve_omni_kv_config_for_stage, +) +from vllm_omni.distributed.omni_coordinator import ( + LeastQueueLengthBalancer, + LoadBalancer, + LoadBalancingPolicy, + RandomBalancer, + RoundRobinBalancer, +) +from vllm_omni.engine.messages import ( + EngineQueueMessage, + RegisterRemoteReplicaMessage, +) +from vllm_omni.engine.stage_client import StageClient, StagePoolClient +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClientBase +from vllm_omni.engine.stage_engine_startup import ( + OmniMasterServer, + StageReplicaResources, + connect_remote_diffusion_proc, + connect_remote_engine_cores, + launch_diffusion_stage_replica, + launch_stage_replica, +) +from vllm_omni.engine.stage_init_utils import ( + LogicalStageInitPlan, + ReplicaInitPlan, + _inject_inferred_kv_tp_topology, + acquire_device_locks, + build_engine_args_dict, + build_llm_stage_output_processor, + build_vllm_config, + compute_replica_layout, + extract_stage_metadata, + get_stage_connector_spec, + inject_kv_stage_info, + inject_omni_kv_connector_config, + load_omni_transfer_config_for_model, + prepare_engine_environment, + release_device_locks, +) +from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.entrypoints.stage_utils import resolve_stage_physical_devices +from vllm_omni.entrypoints.utils import inject_omni_kv_config +from vllm_omni.outputs.output_metadata import FinalOutputModalityType +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + + +@dataclass(frozen=True, slots=True) +class StageRuntimeInfo: + final_output: bool + final_output_type: FinalOutputModalityType | None + stage_type: str + model_stage: str | None = None + + +@dataclass +class StageRemoteFactoryContext: + """Per-stage context for creating head-side clients for remote replicas.""" + + stage_id: int + stage_type: str + stage_cfg: Any + base_metadata: Any + vllm_config: Any | None = None + executor_class: type | None = None + diffusion_batch_size: int = 1 + + +def _build_load_balancer_factory(policy: str) -> Callable[[], LoadBalancer]: + try: + normalized = LoadBalancingPolicy(policy) + except ValueError as exc: + valid = ", ".join(p.value for p in LoadBalancingPolicy) + raise ValueError(f"unknown --omni-lb-policy {policy!r} (valid: {valid})") from exc + if normalized is LoadBalancingPolicy.RANDOM: + return RandomBalancer + if normalized is LoadBalancingPolicy.ROUND_ROBIN: + return RoundRobinBalancer + if normalized is LoadBalancingPolicy.LEAST_QUEUE_LENGTH: + return LeastQueueLengthBalancer + raise ValueError(f"unhandled load balancing policy {normalized!r}") + + +# =========================================================================== +# StageRuntime +# =========================================================================== + + +class StageRuntime: + """Stage runtime for single-node (non-distributed) mode. + + No coordinator, no master server, no hub. Launches stage processes + directly and creates StagePool with static clients. + """ + + def __init__( + self, + stage_configs: list[Any], + model: str, + config_path: str, + *, + stage_init_timeout: int, + diffusion_batch_size: int, + async_chunk: bool, + tokenizer: str | None = None, + ) -> None: + self._stage_configs = stage_configs + self._model = model + self._config_path = config_path + self._stage_init_timeout = stage_init_timeout + self._diffusion_batch_size = diffusion_batch_size + self._async_chunk = async_chunk + self._tokenizer = tokenizer + self._num_stages = len(stage_configs) + + # Populated by initialize() + self.stage_pools: list[StagePool] = [] + self._stage_init_executor: concurrent.futures.ThreadPoolExecutor | None = None + self._spawn_device_lock = threading.Lock() + # Serialize all LLM replica spawning + handshake across device groups + # to prevent ZMQ port-allocation races (get_engine_zmq_addresses) and + # CUDA-context conflicts when multiple engine core subprocesses + # initialize simultaneously on different GPUs. Matches the old + # AsyncOmniEngine._initialize_llm_replica pattern which used a single + # ``llm_stage_launch_lock`` for all replicas. + self._replica_launch_lock = threading.Lock() + self._init_visible_devices_baseline: str | None = None + + @staticmethod + def _client_addresses_from_zmq(addresses: Any) -> dict[str, str]: + client_addresses = { + "input_address": addresses.inputs[0], + "output_address": addresses.outputs[0], + } + if addresses.frontend_stats_publish_address is not None: + client_addresses["stats_update_address"] = addresses.frontend_stats_publish_address + return client_addresses + + @staticmethod + def _cleanup_launched_resources( + *, + stage_id: int, + resources: StageReplicaResources | None = None, + ) -> None: + """Release launch-only resources when client creation never completed.""" + if resources is None: + return + + for resource, resource_name in ( + (resources.manager, "manager"), + (resources.coordinator, "coordinator"), + ): + if resource is None: + continue + try: + resource.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to cleanup launched %s for stage %s: %s", + resource_name, + stage_id, + cleanup_error, + ) + + @staticmethod + def _collect_initialized_clients_for_cleanup( + stage_pools: Sequence[StagePool], + initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]], + ) -> list[StageClient]: + """Collect initialized clients exactly once for failure cleanup.""" + collected: list[StageClient] = [] + seen: set[int] = set() + + def _add_client(client: StageClient | None) -> None: + if client is None: + return + client_id = id(client) + if client_id in seen: + return + seen.add(client_id) + collected.append(client) + + for pool in stage_pools: + for client in getattr(pool, "clients", ()): + _add_client(client) + + for clients in initialized_clients_by_stage.values(): + for client in clients: + _add_client(client) + + return collected + + @staticmethod + def _shutdown_initialized_clients(clients: Sequence[StageClient]) -> None: + """Best-effort shutdown for attached clients after init failure.""" + for client in reversed(list(clients)): + if client is None: + continue + try: + client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to shutdown initialized client after init failure: %s", + cleanup_error, + ) + + def initialize(self) -> None: + """Run the full stage initialization sequence.""" + stage_plans = self._prepare_stage_plans() + initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { + plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans + } + try: + self._before_initialize_stage_replicas(stage_plans) + initialized_clients = self._initialize_stage_replicas(stage_plans, self._stage_init_timeout) + initialized_clients_by_stage = initialized_clients + self._finalize_initialized_stages(stage_plans, initialized_clients) + except Exception as exc: + initialized_clients_by_stage = getattr( + exc, + "_initialized_clients_by_stage", + initialized_clients_by_stage, + ) + cleanup_clients = self._collect_initialized_clients_for_cleanup( + self.stage_pools, + initialized_clients_by_stage, + ) + logger.exception( + "[StageRuntime] Stage initialization failed; shutting down %s initialized client(s)", + len(cleanup_clients), + ) + self._shutdown_initialized_clients(cleanup_clients) + self._cleanup_after_initialize_failure() + raise exc + + def shutdown(self) -> None: + for pool in self.stage_pools: + for client in pool.clients: + if client is not None and hasattr(client, "shutdown"): + try: + client.shutdown() + except Exception: + logger.warning("[StageRuntime] client shutdown failed", exc_info=True) + if self._stage_init_executor is not None: + self._stage_init_executor.shutdown(wait=True, cancel_futures=True) + self._stage_init_executor = None + + def create_membership_controller(self) -> Any | None: + """Return a distributed membership controller, if this runtime needs one.""" + return None + + def _prepare_stage_plans(self) -> list[LogicalStageInitPlan]: + """Build logical stage plans and cache prompt expansion metadata.""" + replicas_per_stage, replica_devices_map = compute_replica_layout(self._stage_configs) + prepare_engine_environment() + omni_transfer_config = load_omni_transfer_config_for_model(self._model, self._config_path) + stage_plans = self._build_logical_stage_init_plans( + omni_transfer_config, + replicas_per_stage, + replica_devices_map, + ) + return stage_plans + + def _finalize_initialized_stages( + self, + stage_plans: Sequence[LogicalStageInitPlan], + initialized_clients: Mapping[int, Sequence[StagePoolClient | None]], + ) -> None: + """Populate runtime fields after replica initialization succeeds.""" + self.stage_pools = self._assemble_stage_pools(stage_plans, initialized_clients) + + def _before_initialize_stage_replicas(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Hook for runtimes that need infrastructure before replica init.""" + return None + + def _cleanup_after_initialize_failure(self) -> None: + """Hook for runtimes that own extra infrastructure during init.""" + return None + + @contextmanager + def _scoped_spawn_device_env(self, physical_devices: str | None) -> Iterator[None]: + """Briefly scope device visibility for spawn-sensitive setup steps.""" + from vllm_omni.engine.stage_engine_startup import scoped_spawn_device_env + + with scoped_spawn_device_env( + physical_devices, + self._spawn_device_lock, + ): + yield + + def _resolve_replica_physical_devices(self, stage_id: int, runtime_cfg: Any) -> str | None: + if runtime_cfg is None: + runtime_cfg = {} + devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + return resolve_stage_physical_devices( + stage_id, + devices, + visible_baseline=self._init_visible_devices_baseline, + ) + + @contextmanager + def _stage_device_scope(self, stage_id: int, runtime_cfg: Any) -> Iterator[None]: + """Temporarily apply the stage device env while launching a replica.""" + physical_devices = self._resolve_replica_physical_devices(stage_id, runtime_cfg) + with self._scoped_spawn_device_env(physical_devices): + yield + + # ---- Internal methods ---- + + def _build_logical_stage_init_plans( + self, + omni_transfer_config: Any, + replicas_per_stage: Sequence[int], + replica_devices_map: Mapping[int, Sequence[str]], + ) -> list[LogicalStageInitPlan]: + """Build startup plans for every logical stage and replica.""" + stage_plans: list[LogicalStageInitPlan] = [] + + for stage_idx, stage_cfg in enumerate(self._stage_configs): + base_metadata = extract_stage_metadata(stage_cfg) + stage_id = int(base_metadata.stage_id) + if stage_id != stage_idx: + raise ValueError( + "stage_id must match its position in stage_configs: " + f"stage_configs[{stage_idx}] has stage_id={stage_id}. " + "Stage ids must be contiguous and zero-based because the " + "orchestrator indexes stage pools by stage_id." + ) + + stage_connector_spec = get_stage_connector_spec( + omni_transfer_config=omni_transfer_config, + stage_id=stage_id, + async_chunk=self._async_chunk, + ) + omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) + num_replicas = replicas_per_stage[stage_idx] + launch_mode = self._get_launch_mode(stage_id) + + replicas: list[ReplicaInitPlan] = [] + stage_vllm_config = None + executor_class = None + engine_args_dict = None + if base_metadata.stage_type != "diffusion": + engine_args_dict = build_engine_args_dict( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + cli_tokenizer=self._tokenizer, + ) + inject_omni_kv_connector_config( + engine_args_dict, + omni_kv_connector, + stage_id, + ) + _inject_inferred_kv_tp_topology( + engine_args_dict.get("omni_kv_config"), + stage_id, + self._stage_configs, + ) + stage_vllm_config, executor_class = build_vllm_config( + stage_cfg, + self._model, + stage_connector_spec=stage_connector_spec, + engine_args_dict=engine_args_dict, + ) + + for replica_id in range(num_replicas): + replica_cfg = copy.deepcopy(stage_cfg) if replica_id > 0 else stage_cfg + if stage_idx in replica_devices_map: + replica_cfg.runtime.devices = replica_devices_map[stage_idx][replica_id] + + replica_metadata = extract_stage_metadata(replica_cfg) + replica_metadata.replica_id = replica_id + if launch_mode == "remote" and replica_metadata.stage_type != "diffusion": + replica_metadata.runtime_cfg = None + + replicas.append( + ReplicaInitPlan( + replica_id=replica_id, + num_replicas=num_replicas, + launch_mode=launch_mode, + stage_cfg=replica_cfg, + metadata=replica_metadata, + stage_connector_spec=stage_connector_spec, + omni_kv_connector=omni_kv_connector, + stage_vllm_config=stage_vllm_config, + executor_class=executor_class, + engine_args_dict=copy.deepcopy(engine_args_dict) if engine_args_dict is not None else None, + ) + ) + + stage_plans.append( + LogicalStageInitPlan( + stage_idx=stage_idx, + stage_id=stage_id, + replicas=replicas, + ) + ) + + return stage_plans + + def _get_launch_mode(self, stage_id: int) -> str: + """Determine launch mode for a stage. Overridden by DistStageRuntime.""" + return "local" + + def _initialize_stage_replicas( + self, + stage_plans: Sequence[LogicalStageInitPlan], + stage_init_timeout: int, + ) -> dict[int, list[StagePoolClient | None]]: + """Initialize all stage replicas. + + Stages sharing the same GPU are initialized sequentially to avoid + memory profiling interference. Stages on different GPUs are + initialized in parallel. + """ + initialized_clients_by_stage: dict[int, list[StagePoolClient | None]] = { + plan.stage_idx: [None] * len(plan.replicas) for plan in stage_plans + } + primary_exc: Exception | None = None + init_state_lock = threading.Lock() + self._init_visible_devices_baseline = os.environ.get(current_omni_platform.device_control_env_var) + + init_groups: dict[str, list[tuple[int, ReplicaInitPlan]]] = {} + for plan in stage_plans: + for replica in plan.replicas: + init_groups.setdefault(self._replica_init_group_key(replica), []).append((plan.stage_idx, replica)) + + def _init_group(group: list[tuple[int, ReplicaInitPlan]]) -> None: + """Initialize replicas in one scheduling group sequentially.""" + nonlocal primary_exc + for stage_idx, replica in group: + with init_state_lock: + if primary_exc is not None: + return + try: + client = self._initialize_replica( + replica, + stage_init_timeout, + ) + except Exception as exc: + with init_state_lock: + if primary_exc is None: + primary_exc = exc + return + with init_state_lock: + initialized_clients_by_stage[stage_idx][replica.replica_id] = client + + inline_keys = [key for key in init_groups if key.startswith("inline:")] + for key in inline_keys: + _init_group(init_groups.pop(key)) + + if primary_exc is None and init_groups: + if len(init_groups) == 1: + _init_group(next(iter(init_groups.values()))) + else: + future_to_group: dict[concurrent.futures.Future[None], str] = {} + if self._stage_init_executor is None: + self._stage_init_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=len(init_groups), + thread_name_prefix="stage-init", + ) + for group_key, group in init_groups.items(): + future_to_group[self._stage_init_executor.submit(_init_group, group)] = group_key + + for future in concurrent.futures.as_completed(future_to_group): + try: + future.result() + except Exception as exc: + with init_state_lock: + if primary_exc is None: + primary_exc = exc + + if primary_exc is not None: + setattr(primary_exc, "_initialized_clients_by_stage", initialized_clients_by_stage) + raise primary_exc + + return initialized_clients_by_stage + + def _replica_init_group_key(self, replica: ReplicaInitPlan) -> str: + """Return the scheduling group used during replica initialization.""" + if replica.launch_mode == "local" and replica.metadata.stage_type == "diffusion": + # Local diffusion process spawning must stay on the orchestrator + # thread. Keep all local diffusion replicas in one sequential group. + return "inline:diffusion" + if replica.launch_mode == "remote": + return f"remote:{replica.metadata.stage_id}:{replica.replica_id}" + + runtime_cfg = replica.metadata.runtime_cfg or {} + devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) + return f"device:{devices}" + + def _initialize_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StagePoolClient: + if plan.launch_mode == "remote": + return self._initialize_remote_replica(plan, stage_init_timeout) + if plan.metadata.stage_type == "diffusion": + return self._initialize_local_diffusion_replica(plan, stage_init_timeout) + return self._initialize_local_llm_replica(plan, stage_init_timeout) + + def _initialize_remote_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StagePoolClient: + """Initialize a remote replica. Only distributed runtime implements this.""" + raise NotImplementedError("Remote replicas require DistStageRuntime") + + def _initialize_local_llm_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StageEngineCoreClientBase: + """Initialize one local LLM replica using vLLM's launch/attach pattern.""" + resources: StageReplicaResources | None = None + stage_client = None + lock_fds: list[int] = [] + try: + physical_devices = self._resolve_replica_physical_devices( + plan.metadata.stage_id, + plan.metadata.runtime_cfg, + ) + if physical_devices: + logger.info( + "[stage_init] Stage-%s set runtime devices: %s", + plan.metadata.stage_id, + physical_devices, + ) + vllm_config = plan.stage_vllm_config + executor_class = plan.executor_class + if vllm_config is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing vllm_config") + if executor_class is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing executor_class") + if plan.engine_args_dict is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} is missing engine args") + with self._scoped_spawn_device_env(physical_devices): + lock_fds = acquire_device_locks( + plan.metadata.stage_id, + plan.engine_args_dict, + stage_init_timeout, + ) + # Serialize engine-core spawning across all LLM replicas to avoid + # ZMQ port-allocation races and simultaneous CUDA context init. + with self._replica_launch_lock: + with launch_stage_replica( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=False, + stage_id=plan.metadata.stage_id, + replica_id=plan.replica_id, + stage_config=plan.stage_cfg, + omni_master_server=self._get_omni_master_server(), + omni_coordinator_address=self._get_coordinator_address(), + stage_visible_devices=physical_devices, + spawn_device_lock=self._spawn_device_lock, + ) as resources: + pass + + logger.info("[StageRuntime] Stage %s engine startup completed", plan.metadata.stage_id) + if resources is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no resources") + if resources.addresses is None: + raise RuntimeError(f"LLM stage {plan.metadata.stage_id} launcher returned no addresses") + stage_client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=vllm_config, + executor_class=executor_class, + metadata=plan.metadata, + client_addresses=self._client_addresses_from_zmq(resources.addresses), + engine_manager=resources.manager, + coordinator=resources.coordinator, + ) + + logger.info("[StageRuntime] Stage %s initialized", plan.metadata.stage_id) + return stage_client + except Exception: + if stage_client is not None: + try: + stage_client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to cleanup stage %s after init failure: %s", + plan.metadata.stage_id, + cleanup_error, + ) + else: + self._cleanup_launched_resources( + stage_id=plan.metadata.stage_id, + resources=resources, + ) + raise + finally: + if lock_fds: + release_device_locks(lock_fds) + + def _get_coordinator_address(self) -> str | None: + """Return coordinator router address. Overridden by DistStageRuntime.""" + return None + + def _get_omni_master_server(self) -> OmniMasterServer | None: + """Return the master server for local distributed launches, if any.""" + return None + + def _initialize_local_diffusion_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> Any: + """Initialize one local diffusion replica end-to-end.""" + client = None + resources = None + try: + with self._stage_device_scope(plan.metadata.stage_id, plan.metadata.runtime_cfg): + omni_conn_cfg, omni_from, omni_to = plan.omni_kv_connector + if omni_conn_cfg: + inject_omni_kv_config(plan.stage_cfg, omni_conn_cfg, omni_from, omni_to) + inject_kv_stage_info(plan.stage_cfg, plan.metadata.stage_id, self._stage_configs) + client, resources = launch_diffusion_stage_replica( + model=self._model, + stage_config=plan.stage_cfg, + metadata=plan.metadata, + stage_init_timeout=stage_init_timeout, + batch_size=self._diffusion_batch_size, + use_inline=self._num_stages == 1 and plan.num_replicas == 1, + replica_id=plan.replica_id, + omni_master_server=self._get_omni_master_server(), + omni_coordinator_address=self._get_coordinator_address(), + ) + + logger.info( + "[StageRuntime] Stage %s replica %s initialized (diffusion, batch_size=%d)", + plan.metadata.stage_id, + plan.replica_id, + self._diffusion_batch_size, + ) + return client + except Exception: + if client is not None: + try: + client.shutdown() + except Exception as cleanup_error: + logger.warning( + "[StageRuntime] Failed to cleanup stage %s after init failure: %s", + plan.metadata.stage_id, + cleanup_error, + ) + else: + self._cleanup_launched_resources( + stage_id=plan.metadata.stage_id, + resources=resources, + ) + raise + finally: + if resources is not None and resources.lock_fds: + release_device_locks(resources.lock_fds) + + def _assemble_stage_pools( + self, + stage_plans: Sequence[LogicalStageInitPlan], + initialized_clients_by_stage: Mapping[int, Sequence[StagePoolClient | None]], + ) -> list[StagePool]: + """Assemble logical stage pools.""" + stage_pools: list[StagePool] = [] + + for plan in stage_plans: + replica_clients = initialized_clients_by_stage[plan.stage_idx] + first_client = replica_clients[0] if replica_clients else None + if first_client is None: + raise RuntimeError(f"Stage {plan.stage_idx} initialization completed with a missing client") + + clients: list[StagePoolClient] = [client for client in replica_clients if client is not None] + stage_vllm_config = None + output_processor = None + if plan.replicas[0].metadata.stage_type != "diffusion": + stage_vllm_config = plan.replicas[0].stage_vllm_config + if stage_vllm_config is None: + raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config") + output_processor = build_llm_stage_output_processor(plan, stage_vllm_config) + + stage_pools.append( + StagePool( + plan.stage_idx, + clients, + output_processor=output_processor, + stage_vllm_config=stage_vllm_config, + ) + ) + + return stage_pools + + +# =========================================================================== +# DistStageRuntime +# =========================================================================== + + +class DistStageRuntime(StageRuntime): + """Stage runtime for distributed (single_stage_mode) deployment. + + Extends StageRuntime with: + - OmniCoordinatorRuntime (independent process) + - OmniMasterServer for replica registration + - Remote replica support + - Dynamic membership via MembershipController + """ + + def __init__( + self, + stage_configs: list[Any], + model: str, + config_path: str, + *, + stage_init_timeout: int, + diffusion_batch_size: int, + async_chunk: bool, + tokenizer: str | None = None, + single_stage_id_filter: int | None, + omni_master_address: str, + omni_master_port: int, + omni_dp_size_local: int = 1, + omni_heartbeat_timeout: float = 30.0, + omni_lb_policy: str = "random", + request_queue: janus.Queue[EngineQueueMessage] | None = None, + ) -> None: + super().__init__( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=diffusion_batch_size, + async_chunk=async_chunk, + tokenizer=tokenizer, + ) + self._single_stage_id_filter = single_stage_id_filter + self._omni_master_address = omni_master_address + self._omni_master_port = omni_master_port + self._omni_heartbeat_timeout = omni_heartbeat_timeout + self._omni_lb_policy = omni_lb_policy + self._request_queue = request_queue + + self._omni_master_server: OmniMasterServer | None = None + self._coordinator_runtime: Any | None = None + self._omni_dp_size_local = omni_dp_size_local + self._stage_remote_factory_contexts: dict[int, StageRemoteFactoryContext] = {} + + def create_membership_controller(self) -> Any | None: + if self._coordinator_runtime is None: + return None + + from vllm_omni.engine.membership_controller import MembershipController + + return MembershipController( + stage_pools=self.stage_pools, + coordinator_pub_address=self._coordinator_runtime.pub_address, + load_balancer_factory=_build_load_balancer_factory(self._omni_lb_policy), + remote_replica_factory=self._build_remote_replica, + ) + + def _prepare_stage_plans(self) -> list[LogicalStageInitPlan]: + self._validate_single_stage_mode_replica_constraints() + return super()._prepare_stage_plans() + + def _validate_single_stage_mode_replica_constraints(self) -> None: + """Apply --omni-dp-size-local to the local stage's runtime.num_replicas.""" + target_stage_id = self._single_stage_id_filter + if target_stage_id is None: + return + + for idx, stage_cfg in enumerate(self._stage_configs): + stage_id = int(getattr(stage_cfg, "stage_id", idx)) + runtime_cfg = getattr(stage_cfg, "runtime", None) + if runtime_cfg is None: + continue + if stage_id == target_stage_id: + try: + runtime_cfg.num_replicas = self._omni_dp_size_local + except (AttributeError, TypeError): + if hasattr(runtime_cfg, "__setitem__"): + runtime_cfg["num_replicas"] = self._omni_dp_size_local + continue + logger.warning( + "[DistStageRuntime] Failed to apply omni_dp_size_local=%s to stage %s runtime config", + self._omni_dp_size_local, + stage_id, + ) + + def _before_initialize_stage_replicas(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + self._start_omni_master_server(stage_plans) + self._stage_remote_factory_contexts = self._capture_stage_factory_contexts(stage_plans) + + def _cleanup_after_initialize_failure(self) -> None: + self._cleanup_distributed_infra() + + def shutdown(self) -> None: + super().shutdown() + self._cleanup_distributed_infra() + + def _cleanup_distributed_infra(self) -> None: + if self._omni_master_server is not None: + try: + self._omni_master_server.stop() + except Exception: + logger.warning("[DistStageRuntime] master server stop failed", exc_info=True) + self._omni_master_server = None + if self._coordinator_runtime is not None: + try: + self._coordinator_runtime.close() + except Exception: + logger.warning("[DistStageRuntime] coordinator close failed", exc_info=True) + self._coordinator_runtime = None + + # ---- Distributed overrides ---- + + def _get_launch_mode(self, stage_id: int) -> str: + if self._single_stage_id_filter is not None and stage_id != self._single_stage_id_filter: + return "remote" + return "local" + + def _get_coordinator_address(self) -> str | None: + if self._coordinator_runtime is not None: + return self._coordinator_runtime.router_address + return None + + def _get_omni_master_server(self) -> OmniMasterServer | None: + return self._omni_master_server + + def _initialize_remote_replica( + self, + plan: ReplicaInitPlan, + stage_init_timeout: int, + ) -> StagePoolClient: + """Wait for a configured remote replica and create its head-side client.""" + if self._omni_master_server is None: + raise RuntimeError("OmniMasterServer is not running; cannot initialize remote replica") + registered_stage_cfg = self._omni_master_server.get_stage_config( + plan.metadata.stage_id, + timeout_s=stage_init_timeout, + replica_id=plan.replica_id, + ) + if registered_stage_cfg is None: + raise ValueError(f"Remote stage {plan.metadata.stage_id} registered without stage config") + + metadata = ( + extract_stage_metadata(OmegaConf.create(registered_stage_cfg)) + if plan.metadata.stage_type == "diffusion" + else copy.deepcopy(plan.metadata) + ) + metadata.replica_id = plan.replica_id + ctx = StageRemoteFactoryContext( + stage_id=plan.metadata.stage_id, + stage_type=plan.metadata.stage_type, + stage_cfg=plan.stage_cfg, + base_metadata=metadata, + vllm_config=plan.stage_vllm_config, + executor_class=plan.executor_class, + diffusion_batch_size=self._diffusion_batch_size, + ) + return self._create_remote_replica_client(ctx, plan.replica_id) + + # ---- Distributed infrastructure ---- + + def _start_omni_master_server(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None: + """Start OmniMasterServer and OmniCoordinatorRuntime.""" + if not self._omni_master_address or not self._omni_master_port: + raise ValueError( + "AsyncOmniEngine single_stage_mode requires both omni_master_address and omni_master_port to be set." + ) + + from vllm_omni.distributed.omni_coordinator import OmniCoordinatorRuntime + + all_stage_ids: list[int] = [] + stage_replica_counts: dict[int, int] = {} + head_local_replicas: dict[int, list[int]] = {} + seen_stage_ids: set[int] = set() + for plan in stage_plans: + stage_id = plan.stage_id + if stage_id in seen_stage_ids: + raise ValueError(f"Duplicate stage_id {stage_id!r} detected") + seen_stage_ids.add(stage_id) + all_stage_ids.append(stage_id) + stage_replica_counts[stage_id] = len(plan.replicas) + local_rids = [rep.replica_id for rep in plan.replicas if rep.launch_mode == "local"] + if local_rids: + head_local_replicas[stage_id] = local_rids + + self._coordinator_runtime = OmniCoordinatorRuntime( + host=self._omni_master_address, + heartbeat_timeout=self._omni_heartbeat_timeout, + ) + + self._omni_master_server = OmniMasterServer( + master_address=self._omni_master_address, + master_port=self._omni_master_port, + stage_ids=all_stage_ids, + stage_replica_counts=stage_replica_counts, + coordinator_router_address=self._coordinator_runtime.router_address, + on_register=self._dispatch_master_register, + head_local_replicas=head_local_replicas, + ) + self._omni_master_server.start() + logger.info("[DistStageRuntime] OmniMasterServer started for stages %s", all_stage_ids) + + def _capture_stage_factory_contexts( + self, stage_plans: Sequence[LogicalStageInitPlan] + ) -> dict[int, StageRemoteFactoryContext]: + contexts: dict[int, StageRemoteFactoryContext] = {} + for plan in stage_plans: + if not plan.replicas: + continue + template = plan.replicas[0] + stage_id = int(plan.stage_id) + contexts[stage_id] = StageRemoteFactoryContext( + stage_id=stage_id, + stage_type=template.metadata.stage_type or "llm", + stage_cfg=template.stage_cfg, + base_metadata=template.metadata, + vllm_config=template.stage_vllm_config, + executor_class=template.executor_class, + diffusion_batch_size=self._diffusion_batch_size, + ) + return contexts + + def _dispatch_master_register(self, stage_id: int, replica_id: int, alloc: Any) -> None: + """Callback from OmniMasterServer when a headless replica registers.""" + if self._request_queue is None: + logger.warning("[DistStageRuntime] on_register fired but no request_queue wired") + return + try: + self._request_queue.sync_q.put_nowait( + RegisterRemoteReplicaMessage(stage_id=stage_id, replica_id=replica_id) + ) + except Exception: + logger.exception("[DistStageRuntime] Failed to enqueue register message") + + def _build_remote_replica(self, stage_id: int, replica_id: int) -> StagePoolClient: + ctx = self._stage_remote_factory_contexts.get(stage_id) + if ctx is None: + raise ValueError(f"No factory context for stage {stage_id}") + return self._create_remote_replica_client(ctx, replica_id) + + def _create_remote_replica_client( + self, + ctx: StageRemoteFactoryContext, + replica_id: int, + ) -> StagePoolClient: + """Create the head-side client for a remote replica. + + Used by both initial remote slots and dynamic headless registrations. + """ + if self._omni_master_server is None: + raise RuntimeError("OmniMasterServer is not running; cannot create remote replica") + stage_id = ctx.stage_id + metadata = copy.deepcopy(ctx.base_metadata) + metadata.replica_id = replica_id + + if ctx.stage_type == "diffusion": + from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + + resources = None + try: + logger.info( + "[DistStageRuntime] Remote diffusion handshake started stage=%d replica=%d", + stage_id, + replica_id, + ) + with connect_remote_diffusion_proc( + omni_master_server=self._omni_master_server, + stage_id=stage_id, + replica_id=replica_id, + ) as remote_resources: + resources = remote_resources + if resources is None: + raise RuntimeError(f"Remote diffusion stage {stage_id} returned no resources") + if resources.addresses is None: + raise RuntimeError(f"Remote diffusion stage {stage_id} returned no addresses") + except Exception: + self._cleanup_launched_resources( + stage_id=stage_id, + resources=resources, + ) + raise + + client = StageDiffusionClient.from_addresses( + metadata, + request_address=resources.addresses.inputs[0], + response_address=resources.addresses.outputs[0], + batch_size=ctx.diffusion_batch_size, + ) + logger.info( + "[DistStageRuntime] Remote diffusion replica attached stage=%d replica=%d", + stage_id, + replica_id, + ) + return client + + if ctx.vllm_config is None: + raise RuntimeError(f"Remote LLM stage {stage_id} is missing vllm_config") + if ctx.executor_class is None: + raise RuntimeError(f"Remote LLM stage {stage_id} is missing executor_class") + vllm_config = copy.deepcopy(ctx.vllm_config) + vllm_config.parallel_config.data_parallel_size_local = 0 + resources = None + try: + logger.info("[DistStageRuntime] Remote LLM handshake started stage=%d replica=%d", stage_id, replica_id) + with connect_remote_engine_cores( + vllm_config=vllm_config, + omni_master_server=self._omni_master_server, + stage_id=stage_id, + replica_id=replica_id, + ) as remote_resources: + resources = remote_resources + if resources is None: + raise RuntimeError(f"Remote LLM stage {stage_id} returned no resources") + if resources.addresses is None: + raise RuntimeError(f"Remote LLM stage {stage_id} returned no addresses") + client_addresses = self._client_addresses_from_zmq(resources.addresses) + replica_host = self._omni_master_server.get_replica_host(stage_id, replica_id) + if replica_host: + client_addresses["replica_host"] = replica_host + client = StageEngineCoreClientBase.make_async_mp_client( + vllm_config=vllm_config, + executor_class=ctx.executor_class, + metadata=metadata, + client_addresses=client_addresses, + engine_manager=resources.manager, + coordinator=resources.coordinator, + ) + logger.info("[DistStageRuntime] Remote LLM replica attached stage=%d replica=%d", stage_id, replica_id) + return client + except Exception: + self._cleanup_launched_resources( + stage_id=stage_id, + resources=resources, + ) + raise + + +# =========================================================================== +# Factory +# =========================================================================== + + +def create_stage_runtime( + stage_configs: list[Any], + model: str, + config_path: str, + *, + single_stage_mode: bool, + stage_init_timeout: int, + diffusion_batch_size: int, + async_chunk: bool, + tokenizer: str | None = None, + # Distributed-only params: + single_stage_id_filter: int | None = None, + omni_master_address: str | None = None, + omni_master_port: int | None = None, + omni_dp_size_local: int = 1, + omni_heartbeat_timeout: float = 30.0, + omni_lb_policy: str = "random", + request_queue: janus.Queue[EngineQueueMessage] | None = None, +) -> StageRuntime: + """Factory: select StageRuntime or DistStageRuntime.""" + if single_stage_mode: + if not omni_master_address or not omni_master_port: + raise ValueError("Distributed mode requires omni_master_address and omni_master_port") + return DistStageRuntime( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=diffusion_batch_size, + async_chunk=async_chunk, + tokenizer=tokenizer, + single_stage_id_filter=single_stage_id_filter, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_dp_size_local=omni_dp_size_local, + omni_heartbeat_timeout=omni_heartbeat_timeout, + omni_lb_policy=omni_lb_policy, + request_queue=request_queue, + ) + return StageRuntime( + stage_configs=stage_configs, + model=model, + config_path=config_path, + stage_init_timeout=stage_init_timeout, + diffusion_batch_size=diffusion_batch_size, + async_chunk=async_chunk, + tokenizer=tokenizer, + ) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 9fcaf44bda2..74911af9e17 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -26,7 +26,7 @@ from vllm.utils import random_uuid from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask +from vllm_omni.diffusion.data import CuMemTag, OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.engine.messages import ErrorMessage, OutputMessage from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.omni_base import ( @@ -141,7 +141,8 @@ def __init__(self, *args: Any, model: str = "", **kwargs: Any) -> None: OmniBase.__init__(self, model=model, **kwargs) self._pause_cond: asyncio.Condition = asyncio.Condition() self._paused: bool = False - self._is_sleeping: bool = False + self._sleeping_tags: set[str] = set() + self._level2_sleeping: bool = False self.final_output_task: asyncio.Task | None = None self.event_resolver = AsyncEventResolver(orchestrator=self) self.config_path = self.engine.config_path @@ -209,10 +210,13 @@ async def get_vllm_config(self) -> Any: def get_diffusion_od_config(self) -> Any | None: """Return the diffusion-stage config when the pipeline has one.""" + saw_diffusion_stage = False for stage_client in self.engine.stage_clients: if getattr(stage_client, "stage_type", None) != "diffusion": continue + saw_diffusion_stage = True + od_config = getattr(stage_client, "od_config", None) if od_config is not None: return od_config @@ -222,6 +226,11 @@ def get_diffusion_od_config(self) -> Any | None: if od_config is not None: return od_config + # Out-of-process diffusion clients don't carry od_config (it lives in the + # worker); fall back to the engine's model_class_name resolution. + if saw_diffusion_stage: + return self.engine.get_diffusion_od_config() + return None @property @@ -271,15 +280,16 @@ async def generate( through all stages in the pipeline and yields outputs as they become available. - **Batch mode (diffusion only):** - When *prompt* is a ``list``, all prompts are dispatched in a single - ``DiffusionEngine.step()`` call at the diffusion stage. The combined - result is yielded as one ``OmniRequestOutput`` with all generated - images. Only a single *request_id* is used for the whole batch. + **Diffusion batching:** + Diffusion stages accept only a single prompt per request. Passing a + ``list`` of prompts to a diffusion stage will raise ``ValueError``. + To batch multiple diffusion prompts, submit each as an independent + request; the scheduler will automatically co-batch compatible requests. Args: - prompt: A single prompt **or** a list of prompts. A list - triggers batch mode when the diffusion stage is reached. + prompt: A single prompt **or** a list of prompts. For diffusion + stages, only a single prompt is accepted; a list will be + rejected with an error. request_id: Unique identifier for this request. If one is not provided, a random one will be generated. sampling_params_list: List of SamplingParams, one per stage. @@ -289,11 +299,10 @@ async def generate( Yields: OmniRequestOutput objects as they are produced by each stage. - In batch mode the diffusion stage yields one output containing - all generated images. Raises: - ValueError: If sampling_params_list has incorrect length. + ValueError: If sampling_params_list has incorrect length, or + if a list prompt is submitted to a diffusion stage. """ # Append a random UUID suffix to the request_id to ensure it is unique # and non-empty, similar to vLLM's input processor. The suffix is used @@ -307,6 +316,23 @@ async def generate( logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}") + _sleeping_tags = getattr(self, "_sleeping_tags", None) + if _sleeping_tags: + raise RuntimeError( + f"Generation rejected: Engine is partially or fully asleep. " + f"Currently sleeping tags: {list(_sleeping_tags)}. " + f"Please perform a full wake_up before generating." + ) + + # Reject diffusion list-prompt early with a clear API error. + if isinstance(prompt, list) and any( + getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) + ): + raise ValueError( + "Diffusion stages accept only a single prompt per request. " + "Submit multiple independent requests to use scheduler batching." + ) + input_stream_task: asyncio.Task | None = None try: # Start final output dispatcher on the first call to generate() @@ -914,11 +940,32 @@ async def sleep( if ack is not None: await self.event_resolver.resolve(ack) final_acks.append(ack) - self._is_sleeping = True + if not hasattr(self, "_sleeping_tags"): + self._sleeping_tags = set() + self._sleeping_tags.update([CuMemTag.WEIGHTS.value, CuMemTag.KV_CACHE.value]) + if level == 2: + self._level2_sleeping = True return final_acks async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | None = None) -> list[OmniACK]: self._final_output_handler() + + if getattr(self, "_level2_sleeping", False): + raise NotImplementedError( + "wake_up() after sleep(level=2) is not yet implemented: weights were " + "discarded from GPU and reloading from disk is not yet supported. " + "Use sleep(level=1) instead, which offloads weights to CPU RAM " + "and supports fast DMA restore." + ) + _current_tags = getattr(self, "_sleeping_tags", set()) + if tags is None: + requested_tags = list(_current_tags) + else: + requested_tags = [t for t in tags if t in _current_tags] + if not requested_tags: + logger.info(f"[{self._name}] Requested tags {tags} are already warm. Skipping wake_up.") + return [] + if stage_ids is None: stage_ids = list(range(len(self.engine.stage_clients))) total_workers = 0 @@ -932,7 +979,7 @@ async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | No task_id = str(uuid.uuid4()) self.event_resolver.watch_task(task_id, expected_count=total_workers) logger.info(f"[{self._name}] Wake-up initiated (Task: {task_id}). Awaiting {total_workers} ACKs...") - task = OmniWakeTask(tags=tags, task_id=task_id) + task = OmniWakeTask(tags=requested_tags, task_id=task_id) rpc_results = await self.collective_rpc(method="handle_wake_task", args=(task,), stage_ids=stage_ids) final_acks = [] for stage_res in rpc_results: @@ -943,7 +990,14 @@ async def wake_up(self, stage_ids: list[int] | None = None, tags: list[str] | No final_acks.append(ack) current_omni_platform.synchronize() await asyncio.sleep(0.1) - self._is_sleeping = False + + for t in requested_tags: + if hasattr(self, "_sleeping_tags"): + self._sleeping_tags.discard(t) + # Only clear the level-2 flag once all tags are warm, in case partial + # wake support (e.g. tags=["kv_cache"] only) is added in the future. + if not getattr(self, "_sleeping_tags", None): + self._level2_sleeping = False logger.info(f"[{self._name}] All {len(final_acks)}/{total_workers} workers reported WARM for task {task_id}.") return final_acks @@ -953,7 +1007,7 @@ async def is_sleeping(self) -> bool: TODO(AsyncOmni): query the orchestrator once all stage backends expose a real sleeping-state RPC. For now we track the requested state locally. """ - return self._is_sleeping + return bool(getattr(self, "_sleeping_tags", None)) async def add_lora(self, lora_request: LoRARequest) -> bool: """Load a new LoRA adapter into all stages. diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index eed7d4f4e54..70e7342c000 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/vllm_omni/entrypoints/cli/benchmark/serve.py @@ -119,7 +119,7 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: help="Keep synthesized audio as 24 kHz mono PCM for WER (works with " "--backend openai-audio-speech or openai-chat-omni). Scoring follows " "zhaochenyang20/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). " - "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[seed-tts-eval]'. " + "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[dev]'. " "Optional: SEED_TTS_EVAL_DEVICE, SEED_TTS_HF_WHISPER_MODEL.", ) g.add_argument( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 5b6c7d9b34f..71483445abe 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -9,10 +9,7 @@ import json import os import signal -import threading -from multiprocessing import connection from types import FrameType -from typing import Any import uvloop from vllm.entrypoints.cli.types import CLISubcommand @@ -151,7 +148,8 @@ def validate(self, args: argparse.Namespace) -> None: "The following CLI args are not supported under --omni: " f"{', '.join(offenders)}. Configure parallelism through the " "per-stage YAML (`--deploy-config` / `--stage-configs-path`) " - "and replica count via `--omni-dp-size-local`." + "and replica count via the per-stage `num_replicas` config field " + "(single-runtime) or `--omni-dp-size-local` (headless / multi-runtime)." ) # --omni-lb-policy is validated against the LoadBalancingPolicy enum. @@ -219,6 +217,29 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu help="Default task type for TTS models (CustomVoice, VoiceDesign, or Base). " "If not specified, will be inferred from model path.", ) + # Forced aligner / word timestamps. --forced-aligner is the opt-in + # toggle; heavier knobs (gpu_memory_utilization, dtype, max_model_len) + # live in the deploy YAML passed via --forced-aligner-config. + omni_config_group.add_argument( + "--forced-aligner", + type=str, + default=None, + help=( + "Enable streaming TTS word timestamps via a forced aligner. " + "Pass the aligner model path/name, e.g. 'Qwen/Qwen3-ForcedAligner-0.6B'. " + "Disabled when omitted." + ), + ) + omni_config_group.add_argument( + "--forced-aligner-config", + type=str, + default=None, + help=( + "Optional YAML file for forced aligner settings (model, runner, " + "gpu_memory_utilization, dtype, max_model_len). The --forced-aligner " + "flag, when set, overrides the YAML model field." + ), + ) # TODO(@lishunyang12): deprecate once all models migrate to --deploy-config omni_config_group.add_argument( "--stage-configs-path", @@ -234,6 +255,15 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu help="Path to a deploy config YAML (new format with stages/engine_args). " "Mutually exclusive with --stage-configs-path.", ) + omni_config_group.add_argument( + "--strategy-config", + type=str, + default=None, + help="Path to a composable-parallel strategy.yaml. Only applies to " + "registry-based models (e.g. qwen2_5_omni): its derived parallel sizing " + "is overlaid onto the registry-merged stages before per-stage engine " + "args are built.", + ) omni_config_group.add_argument( "--stage-overrides", type=str, @@ -449,7 +479,7 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu default=None, help=( "JSON string for diffusion quantization_config. " - 'Example: \'{"method":"gguf","gguf_model":"/path/to/model.gguf"}\'.' + 'Example: \'{"method":"fp8","activation_scheme":"dynamic"}\'.' ), ) omni_config_group.add_argument( @@ -515,7 +545,7 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu "--cache-backend", type=str, default="none", - help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'", + help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'", ) omni_config_group.add_argument( "--cache-config", @@ -536,6 +566,14 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu action="store_true", help="Enable per-step diffusion execution so running requests can be aborted between denoise steps.", ) + omni_config_group.add_argument( + "--request-batch-max-wait-ms", + type=float, + default=0.0, + help="Request-mode batch admission: max milliseconds to wait for compatible " + "requests to accumulate before scheduling a fused forward wave. " + "0 disables admission (default).", + ) # VAE memory optimization parameters omni_config_group.add_argument( @@ -575,7 +613,6 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu action="store_true", help="Enable layerwise (blockwise) offloading on DiT modules.", ) - # Video model parameters (e.g., Wan2.2) - engine-level omni_config_group.add_argument( "--boundary-ratio", @@ -625,6 +662,18 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu "Distributes VAE decode workload across multiple ranks by splitting the latent spatially. " "Equivalent to setting DiffusionParallelConfig.vae_patch_parallel_size.", ) + omni_config_group.add_argument( + "--vae-parallel-mode", + type=str, + default="tile", + choices=["tile", "spatial_shard_height", "spatial_shard_width"], + help="VAE parallel decode strategy for diffusion models. " + "'tile' (default) uses patch/tile parallel decode; " + "'spatial_shard_height'/'spatial_shard_width' use spatially-sharded decode that splits " + "decoder feature maps along height/width and exchanges halo regions. The " + "'spatial_shard_*' modes require vae_patch_parallel_size to match the DiT group size. " + "Equivalent to setting DiffusionParallelConfig.vae_parallel_mode.", + ) # Default sampling parameters omni_config_group.add_argument( @@ -638,8 +687,17 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu # Diffusion model mixed precision omni_config_group.add_argument( "--max-generated-image-size", + default=7680 * 4320, # 8K resolution type=int, - help="The max size of generate image (height * width).", + help="Maximum generated image size in pixels (height * width).", + ) + # Diffusion model (mainly video generation models) streaming output mode + omni_config_group.add_argument( + "--diffusion-streaming-output", + dest="diffusion_streaming_output", + action="store_true", + default=False, + help="Enable chunked streaming output for diffusion (mainly video generation) models that support it.", ) # TTS-specific parameters @@ -670,6 +728,11 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu action="store_true", help="Enable AR stage profiler to include AR stage timing in stage_durations.", ) + omni_config_group.add_argument( + "--enable-orch-monitor", + action="store_true", + help="Enable orchestrator window monitor and write a JSON log at shutdown.", + ) # Supplementary auxiliary text encoder parameters # (e.g., the meta llama/meta llama-3.1-8b-instrument used by hidream) @@ -696,33 +759,27 @@ def run_headless(args: TrackingNamespace) -> None: headless invocations can coexist) and reports heartbeats to the head's OmniCoordinator. """ - from vllm.v1.engine.coordinator import DPCoordinator from vllm.v1.executor.multiproc_executor import MultiprocExecutor from vllm.version import __version__ as VLLM_VERSION - from vllm_omni.diffusion.stage_diffusion_proc import ( - complete_diffusion_handshake, - spawn_diffusion_proc, - ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_omni_kv_config_for_stage - from vllm_omni.engine.omni_core_engine_proc_manager import OmniCoreEngineProcManager - from vllm_omni.engine.stage_engine_startup import register_stage_with_omni_master + from vllm_omni.engine.stage_engine_startup import ( + get_headless_replica_devices, + launch_headless_diffusion_replicas, + launch_headless_llm_replicas, + ) from vllm_omni.engine.stage_init_utils import ( - build_diffusion_config, build_engine_args_dict, build_vllm_config, - extract_stage_metadata, get_stage_connector_spec, - get_stage_devices_per_replica, - inject_kv_stage_info, + inject_omni_kv_connector_config, load_omni_transfer_config_for_model, prepare_engine_environment, - setup_stage_devices, - split_devices_for_replicas, - terminate_alive_proc, ) - from vllm_omni.entrypoints.utils import inject_omni_kv_config, load_and_resolve_stage_configs - from vllm_omni.platforms import current_omni_platform + from vllm_omni.entrypoints.utils import ( + load_and_resolve_stage_configs, + parse_stage_overrides, + ) model = args.model stage_id: int | None = args.stage_id @@ -758,11 +815,19 @@ def run_headless(args: TrackingNamespace) -> None: args.replica_id, ) - config_path, stage_configs = load_and_resolve_stage_configs( + # Parse --stage-overrides (raw JSON string) exactly like the standard + # engine path (AsyncOmniEngine._resolve_stage_configs) so headless and + # standard launches resolve to the same per-stage device layout. + stage_overrides = parse_stage_overrides(args_dict.get("stage_overrides")) + + config_path, stage_configs, _ = load_and_resolve_stage_configs( model, stage_configs_path, args_dict, + trust_remote_code=bool(getattr(args, "trust_remote_code", False)), deploy_config_path=args_dict.get("deploy_config"), + stage_overrides=stage_overrides, + strategy_config_path=args_dict.get("strategy_config"), ) # Locate the stage config that matches stage_id. @@ -777,155 +842,25 @@ def run_headless(args: TrackingNamespace) -> None: ) prepare_engine_environment() - omni_transfer_config = load_omni_transfer_config_for_model(model, config_path) - omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) - - # When ``--omni-dp-size-local > 1``, slice the YAML's ``devices:`` field - # into per-replica subsets so each subprocess we spawn below sees a - # narrowed ``CUDA_VISIBLE_DEVICES`` and doesn't stack on cuda:0. Mirrors - # the head-side per-replica device application at - # ``async_omni_engine.py`` (setup_stage_devices around each launch). - runtime_cfg = getattr(stage_cfg, "runtime", None) - devices_str: str | None = None - if runtime_cfg is not None: - devices_str = ( - runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None) - ) - devices_per_replica = get_stage_devices_per_replica(stage_cfg) - if devices_str: - # Always remap YAML's logical devices through setup_stage_devices, - # even for omni_dp_size_local==1. The launcher's CUDA_VISIBLE_DEVICES - # is dropped from the engine-subprocess env between vllm-serve and - # OmniCoreEngineProcManager.Process, so the worker would otherwise - # default cuda:0 to physical GPU 0 and collide with a co-located - # head on the same host (see hyi3_multi_host_1 reproducer). - per_replica_devices: list[str | None] = split_devices_for_replicas( - devices_str, omni_dp_size_local, devices_per_replica, stage_id - ) - logger.info( - "[Headless] Stage %d: %d local replicas, devices_per_replica=%d, per-replica devices: %s", - stage_id, - omni_dp_size_local, - devices_per_replica, - per_replica_devices, - ) - else: - per_replica_devices = [None] * omni_dp_size_local - device_control_env = current_omni_platform.device_control_env_var + per_replica_devices = get_headless_replica_devices(stage_cfg, stage_id, omni_dp_size_local) if stage_cfg.stage_type == "diffusion": - metadata = extract_stage_metadata(stage_cfg) - if omni_conn_cfg: - inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - # Headless single-stage launch must still infer cross-stage TP topology - # from the loaded deploy config so heterogeneous KV routing keys match - # the head process (e.g. from_tp=2, to_tp=1). - inject_kv_stage_info(stage_cfg, stage_id, stage_configs) - od_config = build_diffusion_config(model, stage_cfg, metadata) - - logger.info( - "[Headless] Launching %d diffusion replica(s) for stage %d via OmniMasterServer at %s:%d", - omni_dp_size_local, - stage_id, - omni_master_address, - omni_master_port, + launch_headless_diffusion_replicas( + model=model, + stage_cfg=stage_cfg, + stage_configs=stage_configs, + stage_id=stage_id, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + config_path=config_path, + replica_bind_address=omni_replica_address, ) + return - procs: list[Any] = [] - try: - for _rep_idx in range(omni_dp_size_local): - # Always auto-assign: headless processes carry no knowledge - # of their per-replica id and the master server is the sole - # authority on the per-stage id namespace. - response = register_stage_with_omni_master( - omni_master_address=omni_master_address, - omni_master_port=omni_master_port, - omni_stage_id=stage_id, - omni_stage_config=stage_cfg, - replica_id=None, - return_full_response=True, - replica_bind_address=omni_replica_address, - ) - # Apply this replica's CUDA_VISIBLE_DEVICES (only when - # ``--omni-dp-size-local > 1`` and the YAML's stage devices - # field is set). The spawned subprocess inherits the env at - # spawn time; we restore the parent env afterwards so the - # next replica's setup sees the same baseline. - previous_visible_devices = os.environ.get(device_control_env) - try: - if per_replica_devices[_rep_idx] is not None: - setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]}) - # Each StageDiffusionProc starts its own - # torch.distributed group bound to - # ``od_config.master_port``. Without an explicit - # per-replica override all spawned subprocesses - # share the value ``OmniDiffusionConfig.__post_init__`` - # picked once (and the second binder hits EADDRINUSE - # on ``init_process_group``). We can't use - # kernel-ephemeral allocation either, because the - # master server's pre-allocated ZMQ ports (returned - # by ``register_stage_with_omni_master``) also live - # in the ephemeral range and are not actually bound - # until the headless ``_perform_diffusion_handshake`` - # runs — so picking an ephemeral port here can steal - # a port the master server already promised to a - # sibling headless. Use ``settle_port`` from a base - # above the Linux default ephemeral range - # (32768-60999) so torch.distributed master ports - # never overlap with ZMQ allocations. - if omni_dp_size_local > 1: - od_config.master_port = od_config.settle_port( - 61000 + _rep_idx * 100, - port_inc=37, - ) - proc, _, _, _ = spawn_diffusion_proc( - model, - od_config, - handshake_address=response.handshake_address, - request_address=response.input_address, - response_address=response.output_address, - omni_coordinator_address=response.coordinator_router_address, - omni_stage_id=stage_id, - omni_replica_id=response.replica_id, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - complete_diffusion_handshake(proc, response.handshake_address, args.stage_init_timeout) - procs.append(proc) - logger.info( - "[Headless] Diffusion replica id=%d for stage %d is up (coord=%s)", - response.replica_id, - stage_id, - response.coordinator_router_address, - ) - - # Block on the sentinel set so any replica crash is detected - # immediately (the previous per-proc join loop only noticed - # crashes in registration order). Any exit triggers fleet - # shutdown via the finally block; non-zero exits propagate. - sentinel_to_proc = {p.sentinel: p for p in procs} - died = connection.wait(list(sentinel_to_proc.keys())) - first = sentinel_to_proc[died[0]] - logger.info( - "[Headless] Diffusion replica %s exited (code=%s); shutting down stage %d.", - first.name, - first.exitcode, - stage_id, - ) - if first.exitcode not in (None, 0): - raise RuntimeError( - f"Diffusion stage {stage_id} replica {first.name!r} exited with code {first.exitcode}" - ) - return - finally: - logger.info("[Headless] Shutting down %d diffusion replica(s) for stage %d.", len(procs), stage_id) - for p in procs: - if p.is_alive(): - terminate_alive_proc(p) - + omni_transfer_config = load_omni_transfer_config_for_model(model, config_path) + omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, stage_id=stage_id, @@ -943,17 +878,7 @@ def run_headless(args: TrackingNamespace) -> None: cli_tokenizer=getattr(args, "tokenizer", None), ) - # Inject omni KV connector config so the engine runner can initialize the - # correct connector (sender/receiver role, type, addresses, etc.). - if omni_conn_cfg: - omni_kv = engine_args_dict.get("omni_kv_config") or {} - if not isinstance(omni_kv, dict): - omni_kv = dict(omni_kv) - omni_kv["connector_config"] = omni_conn_cfg - omni_kv["omni_from_stage"] = omni_from - omni_kv["omni_to_stage"] = omni_to - omni_kv.setdefault("stage_id", stage_id) - engine_args_dict["omni_kv_config"] = omni_kv + inject_omni_kv_connector_config(engine_args_dict, omni_kv_connector, stage_id) vllm_config, executor_class = build_vllm_config( stage_cfg, @@ -963,10 +888,6 @@ def run_headless(args: TrackingNamespace) -> None: headless=True, ) parallel_config = vllm_config.parallel_config - local_engine_count = parallel_config.data_parallel_size_local - - if local_engine_count <= 0: - raise ValueError("data_parallel_size_local must be > 0 in headless mode") shutdown_requested = False @@ -993,126 +914,22 @@ def signal_handler(signum: int, frame: FrameType | None) -> None: executor.start_worker_monitor(inline=True) return - dp_rank = parallel_config.data_parallel_rank if parallel_config.data_parallel_rank is not None else 0 - coordinator = None - if vllm_config.needs_dp_coordinator and dp_rank == 0: - coordinator = DPCoordinator( - parallel_config, - enable_wave_coordination=vllm_config.model_config.is_moe, - ) - logger.info( - "[Headless] Started DP Coordinator process for stage %d (PID: %d)", - stage_id, - coordinator.proc.pid, - ) - - logger.info( - "[Headless] Launching %d omni replica(s) (vLLM dp_size_local=%d each) for stage %d " - "via OmniMasterServer at %s:%d", - omni_dp_size_local, - local_engine_count, - stage_id, - omni_master_address, - omni_master_port, - ) - - # One OmniMasterServer registration per omni replica; each registration - # yields its own (handshake, input, output) allocation and the head's - # OmniCoordinator ROUTER address. We then spawn one - # OmniCoreEngineProcManager per replica so its subprocess gets the - # right replica id wired into its OmniCoordClientForStage. log_stats = bool(args.log_stats) if args.disable_log_stats: log_stats = False - engine_managers: list[Any] = [] - monitor_threads: list[threading.Thread] = [] - - def _monitor_target(mgr: Any) -> None: - try: - mgr.monitor_engine_liveness() - except Exception: - logger.exception("[Headless] monitor_engine_liveness raised") - - try: - for _rep_idx in range(omni_dp_size_local): - # Always auto-assign: see the diffusion branch comment above - # for the rationale (headless owns no replica-id namespace). - response = register_stage_with_omni_master( - omni_master_address=omni_master_address, - omni_master_port=omni_master_port, - omni_stage_id=stage_id, - omni_stage_config=stage_cfg, - coordinator=coordinator, - replica_id=None, - return_full_response=True, - replica_bind_address=omni_replica_address, - # LLM headless: the head binds *all* three sockets — - # handshake ROUTER (``connect_remote_engine_cores``), - # input ROUTER and output PULL (``CoreClient`` — - # ``make_zmq_socket`` defaults bind=True for PULL). - # The remote LLM worker is purely a connector: it - # opens 3 outbound TCP connections to the master's - # host. So the master must keep every address on - # its own host; rewriting any of them to this - # replica's NIC makes the head's ``bind`` go - # EADDRNOTAVAIL on a cross-host launch. - replica_binds_sockets=False, - ) - # Per-replica CUDA_VISIBLE_DEVICES, same pattern as the diffusion - # branch above. OmniCoreEngineProcManager.__init__ spawns its - # subprocesses via context.Process inside the constructor, so we - # must set the env *before* instantiation and restore after. - previous_visible_devices = os.environ.get(device_control_env) - try: - if per_replica_devices[_rep_idx] is not None: - setup_stage_devices(stage_id, {"devices": per_replica_devices[_rep_idx]}) - mgr = OmniCoreEngineProcManager( - local_engine_count=local_engine_count, - start_index=dp_rank, - local_start_index=0, - vllm_config=vllm_config, - local_client=False, - handshake_address=response.handshake_address, - executor_class=executor_class, - log_stats=log_stats, - omni_stage_id=stage_id, - omni_coordinator_address=response.coordinator_router_address, - omni_replica_base_id=response.replica_id, - ) - finally: - if previous_visible_devices is None: - current_omni_platform.unset_device_control_env_var() - else: - current_omni_platform.set_device_control_env_var(previous_visible_devices) - engine_managers.append(mgr) - logger.info( - "[Headless] Stage %d replica id=%d up (coord=%s)", - stage_id, - response.replica_id, - response.coordinator_router_address, - ) - - # Run all managers' liveness monitors in parallel. Each blocks - # until its own subprocesses exit (or fail). - if len(engine_managers) == 1: - engine_managers[0].monitor_engine_liveness() - else: - for mgr in engine_managers: - t = threading.Thread(target=_monitor_target, args=(mgr,), name=f"omni-replica-monitor-{id(mgr):x}") - t.start() - monitor_threads.append(t) - for t in monitor_threads: - t.join() - finally: - logger.info("[Headless] Shutting down stage %d (%d managers).", stage_id, len(engine_managers)) - for mgr in engine_managers: - try: - mgr.shutdown() - except Exception: - logger.exception("[Headless] engine manager shutdown failed") - if coordinator is not None: - coordinator.shutdown() + launch_headless_llm_replicas( + vllm_config=vllm_config, + executor_class=executor_class, + log_stats=log_stats, + omni_master_address=omni_master_address, + omni_master_port=omni_master_port, + stage_id=stage_id, + stage_config=stage_cfg, + omni_dp_size_local=omni_dp_size_local, + per_replica_devices=per_replica_devices, + replica_bind_address=omni_replica_address, + ) def cmd_init() -> list[CLISubcommand]: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 020b483f0a8..14bdc10be62 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -25,16 +25,25 @@ class Omni(OmniBase): """Synchronous entrypoint for offline generation.""" - def _set_final_only_for_llm_stages( + def _maybe_force_final_only_for_llm_stages( self, sampling_params_list: Sequence[OmniSamplingParams], ) -> list[OmniSamplingParams]: - """Return per-stage params with LLM stages forced to FINAL_ONLY.""" + """Return per-stage params with LLM stages forced to FINAL_ONLY. + + The caller may explicitly request ``output_kind = DELTA`` on a stage to + opt into streaming; such stages are left alone. All other LLM stages + are forced to FINAL_ONLY. + """ effective_params: list[OmniSamplingParams] = [] for stage_id, params in enumerate(sampling_params_list): sp = copy.deepcopy(params) stage_meta = self.engine.get_stage_metadata(stage_id) - if stage_meta.stage_type != "diffusion" and hasattr(sp, "output_kind"): + if ( + stage_meta.stage_type != "diffusion" + and hasattr(sp, "output_kind") + and sp.output_kind != RequestOutputKind.DELTA + ): sp.output_kind = RequestOutputKind.FINAL_ONLY effective_params.append(sp) return effective_params @@ -103,7 +112,7 @@ def _run_generation( use_tqdm: bool | Callable[..., tqdm] = True, ) -> Generator[OmniRequestOutput, None, None]: try: - sampling_params_list = self._set_final_only_for_llm_stages(sampling_params_list) + sampling_params_list = self._maybe_force_final_only_for_llm_stages(sampling_params_list) if isinstance(prompts, str) or not isinstance(prompts, Sequence): request_prompts: list[OmniPromptType] = [prompts] diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 029665f59e0..d8b2c029142 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -144,7 +144,7 @@ def __init__( log_stats = kwargs.pop("log_stats", False) self._enable_ar_profiler = kwargs.pop("enable_ar_profiler", False) # NOTE: read-only lookup — must NOT pop. Popping here drops the key - # before it reaches ``StageConfigFactory._create_from_registry``, so + # before it reaches ``StageConfigFactory._create_legacy_from_registry``, so # ``--no-async-chunk`` (``async_chunk=False``) silently fails to # override the deploy YAML's ``async_chunk: true`` default. async_chunk = kwargs.get("async_chunk") @@ -221,6 +221,13 @@ def stage_configs(self) -> list: """Expose engine stage configs for PD disaggregation detection and validation.""" return self.engine.stage_configs + def _consumed_metric_message_ids(self, request_id: str) -> set[int]: + consumed_by_request = getattr(self, "_consumed_metric_messages", None) + if consumed_by_request is None: + consumed_by_request = {} + self._consumed_metric_messages = consumed_by_request + return consumed_by_request.setdefault(request_id, set()) + def _has_dead_stage(self) -> bool: for stage_client in self.engine.stage_clients: if getattr(stage_client, "_engine_dead", False): @@ -310,7 +317,9 @@ def _log_summary_and_cleanup(self, request_id: str) -> None: ) finally: self.request_states.pop(request_id, None) - self._consumed_metric_messages.pop(request_id, None) + consumed_by_request = getattr(self, "_consumed_metric_messages", None) + if consumed_by_request is not None: + consumed_by_request.pop(request_id, None) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. @@ -397,7 +406,7 @@ def _handle_output_message( stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type) msg_id = id(msg) - consumed = self._consumed_metric_messages.setdefault(req_id, set()) + consumed = self._consumed_metric_message_ids(req_id) if msg_id not in consumed: req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type) submit_ts = msg.stage_submit_ts @@ -505,7 +514,11 @@ def _process_single_result( stage_meta = self.engine.get_stage_metadata(stage_id) output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type) if finished and _m is not None: - metrics.on_stage_metrics(stage_id, req_id, _m, output_type) + msg_id = id(result) + consumed = self._consumed_metric_message_ids(req_id) + if msg_id not in consumed: + metrics.on_stage_metrics(stage_id, req_id, _m, output_type) + consumed.add(msg_id) if not stage_meta.final_output: return None @@ -530,6 +543,15 @@ def _process_single_result( finished_reason=fr, ) + # Token counters — aggregate across all stages for this request. + _prompt_tok = 0 + _gen_tok = 0 + for evt in metrics.stage_events.get(rid_key, []): + if evt.stage_id == 0: + _prompt_tok += int(evt.num_tokens_in) + _gen_tok += int(evt.num_tokens_out) + self.prom_metrics.observe_tokens(_prompt_tok, _gen_tok) + # Modality observe inside the same finalize guard so it fires # once per request and inherits the try/except isolation. observe_modality_at_finalize( @@ -578,9 +600,8 @@ def _process_single_result( if current_stage_metrics is not None: response_metrics["stage_id"] = current_stage_metrics["stage_id"] response_metrics["final_output_type"] = current_stage_metrics["final_output_type"] - if current_stage_metrics["final_output_type"] == "text": - response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"] - response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"] + response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"] + response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"] return OmniRequestOutput( request_id=req_id or "", stage_id=stage_id, @@ -595,6 +616,7 @@ def _process_single_result( metrics=response_metrics, stage_durations=stage_durations, peak_memory_mb=peak_memory_mb, + finished=finished, ) def shutdown(self, timeout: float | None = None) -> None: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index cec80751580..680fa6b27d5 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -32,7 +32,7 @@ from starlette.types import ASGIApp, Receive, Scope, Send from vllm.engine.protocol import EngineClient from vllm.entrypoints.anthropic.serving import AnthropicServingMessages -from vllm.entrypoints.chat_utils import load_chat_template +from vllm.entrypoints.chat_utils import ChatTemplateConfig, load_chat_template from vllm.entrypoints.launcher import serve_http, terminate_if_errored from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app @@ -60,13 +60,12 @@ from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding from vllm.entrypoints.pooling.pooling.serving import ServingPooling from vllm.entrypoints.pooling.scoring.serving import ServingScores -from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. # Keep a fallback for older/newer upstream layouts during rebase windows. from vllm.entrypoints.serve.instrumentator.basic import base -from vllm.entrypoints.serve.render.serving import OpenAIServingRender -from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization +from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, process_lora_modules, @@ -85,6 +84,7 @@ OpenAIServingTranslation, ) from vllm.logger import init_logger +from vllm.renderers.online_renderer import OnlineRenderer from vllm.tasks import POOLING_TASKS from vllm.tool_parsers import ToolParserManager from vllm.utils import random_uuid @@ -111,6 +111,7 @@ ImageData, ImageGenerationRequest, ImageGenerationResponse, + ResponseFormat, ) from vllm_omni.entrypoints.openai.protocol.videos import ( SecondStr, @@ -133,18 +134,20 @@ ReferenceImage, ReferenceVideo, ) -from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler +from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler +from vllm_omni.entrypoints.openai.serving_video_stream import create_streaming_video_handler from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, get_default_sampling_params_list, ) -from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER +from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER, FileStorageHandle 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_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.forced_aligner import build_forced_aligner_config from vllm_omni.utils.tracking_parser import TrackingArgumentParser, TrackingNamespace logger = init_logger(__name__) @@ -448,7 +451,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: # Add process-specific prefix to stdout and stderr. decorate_logs("APIServer", skip_if_decorated=True) - listen_address, sock = setup_openai_server(args) + listen_address, sock = setup_openai_server(args, reuse_port=False) # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) @@ -513,6 +516,9 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, else: logger.warning("engine client has no endpoint restrictions attribute") + # Start background processes + await STORAGE_MANAGER.start() + # Conditionally register profiler endpoints based on stage YAML configs stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None if _should_enable_profiler_endpoints(stage_configs): @@ -576,15 +582,14 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: **uvicorn_kwargs, ) - # NB: Await server shutdown only after the backend context is exited - try: - await shutdown_task - finally: - state = getattr(app, "state", None) - serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None - if serving_speech is not None: - serving_speech.shutdown() - sock.close() + try: + await shutdown_task + finally: + state = getattr(app, "state", None) + serving_speech = getattr(state, "openai_serving_speech", None) if state is not None else None + if serving_speech is not None: + serving_speech.shutdown() + sock.close() @asynccontextmanager @@ -685,6 +690,7 @@ async def build_async_omni_from_stage_config( try: kwargs = args.get_explicit_kwargs_dict() model = kwargs.pop("model", None) or args.model + kwargs.setdefault("log_stats", not args.disable_log_stats) async_omni = AsyncOmni(model=model, **kwargs) # # Don't keep the dummy data in memory @@ -752,7 +758,7 @@ async def omni_init_app_state( state.diffusion_engine = engine_client state.openai_serving_models = _DiffusionServingModels(base_model_paths) # OMNI: tokenization endpoints are not supported in pure diffusion mode. - state.openai_serving_tokenization = None + state.serving_tokenization = None # Use for_diffusion method to create chat handler state.openai_serving_chat = OmniOpenAIServingChat.for_diffusion( @@ -776,6 +782,11 @@ async def omni_init_app_state( model_name=model_name, stage_configs=diffusion_stage_configs, ) + state.openai_streaming_video_output = OmniStreamingVideoOutputHandler( + engine_client=engine_client, + model_name=model_name, + stage_configs=diffusion_stage_configs, + ) state.openai_serving_speech = OmniOpenAIServingSpeech.for_diffusion( diffusion_engine=engine_client, @@ -803,19 +814,6 @@ async def omni_init_app_state( state.vllm_config = vllm_config - # Propagate enable_in_reasoning to the API-server process. The engine core - # runs in a separate process, so the contextvar that backs - # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers - # call `get_enable_structured_outputs_in_reasoning()` during request - # handling and need to see the real flag, otherwise they silently fall - # back to False and mismatch the engine-side bitmask gating. - if vllm_config is not None: - from vllm.tool_parsers.structural_tag_registry import ( - set_enable_structured_outputs_in_reasoning, - ) - - set_enable_structured_outputs_in_reasoning(vllm_config.structured_outputs_config.enable_in_reasoning) - # Get supported tasks supported_tasks: set[str] = {"generate"} if hasattr(engine_client, "get_supported_tasks"): @@ -832,6 +830,12 @@ async def omni_init_app_state( if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: resolved_chat_template = _load_model_chat_template_json(args.model) + chat_template_config = ChatTemplateConfig( + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + ) + if args.tool_server == "demo": tool_server: ToolServer | None = DemoToolServer() assert isinstance(tool_server, DemoToolServer) @@ -896,16 +900,11 @@ async def omni_init_app_state( ) await state.openai_serving_models.init_static_loras() - # NOTE: kept aligned with vllm 0.20 `init_app_state`: - # - dropped the `io_processor` kwarg (no longer accepted by 0.20); - # io_processor stays on `engine_client` and downstream serving classes - # read it from there. - # - pass `reasoning_parser` so render-time `adjust_request` runs for - # reasoning models (matches `vllm.entrypoints.openai.api_server`). - state.openai_serving_render = OpenAIServingRender( + # NOTE: kept aligned with upstream `init_app_state`: + # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream). + state.online_renderer = OnlineRenderer( model_config=engine_client.model_config, renderer=engine_client.renderer, - model_registry=state.openai_serving_models.registry, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -922,7 +921,7 @@ async def omni_init_app_state( OpenAIServingResponses( engine_client, state.openai_serving_models, - openai_serving_render=state.openai_serving_render, + state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -943,7 +942,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, args.response_role, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -970,7 +969,7 @@ async def omni_init_app_state( OpenAIServingCompletion( engine_client, state.openai_serving_models, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, @@ -985,9 +984,7 @@ async def omni_init_app_state( state.openai_serving_models, supported_tasks=tuple(supported_tasks), request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if any(task in POOLING_TASKS for task in supported_tasks) else None @@ -997,9 +994,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if "embed" in supported_tasks else None @@ -1009,9 +1004,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if "classify" in supported_tasks else None @@ -1020,17 +1013,17 @@ async def omni_init_app_state( ServingScores( engine_client, state.openai_serving_models, + supported_tasks=tuple(supported_tasks), request_logger=request_logger, - score_template=resolved_chat_template, + chat_template_config=chat_template_config, log_error_stack=args.log_error_stack, ) if any(t in supported_tasks for t in ("embed", "score", "token_embed")) else None ) - state.openai_serving_tokenization = OpenAIServingTokenization( - engine_client, + state.serving_tokenization = ServingTokenization( state.openai_serving_models, - state.openai_serving_render, + state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -1062,7 +1055,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, args.response_role, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -1081,7 +1074,7 @@ async def omni_init_app_state( ServingTokens( engine_client, state.openai_serving_models, - state.openai_serving_render, + state.online_renderer, request_logger=request_logger, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, @@ -1093,7 +1086,11 @@ async def omni_init_app_state( ) state.openai_serving_speech = OmniOpenAIServingSpeech( - engine_client, state.openai_serving_models, request_logger=request_logger, model_name=model_name + engine_client, + state.openai_serving_models, + request_logger=request_logger, + model_name=model_name, + forced_aligner_config=build_forced_aligner_config(args), ) # Warm up speech pipeline (CUDA Graph capture, torch.compile) so the first @@ -1108,7 +1105,7 @@ async def omni_init_app_state( speech_service=state.openai_serving_speech, ) state.openai_streaming_video = ( - OmniStreamingVideoHandler( + create_streaming_video_handler( chat_service=state.openai_serving_chat, engine_client=engine_client, ) @@ -1167,7 +1164,8 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re if not isinstance(modalities, list) or not all(isinstance(m, str) for m in modalities): raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, - detail='modalities must be a list of strings, e.g. ["text", "audio", "image"]', + detail="modalities must be a list of strings, e.g. " + '[\"text\", \"audio\", \"image\"]', ) if "logprobs" in raw_body and not isinstance(raw_body["logprobs"], bool): raise HTTPException( @@ -1177,7 +1175,7 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "") handler = Omnichat(raw_request) if handler is None: - base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None) + base_server = getattr(raw_request.app.state, "serving_tokenization", None) if base_server is None: raise HTTPException( status_code=HTTPStatus.NOT_FOUND.value, @@ -1243,7 +1241,7 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re "/v1/audio/speech", dependencies=[Depends(validate_json_request)], responses={ - HTTPStatus.OK.value: {"content": {"audio/*": {}}}, + HTTPStatus.OK.value: {"content": {"audio/*": {}, "text/event-stream": {}}}, HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse}, HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse}, @@ -1268,7 +1266,7 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request """ handler = Omnispeech(raw_request) if handler is None: - base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None) + base_server = getattr(raw_request.app.state, "serving_tokenization", None) if base_server is None: raise HTTPException( status_code=HTTPStatus.NOT_FOUND.value, @@ -1306,7 +1304,7 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request): handler = Omnispeech(raw_request) if handler is None: - base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None) + base_server = getattr(raw_request.app.state, "serving_tokenization", None) if base_server is None: raise HTTPException( status_code=HTTPStatus.NOT_FOUND.value, @@ -1322,7 +1320,10 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request) result = await handler.create_speech_batch(request) if isinstance(result, ErrorResponse): return _error_response_to_json_response(result) - return JSONResponse(content=result.model_dump()) + # exclude_none so optional per-item fields are omitted rather than + # serialized as null: errored items drop `usage`/`audio_data`/`media_type`, + # successful items drop `error`. Matches the documented batch response shape. + return JSONResponse(content=result.model_dump(exclude_none=True)) except (EngineGenerateError, EngineDeadError) as exc: return _create_engine_error_json_response(raw_request, exc) except ValueError as e: @@ -1346,7 +1347,7 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request) async def create_audio_generate(request: OpenAICreateAudioGenerateRequest, raw_request: Request): handler = OmniAudioGenerate(raw_request) if handler is None: - base_server = getattr(raw_request.app.state, "openai_serving_tokenization", None) + base_server = getattr(raw_request.app.state, "serving_tokenization", None) if base_server is None: raise HTTPException( status_code=HTTPStatus.NOT_FOUND.value, @@ -1605,6 +1606,23 @@ async def streaming_video_chat(websocket: WebSocket): await handler.handle_session(websocket) +@router.websocket("/v1/realtime/video") +async def streaming_video_output(websocket: WebSocket): + """WebSocket endpoint for streaming generated video output chunks.""" + handler = getattr(websocket.app.state, "openai_streaming_video_output", None) + if handler is None: + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "message": "Streaming video generation is not available", + } + ) + await websocket.close() + return + await handler.handle_session(websocket) + + @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" @@ -1694,6 +1712,7 @@ async def show_available_models(raw_request: Request) -> JSONResponse: @router.post( "/v1/images/generations", dependencies=[Depends(validate_json_request)], + response_model=None, responses={ HTTPStatus.OK.value: {"model": ImageGenerationResponse}, HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse}, @@ -1702,7 +1721,9 @@ async def show_available_models(raw_request: Request) -> JSONResponse: }, ) @with_cancellation -async def generate_images(request: ImageGenerationRequest, raw_request: Request): +async def generate_images( + request: ImageGenerationRequest, raw_request: Request +) -> ImageGenerationResponse | StreamingResponse: """Generate images from text prompts using diffusion models. OpenAI DALL-E compatible endpoint for text-to-image generation. @@ -1894,7 +1915,10 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) } if request.size: response_kwargs["size"] = size_str - return ImageGenerationResponse(**response_kwargs) + response = ImageGenerationResponse(**response_kwargs) + if request.response_format != ResponseFormat.FILE: + return response + return response.stream_response() except (EngineGenerateError, EngineDeadError) as exc: return _create_engine_error_json_response(raw_request, exc) @@ -2420,7 +2444,8 @@ def _check_max_generated_image_size( raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, detail=f"Requested image size {width}x{height} exceeds the maximum allowed " - f"size of {max_generated_image_size} pixels.", + f"size of {max_generated_image_size} pixels. You can reduce the requested size " + f"or increase the server's --max-generated-image-size limit.", ) elif resolution is not None: # When resolution is set, the output size is resolution * resolution @@ -2428,7 +2453,8 @@ def _check_max_generated_image_size( raise HTTPException( status_code=HTTPStatus.BAD_REQUEST.value, detail=f"Requested resolution {resolution} (max {resolution}x{resolution} pixels) " - f"exceeds the maximum allowed size of {max_generated_image_size} pixels.", + f"exceeds the maximum allowed size of {max_generated_image_size} pixels. " + f"You can reduce the requested size or increase the server's --max-generated-image-size limit.", ) @@ -2774,12 +2800,11 @@ def _video_error_from_exception(exc: Exception) -> VideoError: ) -def _cleanup_video(video_id: str, output_path: str | None): +async def _cleanup_video(video_id: str): try: - if output_path is not None: - os.remove(output_path) - except OSError: - logger.warning("Failed to cleanup partial video file '%s' for id=%s", output_path, video_id) + await STORAGE_MANAGER.delete(video_id) + except Exception: + logger.warning("Failed to cleanup partial video file '%s'", video_id) async def _run_video_generation_job( @@ -2798,7 +2823,6 @@ async def _run_video_generation_job( await VIDEO_STORE.update_fields(video_id, {"status": VideoGenerationStatus.IN_PROGRESS}) started_at = time.perf_counter() - output_path = None try: video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes( request, @@ -2808,27 +2832,27 @@ async def _run_video_generation_job( reference_audio=reference_audio, ) - file_name = f"{video_id}.{job.file_extension}" - output_path = await STORAGE_MANAGER.save(video_bytes, file_name) - logger.info("Video request %s persisted %s output file.", video_id, output_path) + save_context = await STORAGE_MANAGER.save(video_bytes, video_id) + logger.info("Video request %s persisted %s output file.", video_id, save_context.key) - await VIDEO_STORE.update_fields( - video_id, - { - "status": VideoGenerationStatus.COMPLETED, - "progress": 100, - "file_name": file_name, - "completed_at": int(time.time()), - "inference_time_s": time.perf_counter() - started_at, - "stage_durations": stage_durations, - "peak_memory_mb": peak_memory_mb, - "action": action, - }, - ) + updated_fields = { + "status": VideoGenerationStatus.COMPLETED, + "progress": 100, + "file_name": f"{video_id}.{job.file_extension}", + "completed_at": save_context.created_at, + "inference_time_s": time.perf_counter() - started_at, + "stage_durations": stage_durations, + "peak_memory_mb": peak_memory_mb, + "action": action, + } + if save_context.expires_at is not None: + updated_fields["expires_at"] = save_context.expires_at + + await VIDEO_STORE.update_fields(video_id, updated_fields) except (EngineGenerateError, EngineDeadError) as exc: logger.exception("Video generation failed (engine error) for id=%s", video_id) - _cleanup_video(video_id, output_path) + await _cleanup_video(video_id) await VIDEO_STORE.update_fields( video_id, { @@ -2848,7 +2872,7 @@ async def _run_video_generation_job( except Exception as exc: logger.exception("Video generation failed for id=%s", video_id) - _cleanup_video(video_id, output_path) + await _cleanup_video(video_id) await VIDEO_STORE.update_fields( video_id, { @@ -2859,7 +2883,7 @@ async def _run_video_generation_job( }, ) except asyncio.CancelledError: - _cleanup_video(video_id, output_path) + await _cleanup_video(video_id) await VIDEO_STORE.pop(video_id) raise finally: @@ -3247,7 +3271,7 @@ async def delete_video(video_id: str) -> VideoDeleteResponse: elif job.status is VideoGenerationStatus.FAILED: if job.file_name is not None: try: - await STORAGE_MANAGER.delete(job.file_name) + await STORAGE_MANAGER.delete(video_id) except Exception: logger.warning("Failed to delete stored artifact for failed video job %s", video_id, exc_info=True) @@ -3257,13 +3281,13 @@ async def delete_video(video_id: str) -> VideoDeleteResponse: if job.file_name is None: raise HTTPException(status_code=409, detail="Video output not yet available. Please try again later.") - await STORAGE_MANAGER.delete(job.file_name) + await STORAGE_MANAGER.delete(video_id) await VIDEO_STORE.pop(video_id) return VideoDeleteResponse(id=job.id, deleted=True) @router.get("/v1/videos/{video_id}/content") -async def download_video(video_id: str) -> FileResponse: +async def download_video(video_id: str) -> Response: """Download the generated file for a completed video job. Args: @@ -3286,11 +3310,19 @@ async def download_video(video_id: str) -> FileResponse: if not job.file_name: raise HTTPException(status_code=404, detail="Generation is still in-progress") - full_path = STORAGE_MANAGER.get_full_file_path(job.file_name) - if not os.path.exists(full_path): + file_handle = await STORAGE_MANAGER.open(video_id) + if file_handle is None: raise HTTPException(status_code=404, detail="Generated video file not found on disk") - return FileResponse(path=full_path, media_type=job.media_type, filename=job.file_name) + file_name = job.file_name or f"{video_id}.{job.file_extension}" + if isinstance(file_handle, FileStorageHandle): + response = FileResponse(path=file_handle.path, media_type=job.media_type, filename=file_name) + else: + raise HTTPException( + status_code=500, detail=f"Server generated an unsupported file storage handle for file id {video_id}" + ) + + return response @profiler_router.post("/start_profile") @@ -3309,9 +3341,9 @@ async def start_profile(raw_request: Request, request: ProfileRequest | None = N stages = request.stages if request else None logger.info("Starting profiler for stages: %s", stages if stages else "all") engine_client = raw_request.app.state.engine_client - result = await engine_client.start_profile(stages=stages) + await engine_client.start_profile(stages=stages) logger.info("Profiler started.") - return JSONResponse(content=result) + return JSONResponse(content={"status": "SUCCESS"}) except Exception as e: logger.exception("Failed to start profiler: %s", e) raise HTTPException( @@ -3335,9 +3367,9 @@ async def stop_profile(raw_request: Request, request: ProfileRequest | None = No stages = request.stages if request else None logger.info("Stopping profiler for stages: %s", stages if stages else "all") engine_client = raw_request.app.state.engine_client - result = await engine_client.stop_profile(stages=stages) + await engine_client.stop_profile(stages=stages) logger.info("Profiler stopped.") - return JSONResponse(content=result) + return JSONResponse(content={"status": "SUCCESS"}) except Exception as e: logger.exception("Failed to stop profiler: %s", e) raise HTTPException( diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index a333d58af67..6c704320175 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -24,13 +24,9 @@ def create_audio(self, audio_obj: CreateAudio) -> AudioResponse: audio_tensor = audio_obj.audio_tensor sample_rate = audio_obj.sample_rate response_format = audio_obj.response_format.lower() - stream_format = audio_obj.stream_format base64_encode = audio_obj.base64_encode speed = audio_obj.speed - if stream_format != "audio": - raise ValueError(f"Unsupported stream format: {stream_format}") - if soundfile is None: raise ImportError( "soundfile is required for audio generation. Please install it with: pip install soundfile" @@ -95,12 +91,16 @@ def _apply_speed_adjustment(self, audio_tensor: np.ndarray, speed: float, sample else: waveform = torch.from_numpy(audio_tensor).unsqueeze(0) - # Match librosa.stft defaults: n_fft=2048, hop_length=n_fft//4 - n_fft = 2048 + # Use a speech-sized analysis window. The previous 2048-sample + # window is tuned for music and can smear short consonants after + # aggressive compression, which makes ASR transcript checks flaky. + n_fft = 768 hop_length = n_fft // 4 + window = torch.hann_window(n_fft, device=waveform.device, dtype=waveform.dtype) to_spec = torchaudio.transforms.Spectrogram( n_fft=n_fft, hop_length=hop_length, + window_fn=lambda *_args, **_kwargs: window, power=None, ) stretch = torchaudio.transforms.TimeStretch( @@ -110,6 +110,7 @@ def _apply_speed_adjustment(self, audio_tensor: np.ndarray, speed: float, sample to_wave = torchaudio.transforms.InverseSpectrogram( n_fft=n_fft, hop_length=hop_length, + window_fn=lambda *_args, **_kwargs: window, ) spec = to_spec(waveform) diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index b958d786a20..5bd79e9362f 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/vllm_omni/entrypoints/openai/image_api_utils.py @@ -14,7 +14,7 @@ import PIL.Image SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024) -SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11) +SUPPORTED_LAYERED_LAYERS_RANGE = range(2, 11) def parse_size(size_str: str) -> tuple[int, int]: diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 1ca47d9e0b0..98596eb970b 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -69,12 +69,20 @@ class OpenAICreateSpeechRequest(BaseModel): ge=0.25, le=4.0, ) - stream_format: Literal["sse", "audio"] | None = "audio" + stream_format: Literal["sse", "audio"] | None = Field( + default=None, + description=( + "Streaming output format. 'audio' streams raw pcm/wav bytes; " + "'sse' streams OpenAI speech.audio.* SSE events. If omitted, stream=true " + "selects SSE and stream=false remains non-streaming." + ), + ) stream: bool = Field( default=False, description=( - "If true, stream raw PCM audio chunks as they are decoded. " - "Requires response_format='pcm'. Speed adjustment is not supported when streaming." + "Streaming switch; defaults to OpenAI speech.audio.* SSE events. " + "Set stream_format='audio' to opt into raw pcm/wav byte streaming. " + "Requires response_format='pcm' or 'wav'. Speed adjustment is not supported when streaming." ), ) @@ -149,6 +157,15 @@ class OpenAICreateSpeechRequest(BaseModel): default=None, description=("Optional model-specific parameters passed directly to the model's extra_args."), ) + word_timestamps: bool = Field( + default=False, + description=( + "When true, the server runs a shared forced aligner alongside the streamed " + "audio and emits per-chunk word timestamps. Requires the server to be " + "launched with --forced-aligner pointing at an aligner model. No effect " + "when streaming is off." + ), + ) @field_validator("voice") @classmethod @@ -160,8 +177,6 @@ def validate_voice_not_empty(cls, v: str | None) -> str | None: @field_validator("stream_format") @classmethod def validate_stream_format(cls, v: str) -> str: - if v == "sse": - raise ValueError("'sse' is not a supported stream_format yet. Please use 'audio'.") return v @field_validator("ref_audio", mode="before") @@ -289,20 +304,28 @@ def validate_embedding_constraints(self) -> "OpenAICreateSpeechRequest": raise ValueError("'speaker_embedding' and 'ref_audio' are mutually exclusive") return self + def is_raw_audio_stream(self) -> bool: + return self.stream_format == "audio" + + def is_sse_stream(self) -> bool: + return (self.stream or self.stream_format == "sse") and not self.is_raw_audio_stream() + + def is_streaming(self) -> bool: + return self.is_raw_audio_stream() or self.is_sse_stream() + @model_validator(mode="after") def validate_streaming_constraints(self) -> "OpenAICreateSpeechRequest": - if self.stream: + if self.is_streaming(): if self.response_format not in ("pcm", "wav"): raise ValueError( - "Streaming (stream=true) requires response_format='pcm' or 'wav'. " + "Streaming (stream=true, stream_format='audio', or stream_format='sse') " + "requires response_format='pcm' or 'wav'. " f"Got response_format='{self.response_format}'." ) if self.speed is None: self.speed = 1.0 elif self.speed != 1.0: - raise ValueError( - "Speed adjustment is not supported when streaming (stream=true). Set speed=1.0 or omit it." - ) + raise ValueError("Speed adjustment is not supported when streaming. Set speed=1.0 or omit it.") return self @@ -371,7 +394,6 @@ class CreateAudio(BaseModel): sample_rate: int = 24000 response_format: str = "wav" speed: float = 1.0 - stream_format: Literal["sse", "audio"] | None = "audio" base64_encode: bool = True class Config: @@ -506,12 +528,65 @@ def validate_ref_audio(cls, v: str | None) -> str | None: return _validate_ref_audio_uri(v) +class SpeechInputTokenDetails(BaseModel): + """Per-modality breakdown of the speech request's *input* tokens. + + The aggregate ``input_tokens`` on :class:`SpeechTokenUsage` is the sum of + these. We surface the split (rather than one opaque number) for the same + reason OpenAI's realtime/chat usage does: text and audio inputs are billed + and reasoned about differently, and folding them together is misleading. + + Fields: + text_tokens: Tokens of the text to synthesize. This is ``input`` plus + ``instructions`` (style/emotion prompt), because both are tokenized + into the model prefill. This is the number that should scale with + how much text the caller asked to speak. + audio_tokens: Reference-audio codec frames used as voice-cloning + conditioning. NON-ZERO only when in-context voice cloning is + actually active (Qwen3-TTS ``task_type='Base'`` ICL). It is 0 for + CustomVoice/VoiceDesign and for x-vector-only cloning, because those + paths put no reference codec frames into the prefill. See issue + #4646: this is the value that previously leaked into ``prompt_tokens`` + and made Base usage look independent of the input text. + """ + + text_tokens: int = 0 + audio_tokens: int = 0 + + +class SpeechTokenUsage(BaseModel): + """Token usage for a speech (TTS) request. + + Field naming follows OpenAI's documented ``speech.audio.done`` event + (``input_tokens``/``output_tokens``/``total_tokens``), NOT chat's + ``prompt_tokens``/``completion_tokens``. + + input_tokens = text_tokens + audio_tokens (see SpeechInputTokenDetails) + output_tokens = generated codec/audio tokens (stage-0 decode steps) + total_tokens = input_tokens + output_tokens + + IMPORTANT: ``input_tokens`` is computed from the *semantic* request inputs + (tokenized text + reference-audio frames), NOT from ``len(prompt_token_ids)``. + For staged TTS models that engine prompt is a ``[1] * prefill_len`` + placeholder whose length mirrors the model prefill, so it is not a faithful + input-token count (issue #4646). + """ + + input_tokens: int = 0 + output_tokens: int = 0 + total_tokens: int = 0 + input_token_details: SpeechInputTokenDetails = Field(default_factory=SpeechInputTokenDetails) + + class SpeechBatchItemResult(BaseModel): index: int status: Literal["success", "error"] audio_data: str | None = None media_type: str | None = None error: str | None = None + # Per-item token usage (input text + reference-audio conditioning, and + # generated audio tokens). None when the item errored before generation. + usage: SpeechTokenUsage | None = None class BatchSpeechResponse(BaseModel): @@ -561,11 +636,13 @@ class StreamingSpeechSessionConfig(BaseModel): "Requires response_format='pcm'. Speed adjustment is not supported when streaming." ), ) - split_granularity: Literal["sentence", "clause"] = Field( - default="sentence", + word_timestamps: bool = Field( + default=False, description=( - "Text splitting granularity: 'sentence' splits on .!?。！？, " - "'clause' also splits on CJK commas ， and semicolons ；." + "When true, audio chunks are wrapped in JSON 'audio.chunk' frames carrying " + "base64-encoded PCM plus aligned word timestamps. Requires the server to be " + "launched with --forced-aligner. When false, audio is sent as raw binary " + "frames (existing behavior)." ), ) diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index ffd229619a3..5e5d898f3d6 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -7,9 +7,16 @@ for text-to-image generation, with vllm-omni specific extensions. """ +import base64 +import io +import uuid +import zipfile from enum import Enum +from http import HTTPStatus from typing import Any, Literal +from fastapi import HTTPException +from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field, field_validator from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers @@ -20,6 +27,7 @@ class ResponseFormat(str, Enum): B64_JSON = "b64_json" URL = "url" # Not implemented in PoC + FILE = "file" # file response class ImageGenerationRequest(BaseModel): @@ -52,7 +60,7 @@ class ImageGenerationRequest(BaseModel): user: str | None = Field(default=None, description="User identifier for tracking") layers: int | None = Field( default=None, - description="Number of output layers for layered image models. Supported range: 3-10.", + description="Number of output layers for layered image models. Supported range: 2-10.", ) @field_validator("prompt") @@ -80,9 +88,9 @@ def validate_size(cls, v): @field_validator("response_format") @classmethod def validate_response_format(cls, v): - """Validate response format - only b64_json is supported.""" - if v is not None and v != ResponseFormat.B64_JSON: - raise ValueError(f"Only 'b64_json' response format is supported, got: {v}") + """Validate response format - only b64_json and file are supported.""" + if v is not None and v not in (ResponseFormat.B64_JSON, ResponseFormat.FILE): + raise ValueError(f"Only 'b64_json' or 'file' response format is supported, got: {v}") return v @field_validator("layers") @@ -205,6 +213,40 @@ class ImageGenerationResponse(BaseModel): "Only present for image editing (IT2I) with CoT-enabled models.", ) + def stream_response(self) -> StreamingResponse: + if not self.data or not self.data[0].b64_json: + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + detail="No image data available for file response.", + ) + if len(self.data) == 1: + image_bytes = base64.b64decode(self.data[0].b64_json) + filename = f"image_{uuid.uuid4().hex[:8]}.png" + return StreamingResponse( + io.BytesIO(image_bytes), + media_type="image/png", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Content-Length": str(len(image_bytes)), + }, + ) + else: + zip_buffer = io.BytesIO() + with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf: + for idx, item in enumerate(self.data): + if item.b64_json: + zf.writestr(f"image_{idx}.png", base64.b64decode(item.b64_json)) + zip_bytes = zip_buffer.getvalue() + filename = f"images_{uuid.uuid4().hex[:8]}.zip" + return StreamingResponse( + io.BytesIO(zip_bytes), + media_type="application/zip", + headers={ + "Content-Disposition": f'attachment; filename="{filename}"', + "Content-Length": str(len(zip_bytes)), + }, + ) + class ImageEditARDeltaChunk(BaseModel): """Streaming chunk carrying a text delta from the image-edit AR stage.""" diff --git a/vllm_omni/entrypoints/openai/serving_audio_generate.py b/vllm_omni/entrypoints/openai/serving_audio_generate.py index a27f87c6d86..f8fedd33f95 100644 --- a/vllm_omni/entrypoints/openai/serving_audio_generate.py +++ b/vllm_omni/entrypoints/openai/serving_audio_generate.py @@ -3,7 +3,7 @@ import torch from fastapi import Request from fastapi.responses import Response -from vllm.entrypoints.openai.engine.serving import OpenAIServing +from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing from vllm.logger import init_logger from vllm.utils import random_uuid @@ -152,7 +152,6 @@ async def create_audio_generate( sample_rate=int(sample_rate), response_format=request.response_format or "wav", speed=request.speed or 1.0, - stream_format=request.stream_format, base64_encode=False, ) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 8232265014e..e6751fb967a 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -23,7 +23,7 @@ make_tool_call_id, ) -from vllm_omni.diffusion.diffusion_engine import get_extra_body_params, get_extra_output_params +from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse from vllm_omni.entrypoints.utils import coerce_param_message_types @@ -33,6 +33,7 @@ observe_audio_first_packet, observe_audio_streaming_finalize, ) +from vllm_omni.model_extras import get_extra_body_params, get_extra_output_params try: import soundfile @@ -40,6 +41,7 @@ soundfile = None +from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs from vllm.entrypoints.launcher import terminate_if_errored from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -63,16 +65,16 @@ ToolCall, UsageInfo, ) -from vllm.entrypoints.openai.engine.serving import ChatLikeRequest, clamp_prompt_logprobs from vllm.entrypoints.openai.parser.harmony_utils import ( get_streamable_parser_for_assistant, - parse_chat_output, ) from vllm.entrypoints.openai.responses.protocol import ResponsesRequest +from vllm.entrypoints.serve.engine.typing import ChatLikeRequest from vllm.entrypoints.serve.utils.api_utils import should_include_usage from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls from vllm.inputs import PromptType from vllm.logger import init_logger +from vllm.multimodal.media.connector import MediaConnector from vllm.outputs import RequestOutput from vllm.reasoning import ReasoningParser from vllm.renderers import BaseRenderer, merge_kwargs @@ -88,6 +90,7 @@ ) from vllm.tool_parsers import ToolParser from vllm.tool_parsers.mistral_tool_parser import MistralToolCall +from vllm.tool_parsers.streaming import extract_required_tool_call_streaming from vllm.utils.collection_utils import as_list from vllm.v1.engine.exceptions import EngineDeadError @@ -108,6 +111,7 @@ ) from vllm_omni.entrypoints.openai.stage_params import ( build_stage_sampling_params_list, + clone_sampling_params, get_default_sampling_params_list, ) from vllm_omni.entrypoints.openai.utils import ( @@ -121,11 +125,16 @@ from vllm_omni.errors import OmniClientError from vllm_omni.lora.request import LoRARequest from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.output_metadata import DiffusionMetadataMapping, DiffusionMetadataValue from vllm_omni.utils.audio import audio_chunk_pcm_bytes, audio_chunk_sample_rate logger = init_logger(__name__) +async def _identity_async(value: Any) -> Any: + return value + + class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): """OpenAI-compatible chat serving for both LLM and Diffusion models. @@ -144,6 +153,40 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): _diffusion_extra_body_params: frozenset[str] | None = None _diffusion_extra_output_params: frozenset[str] | None = None + # Harmony flag (always False for vllm-omni models) + use_harmony: bool = False + + @property + def tool_call_id_type(self) -> str: + """Return the tool call ID type, delegating to model config. + + Upstream vLLM removed the stored ``tool_call_id_type`` attribute + from ``OpenAIServingChat`` after the ParserManager refactor; the + field is now resolved on demand via ``get_tool_call_id_type``. + """ + try: + from vllm.entrypoints.chat_utils import get_tool_call_id_type + + return get_tool_call_id_type(self.model_config) + except Exception: + return "random" + + def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool: + """Check if streamed tokens should go through the tool-call parser. + + We only want to do this IF user-provided tools are set, a tool parser + is configured, "auto" tool choice is enabled, and the request's tool + choice field indicates that "auto" tool choice should be used. + + This method existed in upstream vLLM commit 91df0fad4 (OpenAIServingChat) + but was removed in the Harmony refactoring (PR #45171, #45104). Omni's + independently-maintained ``chat_completion_stream_generator`` still calls + it, so we keep a local copy. + """ + # parser_cls may be None (no tool parser configured); guard accordingly + tp_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None + return request.tools and tp_cls is not None and self.enable_auto_tools and request.tool_choice in ["auto", None] + @classmethod def for_diffusion( cls, @@ -167,12 +210,16 @@ def for_diffusion( instance._diffusion_mode = True instance._diffusion_engine = diffusion_engine instance._diffusion_model_name = model_name + instance._diffusion_extra_body_params = None + instance._diffusion_extra_output_params = None instance.engine_client = None instance.has_kv_connector = False + # Extra body/output params are resolved lazily on first use; see + # _get_diffusion_extra_body_params / _get_diffusion_extra_output_params. return instance def _get_diffusion_extra_body_params(self) -> frozenset[str]: - """Return the pipeline-declared ``EXTRA_BODY_PARAMS`` (cached).""" + """Return model-specific extra_body params from the extra registry.""" if self._diffusion_extra_body_params is not None: return self._diffusion_extra_body_params @@ -182,18 +229,22 @@ def _get_diffusion_extra_body_params(self) -> frozenset[str]: if od_config is not None and getattr(od_config, "model_class_name", None): params = get_extra_body_params(od_config.model_class_name) except Exception as e: - logger.warning("Failed to read EXTRA_BODY_PARAMS from pipeline: %s", e) + logger.warning("Failed to read model extra_body params: %s", e) self._diffusion_extra_body_params = params return params def _get_diffusion_extra_output_params( self, - custom_output: dict[str, Any] | None, - ) -> dict[str, Any] | None: - """Pick pipeline-declared ``EXTRA_OUTPUT_PARAMS`` keys from *custom_output*.""" - if not custom_output: - return None + output: object, + ) -> dict[str, DiffusionMetadataValue] | None: + """Pick model-specific extra output keys from diffusion metadata.""" + metadata: DiffusionMetadataMapping = {} + mm_output = getattr(output, "multimodal_output", None) + if isinstance(mm_output, dict): + raw_metadata = mm_output.get("metadata") + if isinstance(raw_metadata, dict): + metadata = raw_metadata if self._diffusion_extra_output_params is None: params: frozenset[str] = frozenset() @@ -202,14 +253,25 @@ def _get_diffusion_extra_output_params( if od_config is not None and getattr(od_config, "model_class_name", None): params = get_extra_output_params(od_config.model_class_name) except Exception as e: - logger.warning("Failed to read EXTRA_OUTPUT_PARAMS from pipeline: %s", e) + logger.warning("Failed to read model extra output params: %s", e) self._diffusion_extra_output_params = params if not self._diffusion_extra_output_params: return None - out = {k: custom_output[k] for k in self._diffusion_extra_output_params if k in custom_output} + flat_metadata: dict[str, DiffusionMetadataValue] = {} + for section in metadata.values(): + if isinstance(section, dict): + flat_metadata.update(section) + out = {k: flat_metadata[k] for k in self._diffusion_extra_output_params if k in flat_metadata} return out or None + @staticmethod + def _get_diffusion_text_output(output: object) -> str: + mm_output = getattr(output, "multimodal_output", None) + if isinstance(mm_output, dict) and mm_output.get("text") is not None: + return str(mm_output["text"]) + return "" + def _get_supported_speakers(self) -> set[str]: """Load supported speakers from model config (cached).""" if self._supported_speakers is not None: @@ -297,17 +359,14 @@ async def _create_chat_completion( tokenizer = await self.engine_client.get_tokenizer() reasoning_parser: ReasoningParser | None = None - if self.reasoning_parser_cls: - chat_template_kwargs = self._prepare_extra_chat_template_kwargs( - request.chat_template_kwargs, - self.default_chat_template_kwargs, - ) - reasoning_parser = self.reasoning_parser_cls( + if self.parser_cls is not None and self.parser_cls.reasoning_parser_cls is not None: + chat_template_kwargs = self._effective_chat_template_kwargs(request) + reasoning_parser = self.parser_cls.reasoning_parser_cls( tokenizer, chat_template_kwargs=chat_template_kwargs, # type: ignore[call-arg] ) - tool_parser = self.tool_parser + tool_parser = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None if isinstance(tokenizer, MistralTokenizer): # because of issues with pydantic we need to potentially @@ -345,7 +404,7 @@ async def _create_chat_completion( tool_dicts = [tool.model_dump() for tool in request.tools] if not self.use_harmony: - error_check_ret = self._validate_chat_template( + error_check_ret = self.online_renderer.validate_chat_template( request_chat_template=request.chat_template, chat_template_kwargs=request.chat_template_kwargs, trust_request_chat_template=self.trust_request_chat_template, @@ -353,14 +412,9 @@ async def _create_chat_completion( if error_check_ret is not None: return error_check_ret - chat_template_kwargs = request.chat_template_kwargs or {} - chat_template_kwargs.update(reasoning_effort=request.reasoning_effort) - - # Merge chat_template_kwargs with defaults - merged_template_kwargs = self._prepare_extra_chat_template_kwargs( - chat_template_kwargs, - self.default_chat_template_kwargs, - ) + # Effective kwargs fold request.chat_template_kwargs, reasoning_effort, + # and server defaults — mirrors OpenAIServingChat._effective_chat_template_kwargs. + merged_template_kwargs = self._effective_chat_template_kwargs(request) conversation, engine_prompts = await self._preprocess_chat( request, request.messages, @@ -378,7 +432,7 @@ async def _create_chat_completion( ) else: should_include_tools = tool_dicts is not None - conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony( + conversation, engine_prompts = self.online_renderer._make_request_with_harmony( request, should_include_tools ) @@ -413,9 +467,13 @@ async def _create_chat_completion( f"Supported modalities: {', '.join(sorted(allowed_modalities))}", ) + if request.modalities and "audio" in request.modalities: + audio_format_check = self._resolve_audio_format(request) + if isinstance(audio_format_check, ErrorResponse): + return audio_format_check + num_inference_steps = None - cfg_text_scale = None - cfg_img_scale = None + extra_body: dict[str, Any] = {} # Omni multistage image generation: Stage-0 (AR) should receive a clean # text prompt (and optional conditioning image/size) so the model's own # processor can construct the correct inputs. @@ -433,9 +491,7 @@ async def _create_chat_completion( # `extra_body` is flattented and merged into the payload's root. # These extra fields are accessible via `model_extra` property (from Pydantic base class). # When sending raw request with curl, no flattening happens. Directly read the `extra_body` dict. - extra_body = getattr(request, "extra_body", None) - if not extra_body: - extra_body = request.model_extra or {} + extra_body = getattr(request, "extra_body", None) or request.model_extra or {} height, width = self._resolve_height_width_from_extra_body(extra_body) @@ -447,9 +503,6 @@ async def _create_chat_completion( num_inference_steps = None negative_prompt = extra_body.get("negative_prompt") - cfg_text_scale = extra_body.get("cfg_text_scale") - cfg_img_scale = extra_body.get("cfg_img_scale") - engine_prompt_image: dict[str, Any] | None = None if reference_images: # Best-effort decode first reference image for i2i. @@ -498,21 +551,6 @@ async def _create_chat_completion( for k, v in engine_prompt_image.items() } - # Forward image_gen extra_args from diffusion-stage sampling - # params into the prompt dict so that thinker2imagegen can - # access them (e.g. byte5_text for ByT5 glyph rendering). - if hasattr(request, "sampling_params_list") and request.sampling_params_list: - for sp_raw in request.sampling_params_list: - ea = ( - (sp_raw.get("extra_args") or {}) - if isinstance(sp_raw, dict) - else (getattr(sp_raw, "extra_args", None) or {}) - ) - ig = ea.get("image_gen") or {} - if ig: - tprompt["image_gen_extra_args"] = ig - break - engine_prompts = [tprompt] # Store height/width for applying to diffusion stage sampling params later _image_gen_height = height @@ -576,15 +614,7 @@ async def _create_chat_completion( sp.width = _image_gen_width if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: sp.num_inference_steps = num_inference_steps - if hasattr(sp, "extra_args") and sp.extra_args is not None: - if cfg_text_scale is not None: - sp.extra_args["cfg_text_scale"] = cfg_text_scale - if cfg_img_scale is not None: - sp.extra_args["cfg_img_scale"] = cfg_img_scale - for _key in self._get_diffusion_extra_body_params(): - _val = extra_body.get(_key) - if _val is not None: - sp.extra_args[_key] = _val + apply_declared_extra_args(sp, self._get_diffusion_extra_body_params(), extra_body) self._log_inputs( request_id, @@ -680,6 +710,13 @@ async def _preprocess_chat( default_mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) + deferred_multi_modal_data: dict[str, Any] | None = None + if self._needs_multistage_multimodal_split(): + messages, deferred_multi_modal_data = await self._prepare_multistage_multimodal_inputs( + messages, + request, + ) + (conversation,), (engine_prompt,) = await renderer.render_chat_async( [messages], chat_params, @@ -724,18 +761,25 @@ async def _preprocess_chat( if hasattr(request, "cache_salt") and request.cache_salt is not None: engine_prompt["cache_salt"] = request.cache_salt + additional_information = getattr(request, "additional_information", None) + if isinstance(additional_information, dict): + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information.update(additional_information) + + if deferred_multi_modal_data: + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["deferred_multi_modal_data"] = deferred_multi_modal_data + speaker = getattr(request, "voice", None) or getattr(request, "speaker", None) normalized = validate_requested_speaker(speaker, self._get_supported_speakers()) if normalized is not None: - if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: - engine_prompt["additional_information"] = {} - engine_prompt["additional_information"]["speaker"] = [normalized] + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["speaker"] = [normalized] language = getattr(request, "language", None) if language is not None and isinstance(language, str) and language.strip(): - if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: - engine_prompt["additional_information"] = {} - engine_prompt["additional_information"]["language"] = [language.strip()] + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["language"] = [language.strip()] # Style instruction — used by Ming-flash-omni instruct TTS path # (ming_task="instruct"). For the omni speech path the thinker2talker @@ -743,20 +787,263 @@ async def _preprocess_chat( # which hardcodes instruction=None. instructions = getattr(request, "instructions", None) if instructions is not None and isinstance(instructions, str) and instructions.strip(): - if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: - engine_prompt["additional_information"] = {} - engine_prompt["additional_information"]["instruction"] = instructions.strip() + prompt_additional_information = self._ensure_prompt_additional_information(engine_prompt) + prompt_additional_information["instruction"] = instructions.strip() return conversation, [engine_prompt] + @staticmethod + def _ensure_prompt_additional_information(engine_prompt: dict[str, Any]) -> dict[str, Any]: + additional_information = engine_prompt.get("additional_information") + if not isinstance(additional_information, dict): + additional_information = {} + engine_prompt["additional_information"] = additional_information + return additional_information + + def _needs_multistage_multimodal_split(self) -> bool: + return bool(self._deferred_multimodal_modalities()) + + def _deferred_multimodal_modalities(self) -> set[str]: + stage_configs = list(getattr(self.engine_client, "stage_configs", []) or []) + if len(stage_configs) < 2: + return set() + + first_stage_modalities = self._stage_input_modalities(stage_configs[0]) + if not first_stage_modalities: + return set() + + downstream_modalities: set[str] = set() + for stage in stage_configs[1:]: + downstream_modalities.update(self._stage_input_modalities(stage)) + return downstream_modalities - first_stage_modalities + + @staticmethod + def _stage_input_modalities(stage: Any) -> set[str]: + engine_args = getattr(stage, "engine_args", None) + explicit = ( + getattr(stage, "input_modalities", None) + or getattr(stage, "modalities", None) + or getattr(engine_args, "input_modalities", None) + or getattr(engine_args, "modalities", None) + ) + if explicit: + return {str(modality) for modality in as_list(explicit)} + + model_stage = str(getattr(engine_args, "model_stage", None) or getattr(stage, "model_stage", "")).lower() + if model_stage in {"asr", "stt"} or model_stage.endswith("_asr"): + return {"audio"} + if any(name in model_stage for name in ("vision", "vl", "aura")): + return {"image", "video"} + if any(name in model_stage for name in ("tts", "talker", "code2wav")): + return set() + + if getattr(stage, "requires_multimodal_data", False): + return {"audio", "image", "video"} + return set() + + async def _prepare_multistage_multimodal_inputs( + self, + messages: list[ChatCompletionMessageParam], + request: ChatLikeRequest | ResponsesRequest, + ) -> tuple[list[ChatCompletionMessageParam], dict[str, Any] | None]: + """Hide modalities unsupported by stage 0 and stash them for downstream stages.""" + deferred_modalities = self._deferred_multimodal_modalities() + deferred_parts: dict[str, list[Any]] = {modality: [] for modality in deferred_modalities} + stripped_messages: list[ChatCompletionMessageParam] = [] + + for message in messages: + content = message.get("content") if isinstance(message, dict) else None + if not isinstance(content, list): + stripped_messages.append(message) + continue + + stripped_content: list[Any] = [] + changed = False + for part in content: + modality, payload = self._deferred_multimodal_part(part, deferred_modalities) + if modality is not None: + if payload is not None: + deferred_parts.setdefault(modality, []).append(payload) + changed = True + continue + stripped_content.append(part) + + if changed: + stripped_message = dict(message) + stripped_message["content"] = stripped_content + stripped_messages.append(cast(ChatCompletionMessageParam, stripped_message)) + else: + stripped_messages.append(message) + + deferred_parts = {modality: parts for modality, parts in deferred_parts.items() if parts} + if not deferred_parts: + return messages, None + + media_connector = MediaConnector( + media_io_kwargs=getattr(request, "media_io_kwargs", None), + allowed_local_media_path=getattr(self.model_config, "allowed_local_media_path", "") or "", + allowed_media_domains=getattr(self.model_config, "allowed_media_domains", None), + ) + multi_modal_data: dict[str, Any] = {} + for modality, parts in deferred_parts.items(): + multi_modal_data[modality] = await self._materialize_deferred_multimodal_parts( + media_connector, + modality, + parts, + ) + return stripped_messages, multi_modal_data + + @staticmethod + def _deferred_multimodal_part(part: Any, deferred_modalities: set[str]) -> tuple[str | None, Any | None]: + if not isinstance(part, dict): + return None, None + part_type = part.get("type") + if part_type in {"video_url", "video"} and "video" in deferred_modalities: + video = part.get("video_url", part.get("video")) + if isinstance(video, dict): + video = video.get("url") + return "video", video + if part_type in {"image_url", "image_pil", "image"} and "image" in deferred_modalities: + image = part.get("image_pil", part.get("image_url", part.get("image"))) + if isinstance(image, dict): + image = image.get("url") + return "image", image + if part_type in {"audio_url", "input_audio", "audio"} and "audio" in deferred_modalities: + audio = part.get("audio_url", part.get("input_audio", part.get("audio"))) + if isinstance(audio, dict): + audio = audio.get("url") or audio.get("data") + return "audio", audio + return None, None + + @staticmethod + async def _materialize_deferred_multimodal_parts( + media_connector: MediaConnector, + modality: str, + parts: list[Any], + ) -> list[Any]: + if modality == "video": + return list(await asyncio.gather(*(media_connector.fetch_video_async(part) for part in parts))) + if modality == "image": + fetch_image = getattr(media_connector, "fetch_image_async", None) + if fetch_image is None: + return parts + return list( + await asyncio.gather( + *(fetch_image(part) if isinstance(part, str) else _identity_async(part) for part in parts) + ) + ) + if modality == "audio": + fetch_audio = getattr(media_connector, "fetch_audio_async", None) + if fetch_audio is None: + return parts + return list( + await asyncio.gather( + *(fetch_audio(part) if isinstance(part, str) else _identity_async(part) for part in parts) + ) + ) + return parts + + 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("<I", 36 + data_size)) + buf.write(b"WAVE") + buf.write(b"fmt ") + buf.write( + struct.pack( + "<IHHIIHH", + 16, + 3, + num_channels, + sr, + sr * num_channels * (bits_per_sample // 8), + num_channels * (bits_per_sample // 8), + bits_per_sample, + ) + ) + buf.write(b"data") + buf.write(struct.pack("<I", data_size)) + buf.write(audio_np.tobytes()) + + audio_b64 = base64.b64encode(buf.getvalue()).decode() + audio_items.append( + { + "type": "audio_url", + "audio_url": {"url": f"data:audio/wav;base64,{audio_b64}"}, + } + ) + + new_content = list(content) + audio_items + if isinstance(msg, dict): + new_msg = {**msg, "content": new_content} + else: + new_msg = msg.model_copy(update={"content": new_content}) + new_messages.append(new_msg) + + return new_messages + def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any]: """Convert request dicts to stage-typed sampling params objects. For diffusion stages, build ``OmniDiffusionSamplingParams`` so downstream ``StageDiffusionClient._sampling_params_to_dict`` (which requires a dataclass) works. For LLM stages build ``SamplingParams``. + If callers provide params for fewer stages than the native pipeline has + (for example AURA has three semantic models but four engine stages), + append cloned deploy defaults for the omitted tail stages. """ stage_configs = list(getattr(self.engine_client, "stage_configs", []) or []) + default_params_list = list(getattr(self.engine_client, "default_sampling_params_list", []) or []) final_sampling_params_list: list[Any] = [] for idx, sampling_params in enumerate(sampling_params_list): stage_type = get_stage_type(stage_configs[idx]) if idx < len(stage_configs) else "llm" @@ -776,6 +1063,11 @@ def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[Any final_sampling_params_list.append(target_cls(**as_dict)) else: raise ValueError(f"Invalid sampling params: {sampling_params}") + for idx in range(len(final_sampling_params_list), len(stage_configs)): + if idx < len(default_params_list): + final_sampling_params_list.append(clone_sampling_params(default_params_list[idx])) + else: + final_sampling_params_list.append(SamplingParams()) return final_sampling_params_list def _get_comprehension_stage_index(self) -> int: @@ -876,7 +1168,8 @@ def _should_check_for_unstreamed_tool_arg_tokens( return ( output.finish_reason is not None and self.enable_auto_tools - and self.tool_parser is not None + and self.parser_cls is not None + and self.parser_cls.tool_parser_cls is not None and delta_message is not None and delta_message.tool_calls ) @@ -999,8 +1292,11 @@ async def chat_completion_stream_generator( all_previous_token_ids = None # Prepare the tool parser if it's needed try: - if tool_choice_auto and self.tool_parser: - tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer, request.tools)] * num_choices + tool_parser_cls = self.parser_cls.tool_parser_cls if self.parser_cls is not None else None + if tool_choice_auto and tool_parser_cls is not None: + tool_parsers: list[ToolParser | None] = [ + tool_parser_cls(tokenizer, request.tools) for _ in range(num_choices) + ] else: tool_parsers = [None] * num_choices except Exception as e: @@ -1328,12 +1624,13 @@ async def chat_completion_stream_generator( # either finished reasoning or no reasoning at all content = current_text - delta_message, function_name_returned[i] = self.extract_tool_call_required_streaming( + delta_message, function_name_returned[i] = extract_required_tool_call_streaming( previous_text=previous_text, current_text=content, delta_text=delta_text, function_name_returned=fn_name_returned, tool_call_idx=history_tool_call_cnt, + tool_call_id_type=self.tool_call_id_type, ) if ( delta_message @@ -1911,8 +2208,8 @@ async def chat_completion_full_generator( else: # Diffusion pipeline text output (e.g. single-stage # img2text / text2text) — no AR request_output, so build - # a simple text choice from custom_output. - text_body = (omni_outputs.custom_output or {}).get("text_output", "") + # a simple text choice from diffusion multimodal output. + text_body = self._get_diffusion_text_output(omni_outputs) message = ChatMessage(role=role, content=text_body) text_choices.append( ChatCompletionResponseChoice( @@ -1938,7 +2235,7 @@ async def chat_completion_full_generator( response_metrics = {} response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {}) response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0)) - extra = self._get_diffusion_extra_output_params(omni_outputs.custom_output) + extra = self._get_diffusion_extra_output_params(omni_outputs) if extra: response_metrics.update(extra) @@ -2044,12 +2341,24 @@ def _create_text_choice( logprobs = None if self.use_harmony: - reasoning, content, _ = parse_chat_output(token_ids) + # Use upstream's HarmonyParser (inherited as self.parser_cls for + # gpt_oss/harmony models) instead of vendoring a copy of the + # pre-refactor parse_chat_output. We take only reasoning+content; + # tool calls are extracted by omni's own path below. + if self.parser_cls is not None: + parser = self.parser_cls(tokenizer, request.tools) + reasoning, content, _ = parser.parse( + "", + request, + model_output_token_ids=token_ids, + ) + else: + reasoning, content = None, None if not request.include_reasoning: reasoning = None - if self.tool_parser is not None: - tool_parser = self.tool_parser(tokenizer, request.tools) + if self.parser_cls is not None and self.parser_cls.tool_parser_cls is not None: + tool_parser = self.parser_cls.tool_parser_cls(tokenizer, request.tools) # NOTE: We use token_ids for openai tool parser tool_call_info = tool_parser.extract_tool_calls( "", @@ -2097,7 +2406,7 @@ def _create_text_choice( auto_tools_called = False # if auto tools are not enabled, and a named tool choice using # outlines is not being used - if (not self.enable_auto_tools or not self.tool_parser) and ( + if (not self.enable_auto_tools or self.parser_cls is None or self.parser_cls.tool_parser_cls is None) and ( not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) and request.tool_choice != "required" ): @@ -2163,10 +2472,11 @@ def _create_text_choice( request.tools and (request.tool_choice == "auto" or request.tool_choice is None) and self.enable_auto_tools - and self.tool_parser + and self.parser_cls is not None + and self.parser_cls.tool_parser_cls is not None ): try: - tool_parser = self.tool_parser(tokenizer, request.tools) + tool_parser = self.parser_cls.tool_parser_cls(tokenizer, request.tools) except RuntimeError as e: logger.exception("Error in tool parser creation.") return self.create_error_response(e) @@ -2297,7 +2607,6 @@ def _create_audio_choice( sample_rate=sample_rate, response_format=audio_format, speed=1.0, - stream_format="audio", base64_encode=True, ) @@ -2632,6 +2941,11 @@ def _build_multistage_generation_inputs( layers=layers, resolution=resolution, ) + apply_declared_extra_args( + default_stage_params, + self._get_diffusion_extra_body_params(), + extra_body, + ) if lora_body and isinstance(lora_body, dict): try: lora_req, lora_scale = parse_lora_request(lora_body) @@ -3026,8 +3340,6 @@ async def _create_diffusion_chat_completion( num_inference_steps = extra_body.get("num_inference_steps") guidance_scale = extra_body.get("guidance_scale") true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale") - cfg_text_scale = extra_body.get("cfg_text_scale") - cfg_img_scale = extra_body.get("cfg_img_scale") seed = extra_body.get("seed") if seed is None: seed = getattr(request, "seed", None) @@ -3085,14 +3397,7 @@ async def _create_diffusion_chat_completion( gen_params.guidance_scale = guidance_scale if true_cfg_scale is not None: gen_params.true_cfg_scale = true_cfg_scale - if cfg_text_scale is not None: - gen_params.extra_args["cfg_text_scale"] = cfg_text_scale - if cfg_img_scale is not None: - gen_params.extra_args["cfg_img_scale"] = cfg_img_scale - for _key in self._get_diffusion_extra_body_params(): - _val = extra_body.get(_key) - if _val is not None: - gen_params.extra_args[_key] = _val + apply_declared_extra_args(gen_params, self._get_diffusion_extra_body_params(), extra_body) if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: @@ -3180,7 +3485,7 @@ async def _create_diffusion_chat_completion( # Text output path (img2text / text2text) if is_text_request and result.final_output_type == "text": - text_body = (result.custom_output or {}).get("text_output", "") + text_body = self._get_diffusion_text_output(result) message = ChatMessage(role="assistant", content=text_body) choice = ChatCompletionResponseChoice( index=0, @@ -3199,7 +3504,7 @@ async def _create_diffusion_chat_completion( completion_tokens=len(text_body.split()), total_tokens=len(prompt.split()) + len(text_body.split()), ), - metrics=self._get_diffusion_extra_output_params(result.custom_output), + metrics=self._get_diffusion_extra_output_params(result), ) logger.info( "Diffusion chat completed for request %s: text output (%d chars)", @@ -3260,7 +3565,6 @@ async def _create_diffusion_chat_completion( sample_rate=sample_rate, response_format=audio_format, speed=1.0, - stream_format="audio", base64_encode=True, ) audio_response: AudioResponse = self.create_audio(audio_obj) @@ -3337,7 +3641,7 @@ async def _create_diffusion_chat_completion( completion_tokens=1, total_tokens=len(prompt.split()) + 1, ), - metrics=self._get_diffusion_extra_output_params(result.custom_output), + metrics=self._get_diffusion_extra_output_params(result), ) logger.info( diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 198d77814d0..3cbd51c499a 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -21,12 +21,12 @@ from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse from transformers.utils.hub import cached_file +from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing from vllm.entrypoints.launcher import terminate_if_errored from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, RequestResponseMetadata, ) -from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.multimodal.media import MediaConnector @@ -43,6 +43,13 @@ OpenAICreateSpeechRequest, SpeechBatchItem, SpeechBatchItemResult, + SpeechInputTokenDetails, + SpeechTokenUsage, +) +from vllm_omni.entrypoints.openai.speech_usage import ( + SpeechOutputTokenCounter, + build_speech_usage, + qwen3_tts_input_token_details, ) from vllm_omni.entrypoints.openai.tts_adapters import ( SpeechServingContext, @@ -75,16 +82,23 @@ _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"} +# CosyVoice3 talker expects its reference transcript wrapped in the model +# instruction template; without the delimiter the talker re-speaks the +# reference (issue #4644). Matches the offline example/test and upstream demo. +_COSYVOICE3_PROMPT_DELIMITER = "<|endofprompt|>" +_COSYVOICE3_PROMPT_PREFIX = f"You are a helpful assistant.{_COSYVOICE3_PROMPT_DELIMITER}" _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} _COVO_AUDIO_MODEL_STAGES = {"fused_thinker_talker"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} +_MOSS_TTS_LOCAL_MODEL_STAGES = {"moss_tts_local", "moss_tts_local_codec"} _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} _HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} _GLM_TTS_MODEL_STAGES = {"glm_tts"} _STEP_AUDIO2_TTS_MODEL_STAGES = {"step_audio2_thinker"} +_INDEXTTS2_TTS_MODEL_STAGES = {"indextts2_talker"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES @@ -98,8 +112,10 @@ | _MING_TTS_MODEL_STAGES | _MOSS_TTS_MODEL_STAGES | _MOSS_TTS_FULL_MODEL_STAGES + | _MOSS_TTS_LOCAL_MODEL_STAGES | _GLM_TTS_MODEL_STAGES | _STEP_AUDIO2_TTS_MODEL_STAGES + | _INDEXTTS2_TTS_MODEL_STAGES ) _SAMPLING_MAX_TOKENS_TTS_MODEL_TYPES = { "fish_tts", @@ -109,6 +125,7 @@ "voxcpm2", "higgs_audio_v2", "higgs_audio_v3", + "indextts2", } _TTS_LANGUAGES = frozenset( { @@ -133,7 +150,6 @@ _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 _VOICE_UPLOAD_MAX_CONSENT_LEN = 1024 _VOICE_UPLOAD_MAX_REF_TEXT_CHARS = 8192 @@ -182,6 +198,18 @@ def _create_wav_header(sample_rate: int, num_channels: int = 1, bits_per_sample: return header +def _infer_audio_num_channels(audio: np.ndarray) -> int: + """Infer channel count before streaming PCM bytes are wrapped as WAV.""" + if audio.ndim == 3 and audio.shape[0] == 1: + audio = audio[0] + if audio.ndim == 2: + if audio.shape[0] in (1, 2): + return int(audio.shape[0]) + if audio.shape[1] in (1, 2): + return int(audio.shape[1]) + return 1 + + def _sanitize_filename(filename: str) -> str: """Sanitize filename to prevent path traversal attacks. @@ -297,8 +325,11 @@ def _init_speaker_storage(self) -> None: self._ref_audio_resolve_cache_bytes = 0 self._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES 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] = {} + # Readiness is keyed by (artifact_key, x_vector_only). An x-vector-only + # request caches a speaker embedding but no ref_code, so its artifact + # must not satisfy a later ICL request that needs ref_code (#5049). + self._ref_audio_model_artifact_ready: set[tuple[str, bool]] = set() + self._request_ref_audio_artifact_keys: dict[str, tuple[str, bool]] = {} 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]] = {} @@ -414,6 +445,7 @@ def for_diffusion( def __init__(self, *args, **kwargs): self.model_name = kwargs.pop("model_name", None) + self.forced_aligner_config: Any | None = kwargs.pop("forced_aligner_config", None) super().__init__(*args, **kwargs) self._init_speaker_storage() @@ -675,7 +707,6 @@ def _detect_tts_model_type(self) -> str | None: if model_stage in _OMNIVOICE_TTS_MODEL_STAGES: return "omnivoice" if model_stage in _COVO_AUDIO_MODEL_STAGES: - model_arch = getattr(self._tts_stage.engine_args, "model_arch", None) if model_arch and "CovoAudio" in model_arch: return "covo_audio" if model_stage in _VOXCPM2_TTS_MODEL_STAGES: @@ -688,6 +719,8 @@ def _detect_tts_model_type(self) -> str | None: return "moss_tts_nano" if model_stage in _MOSS_TTS_FULL_MODEL_STAGES: return "moss_tts" + if model_stage in _MOSS_TTS_LOCAL_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: @@ -696,6 +729,8 @@ def _detect_tts_model_type(self) -> str | None: return "glm_tts" if model_stage in _STEP_AUDIO2_TTS_MODEL_STAGES: return "step_audio2" + if model_stage in _INDEXTTS2_TTS_MODEL_STAGES: + return "indextts2" return None def _get_custom_voice_dir(self) -> str | None: @@ -877,6 +912,69 @@ def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e) return 2048 + def _get_usage_text_tokenizer(self): + """Return a text tokenizer for counting input-text usage tokens. + + Prefer the per-model tokenizer already loaded for prompt-length + estimation (`_tts_tokenizer`, which is the *correct* tokenizer for the + active model). Fall back to a lazily-loaded, cached generic tokenizer + for models that never populate `_tts_tokenizer`. Returns None if no + tokenizer can be obtained (usage then reports text_tokens=0). + """ + if self._tts_tokenizer is not None: + return self._tts_tokenizer + if getattr(self, "_usage_text_tokenizer", None) is None: + try: + from transformers import AutoTokenizer + + self._usage_text_tokenizer = AutoTokenizer.from_pretrained( + self.engine_client.model_config.model, trust_remote_code=True + ) + except Exception as e: # pragma: no cover - environment dependent + logger.warning("Usage: could not load a text tokenizer (%s); text_tokens will be 0", e) + self._usage_text_tokenizer = None + return self._usage_text_tokenizer + + def _count_usage_text_tokens(self, text: str) -> int: + """Token count of `text` using the model's text tokenizer (0 on failure).""" + if not text: + return 0 + tok = self._get_usage_text_tokenizer() + if tok is None: + return 0 + try: + return len(tok(text, padding=False)["input_ids"]) + except Exception: + return 0 + + def _compute_speech_input_details( + self, request: OpenAICreateSpeechRequest, tts_params: dict[str, Any] + ) -> SpeechInputTokenDetails: + """Input-token breakdown (text + reference-audio) for a speech request. + + Counts `input` (+ `instructions`) as text tokens, and reference-audio + codec frames as audio tokens *only* when in-context voice cloning is + active (see `qwen3_tts_input_token_details` / `gate_audio_tokens`). The + audio gating reads Qwen3-TTS `tts_params` conventions; other TTS models + do not set those keys, so they degrade cleanly to text-only counts. + """ + return qwen3_tts_input_token_details( + input_text=request.input, + instructions=request.instructions, + tts_params=tts_params or {}, + count_text_tokens=self._count_usage_text_tokens, + ) + + def _build_speech_usage( + self, + request: OpenAICreateSpeechRequest, + tts_params: dict[str, Any], + output_tokens: int, + ) -> SpeechTokenUsage: + """Assemble the full usage object (input breakdown + generated tokens).""" + details = self._compute_speech_input_details(request, tts_params) + return build_speech_usage(details, output_tokens) + def _estimate_fish_ref_code_len(self, ref_audio: object) -> int | None: """Estimate Fish Speech semantic token length from raw reference audio.""" from vllm_omni.model_executor.models.fish_speech.dac_utils import ( @@ -1319,7 +1417,7 @@ async def upload_voice_embedding(self, embedding_json: str, consent: str, name: emb_dim = len(embedding) if self._tts_model_type == "ming_tts": if emb_dim != 192: - logger.warning("speaker_embedding has %d dimensions; Ming dense expects 192", emb_dim) + raise ValueError(f"Ming speaker embedding must have 192 dims, got {emb_dim}") else: dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_dim) if dim_err is not None: @@ -1414,48 +1512,22 @@ def _is_tts_model(self) -> bool: def _validate_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate TTS request parameters. Returns error message or None.""" - if self._tts_model_type == "step_audio2": - # StepAudio2 only requires non-empty input text - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - return None - if self._tts_model_type == "voxtral_tts": - return self._validate_voxtral_tts_request(request) - if self._tts_model_type == "fish_tts": - return self._validate_fish_tts_request(request) - if self._tts_model_type == "cosyvoice3": - return self._validate_cosyvoice3_request(request) - 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) - if self._tts_model_type == "glm_tts": - return self._validate_glm_tts_request(request) - return self._validate_qwen_tts_request(request) - - def _validate_voxcpm2_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate VoxCPM2 request parameters. Returns error message or None.""" - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - if request.voice is not None: - request.voice = request.voice.lower() - available_voices = set(self.uploaded_speakers) | set(self.precomputed_speakers) | {"default"} - if request.voice not in available_voices: - supported = ", ".join(sorted(available_voices)) or "none" - return f"Invalid voice '{request.voice}'. Supported: {supported}" + adapter = self._get_tts_adapter() + if adapter is not None: + return adapter.validate(request) - 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}" + adapter_cls = resolve_adapter("qwen3_tts") + if adapter_cls is None: + raise ValueError("Qwen3-TTS adapter is not registered") - return None + ctx = SpeechServingContext( + server=self, + engine_client=self.engine_client, + ) + return adapter_cls(ctx).validate(request) def _voxcpm2_encode(self, text: str) -> list[int]: """Tokenize text for VoxCPM2, splitting multichar Chinese tokens.""" @@ -1523,158 +1595,6 @@ def _validate_ref_audio_format(self, ref_audio: str) -> str | None: return "ref_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" return None - def _validate_voxtral_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate Voxtral TTS request parameters. Returns error message or None.""" - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - - # Voxtral TTS requires either a preset voice or ref_audio for voice cloning. - if request.voice is None and request.ref_audio is None: - return "Either 'voice' (preset speaker) or 'ref_audio' (voice cloning) must be provided" - - 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.voice is not None: - request.voice = request.voice.lower() - if self.supported_speakers and request.voice not in self.supported_speakers: - return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(self.supported_speakers))}" - - 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_qwen_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate Qwen TTS request parameters. Returns error message or None.""" - # Infer Base task when ref_audio or ref_text is provided without explicit task_type. - if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None): - request.task_type = "Base" - - # Normalize voice to lowercase for case-insensitive matching - if request.voice is not None: - request.voice = request.voice.lower() - if request.task_type is None and request.voice in self.precomputed_speakers: - request.task_type = "Base" - task_type = request.task_type or "CustomVoice" - - # Validate input is not empty - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - - # 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": - if not self.supported_speakers: - return ( - "This model does not support CustomVoice task (no speakers configured). " - "Use task_type='Base' with ref_audio/ref_text for voice cloning, " - "or use a CustomVoice model." - ) - if request.voice is not None and request.voice not in self.supported_speakers: - return f"Invalid voice '{request.voice}'. Supported: {', '.join(sorted(self.supported_speakers))}" - - # Validate speaker_embedding constraints - if request.speaker_embedding is not None: - if task_type != "Base": - return "'speaker_embedding' is only valid for Base task" - if not request.speaker_embedding: - return "'speaker_embedding' must be a non-empty list of floats" - # speaker_embedding implies x_vector_only_mode — set it before - # Base task validation so callers don't need to pass it explicitly. - request.x_vector_only_mode = True - emb_len = len(request.speaker_embedding) - dim_err = self._validate_qwen_tts_speaker_embedding_dim(emb_len) - if dim_err is not None: - return dim_err - # Validate Base task requirements - if task_type == "Base": - if request.voice is None: - # 1. Ensure a voice source is provided - if request.ref_audio is None and getattr(request, "speaker_embedding", None) is None: - return "Base task requires 'ref_audio' or 'speaker_embedding' for voice cloning" - # 2. Validate ref_audio format if it exists (using the helper from main) - if request.ref_audio is not None: - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - # 3. Validate text requirements based on the mode - if not getattr(request, "x_vector_only_mode", False): - if not request.ref_text or not request.ref_text.strip(): - return ( - "Base task requires non-empty 'ref_text' (transcript of " - "the reference audio) unless 'x_vector_only_mode' is enabled" - ) - else: - voice_lower = request.voice.lower() - if voice_lower in self.uploaded_speakers: - # Check if data file exists for uploaded speaker - speaker_info = self.uploaded_speakers[voice_lower] - file_path = Path(speaker_info["file_path"]) - if not file_path.exists(): - return f"Data file for uploaded speaker '{request.voice}' not found on disk" - elif voice_lower in self.precomputed_speakers: - profile = self.precomputed_speakers[voice_lower] - mode = str(profile.get("mode") or "xvec").lower() - ref_text = request.ref_text or profile.get("ref_text") - if mode == "icl" and (not isinstance(ref_text, str) or not ref_text.strip()): - return ( - f"Precomputed voice '{request.voice}' uses ICL mode but has no ref_text in " - "the request or custom voice manifest" - ) - else: - # need ref_audio for built-in speaker - if request.ref_audio is None: - return ( - f"Base task with built-in speaker '{request.voice}' requires 'ref_audio' for voice cloning" - ) - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - if not getattr(request, "x_vector_only_mode", False) and ( - not request.ref_text or not request.ref_text.strip() - ): - return ( - "Base task requires non-empty 'ref_text' (transcript of " - "the reference audio) unless 'x_vector_only_mode' is enabled" - ) - - # Validate cross-parameter dependencies - if task_type != "Base": - if request.ref_text is not None: - return "'ref_text' is only valid for Base task" - if request.x_vector_only_mode is not None: - return "'x_vector_only_mode' is only valid for Base task" - - # Validate VoiceDesign task requirements - if task_type == "VoiceDesign" and not request.instructions: - return "VoiceDesign task requires 'instructions' to describe the voice" - - # Validate instructions length (using cached value from initialization) - if request.instructions and len(request.instructions) > self._max_instructions_length: - return f"Instructions too long (max {self._max_instructions_length} characters)" - - # Validate max_new_tokens range - 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 _detect_moss_variant(self) -> str: """Sub-classify a ``moss_tts``-stage server into the actual MOSS-TTS variant family (tts, ttsd, sound_effect, voice_generator, realtime). @@ -1690,6 +1610,8 @@ def _detect_moss_variant(self) -> str: name = "" if "realtime" in name: return "realtime" + if "local" in name: + return "local" if "ttsd" in name: return "ttsd" if "soundeffect" in name: @@ -1698,58 +1620,6 @@ def _detect_moss_variant(self) -> str: return "voice_generator" return "tts" - def _validate_moss_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate any MOSS-TTS-family request (nano + 5 full variants). - - Dispatches by ``self._moss_variant``: - - ``tts``/``realtime``: require ``ref_audio`` (voice cloning). - - ``ttsd``: require ``ref_audio`` (speaker 1); ``ref_audio_2`` - optional (defaults to the same ref for both speakers). - - ``sound_effect``: require ``ambient_sound`` (no ref_audio). - - ``voice_generator``: require ``instructions`` (no ref_audio). - - For the legacy moss_tts_nano model_type the variant is None and - we fall through to the original nano contract (ref_audio only). - """ - if not request.input or not request.input.strip(): - # SoundEffect can legitimately have empty input (just ambient_sound). - if self._moss_variant != "sound_effect": - return "Input text cannot be empty" - - v = self._moss_variant - if v in (None, "tts", "realtime"): - if request.ref_audio is None: - label = "MOSS-TTS-Nano" if v is None else ("MOSS-TTS-Realtime" if v == "realtime" else "MOSS-TTS") - return f"{label} requires 'ref_audio' (reference audio for voice cloning)." - return self._validate_ref_audio_format(request.ref_audio) - - if v == "ttsd": - if request.ref_audio is None: - return "MOSS-TTSD requires 'ref_audio' (speaker 1 reference)." - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - if request.ref_audio_2 is not None: - return self._validate_ref_audio_format(request.ref_audio_2) - return None - - if v == "sound_effect": - if not request.ambient_sound or not request.ambient_sound.strip(): - return ( - "MOSS-SoundEffect requires 'ambient_sound' (natural language " - "description of the sound effect to synthesise)." - ) - return None - - if v == "voice_generator": - if not request.instructions or not request.instructions.strip(): - return ( - "MOSS-VoiceGenerator requires 'instructions' (natural language " - "voice description, e.g. 'a warm female voice with an American accent')." - ) - return None - - return None # unreachable - def _get_moss_processor(self): """Lazily load the upstream MOSS-TTS processor once per server. @@ -1821,11 +1691,26 @@ async def _build_moss_tts_params(self, request: OpenAICreateSpeechRequest) -> di params["prompt_audio_array"] = [[wav_list, sr]] return params - # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator): - # call the upstream processor server-side to produce unified codes ---- + # ---- MossTTSDelay family (tts/ttsd/sound_effect/voice_generator) + # and MOSS-TTS-Local-Transformer-v1.5: call the upstream processor + # server-side to produce unified codes. Local-v1.5 ships its own + # AutoProcessor (processor_config.json + processing_moss_tts.py) and + # reuses this exact build_user_message/encode_audios_from_wav path in + # the offline example (examples/.../moss_tts/end2end.py: + # _build_unified_codes) -- it is NOT in the same boat as Realtime + # (no processor_config.json there), so it must not fall back to the + # prompt_audio_array path above (which the talker's preprocess() + # never reads -- info_dict["codes"]["ref"] is the only thing it + # consumes, so skipping this path silently drops all voice-clone + # conditioning and produces unconditioned/garbage audio online). ---- proc = self._get_moss_processor() n_vq = int(getattr(proc.model_config, "n_vq", 32)) - sr_target = int(getattr(proc.model_config, "sampling_rate", 24000)) + # Local-v1.5 encodes reference audio at a fixed 24 kHz working rate + # regardless of its 48 kHz stereo *output* codec -- mirrors the + # offline example's hardcoded encode_audios_from_wav(sampling_rate=24000) + # for this variant; proc.model_config.sampling_rate there is the + # output rate (48000), the wrong value to resample the reference into. + sr_target = 24000 if v == "local" else int(getattr(proc.model_config, "sampling_rate", 24000)) # Reference-audio encoding + speaker caching lives in the model package # (moss_tts.reference_encoder), mirroring Fish Speech / CosyVoice3 / @@ -1852,7 +1737,7 @@ async def _encode_ref(ref_str: str) -> torch.Tensor: ) user_kwargs: dict[str, Any] = {"text": request.input or ""} - if v in ("tts", "realtime"): + if v in ("tts", "local"): user_kwargs["reference"] = [await _encode_ref(request.ref_audio)] elif v == "ttsd": refs = [await _encode_ref(request.ref_audio)] @@ -1893,93 +1778,6 @@ async def _encode_ref(ref_str: str) -> torch.Tensor: params["max_new_frames"] = [request.max_new_tokens] return params - def _validate_higgs_audio_v2_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate higgs_audio_v2 request parameters. Returns error message or None. - - Accepted: plain text -> speech, or shallow voice clone via ``ref_audio`` - + ``ref_text`` (both required together). Still out of scope: preset - ``voice``/``speaker`` selection, ``x_vector_only_mode`` / - ``speaker_embedding`` helpers, ``task_type``/``language``/ - ``instructions``/``speed`` overrides, and multi-speaker ``[SPEAKERn]`` - tags inside the input body. - """ - from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( - MULTI_SPEAKER_TAG_PATTERN, - ) - - if not request.input or not request.input.strip(): - return "higgs_audio_v2: input text cannot be empty" - - # Voice clone: ref_audio and ref_text must come together. - if request.ref_audio is not None and not request.ref_text: - return ( - "higgs_audio_v2 voice clone requires both 'ref_audio' and " - "'ref_text'; received ref_audio without ref_text" - ) - if request.ref_text and request.ref_audio is None: - return ( - "higgs_audio_v2 voice clone requires both 'ref_audio' and " - "'ref_text'; received ref_text without ref_audio" - ) - - if request.x_vector_only_mode is not None: - return "higgs_audio_v2 v1 does not support 'x_vector_only_mode' (voice-cloning helper field)" - if request.speaker_embedding is not None: - return "higgs_audio_v2 v1 does not support 'speaker_embedding' (voice-cloning helper field)" - if request.voice and request.ref_audio is None: - # _apply_uploaded_speaker runs before this validator; if voice was - # an uploaded speaker, ref_audio is now populated and ref_text is - # backfilled from the speaker entry. A bare voice= with no - # ref_audio means the name didn't resolve to an uploaded speaker - # (and higgs has no built-in preset voices). - return ( - "higgs_audio_v2 v1 does not support 'voice'/'speaker' selection for built-in voices; " - f"upload a voice first via POST /v1/audio/voices, or use ref_audio + ref_text. " - f"Got voice={request.voice!r}" - ) - if request.instructions: - return ( - "higgs_audio_v2 v1 does not support 'instructions' (voice " - "style/emotion control); supply plain text instead" - ) - if request.task_type is not None: - return "higgs_audio_v2 v1 does not support 'task_type'; the model is single-mode plain text -> speech" - if request.language is not None: - return ( - "higgs_audio_v2 v1 does not accept 'language' overrides; the model infers language from the input text" - ) - if request.speed is not None and request.speed != 1.0: - return ( - "higgs_audio_v2 v1 does not support 'speed' adjustments; the audio is rendered at native rate (24 kHz)" - ) - - if MULTI_SPEAKER_TAG_PATTERN.search(request.input): - return "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; remove the tag from the input" - - 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}" - - def _validate_glm_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate GLM-TTS request — requires ref_audio for voice cloning.""" - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - - if request.ref_audio is None: - return "GLM-TTS requires 'ref_audio' for zero-shot voice cloning" - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - if not request.ref_text or not request.ref_text.strip(): - return "GLM-TTS voice cloning requires 'ref_text' (transcript of the reference audio)" - - 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 >= {_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 _build_higgs_audio_v2_params(self, request: OpenAICreateSpeechRequest): """Build prompt_token_ids for higgs_audio_v2 via the upstream processor. @@ -2056,18 +1854,6 @@ async def _resolve_higgs_audio_v2_processor(self): # ---- 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. @@ -2201,143 +1987,6 @@ async def _resolve_higgs_audio_v3_adapter(self): 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(): - return "Input text cannot be empty" - - if request.ref_audio is not None: - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - if not request.ref_text or not request.ref_text.strip(): - return "Voice cloning requires 'ref_text' (transcript of the reference audio)" - - 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_cosyvoice3_request(self, request: OpenAICreateSpeechRequest) -> str | None: - """Validate CosyVoice3 request parameters. Returns error message or None.""" - if not request.input or not request.input.strip(): - return "Input text cannot be empty" - - # CosyVoice3 requires reference audio for voice cloning - if request.ref_audio is None: - return "CosyVoice3 requires 'ref_audio' (reference audio for voice cloning)" - - fmt_err = self._validate_ref_audio_format(request.ref_audio) - if fmt_err: - return fmt_err - - if not request.ref_text or not request.ref_text.strip(): - return "CosyVoice3 requires 'ref_text' (transcript of the reference audio)" - - 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_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). @@ -2443,23 +2092,35 @@ def _put_resolved_ref_audio(self, cache_key: str, wav_list: list[float], sr: int def _discard_ref_audio_artifact_ready_if_unreferenced(self, artifact_key: str) -> None: if artifact_key and all(entry[3] != artifact_key for entry in self._ref_audio_resolve_cache.values()): - self._ref_audio_model_artifact_ready.discard(artifact_key) + self._ref_audio_model_artifact_ready = { + (key, mode) for (key, mode) in self._ref_audio_model_artifact_ready if key != artifact_key + } + + @staticmethod + def _tts_x_vector_only(tts_params: dict[str, Any]) -> bool: + return bool((tts_params.get("x_vector_only_mode") or [False])[0]) def _qwen3_tts_can_use_ref_audio_artifact_only(self, tts_params: dict[str, Any], artifact_key: str | None) -> bool: if self._tts_model_type != "qwen3_tts": return False - if not artifact_key or artifact_key not in self._ref_audio_model_artifact_ready: + x_vector_only = self._tts_x_vector_only(tts_params) + if not artifact_key or (artifact_key, x_vector_only) not in self._ref_audio_model_artifact_ready: return False return (tts_params.get("task_type") or ["CustomVoice"])[0] == "Base" - def _track_ref_audio_artifact_warmup(self, request_id: str, artifact_key: str | None) -> None: + def _track_ref_audio_artifact_warmup( + self, request_id: str, artifact_key: str | None, x_vector_only: bool = False + ) -> None: if artifact_key: - self._request_ref_audio_artifact_keys[request_id] = artifact_key + self._request_ref_audio_artifact_keys[request_id] = (artifact_key, bool(x_vector_only)) def _mark_ref_audio_artifact_ready_for_request(self, request_id: str) -> None: - artifact_key = self._request_ref_audio_artifact_keys.pop(request_id, None) + tracked = self._request_ref_audio_artifact_keys.pop(request_id, None) + if tracked is None: + return + artifact_key, x_vector_only = tracked if artifact_key and any(entry[3] == artifact_key for entry in self._ref_audio_resolve_cache.values()): - self._ref_audio_model_artifact_ready.add(artifact_key) + self._ref_audio_model_artifact_ready.add((artifact_key, x_vector_only)) def _discard_ref_audio_artifact_warmup(self, request_id: str) -> None: self._request_ref_audio_artifact_keys.pop(request_id, None) @@ -2472,9 +2133,6 @@ async def _resolve_ref_audio_many(self, ref_audio_list: list[str]) -> list[tuple # ---- 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 @@ -2519,16 +2177,12 @@ 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 @@ -2551,7 +2205,6 @@ 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, ) @@ -2570,7 +2223,8 @@ def _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) + trust_remote_code = bool(getattr(self.engine_client.model_config, "trust_remote_code", False)) + self._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=trust_remote_code) ref_text = request.ref_text prompt_waveform = self._build_ming_prompt_waveform(ref_audio_data) if ref_text is not None else None @@ -2607,6 +2261,8 @@ async def _generate_audio_chunks( response_format: str = "pcm", raw_request: Request | None = None, request_start_s: float | None = None, + include_sample_rate: bool = False, + usage_acc: SpeechOutputTokenCounter | None = None, ): """Generate audio chunks for streaming response. @@ -2633,6 +2289,10 @@ async def _generate_audio_chunks( try: async for res in generator: + # Tally generated codec tokens for usage (reads per-stage metrics + # off the final output; a cheap early-return on every other res). + if usage_acc is not None: + usage_acc.observe(res) audio_output, audio_key = self._extract_audio_output(res) if audio_key is None: continue @@ -2668,7 +2328,12 @@ async def _generate_audio_chunks( assert sr_raw is not None, ( "First audio chunk must include sample rate metadata for WAV streaming" ) - wav_header = _create_wav_header(sample_rate=sample_rate_val, num_channels=1, bits_per_sample=16) + num_channels = _infer_audio_num_channels(np.asarray(chunk_np)) + wav_header = _create_wav_header( + sample_rate=sample_rate_val, + num_channels=num_channels, + bits_per_sample=16, + ) yield wav_header first_chunk = False @@ -2678,12 +2343,15 @@ async def _generate_audio_chunks( sample_rate=sample_rate_val, response_format="pcm", speed=1.0, - stream_format="audio", base64_encode=False, ) if first_audio_chunk_s is None: first_audio_chunk_s = time.perf_counter() - yield self.create_audio(audio_obj).audio_data + audio_bytes = self.create_audio(audio_obj).audio_data + if include_sample_rate: + yield audio_bytes, sample_rate_val + else: + yield audio_bytes self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True total_ms = (time.perf_counter() - stream_start_s) * 1000.0 @@ -2743,6 +2411,69 @@ async def _generate_audio_chunks( if not artifact_ready: self._discard_ref_audio_artifact_warmup(request_id) + async def _generate_audio_sse_events( + self, + generator, + request_id: str, + response_format: str = "pcm", + raw_request: Request | None = None, + request_start_s: float | None = None, + request: OpenAICreateSpeechRequest | None = None, + tts_params: dict[str, Any] | None = None, + ): + """Generate OpenAI-style SSE events with base64 audio deltas. + + Field naming follows the OpenAI ``speech.audio.delta`` schema, which + carries the base64 chunk in ``audio`` (not ``delta`` — that is the + Realtime API ``response.audio.delta`` convention, a different event). + See https://platform.openai.com/docs/api-reference/audio-streaming. + + The terminal ``speech.audio.done`` event carries a ``usage`` object + (``input_tokens``/``output_tokens``/``total_tokens`` + a per-modality + ``input_token_details`` breakdown), matching OpenAI's documented + ``speech.audio.done`` schema. ``output_tokens`` is accumulated from the + stage-0 deltas as they stream (see ``SpeechOutputTokenCounter``); + ``input_tokens`` is computed from the request text + reference audio. + """ + usage_acc = SpeechOutputTokenCounter() + try: + async for chunk in self._generate_audio_chunks( + generator, + request_id, + response_format, + raw_request=raw_request, + request_start_s=request_start_s, + usage_acc=usage_acc, + ): + payload = { + "type": "speech.audio.delta", + "audio": base64.b64encode(chunk).decode("ascii"), + "response_format": response_format, + } + data = json.dumps(payload, separators=(",", ":")) + yield f"event: speech.audio.delta\ndata: {data}\n\n" + done_payload: dict[str, Any] = {"type": "speech.audio.done"} + if request is not None: + # Streaming path: output_tokens = sum of stage-0 deltas. + usage = self._build_speech_usage(request, tts_params or {}, usage_acc.total()) + done_payload["usage"] = usage.model_dump() + done = json.dumps(done_payload, separators=(",", ":")) + yield f"event: speech.audio.done\ndata: {done}\n\n" + except asyncio.CancelledError: + raise + except Exception as e: + payload = { + "type": "speech.audio.error", + "error": { + "message": str(e), + "type": "server_error", + "param": None, + "code": HTTPStatus.INTERNAL_SERVER_ERROR.value, + }, + } + data = json.dumps(payload, separators=(",", ":")) + yield f"event: speech.audio.error\ndata: {data}\n\n" + @staticmethod def _extract_audio_output(res) -> tuple[dict | None, str | None]: """Return (audio_output dict, audio key) or (None, None). @@ -3034,8 +2765,14 @@ async def _build_cosyvoice3_prompt( wav_samples, sr = await self._resolve_ref_audio(request.ref_audio) audio_data = (np.asarray(wav_samples, dtype=np.float32), sr) + # Wrap the reference transcript in the CosyVoice3 instruction template + # so the talker emits target-only speech (see _COSYVOICE3_PROMPT_PREFIX). + # Skip if the caller already supplied a formatted prompt_text. + ref_text = request.ref_text or "" + if _COSYVOICE3_PROMPT_DELIMITER not in ref_text: + ref_text = f"{_COSYVOICE3_PROMPT_PREFIX}{ref_text}" mm_kwargs: dict[str, Any] = { - "prompt_text": request.ref_text, + "prompt_text": ref_text, "sample_rate": sr, } # Pass voice metadata for caching in the processor @@ -3308,6 +3045,7 @@ async def _prepare_speech_generation( self, request: OpenAICreateSpeechRequest, request_id: str | None = None, + has_inline_ref_audio: bool | None = None, ) -> tuple[str, Any, dict[str, Any]]: if self.engine_client.errored: raise self.engine_client.dead_error @@ -3315,12 +3053,19 @@ async def _prepare_speech_generation( request_id = request_id or f"speech-{random_uuid()}" qwen3_ref_audio_warmup_artifact_key: str | None = None - # If this is a streaming request, we need to coerce - # cumulative outputs to delta outputs; this ensures - # we don't emit redundant MM data & drain after emitting. + # If this is a streaming request with real async chunks, we need to + # coerce cumulative outputs to delta outputs; this ensures we don't + # emit redundant MM data & drain after emitting. Qwen3-TTS full-payload + # (async_chunk=False) has no incremental audio chunks, so keep + # FINAL_ONLY semantics and let the streaming response send the final + # waveform once. Scoped to qwen3_tts: other async_chunk=False models + # keep the DELTA coercion they stream with today. # list() makes a copy to avoid mutating the params. sampling_params_list = list(self.engine_client.default_sampling_params_list) - sampling_params_list = coerce_param_message_types(sampling_params_list, request.stream) + async_chunk = getattr(self.model_config, "async_chunk", True) + qwen3_full_payload = self._tts_model_type == "qwen3_tts" and not bool(async_chunk) + is_streaming_request = request.is_streaming() and not qwen3_full_payload + sampling_params_list = coerce_param_message_types(sampling_params_list, is_streaming_request) # Build prompt + tts_params via the per-model adapter (RFC #4327). Every # dedicated TTS model resolves to an adapter that owns its validation, @@ -3332,11 +3077,13 @@ async def _prepare_speech_generation( # apply uploaded speakers inside validate(), which sets request.ref_audio # in place. The builders need to know whether the caller supplied audio # inline vs. via an uploaded voice. - has_inline_ref_audio = request.ref_audio is not None + model_type: str | None = None + has_inline_ref_audio = (request.ref_audio is not None) if has_inline_ref_audio is None else has_inline_ref_audio if self._tts_model_type == "ming_flash_omni_tts": # ming_flash_omni is intentionally NOT migrated onto the adapter # framework in this PR (it has no registered adapter); keep it on the # legacy inline dispatch so serving still works. + model_type = "ming_flash_omni_tts" validation_error = self._validate_ming_flash_omni_tts_request(request) if validation_error: raise ValueError(validation_error) @@ -3350,6 +3097,7 @@ async def _prepare_speech_generation( prepared = await adapter.build(request, sampling_params_list, has_inline_ref_audio) prompt = prepared.prompt tts_params = prepared.tts_params + model_type = prepared.model_type qwen3_ref_audio_warmup_artifact_key = prepared.warmup_artifact_key else: # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker" @@ -3373,34 +3121,11 @@ async def _prepare_speech_generation( tts_params = {} prompt = {"prompt": request.input} - if self._tts_model_type == "step_audio2": - model_type = "step_audio2" - elif self._is_fish_speech: - model_type = "fish_speech" - elif self._tts_model_type == "covo_audio": - model_type = "covo_audio" - elif self._tts_model_type == "voxtral_tts": - model_type = "voxtral_tts" - elif self._tts_model_type == "cosyvoice3": - model_type = "cosyvoice3" - elif self._tts_model_type == "voxcpm2": - 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": - model_type = "moss_tts" - elif self._tts_model_type == "higgs_audio_v2": - model_type = "higgs_audio_v2" - elif self._tts_model_type == "glm_tts": - model_type = "glm_tts" - elif self._is_tts: - model_type = tts_params.get("task_type", ["unknown"])[0] - else: - model_type = "generic" + if model_type is None: + if self._is_tts: + model_type = tts_params.get("task_type", ["unknown"])[0] + else: + model_type = "generic" logger.info( "TTS speech request %s: text=%r, model=%s", request_id, @@ -3500,25 +3225,43 @@ async def _prepare_speech_generation( 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 + from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + MOE_TEXT_EOS_TOKEN_ID, + TEXT_EOS_TOKEN_ID, + ) + + hf_config = self.engine_client.model_config.hf_config + is_moe = getattr(hf_config, "model_type", "") == "bailingmm" + stop_token_id = MOE_TEXT_EOS_TOKEN_ID if is_moe else TEXT_EOS_TOKEN_ID sampling_params_list = copy.deepcopy(sampling_params_list) - sampling_params_list[0].stop_token_ids = [int(TEXT_EOS_TOKEN_ID)] + sampling_params_list[0].stop_token_ids = [int(stop_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: + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + stage0_params = sampling_params_list[0] + stage0_params.seed = request.seed + if stage0_params.extra_args is None: + stage0_params.extra_args = {} + stage0_params.extra_args["tts_local_seed"] = request.seed + + if self._tts_model_type == "qwen3_tts" and sampling_params_list: + stage0_params = sampling_params_list[0] + default_seed = getattr(stage0_params, "seed", None) + if default_seed is not None: import copy sampling_params_list = copy.deepcopy(sampling_params_list) - sampling_params_list[0].seed = request.seed - if self._tts_model_type == "qwen3_tts": - if sampling_params_list[0].extra_args is None: - sampling_params_list[0].extra_args = {} - sampling_params_list[0].extra_args["qwen3_tts_request_seed"] = request.seed + stage0_params = sampling_params_list[0] + if stage0_params.extra_args is None: + stage0_params.extra_args = {} + stage0_params.extra_args.setdefault("tts_local_seed", int(default_seed)) generator = self.engine_client.generate( prompt=prompt, @@ -3526,16 +3269,25 @@ async def _prepare_speech_generation( sampling_params_list=sampling_params_list, output_modalities=["audio"], ) - self._track_ref_audio_artifact_warmup(request_id, qwen3_ref_audio_warmup_artifact_key) + self._track_ref_audio_artifact_warmup( + request_id, + qwen3_ref_audio_warmup_artifact_key, + x_vector_only=self._tts_x_vector_only(tts_params), + ) return request_id, generator, tts_params - async def _generate_pcm_chunks(self, generator, request_id: str): + async def _generate_pcm_chunks(self, generator, request_id: str, *, include_sample_rate: bool = False): """Yield raw PCM byte chunks from the engine generator. Delegates to ``_generate_audio_chunks`` with ``response_format="pcm"``. Used by the WebSocket streaming handler and ``_iter_pcm_audio_bytes``. """ - async for chunk in self._generate_audio_chunks(generator, request_id, response_format="pcm"): + async for chunk in self._generate_audio_chunks( + generator, + request_id, + response_format="pcm", + include_sample_rate=include_sample_rate, + ): yield chunk async def _iter_pcm_audio_bytes(self, request: OpenAICreateSpeechRequest): @@ -3552,8 +3304,16 @@ async def _generate_audio_bytes( request: OpenAICreateSpeechRequest, base64_encode: bool = False, request_id: str | None = None, + usage_out: list[SpeechTokenUsage] | None = None, + has_inline_ref_audio: bool | None = None, ) -> tuple[bytes | str, str]: - request_id, generator, _ = await self._prepare_speech_generation(request, request_id=request_id) + # ``usage_out`` is an opt-in output channel: when a list is passed, the + # computed SpeechTokenUsage is appended to it. The return stays a + # 2-tuple so existing callers (and their test mocks) are unaffected; + # only the batch path, which surfaces per-item usage, opts in. + request_id, generator, bytes_tts_params = await self._prepare_speech_generation( + request, request_id=request_id, has_inline_ref_audio=has_inline_ref_audio + ) artifact_ready = False try: @@ -3566,8 +3326,12 @@ async def _generate_audio_bytes( moss_sample_rate: int | None = None final_output: OmniRequestOutput | None = None + # Non-streaming is FINAL_ONLY, so the stage-0 output carries the full + # token sequence; the counter records its length for output_tokens. + usage_acc = SpeechOutputTokenCounter() async for res in generator: final_output = res + usage_acc.observe(res) if not is_moss: continue try: @@ -3646,12 +3410,13 @@ async def _generate_audio_bytes( sample_rate=sample_rate, response_format=request.response_format or "wav", speed=request.speed or 1.0, - stream_format=request.stream_format, base64_encode=base64_encode, ) audio_response: AudioResponse = self.create_audio(audio_obj) self._mark_ref_audio_artifact_ready_for_request(request_id) artifact_ready = True + if usage_out is not None: + usage_out.append(self._build_speech_usage(request, bytes_tts_params or {}, usage_acc.total())) return audio_response.audio_data, audio_response.media_type finally: if not artifact_ready: @@ -3759,7 +3524,6 @@ async def _create_diffusion_speech( sample_rate=sample_rate, response_format=request.response_format or "wav", speed=request.speed or 1.0, - stream_format=request.stream_format, base64_encode=False, ) audio_response: AudioResponse = self.create_audio(audio_obj) @@ -3789,6 +3553,25 @@ def _diffusion_error_response(message: str, status_code: int = 500) -> Response: error_body = json.dumps({"error": {"message": message, "type": err_type, "param": None, "code": status_code}}) return Response(content=error_body, media_type="application/json", status_code=status_code) + def _validate_speech_streaming_request( + self, + request: OpenAICreateSpeechRequest, + *, + mode_label: str, + ) -> tuple[str, Response | None]: + """Validate pcm/wav + speed constraints for streaming speech responses.""" + response_format = (request.response_format or "wav").lower() + if response_format not in ("pcm", "wav"): + return response_format, self.create_error_response( + f"{mode_label} is only supported for 'pcm' and 'wav' formats. Got '{response_format}'." + ) + if request.speed is not None and request.speed != 1.0: + return response_format, self.create_error_response( + f"{mode_label} is not supported with speed adjustment. " + "Use a non-streaming request or remove the speed parameter." + ) + return response_format, None + async def create_speech( self, request: OpenAICreateSpeechRequest, @@ -3810,9 +3593,11 @@ async def create_speech( - ref_text: Transcript of reference audio (Base task) - x_vector_only_mode: Use speaker embedding only (Base task) - Streaming is supported via stream=True with response_format='pcm' or 'wav'. - Each Code2Wav chunk is yielded as raw audio bytes as soon as it is decoded. - For WAV format, a header with placeholder size values is emitted first. + Streaming is supported via the ``stream=True`` switch or ``stream_format='sse'``, + which return OpenAI ``speech.audio.*`` SSE events. ``stream_format='audio'`` + opts into raw audio streaming with ``response_format='pcm'`` or ``'wav'``. + Raw audio streaming yields each Code2Wav chunk as raw bytes as soon as it is + decoded. Raw WAV streaming emits a header with placeholder size values first. """ if self._diffusion_mode: return await self._create_diffusion_speech(request) @@ -3830,23 +3615,20 @@ async def create_speech( ) try: - if request.stream: - # Determine response format and media type for streaming - response_format = (request.response_format or "wav").lower() - - # Only pcm and wav support streaming without post-processing - if response_format not in ["pcm", "wav"]: - return self.create_error_response( - f"Streaming is only supported for 'pcm' and 'wav' formats. " - f"Got '{response_format}'. For other formats, use stream=False." - ) + if request.is_streaming() and request.word_timestamps: + return self.create_error_response( + "word_timestamps=true is currently supported by the WebSocket " + "/v1/audio/speech/stream path. Use session.config with " + "stream_audio=true and response_format='pcm'." + ) - # Check if speed adjustment is requested (not compatible with streaming) - if request.speed is not None and request.speed != 1.0: - return self.create_error_response( - "Streaming is not supported with speed adjustment. " - "Use stream=False or remove the speed parameter." - ) + if request.is_raw_audio_stream(): + response_format, error = self._validate_speech_streaming_request( + request, + mode_label="Streaming", + ) + if error is not None: + return error media_type = "audio/wav" if response_format == "wav" else "audio/pcm" _, generator, _ = await self._prepare_speech_generation(request, request_id=request_id) @@ -3861,6 +3643,28 @@ async def create_speech( media_type=media_type, ) + if request.is_sse_stream(): + response_format, error = self._validate_speech_streaming_request( + request, + mode_label="SSE streaming", + ) + if error is not None: + return error + + _, generator, sse_tts_params = await self._prepare_speech_generation(request, request_id=request_id) + return StreamingResponse( + self._generate_audio_sse_events( + generator, + request_id, + response_format, + raw_request=raw_request, + request_start_s=request_start_s, + request=request, + tts_params=sse_tts_params, + ), + media_type="text/event-stream", + ) + audio_bytes, media_type = await self._generate_audio_bytes(request, request_id=request_id) total_ms = (time.perf_counter() - request_start_s) * 1000.0 logger.info( @@ -4016,16 +3820,15 @@ async def create_speech_batch( merged_requests = [self._merge_batch_item(batch_request, item) for item in batch_request.items] async def _run_item(idx: int, req: OpenAICreateSpeechRequest) -> SpeechBatchItemResult: - # Batch validation still goes through _validate_tts_request directly - # (not the adapter). The single-request path validates via the - # adapter; both ultimately call the same per-model validators, so the - # surfaces stay in sync. Routing batch through the adapter (and its - # uploaded-speaker handling) is a follow-up (RFC #4327). + has_inline_ref_audio = req.ref_audio is not None validation_error = self._validate_tts_request(req) if validation_error is not None: return SpeechBatchItemResult(index=idx, status="error", error=validation_error) + usage_box: list[SpeechTokenUsage] = [] try: - audio_data, media_type = await self._generate_audio_bytes(req, base64_encode=True) + audio_data, media_type = await self._generate_audio_bytes( + req, base64_encode=True, usage_out=usage_box, has_inline_ref_audio=has_inline_ref_audio + ) except Exception as e: logger.exception("Batch item %d failed: %s", idx, e) return SpeechBatchItemResult(index=idx, status="error", error=str(e)) @@ -4034,6 +3837,7 @@ async def _run_item(idx: int, req: OpenAICreateSpeechRequest) -> SpeechBatchItem status="success", audio_data=audio_data, media_type=media_type, + usage=usage_box[0] if usage_box else None, ) results = await asyncio.gather( diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 695e87d6eef..28acf8679d0 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -1,7 +1,7 @@ """WebSocket handler for streaming text input TTS. -Accepts text incrementally via WebSocket, buffers and splits at sentence -boundaries, and generates audio per sentence using the existing TTS pipeline. +Accepts text incrementally via WebSocket, buffers it until input.done, and +generates audio once for the buffered input using the existing TTS pipeline. Protocol: Client -> Server: @@ -9,15 +9,26 @@ {"type": "input.text", "text": "..."} # Text chunks {"type": "input.done"} # End of input - Server -> Client: + Server -> Client (default, word_timestamps=false): {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "wav"} <binary frame: audio bytes> + ... {"type": "audio.done", "sentence_index": 0} {"type": "session.done", "total_sentences": N} {"type": "error", "message": "..."} + + Server -> Client (when word_timestamps=true): + {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm"} + {"type": "audio.chunk", "sentence_index": 0, "chunk_id": 0, "audio_b64": "<base64 PCM>", "timestamps": null} + ... + {"type": "audio.chunk", "audio_b64": "", "timestamps": [{"word", "start_ms", "end_ms"}, ...]} + {"type": "audio.done", "sentence_index": 0} + # Audio is JSON base64 PCM (not binary). A trailing empty-audio chunk carries the + # full sentence-relative alignment. timestamps: list = aligned, [] = silence, null = failed. """ import asyncio +import base64 import json from contextlib import aclosing @@ -30,17 +41,15 @@ StreamingSpeechSessionConfig, ) from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech -from vllm_omni.entrypoints.openai.text_splitter import ( - SPLIT_CLAUSE, - SPLIT_SENTENCE, - SentenceSplitter, -) +from vllm_omni.utils.forced_aligner import ForcedAlignerLoadError +from vllm_omni.utils.forced_aligner import align as forced_align logger = init_logger(__name__) _DEFAULT_IDLE_TIMEOUT = 30.0 # seconds _DEFAULT_CONFIG_TIMEOUT = 10.0 # seconds _PCM_SAMPLE_RATE = 24000 +_BYTES_PER_SAMPLE = 2 # 16-bit mono PCM _MAX_CONFIG_MESSAGE_SIZE = 4 * 1024 * 1024 # allow large ref_audio payloads _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 @@ -49,8 +58,8 @@ class OmniStreamingSpeechHandler: """Handles WebSocket sessions for streaming text-input TTS. Each WebSocket connection is an independent session. Text arrives - incrementally, is split at sentence boundaries, and audio is generated - per sentence using the existing OmniOpenAIServingSpeech pipeline. + incrementally, is buffered until input.done, and audio is generated once + for the buffered input using the existing OmniOpenAIServingSpeech pipeline. Args: speech_service: The existing TTS serving instance (reused for @@ -88,9 +97,7 @@ async def handle_session(self, websocket: WebSocket) -> None: await self._send_error(websocket, str(error)) return - boundary_re = SPLIT_CLAUSE if config.split_granularity == "clause" else SPLIT_SENTENCE - splitter = SentenceSplitter(boundary_re=boundary_re) - sentence_index = 0 + text_parts: list[str] = [] # 2. Receive text chunks until input.done while True: @@ -124,23 +131,19 @@ async def handle_session(self, websocket: WebSocket) -> None: if not isinstance(text, str): await self._send_error(websocket, "input.text requires a string value") continue - sentences = splitter.add_text(text) - for sentence in sentences: - await self._generate_and_send(websocket, config, sentence, sentence_index) - sentence_index += 1 + text_parts.append(text) elif msg_type == "input.done": - # Flush remaining buffer - remaining = splitter.flush() - if remaining: - await self._generate_and_send(websocket, config, remaining, sentence_index) - sentence_index += 1 + full_text = "".join(text_parts).strip() + total_sentences = 0 + if full_text: + await self._generate_and_send(websocket, config, full_text, 0) + total_sentences = 1 - # Send session.done await websocket.send_json( { "type": "session.done", - "total_sentences": sentence_index, + "total_sentences": total_sentences, } ) return @@ -210,6 +213,24 @@ async def _generate_and_send( """Generate audio for a single sentence and send it over WebSocket.""" response_format = config.response_format or "wav" + # Reject unmet word-timestamps preconditions early with a clear reason. + if config.word_timestamps: + if self._speech_service.forced_aligner_config is None: + await self._send_error( + websocket, + "word_timestamps=true but the server was launched without " + "--forced-aligner; either restart the server with that flag " + "or set word_timestamps=false in session.config.", + ) + return + if not (config.stream_audio and response_format == "pcm"): + await self._send_error( + websocket, + "word_timestamps=true requires stream_audio=true and " + "response_format='pcm' (the aligner consumes raw PCM).", + ) + return + request = OpenAICreateSpeechRequest( input=sentence_text, model=config.model, @@ -227,6 +248,7 @@ async def _generate_and_send( x_vector_only_mode=config.x_vector_only_mode, speaker_embedding=config.speaker_embedding, stream=config.stream_audio, + word_timestamps=config.word_timestamps, ) start_payload = { @@ -236,7 +258,11 @@ async def _generate_and_send( "format": response_format, } if config.stream_audio and response_format == "pcm": + # Nominal stream rate; each audio.chunk carries the authoritative + # per-chunk sample_rate. start_payload["sample_rate"] = _PCM_SAMPLE_RATE + if config.word_timestamps: + start_payload["word_timestamps"] = True await websocket.send_json(start_payload) total_bytes = 0 @@ -245,10 +271,20 @@ async def _generate_and_send( try: if config.stream_audio: request_id, generator, _ = await self._speech_service._prepare_speech_generation(request) - async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream: - async for chunk in stream: - total_bytes += len(chunk) - await websocket.send_bytes(chunk) + if config.word_timestamps: + total_bytes = await self._stream_audio_with_alignments( + websocket=websocket, + request_id=request_id, + generator=generator, + sentence_text=sentence_text, + sentence_index=sentence_index, + language=config.language, + ) + else: + async with aclosing(self._speech_service._generate_pcm_chunks(generator, request_id)) as stream: + async for chunk in stream: + total_bytes += len(chunk) + await websocket.send_bytes(chunk) else: audio_bytes, _ = await self._speech_service._generate_audio_bytes(request) total_bytes = len(audio_bytes) @@ -277,6 +313,115 @@ async def _generate_and_send( except Exception: logger.debug("Failed to send audio.done for sentence %d", sentence_index, exc_info=True) + async def _stream_audio_with_alignments( + self, + *, + websocket: WebSocket, + request_id: str, + generator, + sentence_text: str, + sentence_index: int, + language: str | None = None, + ) -> int: + """Stream PCM as JSON ``audio.chunk`` frames, aligned per sentence. + + Forward each PCM chunk live (``timestamps: null``) while buffering the + sentence audio, then run the forced aligner once over the whole + sentence and emit a final empty-audio ``audio.chunk`` with the word + timestamps. On aligner failure timestamps is ``null``; for silence it + is ``[]`` (audio always flows regardless). + """ + aligner_config = self._speech_service.forced_aligner_config + assert aligner_config is not None # gated by the precondition check + + sentence_audio = bytearray() + total_bytes = 0 + sample_rate = _PCM_SAMPLE_RATE + chunk_id = 0 + + async def send_chunk( + chunk: bytes, + chunk_sample_rate: int, + timestamps_payload: list[dict] | None, + chunk_start_ms: int, + chunk_end_ms: int, + ) -> None: + nonlocal chunk_id + await websocket.send_json( + { + "type": "audio.chunk", + "sentence_index": sentence_index, + "chunk_id": chunk_id, + "chunk_start_ms": chunk_start_ms, + "chunk_end_ms": chunk_end_ms, + "sample_rate": chunk_sample_rate, + "audio_b64": base64.b64encode(chunk).decode("ascii"), + "timestamps": timestamps_payload, + } + ) + chunk_id += 1 + + async with aclosing( + self._speech_service._generate_pcm_chunks( + generator, + request_id, + include_sample_rate=True, + ) + ) as stream: + async for chunk, chunk_sample_rate in stream: + sample_rate = chunk_sample_rate + chunk_start_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + sentence_audio.extend(chunk) + chunk_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + total_bytes += len(chunk) + # Audio first, timestamps after the whole sentence is aligned. + await send_chunk(chunk, chunk_sample_rate, None, chunk_start_ms, chunk_end_ms) + + # Single alignment pass over the full sentence, then emit timestamps. + # A load/config failure is permanent, so surface the reason once; audio + # has already streamed, so the trailing frame still carries null. + try: + timestamps_payload = await self._align_sentence( + audio=bytes(sentence_audio), + text=sentence_text, + sample_rate=sample_rate, + config=aligner_config, + language=language, + ) + except ForcedAlignerLoadError as exc: + await self._send_error(websocket, f"forced aligner unavailable: {exc}") + timestamps_payload = None + sentence_end_ms = int(round((len(sentence_audio) / _BYTES_PER_SAMPLE / sample_rate) * 1000.0)) + await send_chunk(b"", sample_rate, timestamps_payload, 0, sentence_end_ms) + + return total_bytes + + @staticmethod + async def _align_sentence( + *, + audio: bytes, + text: str, + sample_rate: int, + config, + language: str | None = None, + ) -> list[dict] | None: + """Convert a sentence alignment into JSON word-timestamp dicts. + + Returns ``None`` on aligner failure, ``[]`` when it ran but produced no + tokens. Monotonic, non-overlapping bounds are guaranteed by the decoder. + ``language`` is forwarded to word segmentation. + """ + aligned = await forced_align( + audio=audio, + text=text, + sample_rate=sample_rate, + config=config, + language=language, + ) + if aligned is None: + return None + return [{"word": ts.word, "start_ms": ts.start_ms, "end_ms": ts.end_ms} for ts in aligned] + @staticmethod async def _send_error(websocket: WebSocket, message: str) -> None: """Send an error message to the client.""" diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 4bdbfce0217..ed846e7e82a 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -29,6 +29,11 @@ from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request from vllm_omni.entrypoints.openai.video_api_utils import _encode_video_bytes, encode_video_base64 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt +from vllm_omni.outputs.output_metadata import ( + DiffusionMetadataMapping, + DiffusionMultimodalOutput, + DiffusionPayloadValue, +) logger = init_logger(__name__) @@ -58,8 +63,8 @@ class ReferenceAudio: class VideoGenerationArtifacts: """Normalized outputs and profiler metadata extracted from one request.""" - videos: list[Any] - audios: list[Any | None] + videos: list[DiffusionPayloadValue] + audios: list[DiffusionPayloadValue | None] actions: list[VideoAction | None] audio_sample_rate: int output_fps: int @@ -128,6 +133,8 @@ async def _run_and_extract( status_code=HTTPStatus.BAD_REQUEST.value, detail="Provide either an image reference or a video reference, not both.", ) + provided_fields = request.model_fields_set + fps_provided = self._request_fps_provided(request) 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) @@ -147,10 +154,10 @@ async def _run_and_extract( gen_params.height = vp.height if vp.num_frames is not None: gen_params.num_frames = vp.num_frames - if vp.fps is not None: + # Leave fps/frame_rate as None when the user did not provide fps. + if fps_provided and vp.fps is not None: gen_params.fps = vp.fps gen_params.frame_rate = float(vp.fps) - provided_fields = request.model_fields_set if "enable_frame_interpolation" in provided_fields: gen_params.enable_frame_interpolation = request.enable_frame_interpolation if "frame_interpolation_exp" in provided_fields: @@ -223,13 +230,17 @@ async def _run_and_extract( ) result = await self._run_generation(prompt, gen_params, reference_id) - custom_output = self._extract_custom_output(result) - action_only = isinstance(custom_output, dict) and bool(custom_output.get("action_only_output")) + multimodal_output = self._extract_multimodal_output(result) + metadata = multimodal_output.get("metadata") if isinstance(multimodal_output, dict) else {} + common_metadata = metadata.get("common") if isinstance(metadata, dict) else {} + action_only = bool(isinstance(common_metadata, dict) and common_metadata.get("action_only_output")) videos = [{"action_only_output": True}] if action_only else self._extract_video_outputs(result) audios = self._extract_audio_outputs(result, expected_count=len(videos)) actions = self._extract_action_outputs(result, expected_count=len(videos)) audio_sample_rate = self._resolve_audio_sample_rate(result) - output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result) + model_fps = self._resolve_fps(result) + output_fps_base = (vp.fps if fps_provided else None) or model_fps or vp.fps or 24 + output_fps = output_fps_base * self._resolve_video_fps_multiplier(result) return VideoGenerationArtifacts( videos=videos, audios=audios, @@ -336,14 +347,25 @@ async def generate_video_bytes( return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb, artifacts.actions[0] @staticmethod - def _resolve_video_fps_multiplier(result: Any) -> int: - custom_output = OmniOpenAIServingVideo._extract_custom_output(result) - if isinstance(custom_output, dict): - multiplier = custom_output.get("video_fps_multiplier") + def _resolve_video_fps_multiplier(result: object) -> int: + multimodal_output = OmniOpenAIServingVideo._extract_multimodal_output(result) + metadata = multimodal_output.get("metadata") + video_metadata = metadata.get("video") if isinstance(metadata, Mapping) else None + if isinstance(video_metadata, Mapping): + multiplier = video_metadata.get("video_fps_multiplier") if multiplier is not None: return int(multiplier) return 1 + @staticmethod + def _request_fps_provided(request: VideoGenerationRequest) -> bool: + if "fps" in request.model_fields_set and request.fps is not None: + return True + video_params = request.video_params + if video_params is None or "video_params" not in request.model_fields_set: + return False + return "fps" in video_params.model_fields_set and video_params.fps is not None + def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams: default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None) if default_sampling_params_list: @@ -377,7 +399,7 @@ async def _run_generation( prompt: OmniTextPrompt, gen_params: OmniDiffusionSamplingParams, request_id: str, - ) -> Any: + ) -> object: stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None) if not stage_configs: @@ -420,7 +442,7 @@ async def _run_generation( return result @staticmethod - def _normalize_video_outputs(videos: Any) -> list[Any]: + def _normalize_video_outputs(videos: DiffusionPayloadValue | None) -> list[DiffusionPayloadValue]: if videos is None: return [] if hasattr(videos, "ndim") and videos.ndim == 5: @@ -430,7 +452,7 @@ def _normalize_video_outputs(videos: Any) -> list[Any]: return [] first = videos[0] if hasattr(first, "ndim") and first.ndim == 5: - flattened: list[Any] = [] + flattened: list[DiffusionPayloadValue] = [] for item in videos: if hasattr(item, "ndim") and item.ndim == 5: flattened.extend([item[i] for i in range(item.shape[0])]) @@ -446,7 +468,7 @@ def _normalize_video_outputs(videos: Any) -> list[Any]: return videos return [videos] - def _extract_video_outputs(self, result: Any) -> list[Any]: + def _extract_video_outputs(self, result: object) -> list[DiffusionPayloadValue]: videos = None if hasattr(result, "images") and result.images: videos = result.images @@ -456,10 +478,14 @@ def _extract_video_outputs(self, result: Any) -> list[Any]: videos = request_output["images"] elif hasattr(request_output, "images") and request_output.images: videos = request_output.images - elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output: - videos = request_output.multimodal_output.get("video") - if videos is None and hasattr(result, "multimodal_output") and result.multimodal_output: - videos = result.multimodal_output.get("video") + else: + request_multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(request_multimodal_output, Mapping): + videos = request_multimodal_output.get("video") + if videos is None: + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, Mapping): + videos = multimodal_output.get("video") normalized = self._normalize_video_outputs(videos) if not normalized: @@ -470,17 +496,21 @@ def _extract_video_outputs(self, result: Any) -> list[Any]: return normalized @staticmethod - def _extract_audio_outputs(result: Any, expected_count: int) -> list[Any | None]: + def _extract_audio_outputs(result: object, expected_count: int) -> list[DiffusionPayloadValue | None]: audio = None - if hasattr(result, "multimodal_output") and result.multimodal_output: - audio = result.multimodal_output.get("audio") + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, Mapping): + audio = multimodal_output.get("audio") elif hasattr(result, "request_output"): request_output = result.request_output if isinstance(request_output, dict) and request_output.get("multimodal_output"): mm_output = request_output.get("multimodal_output") or {} - audio = mm_output.get("audio") - elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output: - audio = request_output.multimodal_output.get("audio") + if isinstance(mm_output, Mapping): + audio = mm_output.get("audio") + else: + request_multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(request_multimodal_output, Mapping): + audio = request_multimodal_output.get("audio") if audio is None: return [None] * expected_count @@ -501,7 +531,7 @@ def _extract_audio_outputs(result: Any, expected_count: int) -> list[Any | None] return [audio] + [None] * max(expected_count - 1, 0) - def _resolve_audio_sample_rate(self, result: Any) -> int: + def _resolve_audio_sample_rate(self, result: object) -> int: result_sample_rate = self._extract_audio_sample_rate_from_result(result) if result_sample_rate is not None: return result_sample_rate @@ -515,38 +545,45 @@ def _resolve_audio_sample_rate(self, result: Any) -> int: return 24000 @classmethod - def _extract_action_outputs(cls, result: Any, expected_count: int) -> list[VideoAction | None]: - custom_output = cls._extract_custom_output(result) - if not custom_output or "action" not in custom_output: - return [None] * expected_count - - action_payload = custom_output.get("actions", custom_output["action"]) - action_items = cls._split_action_payload(action_payload, expected_count) - return [ - cls._make_video_action(action_item, custom_output) if action_item is not None else None - for action_item in action_items - ] + def _extract_action_outputs(cls, result: object, expected_count: int) -> list[VideoAction | None]: + multimodal_output = cls._extract_multimodal_output(result) + if "actions" in multimodal_output: + action_payload = multimodal_output["actions"] + metadata = multimodal_output.get("metadata") + action_metadata = metadata.get("actions") if isinstance(metadata, Mapping) else {} + action_metadata = action_metadata if isinstance(action_metadata, Mapping) else {} + action_items = cls._split_action_payload(action_payload, expected_count) + return [ + cls._make_video_action(action_item, action_metadata) if action_item is not None else None + for action_item in action_items + ] + + return [None] * expected_count @staticmethod - def _extract_custom_output(result: Any) -> dict[str, Any]: - custom_output = getattr(result, "custom_output", None) - if isinstance(custom_output, dict): - return custom_output + def _extract_multimodal_output(result: object) -> DiffusionMultimodalOutput: + multimodal_output = getattr(result, "multimodal_output", None) + if isinstance(multimodal_output, Mapping): + return dict(multimodal_output) request_output = getattr(result, "request_output", None) if isinstance(request_output, dict): - custom_output = request_output.get("custom_output") - if custom_output is None: - custom_output = request_output.get("_custom_output") + multimodal_output = request_output.get("multimodal_output") + if multimodal_output is None: + multimodal_output = request_output.get("_multimodal_output") elif request_output is not None: - custom_output = getattr(request_output, "custom_output", None) - if custom_output is None: - custom_output = getattr(request_output, "_custom_output", None) + multimodal_output = getattr(request_output, "multimodal_output", None) + if multimodal_output is None: + multimodal_output = getattr(request_output, "_multimodal_output", None) - return custom_output if isinstance(custom_output, dict) else {} + return dict(multimodal_output) if isinstance(multimodal_output, Mapping) else {} @classmethod - def _split_action_payload(cls, action: Any, expected_count: int) -> list[Any | None]: + def _split_action_payload( + cls, + action: DiffusionPayloadValue, + expected_count: int, + ) -> list[DiffusionPayloadValue | None]: if expected_count <= 0: return [] @@ -560,30 +597,34 @@ def _split_action_payload(cls, action: Any, expected_count: int) -> list[Any | N return [action] + [None] * (expected_count - 1) @classmethod - def _make_video_action(cls, action: Any, custom_output: dict[str, Any]) -> VideoAction: + def _make_video_action( + cls, + action: DiffusionPayloadValue, + action_metadata: DiffusionMetadataMapping, + ) -> VideoAction: data = cls._to_jsonable(action) if not isinstance(data, list): data = [data] - action_mode = custom_output.get("action_mode") + action_mode = action_metadata.get("action_mode") return VideoAction( data=data, shape=cls._shape_of(action), dtype=cls._dtype_of(action), - raw_action_dim=cls._coerce_optional_int(custom_output.get("raw_action_dim")), + raw_action_dim=cls._coerce_optional_int(action_metadata.get("raw_action_dim")), action_mode=str(action_mode) if action_mode is not None else None, - domain_id=cls._coerce_optional_int(custom_output.get("domain_id")), + domain_id=cls._coerce_optional_int(action_metadata.get("domain_id")), ) @staticmethod - def _index_action(action: Any, index: int) -> Any: + def _index_action(action: DiffusionPayloadValue, index: int) -> DiffusionPayloadValue | None: try: return action[index] except (IndexError, KeyError, TypeError): return None @classmethod - def _to_jsonable(cls, value: Any) -> Any: + def _to_jsonable(cls, value: DiffusionPayloadValue) -> DiffusionPayloadValue: if hasattr(value, "detach"): value = value.detach() if hasattr(value, "cpu"): @@ -600,7 +641,7 @@ def _to_jsonable(cls, value: Any) -> Any: return value @classmethod - def _shape_of(cls, value: Any) -> list[int]: + def _shape_of(cls, value: DiffusionPayloadValue) -> list[int]: shape = getattr(value, "shape", None) if shape is not None: try: @@ -614,12 +655,12 @@ def _shape_of(cls, value: Any) -> list[int]: return [] @staticmethod - def _dtype_of(value: Any) -> str | None: + def _dtype_of(value: DiffusionPayloadValue) -> str | None: dtype = getattr(value, "dtype", None) return str(dtype) if dtype is not None else None @staticmethod - def _coerce_optional_int(value: Any) -> int | None: + def _coerce_optional_int(value: object) -> int | None: if value is None: return None try: @@ -629,7 +670,7 @@ def _coerce_optional_int(value: Any) -> int | None: return None @staticmethod - def _resolve_fps(result: Any) -> int | None: + def _resolve_fps(result: object) -> int | None: """Extract fps from multimodal_output if the model reported it.""" multimodal_output = getattr(result, "multimodal_output", None) if isinstance(multimodal_output, Mapping): @@ -669,7 +710,7 @@ def _resolve_fps(result: Any) -> int | None: return None @classmethod - def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None: + def _extract_audio_sample_rate_from_result(cls, result: object) -> int | None: multimodal_output = getattr(result, "multimodal_output", None) if isinstance(multimodal_output, Mapping): sample_rate = cls._coerce_audio_sample_rate( @@ -704,7 +745,7 @@ def _extract_audio_sample_rate_from_result(cls, result: Any) -> int | None: return None @classmethod - def _extract_audio_sample_rate_from_config(cls, config: Any) -> int | None: + def _extract_audio_sample_rate_from_config(cls, config: object) -> int | None: if config is None: return None @@ -735,7 +776,7 @@ def _extract_audio_sample_rate_from_config(cls, config: Any) -> int | None: return None @staticmethod - def _coerce_audio_sample_rate(value: Any) -> int | None: + def _coerce_audio_sample_rate(value: object) -> int | None: if value is None: return None @@ -748,12 +789,12 @@ def _coerce_audio_sample_rate(value: Any) -> int | None: return sample_rate if sample_rate > 0 else None @staticmethod - def _extract_stage_durations(result: Any) -> dict[str, float]: + def _extract_stage_durations(result: object) -> dict[str, float]: stage_durations = getattr(result, "stage_durations", None) return stage_durations if isinstance(stage_durations, dict) else {} @staticmethod - def _extract_peak_memory_mb(result: Any) -> float: + def _extract_peak_memory_mb(result: object) -> float: peak_memory_mb = getattr(result, "peak_memory_mb", 0.0) try: return float(peak_memory_mb or 0.0) diff --git a/vllm_omni/entrypoints/openai/serving_video_output_stream.py b/vllm_omni/entrypoints/openai/serving_video_output_stream.py new file mode 100644 index 00000000000..f713df07c98 --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_video_output_stream.py @@ -0,0 +1,519 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""WebSocket handler for streaming generated video chunks. + +Protocol: + Client -> Server: + {"type": "session.start", "model": "...", "prompt": "...", "format": "m4s", ...} + {"type": "session.prompt_update", "prompt": "..."} # currently unsupported + {"type": "session.stop"} + {"type": "session.ping"} # optional; refreshes stall clock + + Server -> Client: + {"type": "video.start", "request_id": "...", "config": {...}, "format": "m4s"} + <binary frame: fragmented MP4 bytes> + {"type": "session.done", ...} + {"type": "session.pong"} # reply to session.ping + {"type": "error", "message": "..."} +""" + +from __future__ import annotations + +import asyncio +import copy +import inspect +import json +import time +import uuid +from collections.abc import AsyncGenerator +from http import HTTPStatus +from typing import Any, cast + +from fastapi import HTTPException, WebSocket, WebSocketDisconnect +from pydantic import ValidationError +from vllm.engine.protocol import EngineClient +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoParams +from vllm_omni.entrypoints.openai.stage_params import ( + build_stage_sampling_params_list, + get_default_sampling_params_list, +) +from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request +from vllm_omni.entrypoints.openai.video_api_utils import StreamingVideoFormat, create_streaming_video_encoder +from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt + +logger = init_logger(__name__) + +_DEFAULT_STALL_TIMEOUT = 60.0 +_DEFAULT_START_TIMEOUT = 10.0 +_CONTROL_POLL_INTERVAL = 1.0 +_MAX_START_MESSAGE_SIZE = 4 * 1024 * 1024 +_MAX_CONTROL_MESSAGE_SIZE = 128 * 1024 + + +class _SessionProgress: + """Tracks last server-side progress for stall-timeout enforcement.""" + + def __init__(self) -> None: + self._last_at = time.monotonic() + + def touch(self) -> None: + self._last_at = time.monotonic() + + def stalled_for(self, stall_timeout: float) -> bool: + return time.monotonic() - self._last_at > stall_timeout + + +class OmniStreamingVideoOutputHandler: + """Handles `/v1/realtime/video` sessions for generated video output.""" + + def __init__( + self, + engine_client: EngineClient, + model_name: str | None = None, + stage_configs: list[Any] | None = None, + stall_timeout: float = _DEFAULT_STALL_TIMEOUT, + start_timeout: float = _DEFAULT_START_TIMEOUT, + ) -> None: + self._engine_client = engine_client + self._model_name = model_name + self._stage_configs = stage_configs + self._stall_timeout = stall_timeout + self._start_timeout = start_timeout + + async def handle_session(self, websocket: WebSocket) -> None: + await websocket.accept() + + request_id: str | None = None + stop_event = asyncio.Event() + send_lock = asyncio.Lock() + control_task: asyncio.Task[None] | None = None + stopped = False + chunks_sent = 0 + + try: + start = await self._receive_start(websocket) + if start is None: + return + request, output_format = start + + request_id = f"video_stream-{uuid.uuid4().hex}" + progress = _SessionProgress() + async with send_lock: + await websocket.send_json( + { + "type": "video.start", + "request_id": request_id, + "format": output_format, + "config": request.model_dump(exclude_none=True), + } + ) + progress.touch() + + control_task = asyncio.create_task( + self._control_loop(websocket, request_id, stop_event, send_lock, progress) + ) + async for chunk in self._stream_video_bytes(request, request_id, output_format, progress): + if stop_event.is_set(): + stopped = True + break + async with send_lock: + await websocket.send_bytes(chunk) + chunks_sent += 1 + progress.touch() + + async with send_lock: + await websocket.send_json( + { + "type": "session.done", + "request_id": request_id, + "chunks": chunks_sent, + "stopped": stopped or stop_event.is_set(), + } + ) + except WebSocketDisconnect: + logger.info("Streaming video output: client disconnected") + if request_id is not None: + await self._abort_request(request_id) + except Exception as e: + if stop_event.is_set(): + try: + async with send_lock: + await websocket.send_json( + { + "type": "session.done", + "request_id": request_id, + "chunks": chunks_sent, + "stopped": True, + } + ) + except Exception: + logger.debug("Failed to send stopped session.done", exc_info=True) + return + logger.exception("Streaming video output session error: %s", e) + await self._send_error(websocket, str(e)) + finally: + if control_task is not None: + control_task.cancel() + await asyncio.gather(control_task, return_exceptions=True) + + async def _receive_start(self, websocket: WebSocket) -> tuple[VideoGenerationRequest, StreamingVideoFormat] | None: + try: + raw = await asyncio.wait_for(websocket.receive_text(), timeout=self._start_timeout) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.start") + return None + + if len(raw) > _MAX_START_MESSAGE_SIZE: + await self._send_error(websocket, "session.start message too large") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.start") + return None + + if not isinstance(msg, dict): + await self._send_error(websocket, "session.start must be a JSON object") + return None + + if msg.get("type") != "session.start": + await self._send_error(websocket, f"Expected session.start, got: {msg.get('type')}") + return None + + try: + output_format = self._coerce_output_format(msg.pop("format", "m4s")) + except ValueError as e: + await self._send_error(websocket, str(e)) + return None + + request_data = self._coerce_request_data(msg) + try: + request = VideoGenerationRequest(**request_data) + except ValidationError as e: + await self._send_error(websocket, f"Invalid session.start: {e}") + return None + + configured_model = self._model_name + if request.model is not None and configured_model is not None and request.model != configured_model: + await self._send_error( + websocket, + f"Model mismatch: request specifies '{request.model}' but server is running '{configured_model}'.", + ) + return None + + return request, output_format + + @staticmethod + def _coerce_request_data(msg: dict[str, Any]) -> dict[str, Any]: + request_data = {k: v for k, v in msg.items() if k not in {"type", "format", "sampling_params"}} + sampling_params = msg.get("sampling_params") + if isinstance(sampling_params, dict): + extra_params = dict(request_data.get("extra_params") or {}) + for key, value in sampling_params.items(): + if key in VideoGenerationRequest.model_fields and key not in request_data: + request_data[key] = value + else: + extra_params[key] = value + if extra_params: + request_data["extra_params"] = extra_params + return request_data + + @staticmethod + def _coerce_output_format(raw_format: Any) -> StreamingVideoFormat: + if raw_format == "m4s": + return cast(StreamingVideoFormat, raw_format) + raise ValueError("Invalid session.start format: expected 'm4s'.") + + async def _control_loop( + self, + websocket: WebSocket, + request_id: str, + stop_event: asyncio.Event, + send_lock: asyncio.Lock, + progress: _SessionProgress, + ) -> None: + while not stop_event.is_set(): + if progress.stalled_for(self._stall_timeout): + await self._send_error( + websocket, + "Stall timeout: no generation progress", + send_lock=send_lock, + ) + await self._abort_request(request_id) + stop_event.set() + return + + poll_timeout = min(_CONTROL_POLL_INTERVAL, self._stall_timeout) + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=poll_timeout, + ) + except asyncio.TimeoutError: + continue + + if len(raw) > _MAX_CONTROL_MESSAGE_SIZE: + await self._send_error(websocket, "control message too large", send_lock=send_lock) + continue + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON message", send_lock=send_lock) + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "WebSocket messages must be JSON objects", send_lock=send_lock) + continue + + msg_type = msg.get("type") + if msg_type == "session.stop": + await self._abort_request(request_id) + stop_event.set() + return + if msg_type == "session.ping": + progress.touch() + try: + async with send_lock: + await websocket.send_json({"type": "session.pong"}) + except Exception: + pass + continue + if msg_type == "session.prompt_update": + await self._send_error( + websocket, + "session.prompt_update is not supported by this backend yet", + send_lock=send_lock, + ) + continue + await self._send_error(websocket, f"Unknown message type: {msg_type}", send_lock=send_lock) + + async def _abort_request(self, request_id: str) -> None: + try: + abort = getattr(self._engine_client, "abort", None) + if callable(abort): + result = abort(request_id) + if inspect.isawaitable(result): + await result + except Exception: + logger.debug("Failed to abort streaming video request %s", request_id, exc_info=True) + + async def _stream_video_bytes( + self, + request: VideoGenerationRequest, + request_id: str, + output_format: StreamingVideoFormat, + progress: _SessionProgress, + ) -> AsyncGenerator[bytes, None]: + """Yield encoded video bytes from diffusion streaming outputs.""" + prompt, gen_params, vp = self._build_prompt_and_sampling_params(request) + video_codec_options = {"preset": "ultrafast", "threads": "0", "tune": "zerolatency"} + if isinstance(request.extra_params, dict) and "video_codec_options" in request.extra_params: + video_codec_options = request.extra_params["video_codec_options"] + + output_fps = vp.fps or gen_params.resolved_frame_rate or 16 + encoder = create_streaming_video_encoder( + output_format=output_format, + fps=output_fps, + video_codec_options=video_codec_options, + ) + completed = False + try: + async for result in self._iter_generation_outputs(prompt, gen_params, request_id, progress): + if error := getattr(result, "error", None): + raise RuntimeError(str(error)) + videos = self._extract_video_outputs(result) + for video in videos: + chunk = encoder.encode(video) + if chunk: + yield chunk + + completed = True + final_chunk = encoder.close() + if final_chunk: + yield final_chunk + finally: + if not completed: + encoder.close() + + def _build_prompt_and_sampling_params( + self, + request: VideoGenerationRequest, + ) -> tuple[OmniTextPrompt, OmniDiffusionSamplingParams, VideoParams]: + """Build text-only diffusion inputs for the streaming endpoint.""" + prompt: OmniTextPrompt = OmniTextPrompt(prompt=request.prompt) + if request.negative_prompt is not None: + prompt["negative_prompt"] = request.negative_prompt + + gen_params = self._resolve_default_sampling_params() + vp = request.resolve_video_params() + if vp.width is not None and vp.height is not None: + gen_params.width = vp.width + gen_params.height = vp.height + if vp.num_frames is not None: + gen_params.num_frames = vp.num_frames + if vp.fps is not None: + gen_params.fps = vp.fps + gen_params.frame_rate = float(vp.fps) + + provided_fields = request.model_fields_set + if "num_inference_steps" in provided_fields and request.num_inference_steps is not None: + gen_params.num_inference_steps = request.num_inference_steps + if "guidance_scale" in provided_fields and request.guidance_scale is not None: + gen_params.guidance_scale = request.guidance_scale + if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None: + gen_params.guidance_scale_2 = request.guidance_scale_2 + if "true_cfg_scale" in provided_fields and request.true_cfg_scale is not None: + gen_params.true_cfg_scale = request.true_cfg_scale + if "seed" in provided_fields and request.seed is not None: + gen_params.seed = request.seed + if "boundary_ratio" in provided_fields and request.boundary_ratio is not None: + gen_params.boundary_ratio = request.boundary_ratio + if "flow_shift" in provided_fields and request.flow_shift is not None: + gen_params.extra_args["flow_shift"] = request.flow_shift + + if request.extra_params is not None: + if not isinstance(request.extra_params, dict): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="extra_params must be a JSON object/dict.", + ) + gen_params.extra_args.update(request.extra_params) + + self._apply_lora(request.lora, gen_params) + return prompt, gen_params, vp + + def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams: + default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None) + if default_sampling_params_list: + for params in default_sampling_params_list: + if isinstance(params, OmniDiffusionSamplingParams): + return copy.deepcopy(params) + return OmniDiffusionSamplingParams() + + @staticmethod + def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None: + try: + lora_request, lora_scale = parse_lora_request(lora_body) + except ValueError as e: + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail=str(e), + ) from e + + if lora_request is None: + return + + gen_params.lora_request = lora_request + if lora_scale is not None: + gen_params.lora_scale = lora_scale + + async def _iter_generation_outputs( + self, + prompt: OmniTextPrompt, + gen_params: OmniDiffusionSamplingParams, + request_id: str, + progress: _SessionProgress, + ) -> AsyncGenerator[Any, None]: + stage_configs = self._stage_configs or getattr(self._engine_client, "stage_configs", None) + if not stage_configs: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail="Stage configs not found. Start server with an omni diffusion model.", + ) + + for stage in stage_configs: + stage_type = get_stage_type(stage) + if stage_type != "diffusion": + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE.value, + detail=f"Video generation only supports diffusion stages, found '{stage_type}' stage.", + ) + + engine_client = cast(Any, self._engine_client) + sampling_params_list = build_stage_sampling_params_list( + list(stage_configs), + get_default_sampling_params_list(engine_client), + diffusion_params=gen_params, + replace_diffusion_params=True, + ) + + async for output in engine_client.generate( + prompt=prompt, + request_id=request_id, + sampling_params_list=sampling_params_list, + ): + progress.touch() + yield output + + @staticmethod + def _normalize_video_outputs(videos: Any) -> list[Any]: + if videos is None: + return [] + if hasattr(videos, "ndim") and videos.ndim == 5: + return [videos[i] for i in range(videos.shape[0])] + if isinstance(videos, list): + if not videos: + return [] + first = videos[0] + if hasattr(first, "ndim") and first.ndim == 5: + flattened: list[Any] = [] + for item in videos: + if hasattr(item, "ndim") and item.ndim == 5: + flattened.extend([item[i] for i in range(item.shape[0])]) + else: + flattened.append(item) + return flattened + if isinstance(first, list): + return videos + if hasattr(first, "ndim") and first.ndim == 3: + return [videos] + return videos + return [videos] + + def _extract_video_outputs(self, result: Any) -> list[Any]: + videos = None + if hasattr(result, "images") and result.images: + videos = result.images + elif hasattr(result, "request_output"): + request_output = result.request_output + if isinstance(request_output, dict) and request_output.get("images"): + videos = request_output["images"] + elif not isinstance(request_output, dict): + if hasattr(request_output, "images") and request_output.images: + videos = request_output.images + elif hasattr(request_output, "multimodal_output") and request_output.multimodal_output: + videos = request_output.multimodal_output.get("video") + if videos is None and hasattr(result, "multimodal_output") and result.multimodal_output: + videos = result.multimodal_output.get("video") + + normalized = self._normalize_video_outputs(videos) + if not normalized: + raise HTTPException( + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, + detail="No video outputs found in generation result.", + ) + return normalized + + @staticmethod + async def _send_error( + websocket: WebSocket, + message: str, + code: int | None = None, + *, + send_lock: asyncio.Lock | None = None, + ) -> None: + try: + payload: dict[str, Any] = {"type": "error", "message": message} + if code is not None: + payload["code"] = code + if send_lock is None: + await websocket.send_json(payload) + else: + async with send_lock: + await websocket.send_json(payload) + except Exception: + pass diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py index 0c1fb52e9d2..af0083fcb70 100644 --- a/vllm_omni/entrypoints/openai/serving_video_stream.py +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""WebSocket handler for streaming video input understanding. +"""Qwen-Omni streaming video WebSocket handler. Accepts video frames incrementally via WebSocket, buffers them, and -generates text + optional audio responses using the existing Qwen3-Omni -multi-stage pipeline (thinker -> talker -> code2wav). +generates text + optional audio responses using the Qwen3-Omni multi-stage +pipeline (thinker -> talker -> code2wav). Protocol: Client -> Server: @@ -24,647 +24,35 @@ {"type": "error", "message": "..."} """ -import asyncio -import base64 -import hashlib -import io -import json -import time as _time -import uuid -import wave -from collections.abc import Mapping -from typing import Any +from __future__ import annotations -import torch -from fastapi import WebSocket, WebSocketDisconnect -from PIL import Image -from pydantic import BaseModel, Field, ValidationError -from vllm.logger import init_logger +from typing import Any -from vllm_omni.entrypoints.openai import video_stream_envs -from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter -from vllm_omni.entrypoints.openai.video_stream_context import ( - text_only_message, +from vllm_omni.entrypoints.openai.video_stream_base import ( + _BAD_FRAME, + _DEFAULT_CONFIG_TIMEOUT, + _DEFAULT_IDLE_TIMEOUT, + StreamingVideoSessionConfig, + VideoStreamTurnTrigger, +) +from vllm_omni.entrypoints.openai.video_stream_base import ( + OmniStreamingVideoHandler as OmniStreamingVideoHandlerBase, ) -from vllm_omni.outputs import OmniRequestOutput - -logger = init_logger(__name__) - -_DEFAULT_IDLE_TIMEOUT = 60.0 -_DEFAULT_CONFIG_TIMEOUT = 10.0 -_MAX_FRAME_SIZE = 10 * 1024 * 1024 # 10MB per frame -_MAX_BUFFER_FRAMES = 64 -_MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 -_MAX_MSG_QUEUE = 200 -_CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length -_BAD_FRAME = object() - - -def _decode_frame_bytes(raw_bytes: bytes) -> Any: - return Image.open(io.BytesIO(raw_bytes)).convert("RGB") - - -class StreamingVideoSessionConfig(BaseModel): - """Configuration sent as the first WebSocket message.""" - - model: str | None = None - modalities: list[str] = Field( - default_factory=lambda: ["text", "audio"], - description="Output modalities: 'text', 'audio', or both.", - ) - num_frames: int = Field( - default=4, - ge=1, - le=128, - description="Max frames to sample from buffer for the model.", - ) - max_frames: int = Field( - default=50, - ge=1, - le=256, - description="Max frames to keep in the buffer.", - ) - system_prompt: str | None = Field( - default=None, - description="Custom system prompt.", - ) - use_audio_in_video: bool = Field( - default=True, - description="Interleave audio chunks with video frames when audio input is present.", - ) - sampling_params_list: list[dict[str, Any]] | None = Field( - default=None, - description="Per-stage sampling params [thinker, talker, code2wav].", - ) - enable_frame_filter: bool = Field( - default=True, - description="EVS pixel-similarity pre-filter to drop near-duplicate frames.", - ) - frame_filter_threshold: float = Field( - default=0.95, - ge=0.0, - le=1.0, - description="EVS similarity threshold (higher = keep more frames).", - ) - - -class OmniStreamingVideoHandler: - """Handles WebSocket sessions for streaming video input. - - Supports: - - Concurrent frame reception during query processing (reader/processor split) - - PCM audio input (``audio.chunk``) - - Async-chunk incremental audio output via ``engine_client.generate()`` - - Multi-turn conversation history - - Soft interrupt (new query cancels current generation) - """ - - def __init__( - self, - chat_service: Any, - idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, - config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, - engine_client: Any | None = None, - ) -> None: - self._chat_service = chat_service - self._idle_timeout = idle_timeout - self._config_timeout = config_timeout - self._engine_client = engine_client - - async def handle_session(self, websocket: WebSocket) -> None: - """Main session loop for a single WebSocket connection.""" - await websocket.accept() - - try: - config = await self._receive_config(websocket) - if config is None: - return - - frame_buffer: list[str] = [] # base64-encoded JPEG frames - # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. - frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME - frame_filter = ( - FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None - ) - audio_buffer = bytearray() # raw PCM16 16kHz mono - message_history: list[dict[str, Any]] = [] - active_request_id: str | None = None - prev_request_id: str | None = None # abort target iff prev was interrupted - prev_was_interrupted: bool = False - interrupt_event = asyncio.Event() - prewarm_tasks: set[asyncio.Task[Any]] = set() - query_task: asyncio.Task[Any] | None = None - - msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=_MAX_MSG_QUEUE) - - async def _reader() -> None: - """Receive WebSocket messages and enqueue them.""" - try: - while True: - try: - raw = await asyncio.wait_for( - websocket.receive_text(), - timeout=self._idle_timeout, - ) - except asyncio.TimeoutError: - await self._send_error(websocket, "Idle timeout") - await msg_queue.put(None) - return - - try: - msg = json.loads(raw) - except json.JSONDecodeError: - await self._send_error(websocket, "Invalid JSON") - continue - - if not isinstance(msg, dict): - await self._send_error(websocket, "Messages must be JSON objects") - continue - - msg_type = str(msg.get("type", "")) - if msg_type.startswith("_internal."): - await self._send_error(websocket, f"Unknown type: {msg_type}") - continue - - await msg_queue.put(msg) - if msg.get("type") == "video.done": - return - except WebSocketDisconnect: - await msg_queue.put(None) - except Exception: - await msg_queue.put(None) - raise - - async def _cancel_active_query(*, abort_now: bool = False) -> None: - """Signal soft interrupt for the active query.""" - nonlocal active_request_id, prev_was_interrupted, query_task - if active_request_id is not None: - interrupt_event.set() - prev_was_interrupted = True - logger.info("Interrupt signaled for %s", active_request_id) - if abort_now and self._engine_client: - try: - await self._engine_client.abort(active_request_id) - except Exception: - logger.debug("Abort failed for %s", active_request_id, exc_info=True) - if query_task is not None and not query_task.done(): - query_task.cancel() - await asyncio.gather(query_task, return_exceptions=True) - query_task = None - - async def _processor() -> None: - """Process enqueued messages.""" - nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task - - while True: - msg = await msg_queue.get() - if msg is None: - await _cancel_active_query(abort_now=True) - return - - msg_type = msg.get("type") - - if msg_type == "_internal.frame_decode_failed": - frame_data = msg.get("b64", "") - removed = frame_data in frame_buffer - if removed: - frame_buffer[:] = [f for f in frame_buffer if f != frame_data] - if frame_pil_cache.get(frame_data) is _BAD_FRAME: - frame_pil_cache.pop(frame_data, None) - if removed: - await self._send_error(websocket, "Frame decode failed") - - elif msg_type == "video.frame": - frame_data = msg.get("data", "") - if not frame_data: - continue - if len(frame_data) > _MAX_FRAME_SIZE: - await self._send_error(websocket, "Frame too large") - continue - try: - raw_bytes = base64.b64decode(frame_data, validate=True) - except Exception: - await self._send_error(websocket, "Invalid image data") - continue - if frame_filter is not None: - try: - if not frame_filter.should_retain(raw_bytes): - continue - except Exception: - await self._send_error(websocket, "Invalid image data") - continue - max_buf = config.max_frames - if len(frame_buffer) >= max_buf: - dropped = frame_buffer.pop(0) - frame_pil_cache.pop(dropped, None) - frame_buffer.append(frame_data) - # Prewarm: decode PIL off the event loop so query-time chat_template - # can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames. - if frame_data not in frame_pil_cache: - mm_uuid = hashlib.md5(raw_bytes, usedforsecurity=False).hexdigest() - - async def _prewarm(b64: str, b: bytes, u: str) -> None: - try: - pil = await asyncio.to_thread(_decode_frame_bytes, b) - frame_pil_cache[b64] = (pil, u) - except Exception: - frame_pil_cache[b64] = _BAD_FRAME - logger.warning("prewarm decode failed for frame (len=%d)", len(b)) - try: - msg_queue.put_nowait({"type": "_internal.frame_decode_failed", "b64": b64}) - except asyncio.QueueFull: - logger.warning( - "frame decode failure event dropped because message queue is full" - ) - - task = asyncio.create_task(_prewarm(frame_data, raw_bytes, mm_uuid)) - prewarm_tasks.add(task) - task.add_done_callback(prewarm_tasks.discard) - - elif msg_type == "audio.chunk": - data_b64 = msg.get("data", "") - try: - pcm_bytes = base64.b64decode(data_b64) - except Exception: - continue - if len(audio_buffer) + len(pcm_bytes) > _MAX_AUDIO_BUFFER_BYTES: - await self._send_error(websocket, "Audio buffer overflow") - audio_buffer.clear() - continue - audio_buffer.extend(pcm_bytes) - - elif msg_type == "video.query": - await _cancel_active_query() - - query_text = msg.get("text", "") - audio_data_b64 = msg.get("audio_data") - if audio_data_b64: - try: - decoded = base64.b64decode(audio_data_b64) - if len(audio_buffer) + len(decoded) <= _MAX_AUDIO_BUFFER_BYTES: - audio_buffer.extend(decoded) - else: - await self._send_error(websocket, "Audio buffer overflow") - audio_buffer.clear() - except Exception: - pass - - if not frame_buffer: - await self._send_error(websocket, "No frames buffered") - continue - - # Abort only if the previous turn was interrupted mid-flight. - # A naturally-finished request is already released by the scheduler; - # aborting it again can race with stage-1/2 tear-down and has been - # observed to crash flash_attn with a mixed prefill+decode batch - # (scheduler_metadata shape mismatch) under longer sessions. - if prev_was_interrupted and prev_request_id and self._engine_client: - try: - await self._engine_client.abort(prev_request_id) - except Exception: - pass - await asyncio.sleep(0.1) - prev_was_interrupted = False - - request_id = f"video-{uuid.uuid4().hex[:12]}" - active_request_id = request_id - interrupt_event.clear() - query_frames = list(frame_buffer) - query_audio_buffer = bytearray(audio_buffer) - audio_buffer.clear() - query_prewarmed_frames = dict(frame_pil_cache) - - async def _run_query() -> None: - nonlocal active_request_id, prev_request_id - try: - await self._process_query( - websocket, - config, - query_frames, - query_audio_buffer, - message_history, - query_text, - request_id, - interrupt_event, - query_prewarmed_frames, - ) - finally: - if active_request_id == request_id: - prev_request_id = request_id - active_request_id = None - - query_task = asyncio.create_task(_run_query()) - - elif msg_type == "video.done": - if query_task is not None and not query_task.done(): - await asyncio.gather(query_task, return_exceptions=True) - query_task = None - await websocket.send_json({"type": "session.done"}) - return - - elif msg_type == "ping": - try: - await websocket.send_json({"type": "pong"}) - except Exception: - pass - - else: - await self._send_error(websocket, f"Unknown type: {msg_type}") - - reader_task = asyncio.create_task(_reader()) - try: - await _processor() - finally: - reader_task.cancel() - try: - await reader_task - except (asyncio.CancelledError, Exception): - pass - for t in list(prewarm_tasks): - t.cancel() - if prewarm_tasks: - await asyncio.gather(*prewarm_tasks, return_exceptions=True) - if query_task is not None and not query_task.done(): - await _cancel_active_query(abort_now=True) - - except WebSocketDisconnect: - logger.info("Streaming video: client disconnected") - except Exception as e: - logger.exception("Streaming video session error: %s", e) - try: - await self._send_error(websocket, f"Internal error: {e}") - except Exception: - pass - - async def _receive_config(self, websocket: WebSocket) -> StreamingVideoSessionConfig | None: - """Wait for and validate the session.config message.""" - try: - raw = await asyncio.wait_for( - websocket.receive_text(), - timeout=self._config_timeout, - ) - except asyncio.TimeoutError: - await self._send_error(websocket, "Timeout waiting for session.config") - return None - - try: - msg = json.loads(raw) - except json.JSONDecodeError: - await self._send_error(websocket, "Invalid JSON in session.config") - return None - - if not isinstance(msg, dict) or msg.get("type") != "session.config": - await self._send_error( - websocket, - f"Expected session.config, got: {msg.get('type') if isinstance(msg, dict) else type(msg).__name__}", - ) - return None - - config_data = {k: v for k, v in msg.items() if k != "type"} - alias_map = { - "num_sample_frames": "num_frames", - "evs_enabled": "enable_frame_filter", - "evs_threshold": "frame_filter_threshold", - } - for old_key, new_key in alias_map.items(): - if old_key in config_data and new_key not in config_data: - config_data[new_key] = config_data[old_key] - - try: - config = StreamingVideoSessionConfig(**config_data) - except ValidationError as e: - await self._send_error(websocket, f"Invalid session config: {e}") - return None - - return config - - async def _process_query( - self, - websocket: WebSocket, - config: StreamingVideoSessionConfig, - frame_buffer: list[str], - audio_buffer: bytearray, - message_history: list[dict[str, Any]], - query_text: str, - request_id: str, - interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], - ) -> None: - """Build prompt, run inference, stream text + audio response.""" - - if self._engine_client is None: - await self._send_error(websocket, "Streaming video requires an engine client") - return - - await self._process_query_engine( - websocket, - config, - frame_buffer, - audio_buffer, - message_history, - query_text, - request_id, - interrupt_event, - prewarmed_frames, - ) - - # ------------------------------------------------------------------ - # Engine-client path (async_chunk audio streaming) - # ------------------------------------------------------------------ - - async def _process_query_engine( - self, - websocket: WebSocket, - config: StreamingVideoSessionConfig, - frame_buffer: list[str], - audio_buffer: bytearray, - message_history: list[dict[str, Any]], - query_text: str, - request_id: str, - interrupt_event: asyncio.Event, - prewarmed_frames: dict[str, tuple[Any, str]], - ) -> None: - """Direct engine_client.generate() path for async_chunk audio.""" - from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionRequest, - ) - - messages, user_message = self._build_messages( - config, - frame_buffer, - audio_buffer, - message_history, - query_text, - prewarmed_frames, - ) - - request_kwargs: dict[str, Any] = { - "model": config.model or "default", - "messages": messages, - "stream": True, - "modalities": config.modalities, - "add_generation_prompt": True, - "continue_final_message": False, - "add_special_tokens": False, - } - if config.use_audio_in_video and len(audio_buffer) > 0: - request_kwargs["mm_processor_kwargs"] = { - "use_audio_in_video": True, - } - if config.sampling_params_list: - request_kwargs["sampling_params_list"] = config.sampling_params_list - - try: - chat_request = ChatCompletionRequest(**request_kwargs) - except Exception as e: - await self._send_error(websocket, f"Failed to build request: {e}") - return - - try: - engine_prompt = await self._preprocess_to_engine_prompt(chat_request) - except Exception as e: - await self._send_error(websocket, f"Preprocess failed: {e}") - return - - await websocket.send_json({"type": "response.start"}) - text_parts: list[str] = [] - text_done_sent = False - audio_chunk_count = 0 - # Number of per-step tensors in OmniRequestOutput.audio_data already - # drained. Used by the fast path to skip already-emitted history. - audio_chunks_drained = 0 - previous_text = "" - interrupted = False - t_start = _time.monotonic() - t_first_text = None - t_first_audio = None - - # Wire-level async-chunk switch. "off" means - # buffer all deltas server-side and flush once at the end; the engine - # pipeline still overlaps internally. - async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK - streaming = async_chunk_mode == "on" - audio_tail_tensors: list[Any] = [] - - try: - result_gen = self._engine_client.generate( - prompt=engine_prompt, - request_id=request_id, - output_modalities=config.modalities, - ) - - async for output in result_gen: - # Soft interrupt: drain without sending - if interrupt_event.is_set(): - if not interrupted: - logger.info("Generation interrupted — draining") - interrupted = True - continue - - if not isinstance(output, OmniRequestOutput): - continue - - out_type = getattr(output, "final_output_type", "text") - - if out_type == "audio": - if streaming and not text_done_sent: - full_text = "".join(text_parts) - await websocket.send_json({"type": "response.text.done", "text": full_text}) - text_done_sent = True - - if t_first_audio is None: - t_first_audio = _time.monotonic() - audio_chunk_count += 1 - if streaming: - b64, audio_chunks_drained = self._extract_audio_delta_b64( - output, - audio_chunks_drained, - ) - if b64: - await websocket.send_json( - { - "type": "response.audio.delta", - "data": b64, - "format": "wav", - } - ) - else: - audio_data = self._get_audio_data(output) - if audio_data is not None: - if isinstance(audio_data, list): - audio_tail_tensors = list(audio_data) - else: - audio_tail_tensors = [audio_data] - else: - delta_text, previous_text = self._extract_text_delta( - output, - previous_text, - ) - if delta_text: - if t_first_text is None: - t_first_text = _time.monotonic() - text_parts.append(delta_text) - if streaming: - await websocket.send_json({"type": "response.text.delta", "delta": delta_text}) - - if not text_done_sent: - full_text = "".join(text_parts) - await websocket.send_json({"type": "response.text.done", "text": full_text}) - text_done_sent = True - - if not streaming and audio_tail_tensors: - try: - coalesced = ( - audio_tail_tensors[0] if len(audio_tail_tensors) == 1 else torch.cat(audio_tail_tensors, dim=-1) - ) - tail_np = self._tensor_to_1d_np(coalesced) - b64, _ = self._encode_tail( - tail_np, - 0, - new_drained=len(audio_tail_tensors), - is_first=True, - ) - if b64: - await websocket.send_json( - { - "type": "response.audio.delta", - "data": b64, - "format": "wav", - } - ) - except Exception: - logger.exception("Failed to coalesce off-path audio") - - if audio_chunk_count > 0: - await websocket.send_json({"type": "response.audio.done"}) - - response_text = "".join(text_parts) - message_history.append(user_message) - message_history.append({"role": "assistant", "content": response_text}) - t_end = _time.monotonic() - logger.info( - "[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d", - async_chunk_mode, - t_end - t_start, - (t_first_text - t_start) if t_first_text else -1, - (t_first_audio - t_start) if t_first_audio else -1, - audio_chunk_count, - ) +__all__ = [ + "QwenOmniStreamingVideoHandler", + "StreamingVideoSessionConfig", + "create_streaming_video_handler", +] - except Exception: - logger.exception("Engine query failed") - await self._send_error(websocket, "Query processing failed") - if not text_done_sent: - full_text = "".join(text_parts) - await websocket.send_json({"type": "response.text.done", "text": full_text}) +class QwenOmniStreamingVideoHandler(OmniStreamingVideoHandlerBase): + """Qwen-Omni pipeline: manual ``video.query`` trigger and image_pil prompts.""" - # ------------------------------------------------------------------ - # Message building - # ------------------------------------------------------------------ + def should_trigger_turn(self, trigger: VideoStreamTurnTrigger) -> bool: + return False - def _build_messages( + def build_engine_prompt( self, config: StreamingVideoSessionConfig, frame_buffer: list[str], @@ -673,11 +61,6 @@ def _build_messages( query_text: str, prewarmed_frames: dict[str, tuple[Any, str]], ) -> tuple[list[dict[str, Any]], dict[str, Any]]: - """Build OpenAI-style messages list and the current user message. - - Returns (messages, user_message). - """ - # Stride sampling (index 0 anchor, last slot = newest). Covers full buffer + stable mm_hash. n_buf = len(frame_buffer) if n_buf <= config.num_frames: frames = list(frame_buffer) @@ -686,7 +69,6 @@ def _build_messages( idx = [i * stride for i in range(config.num_frames - 1)] + [n_buf - 1] frames = [frame_buffer[i] for i in idx] - # Prefer prewarmed PIL + uuid so mm_cache can dedupe by hash. prewarmed = prewarmed_frames or {} user_content: list[dict] = [] for frame_b64 in frames: @@ -731,12 +113,6 @@ def _build_messages( if config.system_prompt: messages.append({"role": "system", "content": config.system_prompt}) - # Add text-only history (strip images/audio from old turns). - # Keep only the last turn (2 messages) to keep prompt short - # enough for single-step mm_encoder scheduling. When prompt - # exceeds ~50 tokens, the V1 scheduler splits mm_encoder and - # text prefill, causing incomplete thinker embeddings and - # garbled audio. recent_history = message_history[-2:] if len(message_history) > 2 else message_history for hist_msg in recent_history: messages.append(self._text_only_message(hist_msg)) @@ -745,228 +121,32 @@ def _build_messages( return messages, user_message - # ------------------------------------------------------------------ - # Audio helpers - # ------------------------------------------------------------------ - - @staticmethod - def _pcm_to_wav_b64(pcm_data: bytes, sample_rate: int = 16000) -> str: - """Wrap raw PCM16 mono in a WAV container and return base64.""" - buf = io.BytesIO() - with wave.open(buf, "wb") as wf: - wf.setnchannels(1) - wf.setsampwidth(2) - wf.setframerate(sample_rate) - wf.writeframes(pcm_data) - return base64.b64encode(buf.getvalue()).decode() - - @classmethod - def _extract_audio_delta_b64( - cls, - result: OmniRequestOutput, - chunks_drained: int, - ) -> tuple[str | None, int]: - """Return (base64 WAV of new samples, updated chunks_drained). - - `chunks_drained` is the number of per-step tensors in - ``audio_data`` that have already been emitted. Each engine step appends - one tensor, so new samples are ``audio_data[chunks_drained:]`` — no - matter how many steps accumulated between reads (handles backpressure - cleanly, unlike a simple ``audio_data[-1]``). - - Two paths, selected at runtime by ``VLLM_VIDEO_AUDIO_DELTA_MODE``: - * fast — only D2H the new tail. Per-call cost ∝ new chunks. - * slow — full cat + D2H each call. Per-call cost ∝ total history. - Retained for A/B; remove once downstream callers confirm. - """ - audio_data = cls._get_audio_data(result) - if audio_data is None: - return None, chunks_drained - - if video_stream_envs.VLLM_VIDEO_AUDIO_DELTA_MODE == "slow": - return cls._delta_slow(audio_data, chunks_drained) - return cls._delta_fast(audio_data, chunks_drained) - - @staticmethod - def _get_audio_data(result: OmniRequestOutput): - """Navigate OmniRequestOutput → multimodal_output['audio']. None on miss.""" - request_output = getattr(result, "request_output", None) - if request_output is None: - return None - outputs = getattr(request_output, "outputs", None) - if not isinstance(outputs, list) or not outputs: - return None - mm_output = getattr(outputs[0], "multimodal_output", None) - if not isinstance(mm_output, Mapping): - return None - return mm_output.get("audio") - - @classmethod - def _delta_fast( - cls, - audio_data, - chunks_drained: int, - ) -> tuple[str | None, int]: - """Emit only tensors appended since the last call.""" - # Single tensor: output_processor hands us one tensor before it becomes a - # list (see output_processor.py:89). Treat it as chunk #0. - if not isinstance(audio_data, list): - if chunks_drained >= 1: - return None, chunks_drained - tail_np = cls._tensor_to_1d_np(audio_data) - return cls._encode_tail(tail_np, chunks_drained, new_drained=1, is_first=True) - - n = len(audio_data) - if n <= chunks_drained: - return None, chunks_drained - - new_chunks = audio_data[chunks_drained:] - tail = new_chunks[0] if len(new_chunks) == 1 else torch.cat(new_chunks, dim=-1) - tail_np = cls._tensor_to_1d_np(tail) - return cls._encode_tail(tail_np, chunks_drained, new_drained=n, is_first=(chunks_drained == 0)) - - @classmethod - def _delta_slow( - cls, - audio_data, - chunks_drained: int, - ) -> tuple[str | None, int]: - """Pre-fix behaviour: concat everything each call and slice on CPU.""" - if isinstance(audio_data, list): - if not audio_data: - return None, chunks_drained - audio_tensor = torch.cat(audio_data, dim=-1) - new_drained = len(audio_data) - else: - audio_tensor = audio_data - new_drained = 1 - - full_np = cls._tensor_to_1d_np(audio_tensor) - if full_np is None: - return None, chunks_drained - # chunks_drained doesn't map directly to sample offset without tracking - # per-chunk lengths, so we re-derive: replay the tail that corresponds - # to chunks appended since last call by slicing off the part produced - # by the already-drained prefix. For slow path this is intentionally - # wasteful — the point is to reproduce the pre-fix hot loop. - if chunks_drained == 0: - tail_np = full_np - else: - # Recover prefix length by re-concatenating the already-drained - # prefix tensors (cost intentionally identical to the baseline - # implementation this was lifted from). - if isinstance(audio_data, list) and chunks_drained < len(audio_data): - prefix_len = sum(int(t.shape[-1]) for t in audio_data[:chunks_drained]) - tail_np = full_np[prefix_len:] - else: - tail_np = full_np[0:0] - return cls._encode_tail(tail_np, chunks_drained, new_drained=new_drained, is_first=(chunks_drained == 0)) - - @classmethod - def _encode_tail( - cls, - tail_np, - old_drained: int, - *, - new_drained: int, - is_first: bool, - ) -> tuple[str | None, int]: - """Strip the CausalConv leading artifact on first emit, then b64-encode.""" - if tail_np is None or len(tail_np) == 0: - return None, new_drained - if is_first and len(tail_np) > _CODEC_FRAME_SAMPLES * 2: - tail_np = tail_np[_CODEC_FRAME_SAMPLES:] - if len(tail_np) == 0: - return None, new_drained - try: - return cls._encode_audio_wav_b64(tail_np), new_drained - except Exception: - logger.exception("Failed to encode audio delta WAV") - return None, old_drained - - @staticmethod - def _tensor_to_1d_np(t): - """Tensor → flat float32 numpy on CPU. None on failure.""" - if t is None or not hasattr(t, "float"): - return None - arr = t.float().detach().cpu().numpy() - if arr.ndim > 1: - arr = arr.flatten() - return arr - - @staticmethod - def _encode_audio_wav_b64(audio_np) -> str: - """Encode numpy float32 audio to base64 WAV (24kHz).""" - from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin - from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio - - audio_obj = CreateAudio( - audio_tensor=audio_np, - sample_rate=24000, - response_format="wav", - speed=1.0, - stream_format="audio", - base64_encode=True, - ) - mixin = AudioMixin() - resp = mixin.create_audio(audio_obj) - return resp.audio_data - - @staticmethod - def _extract_text_delta( - result: OmniRequestOutput, - previous_text: str, - ) -> tuple[str, str]: - """Extract incremental text delta from OmniRequestOutput.""" - if result.final_output_type != "text": - return "", previous_text - - request_output = getattr(result, "request_output", None) - if request_output is None: - return "", previous_text - - outputs = getattr(request_output, "outputs", None) - if not isinstance(outputs, list) or not outputs: - return "", previous_text - - text = getattr(outputs[0], "text", None) - if not isinstance(text, str) or not text: - return "", previous_text - - if text.startswith(previous_text): - return text[len(previous_text) :], text - return text, text - - # ------------------------------------------------------------------ - # Preprocessing - # ------------------------------------------------------------------ - - async def _preprocess_to_engine_prompt(self, request) -> Any: - """Use the chat handler's preprocessing to build an engine prompt.""" - handler = self._chat_service - renderer = handler.renderer + def on_turn_complete( + self, + message_history: list[dict[str, Any]], + user_message: dict[str, Any], + response_text: str, + ) -> None: + message_history.append(user_message) + message_history.append({"role": "assistant", "content": response_text}) - _conversation, engine_prompts = await handler._preprocess_chat( - request, - request.messages, - default_template=getattr(request, "chat_template", None) or handler.chat_template, - default_template_content_format=handler.chat_template_content_format, - renderer=renderer, - add_generation_prompt=request.add_generation_prompt, - continue_final_message=request.continue_final_message, - add_special_tokens=request.add_special_tokens, - ) - return engine_prompts[0] + _build_messages = build_engine_prompt - # ------------------------------------------------------------------ - # Utilities - # ------------------------------------------------------------------ - _text_only_message = staticmethod(text_only_message) +def create_streaming_video_handler( + chat_service: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + engine_client: Any | None = None, +) -> OmniStreamingVideoHandlerBase: + """Create the handler for ``/v1/video/chat/stream``. - async def _send_error(self, websocket: WebSocket, message: str) -> None: - """Send an error message to the client.""" - try: - await websocket.send_json({"type": "error", "message": message}) - except Exception: - pass + Returns :class:`QwenOmniStreamingVideoHandler` today. Additional pipelines + can be selected here in follow-up PRs. + """ + return QwenOmniStreamingVideoHandler( + chat_service=chat_service, + idle_timeout=idle_timeout, + config_timeout=config_timeout, + engine_client=engine_client, + ) diff --git a/vllm_omni/entrypoints/openai/speech_usage.py b/vllm_omni/entrypoints/openai/speech_usage.py new file mode 100644 index 00000000000..5e38405fd31 --- /dev/null +++ b/vllm_omni/entrypoints/openai/speech_usage.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Token-usage accounting for the Speech (``/v1/audio/speech``) API. + +Why this module exists (issue #4646) +------------------------------------ +For staged TTS models the engine prompt is a PLACEHOLDER: the serving layer +builds ``prompt_token_ids = [1] * prefill_len`` and lets the model rebuild the +real conditioning (text / ref_audio / ref_text) into ``inputs_embeds`` later. +So ``len(prompt_token_ids)`` is NOT a faithful count of what the caller sent. + +For Qwen3-TTS specifically the placeholder length mirrors the model prefill: + * CustomVoice/VoiceDesign: the full input text is embedded in the prefill, so + the placeholder length scales with the input text. + * Base in-context voice cloning: the prefill embeds only ``codec_bos`` + the + reference-audio codec frames; the input text is consumed incrementally + during DECODE, not prefill. So the placeholder length tracks the *reference + audio* and is independent of the input text. + +That is why ``usage.prompt_tokens`` (== ``len(prompt_token_ids)``) looked wrong +for Base: it counted reference-audio frames instead of the synthesized text. + +This module computes usage from the *semantic* inputs instead: + + input_tokens = text_tokens + audio_tokens + text_tokens -> tokens of ``input`` (+ ``instructions``); the text to speak + audio_tokens -> reference-audio codec frames used as voice-clone + conditioning, counted ONLY when in-context cloning is + actually active (see ``gate_audio_tokens``) + output_tokens = generated codec/audio tokens (stage-0 decode steps) + total_tokens = input_tokens + output_tokens + +Naming (``input_tokens``/``output_tokens``) follows OpenAI's ``speech.audio.done`` +event; the ``input_token_details`` breakdown follows OpenAI's realtime/chat +convention of never folding audio into an opaque text count. +""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from vllm_omni.entrypoints.openai.protocol.audio import ( + SpeechInputTokenDetails, + SpeechTokenUsage, +) + + +def _first(value: Any, default: Any = None) -> Any: + """Unwrap the singleton-list convention used by ``tts_params``. + + ``tts_params`` wraps scalars in 1-element lists (e.g. ``task_type=["Base"]``) + because the model side batches per request. This returns the inner scalar. + """ + if isinstance(value, (list, tuple)): + return value[0] if value else default + return value if value is not None else default + + +def gate_audio_tokens( + *, + task_type: str | None, + x_vector_only_mode: bool, + icl_mode_override: bool | None, + ref_code_length: Any, +) -> int: + """Reference-audio codec frames that actually enter the prefill. + + Mirrors the ``in_context_mode`` decision in + ``Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information``: + reference codec frames are prefilled ONLY for Base in-context voice cloning. + They are NOT prefilled for: + * CustomVoice / VoiceDesign (no reference audio at all), or + * Base ``x_vector_only_mode`` (the reference audio is reduced to a single + speaker embedding vector; no per-frame codec context is inserted). + + Counting ``ref_code_length`` outside the in-context path would re-introduce + the issue #4646 inconsistency in the other direction, so we gate it here. + """ + in_context = (task_type == "Base") and not x_vector_only_mode + # An explicit per-request ``voice_clone_prompt.icl_mode`` wins if present, + # but ONLY for the Base task -- the estimator consults ``icl_mode`` inside + # its ``task_type == "Base"`` branch only, so CustomVoice / VoiceDesign must + # never have audio counted even if such a flag leaks into their params. + if task_type == "Base" and icl_mode_override is not None: + in_context = bool(icl_mode_override) + if not in_context: + return 0 + try: + return max(0, int(ref_code_length)) if ref_code_length is not None else 0 + except (TypeError, ValueError): + return 0 + + +def qwen3_tts_input_token_details( + *, + input_text: str, + instructions: str | None, + tts_params: dict[str, Any], + count_text_tokens: Callable[[str], int], +) -> SpeechInputTokenDetails: + """Compute the input-token breakdown for a Qwen3-TTS request. + + ``count_text_tokens`` tokenizes a string with the model's *text* tokenizer + and returns the token count. ``tts_params`` is the finalized param dict from + the adapter build (it already carries ``ref_code_length`` for ICL clones and + the resolved ``task_type`` / ``x_vector_only_mode``). + """ + # Text tokens = the text to synthesize plus the style/emotion instructions, + # because Qwen3-TTS tokenizes and prepends the instruction block too. + text = input_text or "" + instr = instructions or "" + text_tokens = count_text_tokens(text) if text else 0 + if instr.strip(): + text_tokens += count_text_tokens(instr) + + # Audio tokens = reference codec frames, gated to the in-context clone path. + voice_clone_prompt = _first(tts_params.get("voice_clone_prompt"), None) + icl_override = None + if isinstance(voice_clone_prompt, dict): + icl_flag = voice_clone_prompt.get("icl_mode") + if isinstance(icl_flag, bool): + icl_override = icl_flag + audio_tokens = gate_audio_tokens( + task_type=_first(tts_params.get("task_type"), "CustomVoice"), + x_vector_only_mode=bool(_first(tts_params.get("x_vector_only_mode"), False)), + icl_mode_override=icl_override, + ref_code_length=_first(tts_params.get("ref_code_length"), None), + ) + return SpeechInputTokenDetails(text_tokens=int(text_tokens), audio_tokens=int(audio_tokens)) + + +def build_speech_usage(details: SpeechInputTokenDetails, output_tokens: int) -> SpeechTokenUsage: + """Assemble the final usage object from the input breakdown + output count.""" + input_tokens = int(details.text_tokens) + int(details.audio_tokens) + out = max(0, int(output_tokens)) + return SpeechTokenUsage( + input_tokens=input_tokens, + output_tokens=out, + total_tokens=input_tokens + out, + input_token_details=details, + ) + + +@dataclass +class SpeechOutputTokenCounter: + """Accumulates the count of generated output (codec) tokens off the stream. + + Why not count ``res.outputs[0].token_ids``? In the staged TTS pipeline the + token-generating stage (stage 0, the AR talker) is NOT surfaced to the audio + consumer loop — only the downstream code2wav (audio) outputs reach it, and + those carry an empty ``token_ids`` (their payload is a waveform, not tokens). + + The generated codec-token count is instead reported by the pipeline itself + in per-stage metrics on the final (finished) output: + ``res.metrics['stage_metrics'][<stage>]['num_tokens_out']``. Only the AR + stage has a non-zero ``num_tokens_out`` (code2wav reports 0), so we take the + max across stages. Tracking the max is safe because the metrics appear on + the terminal output (intermediate outputs carry empty metrics). + """ + + output_tokens: int = 0 + + def observe(self, res: Any) -> None: + metrics = getattr(res, "metrics", None) + if not isinstance(metrics, dict): + return + stage_metrics = metrics.get("stage_metrics") + if not isinstance(stage_metrics, dict): + return + for stage in stage_metrics.values(): + if not isinstance(stage, dict): + continue + n = stage.get("num_tokens_out") + if isinstance(n, int) and n > self.output_tokens: + self.output_tokens = n + + def total(self) -> int: + """Generated codec/audio token count (max stage ``num_tokens_out``).""" + return self.output_tokens diff --git a/vllm_omni/entrypoints/openai/storage.py b/vllm_omni/entrypoints/openai/storage.py index 52aef8948f1..a36688956e1 100644 --- a/vllm_omni/entrypoints/openai/storage.py +++ b/vllm_omni/entrypoints/openai/storage.py @@ -1,32 +1,78 @@ import asyncio +import contextlib import os +import stat +import time from abc import ABC, abstractmethod +from dataclasses import dataclass, field from tempfile import NamedTemporaryFile +from typing import Generic, Literal, TypeVar +from vllm.logger import init_logger -class StorageBaseManager(ABC): +from vllm_omni.config.server_settings import SERVER_SETTINGS_CONFIG, FileBackend + +logger = init_logger(__name__) + + +@dataclass +class SaveContext: + key: str + created_at: int + expires_at: int | None = None + + +@dataclass(frozen=True) +class BaseStorageHandle: + kind: str + + +@dataclass(frozen=True) +class FileStorageHandle(BaseStorageHandle): + path: str + kind: Literal["path"] = field(default="path", init=False) + + +K = TypeVar("K", bound=BaseStorageHandle, covariant=True) + + +class StorageBaseManager(Generic[K], ABC): @abstractmethod - async def save(self, *args, **kwargs) -> str: + async def save(self, *args, **kwargs) -> SaveContext: pass @abstractmethod async def delete(self, *args, **kwargs) -> bool: pass + async def start(self, *args, **kwargs) -> None: + pass + + async def stop(self, *args, **kwargs) -> None: + pass + + @abstractmethod + async def open(self, storage_key: str) -> K | None: + pass + -class LocalStorageManager(StorageBaseManager): - def __init__(self, storage_path: str | None = None): - if storage_path is None: - storage_path = os.getenv("VLLM_OMNI_STORAGE_PATH", "/tmp/storage") - self.storage_path = storage_path +class LocalStorageManager(StorageBaseManager[FileStorageHandle]): + def __init__(self, storage_path: str, max_concurrency: int = 4): + self.storage_path = os.path.realpath(storage_path) os.makedirs(self.storage_path, exist_ok=True) - max_concurrency = int(os.getenv("VLLM_OMNI_STORAGE_MAX_CONCURRENCY", "4")) self._io_semaphore = asyncio.Semaphore(max(1, max_concurrency)) - def _save_sync(self, data: bytes, file_name: str) -> str: + async def open(self, storage_key: str) -> FileStorageHandle | None: + local_file = self.get_full_file_path(storage_key) + if not os.path.exists(local_file): + return None + return FileStorageHandle(path=local_file) + + def _save_sync(self, data: bytes, file_name: str) -> SaveContext: filename = self.get_full_file_path(file_name) tmp_name: str | None = None + try: with NamedTemporaryFile("wb", dir=self.storage_path, delete=False) as f: tmp_name = f.name @@ -34,7 +80,8 @@ def _save_sync(self, data: bytes, file_name: str) -> str: f.flush() os.fsync(f.fileno()) os.replace(tmp_name, filename) - return filename + response = SaveContext(key=file_name, created_at=int(time.time())) + return response except Exception: if tmp_name is not None: try: @@ -43,7 +90,7 @@ def _save_sync(self, data: bytes, file_name: str) -> str: pass raise - async def save(self, data: bytes, file_name: str) -> str: + async def save(self, data: bytes, file_name: str) -> SaveContext: async with self._io_semaphore: return await asyncio.to_thread(self._save_sync, data, file_name) @@ -62,9 +109,102 @@ def exists(self, file_name: str) -> bool: return os.path.exists(self.get_full_file_path(file_name)) def get_full_file_path(self, file_name: str) -> str: - return os.path.join(self.storage_path, file_name) + full = os.path.realpath(os.path.join(self.storage_path, file_name)) + + if os.path.commonpath([self.storage_path, full]) != self.storage_path: + raise ValueError(f"Illegal storage key: {file_name!r}") + + return full + + +class LocalStorageTTLManager(LocalStorageManager): + def __init__(self, ttl_seconds: int, sweep_interval_seconds: int, *args, **kwargs): + if ttl_seconds <= 0: + raise ValueError("`ttl_seconds` must be greater than or equal to 1.") + if sweep_interval_seconds <= 0: + raise ValueError("`sweep_interval_seconds` must be greater than or equal to 1.") + + self._ttl_seconds = ttl_seconds + self._sweep_interval_seconds = sweep_interval_seconds + self._sweeper_task: asyncio.Task[None] | None = None + + super().__init__(*args, **kwargs) + + async def save(self, data: bytes, file_name: str) -> SaveContext: + result = await super().save(data, file_name) + result.expires_at = result.created_at + self._ttl_seconds + return result + + async def _sweep_once(self, cutoff: float) -> int: + expired: list[str] = [] + for entry in os.scandir(self.storage_path): + try: + if not entry.is_file(follow_symlinks=False): + continue + if entry.stat(follow_symlinks=False).st_mtime < cutoff: + expired.append(entry.path) + except FileNotFoundError: + logger.debug("TTL sweep skipped %s; file removed during scan", entry.path) + except OSError: + logger.warning("TTL sweep failed to inspect %s", entry.path, exc_info=True) + + deleted = 0 + for path in expired: + try: + async with self._io_semaphore: + st = await asyncio.to_thread(os.stat, path, follow_symlinks=False) + if not stat.S_ISREG(st.st_mode): + continue + if st.st_mtime >= cutoff: + continue + await asyncio.to_thread(os.remove, path) + deleted += 1 + except FileNotFoundError: + logger.debug("TTL sweep skipped %s; file already removed", path) + except OSError: + logger.warning("TTL sweep failed to delete expired file %s", path, exc_info=True) + + return deleted + + async def _sweep_loop(self) -> None: + while True: + try: + cutoff = time.time() - self._ttl_seconds + await self._sweep_once(cutoff) + except Exception: + logger.exception("TTL sweep failed for storage path %s", self.storage_path) + await asyncio.sleep(self._sweep_interval_seconds) + + async def start(self) -> None: + if self._sweeper_task is None or self._sweeper_task.done(): + self._sweeper_task = asyncio.create_task(self._sweep_loop()) + + async def stop(self) -> None: + if self._sweeper_task is None: + return + self._sweeper_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._sweeper_task + self._sweeper_task = None + + +def get_storage_manager(storage_config: FileBackend) -> StorageBaseManager[FileStorageHandle]: + if isinstance(storage_config, FileBackend): + if storage_config.file_ttl is not None and storage_config.ttl_sweep_interval is not None: + manager = LocalStorageTTLManager( + storage_path=storage_config.path, + max_concurrency=storage_config.file_concurrency, + ttl_seconds=storage_config.file_ttl, + sweep_interval_seconds=storage_config.ttl_sweep_interval, + ) + else: + manager = LocalStorageManager( + storage_path=storage_config.path, max_concurrency=storage_config.file_concurrency + ) + else: + raise ValueError("No supported storage managers") + + return manager -# Should implement TTL management on local file storage -# and recovery on restart in the *STORE's -STORAGE_MANAGER = LocalStorageManager() +STORAGE_MANAGER = get_storage_manager(SERVER_SETTINGS_CONFIG.storage) diff --git a/vllm_omni/entrypoints/openai/text_splitter.py b/vllm_omni/entrypoints/openai/text_splitter.py deleted file mode 100644 index 9c02de1396b..00000000000 --- a/vllm_omni/entrypoints/openai/text_splitter.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Multi-language sentence boundary detector for streaming TTS input. - -Buffers incoming text and splits at sentence boundaries (English and CJK), -yielding complete sentences for audio generation. -""" - -import re -from re import Pattern - -# Maximum buffer size (in characters) to prevent unbounded memory growth. -_MAX_BUFFER_SIZE = 100_000 # ~100 KB of text - -# Sentence-level: .!? + CJK sentence-ending 。！？ -# NOTE: English requires trailing whitespace to confirm a boundary — -# end-of-string is NOT treated as a boundary (that is what flush() is for). -SPLIT_SENTENCE = re.compile( - r"(?<=[.!?])\s+" - r"|(?<=[。！？])" -) - -# Clause-level: adds CJK commas ， and semicolons ； -SPLIT_CLAUSE = re.compile( - r"(?<=[.!?])\s+" - r"|(?<=[。！？，；])" -) - -# Default alias -_SENTENCE_BOUNDARY_RE = SPLIT_SENTENCE - - -class SentenceSplitter: - """Incremental sentence splitter for streaming text input. - - Buffers text and yields complete sentences when boundaries are detected. - Designed for TTS pipelines where text arrives incrementally (e.g., from STT). - - Args: - min_sentence_length: Minimum character length for a sentence. - Sentences shorter than this are kept in the buffer to avoid - splitting on abbreviations like "Dr." or "U.S.". - boundary_re: Custom compiled regex for sentence boundaries. - Use ``SPLIT_SENTENCE`` (default) for sentence-level splitting, - ``SPLIT_CLAUSE`` for finer-grained clause-level splitting, - or pass your own ``re.Pattern``. - """ - - def __init__( - self, - min_sentence_length: int = 2, - boundary_re: Pattern[str] | None = None, - ) -> None: - self._buffer: str = "" - self._min_sentence_length = min_sentence_length - self._boundary_re = boundary_re or _SENTENCE_BOUNDARY_RE - - @property - def buffer(self) -> str: - """Current buffered text.""" - return self._buffer - - def add_text(self, text: str) -> list[str]: - """Add text to the buffer and return any complete sentences. - - Args: - text: Incoming text chunk. - - Returns: - List of complete sentences extracted from the buffer. - May be empty if no sentence boundary was found. - - Raises: - ValueError: If the buffer exceeds the maximum size. - """ - if not text: - return [] - - self._buffer += text - if len(self._buffer) > _MAX_BUFFER_SIZE: - raise ValueError( - f"Text buffer exceeded maximum size ({_MAX_BUFFER_SIZE} chars). " - "Consider adding sentence-ending punctuation to your input." - ) - return self._extract_sentences() - - def flush(self) -> str | None: - """Flush remaining buffered text as a final sentence. - - Returns: - The remaining buffered text (stripped), or None if buffer is empty. - """ - remaining = self._buffer.strip() - self._buffer = "" - return remaining if remaining else None - - def _extract_sentences(self) -> list[str]: - """Split buffer at sentence boundaries, keeping incomplete text buffered.""" - parts = self._boundary_re.split(self._buffer) - - if len(parts) <= 1: - # No boundary found — keep everything in buffer - return [] - - sentences: list[str] = [] - carry = "" - # All parts except the last are complete sentences - for i in range(len(parts) - 1): - text = carry + parts[i] - carry = "" - stripped = text.strip() - if len(stripped) >= self._min_sentence_length: - sentences.append(stripped) - elif stripped: - # Too short (e.g. "Dr.") — carry forward to next part - carry = text - # else: empty, skip - - # Last part stays in buffer (may be incomplete) - self._buffer = carry + parts[-1] - - return sentences diff --git a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py index ca7b9612a6e..a2b9cce0915 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/__init__.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/__init__.py @@ -57,6 +57,7 @@ def resolve_adapter(model_type: str | None) -> type[TTSModelAdapter] | None: glm_tts, higgs_audio_v2, higgs_audio_v3, + indextts2, ming_tts, moss_tts, omnivoice, diff --git a/vllm_omni/entrypoints/openai/tts_adapters/base.py b/vllm_omni/entrypoints/openai/tts_adapters/base.py index b6d3d3e9f85..58c135044f7 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/base.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/base.py @@ -104,6 +104,10 @@ class TTSModelAdapter(ABC): #: Serving backend: ``"ar"`` (engine_client) or ``"diffusion"``. backend: ClassVar[str] = "ar" + max_new_tokens_min = 1 + + max_new_tokens_max = 4096 + def __init__(self, ctx: SpeechServingContext) -> None: self.ctx = ctx @@ -194,6 +198,3 @@ def extra_body_params(cls) -> frozenset[str]: "SpeechServingContext", "TTSModelAdapter", ] - -# Type alias kept for readability in adapter signatures. -NormalizeFn = Callable[["OpenAICreateSpeechRequest"], None] diff --git a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py index 6a7a89702f2..85648442f2b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py @@ -16,10 +16,32 @@ class CosyVoice3Adapter(ARTTSAdapter): name = "cosyvoice3" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - err = self.ctx.server._apply_uploaded_speaker(request) + """Validate CosyVoice3 request parameters. Returns error message or None.""" + server = self.ctx.server + err = server._apply_uploaded_speaker(request) if err: return err - return self.ctx.server._validate_cosyvoice3_request(request) + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + # CosyVoice3 requires reference audio for voice cloning + if request.ref_audio is None: + return "CosyVoice3 requires 'ref_audio' (reference audio for voice cloning)" + + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + + if not request.ref_text or not request.ref_text.strip(): + return "CosyVoice3 requires 'ref_text' (transcript of the reference audio)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py b/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py index f98f8961914..70673ac5454 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/covo_audio.py @@ -15,6 +15,11 @@ class CovoAudioAdapter(ARTTSAdapter): stage_keys = frozenset({"fused_thinker_talker"}) name = "covo_audio" + def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None + async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool ) -> PreparedRequest: diff --git a/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py b/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py index d216117b708..e8080a4275b 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/fish_speech.py @@ -16,10 +16,28 @@ class FishSpeechAdapter(ARTTSAdapter): name = "fish_tts" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - err = self.ctx.server._apply_uploaded_speaker(request) + """Validate Fish Speech request parameters. Returns error message or None.""" + server = self.ctx.server + err = server._apply_uploaded_speaker(request) if err: return err - return self.ctx.server._validate_fish_tts_request(request) + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.ref_audio is not None: + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not request.ref_text or not request.ref_text.strip(): + return "Voice cloning requires 'ref_text' (transcript of the reference audio)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py index e9e8700e63d..ddfd194e2fc 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/glm_tts.py @@ -16,10 +16,28 @@ class GlmTTSAdapter(ARTTSAdapter): name = "glm_tts" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - err = self.ctx.server._apply_uploaded_speaker(request) + """Validate GLM-TTS request — requires ref_audio for voice cloning.""" + server = self.ctx.server + err = server._apply_uploaded_speaker(request) if err: return err - return self.ctx.server._validate_glm_tts_request(request) + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.ref_audio is None: + return "GLM-TTS requires 'ref_audio' for zero-shot voice cloning" + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not request.ref_text or not request.ref_text.strip(): + return "GLM-TTS voice cloning requires 'ref_text' (transcript of the reference audio)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be >= {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py index 606b4bca177..99949d8b366 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v2.py @@ -16,10 +16,76 @@ class HiggsAudioV2Adapter(ARTTSAdapter): name = "higgs_audio_v2" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - err = self.ctx.server._apply_uploaded_speaker(request) + """Validate higgs_audio_v2 request parameters. Returns error message or None. + + Accepted: plain text -> speech, or shallow voice clone via ``ref_audio`` + + ``ref_text`` (both required together). Still out of scope: preset + ``voice``/``speaker`` selection, ``x_vector_only_mode`` / + ``speaker_embedding`` helpers, ``task_type``/``language``/ + ``instructions``/``speed`` overrides, and multi-speaker ``[SPEAKERn]`` + tags inside the input body. + """ + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import ( + MULTI_SPEAKER_TAG_PATTERN, + ) + + server = self.ctx.server + err = server._apply_uploaded_speaker(request) if err: return err - return self.ctx.server._validate_higgs_audio_v2_request(request) + + if not request.input or not request.input.strip(): + return "higgs_audio_v2: input text cannot be empty" + + # Voice clone: ref_audio and ref_text must come together. + if request.ref_audio is not None and not request.ref_text: + return ( + "higgs_audio_v2 voice clone requires both 'ref_audio' and " + "'ref_text'; received ref_audio without ref_text" + ) + if request.ref_text and request.ref_audio is None: + return ( + "higgs_audio_v2 voice clone requires both 'ref_audio' and " + "'ref_text'; received ref_text without ref_audio" + ) + + if request.x_vector_only_mode is not None: + return "higgs_audio_v2 v1 does not support 'x_vector_only_mode' (voice-cloning helper field)" + if request.speaker_embedding is not None: + return "higgs_audio_v2 v1 does not support 'speaker_embedding' (voice-cloning helper field)" + if request.voice and request.ref_audio is None: + # _apply_uploaded_speaker runs before this validator; if voice was + # an uploaded speaker, ref_audio is now populated and ref_text is + # backfilled from the speaker entry. A bare voice= with no + # ref_audio means the name didn't resolve to an uploaded speaker + # (and higgs has no built-in preset voices). + return ( + "higgs_audio_v2 v1 does not support 'voice'/'speaker' selection for built-in voices; " + f"upload a voice first via POST /v1/audio/voices, or use ref_audio + ref_text. " + f"Got voice={request.voice!r}" + ) + if request.instructions: + return ( + "higgs_audio_v2 v1 does not support 'instructions' (voice " + "style/emotion control); supply plain text instead" + ) + if request.task_type is not None: + return "higgs_audio_v2 v1 does not support 'task_type'; the model is single-mode plain text -> speech" + if request.language is not None: + return ( + "higgs_audio_v2 v1 does not accept 'language' overrides; the model infers language from the input text" + ) + if request.speed is not None and request.speed != 1.0: + return ( + "higgs_audio_v2 v1 does not support 'speed' adjustments; the audio is rendered at native rate (24 kHz)" + ) + + if MULTI_SPEAKER_TAG_PATTERN.search(request.input): + return "higgs_audio_v2 v1 does not support multi-speaker [SPEAKERn] tags; remove the tag from the input" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py index dde5e0815fa..9cae9aa2d1a 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/higgs_audio_v3.py @@ -16,10 +16,20 @@ class HiggsAudioV3Adapter(ARTTSAdapter): name = "higgs_audio_v3" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - err = self.ctx.server._apply_uploaded_speaker(request) + """Validate higgs_audio_v3 request parameters.""" + server = self.ctx.server + err = server._apply_uploaded_speaker(request) if err: return err - return self.ctx.server._validate_higgs_audio_v3_request(request) + 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 < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py new file mode 100644 index 00000000000..641e380da85 --- /dev/null +++ b/vllm_omni/entrypoints/openai/tts_adapters/indextts2.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +"""IndexTTS2 serving adapter.""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +from vllm.inputs import tokens_input +from vllm.utils import random_uuid + +from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter +from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest + +if TYPE_CHECKING: + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + +_TTS_MAX_NEW_TOKENS_MIN = 1 +_TTS_MAX_NEW_TOKENS_MAX = 4096 +_INDEXTTS2_EMOTION_KEYS = ("emo_audio", "emo_vector", "emo_alpha", "emo_text", "use_emo_text", "use_random") + + +def _update_conditioning_hash(h: hashlib._Hash, value: Any) -> None: + """Hash request conditioning values without serializing huge repr strings.""" + if value is None: + h.update(b"none") + return + if isinstance(value, torch.Tensor): + arr = value.detach().cpu().numpy() + h.update(b"tensor") + h.update(repr((arr.shape, str(arr.dtype))).encode("utf-8")) + h.update(np.ascontiguousarray(arr).tobytes()) + return + if isinstance(value, np.ndarray): + arr = np.ascontiguousarray(value) + h.update(b"ndarray") + h.update(repr((arr.shape, str(arr.dtype))).encode("utf-8")) + h.update(arr.tobytes()) + return + if isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[1], (int, np.integer)): + samples = value[0] + if isinstance(samples, (list, tuple, np.ndarray, torch.Tensor)): + if isinstance(samples, torch.Tensor): + samples_arr = samples.detach().cpu().numpy() + else: + samples_arr = np.asarray(samples, dtype=np.float32) + samples_arr = np.ascontiguousarray(samples_arr, dtype=np.float32) + h.update(b"audio") + h.update(int(value[1]).to_bytes(8, byteorder="little", signed=True)) + h.update(int(samples_arr.size).to_bytes(8, byteorder="little", signed=False)) + h.update(samples_arr.tobytes()) + return + if isinstance(value, Mapping): + h.update(b"mapping") + for key in sorted(value): + h.update(b"\x01") + h.update(repr(key).encode("utf-8")) + h.update(b"\x02") + _update_conditioning_hash(h, value[key]) + return + if isinstance(value, (list, tuple)): + h.update(f"sequence:{len(value)}".encode()) + for item in value: + h.update(b"\x03") + _update_conditioning_hash(h, item) + return + h.update(repr(value).encode("utf-8")) + + +def _first_conditioning_value(value: Any) -> Any: + if isinstance(value, list) and len(value) == 1: + return value[0] + return value + + +def indextts2_conditioning_cache_salt( + request: OpenAICreateSpeechRequest, + tts_params: dict[str, Any], + *, + request_id: str | None = None, +) -> str: + """Stable prefix-cache salt for IndexTTS2 placeholder prompts. + + IndexTTS2 builds prefill embeddings from ``additional_information`` (speaker + audio plus emotion controls), while the scheduler only sees dummy token ids. + Include all supported control inputs so prefix-cache hits are only possible + when text, speaker, and emotion conditioning are identical. + """ + h = hashlib.sha256() + for key in ( + "text", + "voice", + "voice_name", + "voice_created_at", + "emo_audio", + "emo_vector", + "emo_alpha", + "emo_text", + "use_emo_text", + "use_random", + ): + h.update(b"\x00") + h.update(key.encode("utf-8")) + h.update(b"\x00") + _update_conditioning_hash(h, tts_params.get(key)) + + for key, value in ( + ("input", request.input), + ("ref_audio", request.ref_audio), + ): + h.update(b"\x00") + h.update(key.encode("utf-8")) + h.update(b"\x00") + _update_conditioning_hash(h, value) + + use_random = bool(_first_conditioning_value(tts_params.get("use_random"))) + if use_random: + h.update(b"\x00indextts2-use-random-request\x00") + h.update((request_id or random_uuid()).encode("utf-8")) + return h.hexdigest()[:32] + + +@register_tts_adapter +class IndexTTS2Adapter(ARTTSAdapter): + stage_keys = frozenset({"indextts2_talker"}) + name = "indextts2" + + def validate(self, request: OpenAICreateSpeechRequest) -> str | None: + server = self.ctx.server + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + has_uploaded_voice = bool( + voice_lower + and voice_lower in server.uploaded_speakers + and server._get_uploaded_audio_data(voice_lower) is not None + ) + if request.ref_audio is None and not has_uploaded_voice: + return "IndexTTS2 requires 'ref_audio' or an uploaded audio voice for voice cloning" + if request.ref_audio is not None: + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + 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}" + if request.extra_params is not None: + if not isinstance(request.extra_params, dict): + return "extra_params must be a JSON object/dict." + extra_err = self._validate_extra_params(request.extra_params) + if extra_err: + return extra_err + return None + + def _validate_extra_params(self, extras: Mapping[str, Any]) -> str | None: + server = self.ctx.server + if "emo_audio" in extras: + value = extras["emo_audio"] + if not isinstance(value, str): + return "extra_params.emo_audio must be a URL (http/https), base64 data URL (data:...), or file URI (file://...)" + fmt_err = server._validate_ref_audio_format(value) + if fmt_err: + return fmt_err.replace("ref_audio", "extra_params.emo_audio", 1) + + if "emo_vector" in extras: + value = extras["emo_vector"] + if not isinstance(value, list) or len(value) != 8: + return "extra_params.emo_vector must be a list of 8 numbers" + for item in value: + if isinstance(item, bool) or not isinstance(item, (int, float)) or not math.isfinite(float(item)): + return "extra_params.emo_vector must be a list of 8 finite numbers" + if float(item) < 0.0 or float(item) > 1.2: + return "extra_params.emo_vector values must be in [0, 1.2]" + + if "emo_alpha" in extras: + value = extras["emo_alpha"] + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(float(value)): + return "extra_params.emo_alpha must be a finite number" + + if "emo_text" in extras and extras["emo_text"] is not None and not isinstance(extras["emo_text"], str): + return "extra_params.emo_text must be a string" + + for key in ("use_emo_text", "use_random"): + if key in extras and not isinstance(extras[key], bool): + return f"extra_params.{key} must be a boolean" + return None + + async def build( + self, + request: OpenAICreateSpeechRequest, + sampling_params_list: list, + has_inline_ref_audio: bool, + ) -> PreparedRequest: + del sampling_params_list, has_inline_ref_audio + + server = self.ctx.server + tts_params = await self._build_params(request) + from vllm_omni.model_executor.models.indextts2.prompt_utils import ( + estimate_indextts2_prefill_prompt_len, + ) + + ph_len = estimate_indextts2_prefill_prompt_len( + server.engine_client.model_config.model, + request.input, + ) + prompt = tokens_input(prompt_token_ids=[1] * ph_len) + prompt["additional_information"] = tts_params + prompt["cache_salt"] = indextts2_conditioning_cache_salt(request, tts_params) + return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="indextts2") + + async def _build_params(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + server = self.ctx.server + params: dict[str, Any] = {"text": [request.input]} + voice_lower = request.voice.lower() if isinstance(request.voice, str) else None + ref_audio_source = request.ref_audio + using_uploaded_voice = False + if ref_audio_source is None and voice_lower and voice_lower in server.uploaded_speakers: + ref_audio_source = server._get_uploaded_audio_data(voice_lower) + using_uploaded_voice = ref_audio_source is not None + if ref_audio_source is not None and isinstance(ref_audio_source, str): + wav_list, sr = await server._resolve_ref_audio(ref_audio_source) + params["voice"] = [[wav_list, sr]] + if using_uploaded_voice and voice_lower: + params["voice_name"] = [voice_lower] + params["voice_created_at"] = [server._voice_created_at(voice_lower)] + + extras = request.extra_params if isinstance(request.extra_params, dict) else {} + for key in _INDEXTTS2_EMOTION_KEYS: + if key not in extras: + continue + if key == "emo_audio": + wav_list, sr = await server._resolve_ref_audio(extras[key]) + params[key] = [[wav_list, sr]] + else: + params[key] = [extras[key]] + return params diff --git a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py index 37318a521fb..c736afb51a6 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/ming_tts.py @@ -9,6 +9,10 @@ if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm.logger import init_logger + +logger = init_logger(__name__) + @register_tts_adapter class MingTTSAdapter(ARTTSAdapter): @@ -16,7 +20,13 @@ class MingTTSAdapter(ARTTSAdapter): name = "ming_tts" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - return self.ctx.server._validate_ming_tts_request(request) + """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) async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool @@ -25,9 +35,18 @@ async def build( 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 server.uploaded_speakers: - ref_audio_source = server._get_uploaded_audio_data(request.voice) - if request.ref_text is None: - request.ref_text = server.uploaded_speakers[voice_lower].get("ref_text") + speaker_info = server.uploaded_speakers[voice_lower] + if speaker_info.get("embedding_source") == "direct": + if request.speaker_embedding is None: + request.speaker_embedding = server._get_uploaded_speaker_embedding(request.voice) + if request.speaker_embedding is None: + raise ValueError(f"Speaker embedding for uploaded voice '{request.voice}' is missing") + else: + ref_audio_source = server._get_uploaded_audio_data(request.voice) + if not ref_audio_source: + raise ValueError(f"Audio file for uploaded voice '{request.voice}' is missing") + if request.ref_text is None: + request.ref_text = speaker_info.get("ref_text") ref_audio_data = None if isinstance(ref_audio_source, list): ref_audio_data = await server._resolve_ref_audio_many(ref_audio_source) @@ -42,3 +61,82 @@ async def build( tts_params = prompt.get("additional_information", {}) # Ming stop-token / max_tokens sampling stays in the orchestrator tail. return PreparedRequest(prompt=prompt, tts_params=tts_params, model_type="ming_tts") + + def _validate_ming_tts_single_speaker_request(self, request: "OpenAICreateSpeechRequest") -> str | None: + server = self.ctx.server + if request.ref_audio is not None: + fmt_err = server._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 server.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'" + + 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) > server._max_instructions_length: + return f"Instructions too long (max {server._max_instructions_length} characters)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None + + def _validate_ming_tts_podcast_request(self, request: "OpenAICreateSpeechRequest") -> str | None: + server = self.ctx.server + 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 = server._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) > server._max_instructions_length: + return f"Instructions too long (max {server._max_instructions_length} characters)" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None diff --git a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py index 62d0c6568fb..ad43a1c0184 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/moss_tts.py @@ -18,10 +18,69 @@ class _MossTTSAdapterBase(ARTTSAdapter): def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - err = self.ctx.server._apply_uploaded_speaker(request) + """Validate any MOSS-TTS-family request (nano + 5 full variants). + + Dispatches by ``self._moss_variant``: + - ``tts``/``realtime``: require ``ref_audio`` (voice cloning). + - ``ttsd``: require ``ref_audio`` (speaker 1); ``ref_audio_2`` + optional (defaults to the same ref for both speakers). + - ``sound_effect``: require ``ambient_sound`` (no ref_audio). + - ``voice_generator``: require ``instructions`` (no ref_audio). + - For the legacy moss_tts_nano model_type the variant is None and + we fall through to the original nano contract (ref_audio only). + """ + server = self.ctx.server + err = server._apply_uploaded_speaker(request) if err: return err - return self.ctx.server._validate_moss_tts_request(request) + + if not request.input or not request.input.strip(): + # SoundEffect can legitimately have empty input (just ambient_sound). + if server._moss_variant != "sound_effect": + return "Input text cannot be empty" + + v = server._moss_variant + if v in (None, "tts", "realtime", "local"): + if request.ref_audio is None: + label = ( + "MOSS-TTS-Nano" + if v is None + else ( + "MOSS-TTS-Realtime" + if v == "realtime" + else ("MOSS-TTS-Local-Transformer" if v == "local" else "MOSS-TTS") + ) + ) + return f"{label} requires 'ref_audio' (reference audio for voice cloning)." + return server._validate_ref_audio_format(request.ref_audio) + + if v == "ttsd": + if request.ref_audio is None: + return "MOSS-TTSD requires 'ref_audio' (speaker 1 reference)." + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if request.ref_audio_2 is not None: + return server._validate_ref_audio_format(request.ref_audio_2) + return None + + if v == "sound_effect": + if not request.ambient_sound or not request.ambient_sound.strip(): + return ( + "MOSS-SoundEffect requires 'ambient_sound' (natural language " + "description of the sound effect to synthesise)." + ) + return None + + if v == "voice_generator": + if not request.instructions or not request.instructions.strip(): + return ( + "MOSS-VoiceGenerator requires 'instructions' (natural language " + "voice description, e.g. 'a warm female voice with an American accent')." + ) + return None + + return None # unreachable async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool @@ -54,5 +113,5 @@ class MossTTSNanoAdapter(_MossTTSAdapterBase): @register_tts_adapter class MossTTSAdapter(_MossTTSAdapterBase): - stage_keys = frozenset({"moss_tts", "moss_tts_codec"}) + stage_keys = frozenset({"moss_tts", "moss_tts_codec", "moss_tts_local", "moss_tts_local_codec"}) name = "moss_tts" diff --git a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py index 0f62a169ce4..c0be8c6855c 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py @@ -1,18 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -"""Qwen3-TTS serving adapter. - -First model migrated onto the adapter framework (RFC #4327). It is the dispatch -surface for Qwen3-TTS: the serving orchestrator routes prompt/param building -through ``build()`` instead of the inline ``else`` branch in -``_prepare_speech_generation``. - -To stay behaviour-identical and avoid divergence from the batch path (which -shares ``_validate_qwen_tts_request`` and ``_build_tts_params``), the adapter -reuses the single-source helper implementations on the serving instance through -``ctx.server`` rather than copying them. Subsequent PRs relocate those bodies as -more models migrate and the batch path also routes through adapters. -""" +"""Qwen3-TTS serving adapter.""" +from pathlib import Path from typing import TYPE_CHECKING from vllm.logger import init_logger @@ -36,13 +25,133 @@ class Qwen3TTSAdapter(ARTTSAdapter): def normalize(self, request: "OpenAICreateSpeechRequest") -> None: """Qwen3-TTS normalization (Base-task inference, voice lowercasing) is performed inside ``validate`` today; kept fused for a strict behaviour - match. See ``_validate_qwen_tts_request``.""" + match.""" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - # Go through the shared dispatcher (which routes to - # _validate_qwen_tts_request for this model-type) rather than the leaf - # validator directly, matching the pre-adapter qwen3 path. - return self.ctx.server._validate_tts_request(request) + """Validate Qwen TTS request parameters. Returns error message or None.""" + # Infer Base task when ref_audio or ref_text is provided without explicit task_type. + server = self.ctx.server + if request.task_type is None and (request.ref_audio is not None or request.ref_text is not None): + request.task_type = "Base" + + # Normalize voice to lowercase for case-insensitive matching + if request.voice is not None: + request.voice = request.voice.lower() + if request.task_type is None and request.voice in server.precomputed_speakers: + request.task_type = "Base" + task_type = request.task_type or "CustomVoice" + + # Validate input is not empty + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + # 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 server.supported_languages: + return ( + f"Invalid language '{request.language}'. Supported: {', '.join(sorted(server.supported_languages))}" + ) + + # Validate speaker for CustomVoice task + if task_type == "CustomVoice": + if not server.supported_speakers: + return ( + "This model does not support CustomVoice task (no speakers configured). " + "Use task_type='Base' with ref_audio/ref_text for voice cloning, " + "or use a CustomVoice model." + ) + if request.voice is not None and request.voice not in server.supported_speakers: + return f"Invalid voice '{request.voice}'. Supported: {', '.join(sorted(server.supported_speakers))}" + + # Validate speaker_embedding constraints + if request.speaker_embedding is not None: + if task_type != "Base": + return "'speaker_embedding' is only valid for Base task" + if not request.speaker_embedding: + return "'speaker_embedding' must be a non-empty list of floats" + # speaker_embedding implies x_vector_only_mode — set it before + # Base task validation so callers don't need to pass it explicitly. + request.x_vector_only_mode = True + emb_len = len(request.speaker_embedding) + dim_err = server._validate_qwen_tts_speaker_embedding_dim(emb_len) + if dim_err is not None: + return dim_err + # Validate Base task requirements + if task_type == "Base": + if request.voice is None: + # 1. Ensure a voice source is provided + if request.ref_audio is None and getattr(request, "speaker_embedding", None) is None: + return "Base task requires 'ref_audio' or 'speaker_embedding' for voice cloning" + # 2. Validate ref_audio format if it exists (using the helper from main) + if request.ref_audio is not None: + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + # 3. Validate text requirements based on the mode + if not getattr(request, "x_vector_only_mode", False): + if not request.ref_text or not request.ref_text.strip(): + return ( + "Base task requires non-empty 'ref_text' (transcript of " + "the reference audio) unless 'x_vector_only_mode' is enabled" + ) + else: + voice_lower = request.voice.lower() + if voice_lower in server.uploaded_speakers: + # Check if data file exists for uploaded speaker + speaker_info = server.uploaded_speakers[voice_lower] + file_path = Path(speaker_info["file_path"]) + if not file_path.exists(): + return f"Data file for uploaded speaker '{request.voice}' not found on disk" + elif voice_lower in server.precomputed_speakers: + profile = server.precomputed_speakers[voice_lower] + mode = str(profile.get("mode") or "xvec").lower() + ref_text = request.ref_text or profile.get("ref_text") + if mode == "icl" and (not isinstance(ref_text, str) or not ref_text.strip()): + return ( + f"Precomputed voice '{request.voice}' uses ICL mode but has no ref_text in " + "the request or custom voice manifest" + ) + else: + # need ref_audio for built-in speaker + if request.ref_audio is None: + return ( + f"Base task with built-in speaker '{request.voice}' requires 'ref_audio' for voice cloning" + ) + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not getattr(request, "x_vector_only_mode", False) and ( + not request.ref_text or not request.ref_text.strip() + ): + return ( + "Base task requires non-empty 'ref_text' (transcript of " + "the reference audio) unless 'x_vector_only_mode' is enabled" + ) + + # Validate cross-parameter dependencies + if task_type != "Base": + if request.ref_text is not None: + return "'ref_text' is only valid for Base task" + if request.x_vector_only_mode is not None: + return "'x_vector_only_mode' is only valid for Base task" + + # Validate VoiceDesign task requirements + if task_type == "VoiceDesign" and not request.instructions: + return "VoiceDesign task requires 'instructions' to describe the voice" + + # Validate instructions length (using cached value from initialization) + if request.instructions and len(request.instructions) > server._max_instructions_length: + return f"Instructions too long (max {server._max_instructions_length} characters)" + + # Validate max_new_tokens range + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py b/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py index dd1365204c8..2c720d8e1af 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/step_audio2.py @@ -23,8 +23,9 @@ class StepAudio2Adapter(ARTTSAdapter): name = "step_audio2" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - # Dispatcher routes to the step_audio2 case (non-empty input check). - return self.ctx.server._validate_tts_request(request) + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py index 221e2f64bd5..361aa2522d9 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxcpm2.py @@ -19,7 +19,25 @@ class VoxCPM2Adapter(ARTTSAdapter): name = "voxcpm2" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - return self.ctx.server._validate_voxcpm2_request(request) + """Validate VoxCPM2 request parameters. Returns error message or None.""" + server = self.ctx.server + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.voice is not None: + request.voice = request.voice.lower() + available_voices = set(server.uploaded_speakers) | set(server.precomputed_speakers) | {"default"} + if request.voice not in available_voices: + supported = ", ".join(sorted(available_voices)) or "none" + return f"Invalid voice '{request.voice}'. Supported: {supported}" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py index 72e18dc8408..0b2d400b1db 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/voxtral.py @@ -16,7 +16,32 @@ class VoxtralTTSAdapter(ARTTSAdapter): name = "voxtral_tts" def validate(self, request: "OpenAICreateSpeechRequest") -> str | None: - return self.ctx.server._validate_voxtral_tts_request(request) + """Validate Voxtral TTS request parameters. Returns error message or None.""" + server = self.ctx.server + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + # Voxtral TTS requires either a preset voice or ref_audio for voice cloning. + if request.voice is None and request.ref_audio is None: + return "Either 'voice' (preset speaker) or 'ref_audio' (voice cloning) must be provided" + + if request.ref_audio is not None: + fmt_err = server._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + + if request.voice is not None: + request.voice = request.voice.lower() + if server.supported_speakers and request.voice not in server.supported_speakers: + return f"Invalid speaker '{request.voice}'. Supported: {', '.join(sorted(server.supported_speakers))}" + + if request.max_new_tokens is not None: + if request.max_new_tokens < self.max_new_tokens_min: + return f"max_new_tokens must be at least {self.max_new_tokens_min}" + if request.max_new_tokens > self.max_new_tokens_max: + return f"max_new_tokens cannot exceed {self.max_new_tokens_max}" + + return None async def build( self, request: "OpenAICreateSpeechRequest", sampling_params_list: list, has_inline_ref_audio: bool diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index aa4eeaba50b..aa647918989 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -28,6 +28,29 @@ ) +class VideoFrames(list[Image.Image]): + """Decoded video frames plus source metadata.""" + + def __init__(self, frames: list[Image.Image] | None = None, *, fps: float | None = None) -> None: + super().__init__(frames or []) + self.fps = fps + self.frame_rate = fps + + +def positive_float(value: Any) -> float | None: + if value is None: + return None + if hasattr(value, "item") and not isinstance(value, (bytes, str)): + value = value.item() + try: + result = float(value) + except (TypeError, ValueError): + return None + if result <= 0: + return None + return result + + def _decode_image_bytes(image_bytes: bytes, *, source: str) -> Image.Image: try: return Image.open(BytesIO(image_bytes)).convert("RGB") @@ -41,7 +64,7 @@ def _decode_video_bytes( source: str, max_frames: int | None = None, keep: Literal["first", "last"] = "first", -) -> list[Image.Image]: +) -> VideoFrames: try: import av except ImportError as exc: # pragma: no cover - av is a serving dependency via media_utils @@ -56,8 +79,16 @@ def _decode_video_bytes( tail_frames: deque[Image.Image] | None = ( deque(maxlen=max_frames) if keep == "last" and max_frames is not None else None ) + fps: float | None = None try: with av.open(BytesIO(video_bytes)) as container: + video_stream = container.streams.video[0] if container.streams.video else None + if video_stream is not None: + fps = ( + positive_float(getattr(video_stream, "average_rate", None)) + or positive_float(getattr(video_stream, "base_rate", None)) + or positive_float(getattr(video_stream, "guessed_rate", None)) + ) for frame in container.decode(video=0): image = frame.to_image().convert("RGB") if tail_frames is not None: @@ -73,7 +104,7 @@ def _decode_video_bytes( frames = list(tail_frames) if not frames: raise InvalidInputReferenceError(f"Invalid {source}: provided content is not a valid video.") - return frames + return VideoFrames(frames, fps=fps) def _decode_media_bytes( @@ -82,7 +113,7 @@ def _decode_media_bytes( source: str, max_video_frames: int | None = None, video_keep: Literal["first", "last"] = "first", -) -> Image.Image | list[Image.Image]: +) -> Image.Image | VideoFrames: try: return _decode_image_bytes(media_bytes, source=source) except InvalidInputReferenceError: @@ -138,7 +169,7 @@ def _decode_base64_video( source: str, max_frames: int | None = None, keep: Literal["first", "last"] = "first", -) -> list[Image.Image]: +) -> VideoFrames: if video_reference: if video_reference.startswith("data:video"): _, b64_data = video_reference.split(",", 1) @@ -158,7 +189,7 @@ async def decode_video_url( *, max_frames: int | None = None, keep: Literal["first", "last"] = "first", -) -> list[Image.Image]: +) -> VideoFrames: if video_url.startswith("data:video"): return _decode_base64_video( video_url, @@ -240,7 +271,7 @@ async def decode_input_reference( *, max_video_frames: int | None = None, video_keep: Literal["first", "last"] = "first", -) -> Image.Image | list[Image.Image] | None: +) -> Image.Image | VideoFrames | None: """Decode media input from multipart bytes, data URLs, or typed references.""" provided = sum(item is not None for item in (input_reference_bytes, image_reference, video_reference)) @@ -392,16 +423,8 @@ def _coerce_audio_to_numpy(audio: Any) -> np.ndarray: return arr.astype(np.float32) -def _encode_video_bytes( - video: Any, - fps: int, - audio: Any | None = None, - audio_sample_rate: int | None = None, - video_codec_options: dict[str, str] | None = None, -) -> bytes: - """Encode a video payload into MP4 bytes, optionally muxing audio.""" - from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes - +def _coerce_video_to_uint8_frames(video: Any) -> np.ndarray: + """Convert a video payload into contiguous uint8 frames shaped (F, H, W, 3).""" frames = _coerce_video_to_frames(video) if not frames: raise ValueError("No frames found to encode.") @@ -417,13 +440,23 @@ def _encode_video_bytes( frames_np *= 255.0 frames_u8 = np.round(frames_np).astype(np.uint8) - # Ensure contiguous memory layout for faster PyAV muxing - frames_u8 = np.ascontiguousarray(frames_u8) + return np.ascontiguousarray(frames_u8) + + +def _encode_video_bytes( + video: Any, + fps: int, + audio: Any | None = None, + audio_sample_rate: int | None = None, + video_codec_options: dict[str, str] | None = None, +) -> bytes: + """Encode a video payload into MP4 bytes, optionally muxing audio.""" + from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None return mux_video_audio_bytes( - frames_u8, + _coerce_video_to_uint8_frames(video), audio_np, fps=float(fps), audio_sample_rate=audio_sample_rate or 24000, @@ -431,6 +464,55 @@ def _encode_video_bytes( ) +class FragmentedMP4VideoEncoder: + """Normalize video chunks and append them to one fragmented MP4 stream.""" + + def __init__( + self, + *, + fps: int | float, + video_codec_options: dict[str, str] | None = None, + ) -> None: + self._fps = float(fps) + self._video_codec_options = video_codec_options + self._muxer: Any | None = None + + def encode(self, video: Any) -> bytes: + """Encode one generated video chunk and return newly emitted fMP4 bytes.""" + from vllm_omni.diffusion.utils.media_utils import FragmentedMP4Muxer + + frames_u8 = _coerce_video_to_uint8_frames(video) + if self._muxer is None: + self._muxer = FragmentedMP4Muxer( + width=frames_u8.shape[2], + height=frames_u8.shape[1], + fps=self._fps, + video_codec_options=self._video_codec_options, + ) + return self._muxer.mux_video_frames(frames_u8) + + def close(self) -> bytes: + """Close the underlying fMP4 muxer and return trailing bytes, if any.""" + if self._muxer is None: + return b"" + return self._muxer.close() + + +StreamingVideoFormat = Literal["m4s"] + + +def create_streaming_video_encoder( + *, + output_format: StreamingVideoFormat, + fps: int | float, + video_codec_options: dict[str, str] | None = None, +) -> FragmentedMP4VideoEncoder: + """Create an incremental encoder for the requested WebSocket video format.""" + if output_format == "m4s": + return FragmentedMP4VideoEncoder(fps=fps, video_codec_options=video_codec_options) + raise ValueError(f"Unsupported streaming video format: {output_format}") + + def encode_video_base64( video: Any, fps: int, diff --git a/vllm_omni/entrypoints/openai/video_stream_base.py b/vllm_omni/entrypoints/openai/video_stream_base.py new file mode 100644 index 00000000000..460a66d7c6d --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_base.py @@ -0,0 +1,973 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Base WebSocket handler for streaming video input understanding. + +Shared session loop, frame/audio buffering, EVS pre-filter, prewarm, +interrupt handling, and engine ``generate()`` streaming. Pipeline-specific +behavior (trigger rules, prompt shape, history) is supplied by subclasses +via :class:`VideoStreamPipelineHooks`. + +Protocol: + Client -> Server: + {"type": "session.config", ...} # Session config (sent once) + {"type": "video.frame", "data": "..."} # base64 JPEG/PNG frame + {"type": "audio.chunk", "data": "..."} # base64 PCM16 16kHz mono + {"type": "video.query", "text": "..."} # Submit query about buffered frames + {"type": "video.done"} # End of session + + Server -> Client: + {"type": "response.start"} + {"type": "response.text.delta", "delta": "..."} + {"type": "response.text.done", "text": "..."} + {"type": "response.audio.delta", "data": "...", "format": "wav"} + {"type": "response.audio.done"} + {"type": "session.done"} + {"type": "error", "message": "..."} +""" + +import asyncio +import base64 +import hashlib +import io +import json +import time as _time +import uuid +import wave +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any, Protocol, runtime_checkable + +import torch +from fastapi import WebSocket, WebSocketDisconnect +from PIL import Image +from pydantic import BaseModel, Field, ValidationError +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai import video_stream_envs +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter +from vllm_omni.entrypoints.openai.video_stream_context import ( + text_only_message, +) +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + +_DEFAULT_IDLE_TIMEOUT = 60.0 +_DEFAULT_CONFIG_TIMEOUT = 10.0 +_MAX_FRAME_SIZE = 10 * 1024 * 1024 # 10MB per frame +_MAX_BUFFER_FRAMES = 64 +_MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 +_MAX_MSG_QUEUE = 200 +_CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length +_BAD_FRAME = object() + + +def _decode_frame_bytes(raw_bytes: bytes) -> Any: + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + +@runtime_checkable +class VideoStreamPipelineHooks(Protocol): + """Pipeline-specific hooks for streaming video handlers.""" + + def should_trigger_turn(self, trigger: "VideoStreamTurnTrigger") -> bool: + """Return True to auto-start a turn after a new frame (no ``video.query``).""" + ... + + def build_engine_prompt( + self, + config: "StreamingVideoSessionConfig", + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Build OpenAI-style messages and the current user message.""" + ... + + def on_turn_complete( + self, + message_history: list[dict[str, Any]], + user_message: dict[str, Any], + response_text: str, + ) -> None: + """Update session state after a successful turn.""" + ... + + +@dataclass(frozen=True) +class VideoStreamTurnTrigger: + """Snapshot passed to :meth:`OmniStreamingVideoHandler.should_trigger_turn`.""" + + frame_count: int + is_generating: bool + config: "StreamingVideoSessionConfig" + + +class StreamingVideoSessionConfig(BaseModel): + """Configuration sent as the first WebSocket message.""" + + model: str | None = None + modalities: list[str] = Field( + default_factory=lambda: ["text", "audio"], + description="Output modalities: 'text', 'audio', or both.", + ) + num_frames: int = Field( + default=4, + ge=1, + le=128, + description="Max frames to sample from buffer for the model.", + ) + max_frames: int = Field( + default=50, + ge=1, + le=256, + description="Max frames to keep in the buffer.", + ) + system_prompt: str | None = Field( + default=None, + description="Custom system prompt.", + ) + use_audio_in_video: bool = Field( + default=True, + description="Interleave audio chunks with video frames when audio input is present.", + ) + sampling_params_list: list[dict[str, Any]] | None = Field( + default=None, + description="Per-stage sampling params [thinker, talker, code2wav].", + ) + enable_frame_filter: bool = Field( + default=True, + description="EVS pixel-similarity pre-filter to drop near-duplicate frames.", + ) + frame_filter_threshold: float = Field( + default=0.95, + ge=0.0, + le=1.0, + description="EVS similarity threshold (higher = keep more frames).", + ) + + +class OmniStreamingVideoHandler: + """Base handler for WebSocket streaming video sessions. + + Subclasses implement :class:`VideoStreamPipelineHooks` to customize turn + triggering, prompt construction, and history updates. + """ + + def should_trigger_turn(self, trigger: VideoStreamTurnTrigger) -> bool: + """Auto-trigger after ``video.frame`` when True (default: never).""" + return False + + def build_engine_prompt( + self, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + raise NotImplementedError + + def on_turn_complete( + self, + message_history: list[dict[str, Any]], + user_message: dict[str, Any], + response_text: str, + ) -> None: + raise NotImplementedError + + def create_message_history(self, config: StreamingVideoSessionConfig) -> Any: + """Per-session conversation state (default: empty OpenAI-style list).""" + return [] + + def on_frame_buffered( + self, + raw_bytes: bytes, + frame_b64: str, + message_history: Any, + config: StreamingVideoSessionConfig, + ) -> None: + """Hook after a frame is accepted into the session buffer.""" + del raw_bytes, frame_b64, message_history, config + + def __init__( + self, + chat_service: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + engine_client: Any | None = None, + ) -> None: + self._chat_service = chat_service + self._idle_timeout = idle_timeout + self._config_timeout = config_timeout + self._engine_client = engine_client + + async def handle_session(self, websocket: WebSocket) -> None: + """Main session loop for a single WebSocket connection.""" + await websocket.accept() + + try: + config = await self._receive_config(websocket) + if config is None: + return + + frame_buffer: list[str] = [] # base64-encoded JPEG frames + # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. + frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME + frame_filter = ( + FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None + ) + audio_buffer = bytearray() # raw PCM16 16kHz mono + message_history: Any = self.create_message_history(config) + active_request_id: str | None = None + prev_request_id: str | None = None # abort target iff prev was interrupted + prev_was_interrupted: bool = False + interrupt_event = asyncio.Event() + prewarm_tasks: set[asyncio.Task[Any]] = set() + query_task: asyncio.Task[Any] | None = None + + msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=_MAX_MSG_QUEUE) + + async def _reader() -> None: + """Receive WebSocket messages and enqueue them.""" + try: + while True: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._idle_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Idle timeout") + await msg_queue.put(None) + return + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON") + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "Messages must be JSON objects") + continue + + msg_type = str(msg.get("type", "")) + if msg_type.startswith("_internal."): + await self._send_error(websocket, f"Unknown type: {msg_type}") + continue + + await msg_queue.put(msg) + if msg.get("type") == "video.done": + return + except WebSocketDisconnect: + await msg_queue.put(None) + except Exception: + await msg_queue.put(None) + raise + + async def _cancel_active_query(*, abort_now: bool = False) -> None: + """Signal soft interrupt for the active query.""" + nonlocal active_request_id, prev_was_interrupted, query_task + if active_request_id is not None: + interrupt_event.set() + prev_was_interrupted = True + logger.info("Interrupt signaled for %s", active_request_id) + if abort_now and self._engine_client: + try: + await self._engine_client.abort(active_request_id) + except Exception: + logger.debug("Abort failed for %s", active_request_id, exc_info=True) + if query_task is not None and not query_task.done(): + query_task.cancel() + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + + async def _start_query_turn(*, query_text: str) -> None: + """Schedule a new inference turn from the current buffers.""" + nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task + + await _cancel_active_query() + + if not frame_buffer: + await self._send_error(websocket, "No frames buffered") + return + + if prev_was_interrupted and prev_request_id and self._engine_client: + try: + await self._engine_client.abort(prev_request_id) + except Exception: + pass + await asyncio.sleep(0.1) + prev_was_interrupted = False + + request_id = f"video-{uuid.uuid4().hex[:12]}" + active_request_id = request_id + interrupt_event.clear() + query_frames = list(frame_buffer) + query_audio_buffer = bytearray(audio_buffer) + audio_buffer.clear() + query_prewarmed_frames = dict(frame_pil_cache) + + async def _run_query() -> None: + nonlocal active_request_id, prev_request_id + try: + await self._process_query( + websocket, + config, + query_frames, + query_audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + query_prewarmed_frames, + ) + finally: + if active_request_id == request_id: + prev_request_id = request_id + active_request_id = None + + query_task = asyncio.create_task(_run_query()) + + async def _processor() -> None: + """Process enqueued messages.""" + nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task + + while True: + msg = await msg_queue.get() + if msg is None: + await _cancel_active_query(abort_now=True) + return + + msg_type = msg.get("type") + + if msg_type == "_internal.frame_decode_failed": + frame_data = msg.get("b64", "") + removed = frame_data in frame_buffer + if removed: + frame_buffer[:] = [f for f in frame_buffer if f != frame_data] + if frame_pil_cache.get(frame_data) is _BAD_FRAME: + frame_pil_cache.pop(frame_data, None) + if removed: + await self._send_error(websocket, "Frame decode failed") + + elif msg_type == "video.frame": + frame_data = msg.get("data", "") + if not frame_data: + continue + if len(frame_data) > _MAX_FRAME_SIZE: + await self._send_error(websocket, "Frame too large") + continue + try: + raw_bytes = base64.b64decode(frame_data, validate=True) + except Exception: + await self._send_error(websocket, "Invalid image data") + continue + if frame_filter is not None: + try: + if not frame_filter.should_retain(raw_bytes): + continue + except Exception: + await self._send_error(websocket, "Invalid image data") + continue + max_buf = config.max_frames + if len(frame_buffer) >= max_buf: + dropped = frame_buffer.pop(0) + frame_pil_cache.pop(dropped, None) + frame_buffer.append(frame_data) + self.on_frame_buffered(raw_bytes, frame_data, message_history, config) + # Prewarm: decode PIL off the event loop so query-time chat_template + # can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames. + if frame_data not in frame_pil_cache: + mm_uuid = hashlib.md5(raw_bytes, usedforsecurity=False).hexdigest() + + async def _prewarm(b64: str, b: bytes, u: str) -> None: + try: + pil = await asyncio.to_thread(_decode_frame_bytes, b) + frame_pil_cache[b64] = (pil, u) + except Exception: + frame_pil_cache[b64] = _BAD_FRAME + logger.warning("prewarm decode failed for frame (len=%d)", len(b)) + try: + msg_queue.put_nowait({"type": "_internal.frame_decode_failed", "b64": b64}) + except asyncio.QueueFull: + logger.warning( + "frame decode failure event dropped because message queue is full" + ) + + task = asyncio.create_task(_prewarm(frame_data, raw_bytes, mm_uuid)) + prewarm_tasks.add(task) + task.add_done_callback(prewarm_tasks.discard) + + is_generating = active_request_id is not None or ( + query_task is not None and not query_task.done() + ) + if self.should_trigger_turn( + VideoStreamTurnTrigger( + frame_count=len(frame_buffer), + is_generating=is_generating, + config=config, + ) + ): + await _start_query_turn(query_text="") + + elif msg_type == "audio.chunk": + data_b64 = msg.get("data", "") + try: + pcm_bytes = base64.b64decode(data_b64) + except Exception: + continue + if len(audio_buffer) + len(pcm_bytes) > _MAX_AUDIO_BUFFER_BYTES: + await self._send_error(websocket, "Audio buffer overflow") + audio_buffer.clear() + continue + audio_buffer.extend(pcm_bytes) + + elif msg_type == "video.query": + query_text = msg.get("text", "") + audio_data_b64 = msg.get("audio_data") + if audio_data_b64: + try: + decoded = base64.b64decode(audio_data_b64) + if len(audio_buffer) + len(decoded) <= _MAX_AUDIO_BUFFER_BYTES: + audio_buffer.extend(decoded) + else: + await self._send_error(websocket, "Audio buffer overflow") + audio_buffer.clear() + except Exception: + pass + + await _start_query_turn(query_text=query_text) + + elif msg_type == "video.done": + if query_task is not None and not query_task.done(): + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + await websocket.send_json({"type": "session.done"}) + return + + elif msg_type == "ping": + try: + await websocket.send_json({"type": "pong"}) + except Exception: + pass + + else: + await self._send_error(websocket, f"Unknown type: {msg_type}") + + reader_task = asyncio.create_task(_reader()) + try: + await _processor() + finally: + reader_task.cancel() + try: + await reader_task + except (asyncio.CancelledError, Exception): + pass + for t in list(prewarm_tasks): + t.cancel() + if prewarm_tasks: + await asyncio.gather(*prewarm_tasks, return_exceptions=True) + if query_task is not None and not query_task.done(): + await _cancel_active_query(abort_now=True) + + except WebSocketDisconnect: + logger.info("Streaming video: client disconnected") + except Exception as e: + logger.exception("Streaming video session error: %s", e) + try: + await self._send_error(websocket, f"Internal error: {e}") + except Exception: + pass + + async def _receive_config(self, websocket: WebSocket) -> StreamingVideoSessionConfig | None: + """Wait for and validate the session.config message.""" + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._config_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.config") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.config") + return None + + if not isinstance(msg, dict) or msg.get("type") != "session.config": + await self._send_error( + websocket, + f"Expected session.config, got: {msg.get('type') if isinstance(msg, dict) else type(msg).__name__}", + ) + return None + + config_data = {k: v for k, v in msg.items() if k != "type"} + alias_map = { + "num_sample_frames": "num_frames", + "evs_enabled": "enable_frame_filter", + "evs_threshold": "frame_filter_threshold", + } + for old_key, new_key in alias_map.items(): + if old_key in config_data and new_key not in config_data: + config_data[new_key] = config_data[old_key] + + try: + config = StreamingVideoSessionConfig(**config_data) + except ValidationError as e: + await self._send_error(websocket, f"Invalid session config: {e}") + return None + + return config + + async def _process_query( + self, + websocket: WebSocket, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + request_id: str, + interrupt_event: asyncio.Event, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> None: + """Build prompt, run inference, stream text + audio response.""" + + if self._engine_client is None: + await self._send_error(websocket, "Streaming video requires an engine client") + return + + await self._process_query_engine( + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ) + + # ------------------------------------------------------------------ + # Engine-client path (async_chunk audio streaming) + # ------------------------------------------------------------------ + + async def _process_query_engine( + self, + websocket: WebSocket, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + request_id: str, + interrupt_event: asyncio.Event, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> None: + """Direct engine_client.generate() path for async_chunk audio.""" + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + + messages, user_message = self.build_engine_prompt( + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + prewarmed_frames, + ) + + request_kwargs: dict[str, Any] = { + "model": config.model or "default", + "messages": messages, + "stream": True, + "modalities": config.modalities, + "add_generation_prompt": True, + "continue_final_message": False, + "add_special_tokens": False, + } + if config.use_audio_in_video and len(audio_buffer) > 0: + request_kwargs["mm_processor_kwargs"] = { + "use_audio_in_video": True, + } + if config.sampling_params_list: + request_kwargs["sampling_params_list"] = config.sampling_params_list + + try: + chat_request = ChatCompletionRequest(**request_kwargs) + except Exception as e: + await self._send_error(websocket, f"Failed to build request: {e}") + return + + try: + engine_prompt = await self._preprocess_to_engine_prompt(chat_request) + except Exception as e: + await self._send_error(websocket, f"Preprocess failed: {e}") + return + + await websocket.send_json({"type": "response.start"}) + text_parts: list[str] = [] + text_done_sent = False + audio_chunk_count = 0 + # Number of per-step tensors in OmniRequestOutput.audio_data already + # drained. Used by the fast path to skip already-emitted history. + audio_chunks_drained = 0 + previous_text = "" + interrupted = False + t_start = _time.monotonic() + t_first_text = None + t_first_audio = None + + # Wire-level async-chunk switch. "off" means + # buffer all deltas server-side and flush once at the end; the engine + # pipeline still overlaps internally. + async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK + streaming = async_chunk_mode == "on" + audio_tail_tensors: list[Any] = [] + + try: + result_gen = self._engine_client.generate( + prompt=engine_prompt, + request_id=request_id, + output_modalities=config.modalities, + ) + + async for output in result_gen: + # Soft interrupt: drain without sending + if interrupt_event.is_set(): + if not interrupted: + logger.info("Generation interrupted — draining") + interrupted = True + continue + + if not isinstance(output, OmniRequestOutput): + continue + + out_type = getattr(output, "final_output_type", "text") + + if out_type == "audio": + if streaming and not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if t_first_audio is None: + t_first_audio = _time.monotonic() + audio_chunk_count += 1 + if streaming: + b64, audio_chunks_drained = self._extract_audio_delta_b64( + output, + audio_chunks_drained, + ) + if b64: + await websocket.send_json( + { + "type": "response.audio.delta", + "data": b64, + "format": "wav", + } + ) + else: + audio_data = self._get_audio_data(output) + if audio_data is not None: + if isinstance(audio_data, list): + audio_tail_tensors = list(audio_data) + else: + audio_tail_tensors = [audio_data] + else: + delta_text, previous_text = self._extract_text_delta( + output, + previous_text, + ) + if delta_text: + if t_first_text is None: + t_first_text = _time.monotonic() + text_parts.append(delta_text) + if streaming: + await websocket.send_json({"type": "response.text.delta", "delta": delta_text}) + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if not streaming and audio_tail_tensors: + try: + coalesced = ( + audio_tail_tensors[0] if len(audio_tail_tensors) == 1 else torch.cat(audio_tail_tensors, dim=-1) + ) + tail_np = self._tensor_to_1d_np(coalesced) + b64, _ = self._encode_tail( + tail_np, + 0, + new_drained=len(audio_tail_tensors), + is_first=True, + ) + if b64: + await websocket.send_json( + { + "type": "response.audio.delta", + "data": b64, + "format": "wav", + } + ) + except Exception: + logger.exception("Failed to coalesce off-path audio") + + if audio_chunk_count > 0: + await websocket.send_json({"type": "response.audio.done"}) + + response_text = "".join(text_parts) + self.on_turn_complete(message_history, user_message, response_text) + + t_end = _time.monotonic() + logger.info( + "[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d", + async_chunk_mode, + t_end - t_start, + (t_first_text - t_start) if t_first_text else -1, + (t_first_audio - t_start) if t_first_audio else -1, + audio_chunk_count, + ) + + except Exception: + logger.exception("Engine query failed") + await self._send_error(websocket, "Query processing failed") + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + + # ------------------------------------------------------------------ + # Audio helpers + # ------------------------------------------------------------------ + + @staticmethod + def _pcm_to_wav_b64(pcm_data: bytes, sample_rate: int = 16000) -> str: + """Wrap raw PCM16 mono in a WAV container and return base64.""" + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(pcm_data) + return base64.b64encode(buf.getvalue()).decode() + + @classmethod + def _extract_audio_delta_b64( + cls, + result: OmniRequestOutput, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Return (base64 WAV of new samples, updated chunks_drained). + + `chunks_drained` is the number of per-step tensors in + ``audio_data`` that have already been emitted. Each engine step appends + one tensor, so new samples are ``audio_data[chunks_drained:]`` — no + matter how many steps accumulated between reads (handles backpressure + cleanly, unlike a simple ``audio_data[-1]``). + + Two paths, selected at runtime by ``VLLM_VIDEO_AUDIO_DELTA_MODE``: + * fast — only D2H the new tail. Per-call cost ∝ new chunks. + * slow — full cat + D2H each call. Per-call cost ∝ total history. + Retained for A/B; remove once downstream callers confirm. + """ + audio_data = cls._get_audio_data(result) + if audio_data is None: + return None, chunks_drained + + if video_stream_envs.VLLM_VIDEO_AUDIO_DELTA_MODE == "slow": + return cls._delta_slow(audio_data, chunks_drained) + return cls._delta_fast(audio_data, chunks_drained) + + @staticmethod + def _get_audio_data(result: OmniRequestOutput): + """Navigate OmniRequestOutput → multimodal_output['audio']. None on miss.""" + request_output = getattr(result, "request_output", None) + if request_output is None: + return None + outputs = getattr(request_output, "outputs", None) + if not isinstance(outputs, list) or not outputs: + return None + mm_output = getattr(outputs[0], "multimodal_output", None) + if not isinstance(mm_output, Mapping): + return None + return mm_output.get("audio") + + @classmethod + def _delta_fast( + cls, + audio_data, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Emit only tensors appended since the last call.""" + # Single tensor: output_processor hands us one tensor before it becomes a + # list (see output_processor.py:89). Treat it as chunk #0. + if not isinstance(audio_data, list): + if chunks_drained >= 1: + return None, chunks_drained + tail_np = cls._tensor_to_1d_np(audio_data) + return cls._encode_tail(tail_np, chunks_drained, new_drained=1, is_first=True) + + n = len(audio_data) + if n <= chunks_drained: + return None, chunks_drained + + new_chunks = audio_data[chunks_drained:] + tail = new_chunks[0] if len(new_chunks) == 1 else torch.cat(new_chunks, dim=-1) + tail_np = cls._tensor_to_1d_np(tail) + return cls._encode_tail(tail_np, chunks_drained, new_drained=n, is_first=(chunks_drained == 0)) + + @classmethod + def _delta_slow( + cls, + audio_data, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Pre-fix behaviour: concat everything each call and slice on CPU.""" + if isinstance(audio_data, list): + if not audio_data: + return None, chunks_drained + audio_tensor = torch.cat(audio_data, dim=-1) + new_drained = len(audio_data) + else: + audio_tensor = audio_data + new_drained = 1 + + full_np = cls._tensor_to_1d_np(audio_tensor) + if full_np is None: + return None, chunks_drained + # chunks_drained doesn't map directly to sample offset without tracking + # per-chunk lengths, so we re-derive: replay the tail that corresponds + # to chunks appended since last call by slicing off the part produced + # by the already-drained prefix. For slow path this is intentionally + # wasteful — the point is to reproduce the pre-fix hot loop. + if chunks_drained == 0: + tail_np = full_np + else: + # Recover prefix length by re-concatenating the already-drained + # prefix tensors (cost intentionally identical to the baseline + # implementation this was lifted from). + if isinstance(audio_data, list) and chunks_drained < len(audio_data): + prefix_len = sum(int(t.shape[-1]) for t in audio_data[:chunks_drained]) + tail_np = full_np[prefix_len:] + else: + tail_np = full_np[0:0] + return cls._encode_tail(tail_np, chunks_drained, new_drained=new_drained, is_first=(chunks_drained == 0)) + + @classmethod + def _encode_tail( + cls, + tail_np, + old_drained: int, + *, + new_drained: int, + is_first: bool, + ) -> tuple[str | None, int]: + """Strip the CausalConv leading artifact on first emit, then b64-encode.""" + if tail_np is None or len(tail_np) == 0: + return None, new_drained + if is_first and len(tail_np) > _CODEC_FRAME_SAMPLES * 2: + tail_np = tail_np[_CODEC_FRAME_SAMPLES:] + if len(tail_np) == 0: + return None, new_drained + try: + return cls._encode_audio_wav_b64(tail_np), new_drained + except Exception: + logger.exception("Failed to encode audio delta WAV") + return None, old_drained + + @staticmethod + def _tensor_to_1d_np(t): + """Tensor → flat float32 numpy on CPU. None on failure.""" + if t is None or not hasattr(t, "float"): + return None + arr = t.float().detach().cpu().numpy() + if arr.ndim > 1: + arr = arr.flatten() + return arr + + @staticmethod + def _encode_audio_wav_b64(audio_np) -> str: + """Encode numpy float32 audio to base64 WAV (24kHz).""" + from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin + from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio + + audio_obj = CreateAudio( + audio_tensor=audio_np, + sample_rate=24000, + response_format="wav", + speed=1.0, + base64_encode=True, + ) + mixin = AudioMixin() + resp = mixin.create_audio(audio_obj) + return resp.audio_data + + @staticmethod + def _extract_text_delta( + result: OmniRequestOutput, + previous_text: str, + ) -> tuple[str, str]: + """Extract incremental text delta from OmniRequestOutput.""" + if result.final_output_type != "text": + return "", previous_text + + request_output = getattr(result, "request_output", None) + if request_output is None: + return "", previous_text + + outputs = getattr(request_output, "outputs", None) + if not isinstance(outputs, list) or not outputs: + return "", previous_text + + text = getattr(outputs[0], "text", None) + if not isinstance(text, str) or not text: + return "", previous_text + + if text.startswith(previous_text): + return text[len(previous_text) :], text + return text, text + + # ------------------------------------------------------------------ + # Preprocessing + # ------------------------------------------------------------------ + + async def _preprocess_to_engine_prompt(self, request) -> Any: + """Use the chat handler's preprocessing to build an engine prompt.""" + handler = self._chat_service + renderer = handler.renderer + + _conversation, engine_prompts = await handler._preprocess_chat( + request, + request.messages, + default_template=getattr(request, "chat_template", None) or handler.chat_template, + default_template_content_format=handler.chat_template_content_format, + renderer=renderer, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + add_special_tokens=request.add_special_tokens, + ) + return engine_prompts[0] + + # ------------------------------------------------------------------ + # Utilities + # ------------------------------------------------------------------ + + _text_only_message = staticmethod(text_only_message) + + async def _send_error(self, websocket: WebSocket, message: str) -> None: + """Send an error message to the client.""" + try: + await websocket.send_json({"type": "error", "message": message}) + except Exception: + pass diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index 69f722ac824..b6526458363 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -8,7 +8,13 @@ Infer -> client sends msgpack(obs), server sends msgpack(ndarray) Reset -> client sends msgpack({endpoint:reset}), server sends msgpack(status) -NumPy values use the msgpack-numpy marker mapping: +NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the +openpi-client format so official OpenPI clients can decode action responses: + ndarray -> {__ndarray__: true, dtype, shape, data} + scalar -> {__npgeneric__: true, dtype, data} + +Inbound payloads accept both the openpi-client markers above and the legacy +vLLM-native markers: ndarray -> {nd: true, type, kind, shape, data} scalar -> {nd: false, type, kind, data} """ @@ -41,18 +47,16 @@ def _pack_numpy(obj: Any) -> Any: if not obj.flags.c_contiguous: obj = np.ascontiguousarray(obj) return { - b"nd": True, + b"__ndarray__": True, b"data": obj.tobytes(), - b"type": obj.dtype.str, - b"kind": obj.dtype.kind, + b"dtype": obj.dtype.str, b"shape": obj.shape, } if isinstance(obj, np.generic): return { - b"nd": False, - b"data": obj.tobytes(), - b"type": obj.dtype.str, - b"kind": obj.dtype.kind, + b"__npgeneric__": True, + b"data": obj.item(), + b"dtype": obj.dtype.str, } raise TypeError(f"Unsupported type: {type(obj)!r}") @@ -67,27 +71,67 @@ def _decode_marker_text(value: Any) -> str: return str(value) +def _is_truthy_marker(value: Any) -> bool: + return value is True or value == 1 + + +def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any: + nd = _mapping_get(obj, "nd", _MISSING) + dtype = _mapping_get(obj, "type", _MISSING) + kind = _mapping_get(obj, "kind", _MISSING) + data = _mapping_get(obj, "data", _MISSING) + if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING: + return _MISSING + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + kind_text = _decode_marker_text(kind) + if dtype_obj.kind != kind_text: + raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + + array = np.frombuffer(data, dtype=dtype_obj).copy() + if nd: + shape = _mapping_get(obj, "shape", _MISSING) + if shape is _MISSING: + raise ValueError("NumPy ndarray marker is missing shape") + return array.reshape(tuple(shape)) + return array[0] + + +def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any: + if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)): + data = _mapping_get(obj, "data", _MISSING) + dtype = _mapping_get(obj, "dtype", _MISSING) + shape = _mapping_get(obj, "shape", _MISSING) + if data is _MISSING or dtype is _MISSING or shape is _MISSING: + raise ValueError("OpenPI ndarray marker is missing required fields") + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + array = np.frombuffer(data, dtype=dtype_obj).copy() + return array.reshape(tuple(shape)) + + if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)): + data = _mapping_get(obj, "data", _MISSING) + dtype = _mapping_get(obj, "dtype", _MISSING) + if data is _MISSING or dtype is _MISSING: + raise ValueError("OpenPI scalar marker is missing required fields") + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + return dtype_obj.type(data) + + return _MISSING + + def _unpack_numpy(obj: Any) -> Any: if isinstance(obj, dict): - nd = _mapping_get(obj, "nd", _MISSING) - dtype = _mapping_get(obj, "type", _MISSING) - kind = _mapping_get(obj, "kind", _MISSING) - data = _mapping_get(obj, "data", _MISSING) - if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING: - dtype_obj = np.dtype(_decode_marker_text(dtype)) - kind_text = _decode_marker_text(kind) - if dtype_obj.kind != kind_text: - raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") - if dtype_obj.kind in ("V", "O", "c"): - raise ValueError(f"Unsupported dtype: {dtype_obj}") - - array = np.frombuffer(data, dtype=dtype_obj).copy() - if nd: - shape = _mapping_get(obj, "shape", _MISSING) - if shape is _MISSING: - raise ValueError("NumPy ndarray marker is missing shape") - return array.reshape(tuple(shape)) - return array[0] + for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker): + decoded = decoder(obj) + if decoded is not _MISSING: + return decoded return {key: _unpack_numpy(value) for key, value in obj.items()} if isinstance(obj, list): return [_unpack_numpy(value) for value in obj] diff --git a/vllm_omni/entrypoints/openpi/serving.py b/vllm_omni/entrypoints/openpi/serving.py index ec5c88408ea..4f1d5808144 100644 --- a/vllm_omni/entrypoints/openpi/serving.py +++ b/vllm_omni/entrypoints/openpi/serving.py @@ -128,7 +128,7 @@ async def infer(self, obs: dict, *, session_id: str, reset: bool) -> ActionOutpu # exposes an async iterator, so consume it to completion and use the # final output, matching other non-streaming OpenAI serving paths. async for output in self.engine_client.generate( - prompt=request.prompts[0], + prompt=request.prompt, request_id=request.request_id, sampling_params_list=[request.sampling_params], ): @@ -159,7 +159,7 @@ def _build_request(self, obs: dict, *, session_id: str, reset: bool) -> Any: prompt = obs.get("prompt", "") sampling_params = OmniDiffusionSamplingParams(extra_args=extra_args) return OmniDiffusionRequest( - prompts=[prompt], + prompt=prompt, sampling_params=sampling_params, request_id=self._next_request_id(session_id), ) diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 74b533b51f7..c899b269797 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -11,9 +11,39 @@ logger = logging.getLogger(__name__) +def resolve_stage_physical_devices( + stage_id: int, + devices: str | int | None, + *, + visible_baseline: str | None = None, +) -> str | None: + """Map logical stage devices to physical IDs without mutating process env. + + When ``visible_baseline`` is provided it is used for logical-to-physical + mapping instead of the current device-control env var. This keeps parallel + stage initialization correct even if another thread has temporarily + narrowed ``CUDA_VISIBLE_DEVICES`` for a different stage. + """ + if devices in (None, "cpu"): + return None + + if not isinstance(devices, (int, str)): + raise TypeError(f"Expected str or int device IDs for stage initialization, got type {type(devices)}") + + device_list = _parse_device_list(devices) + env_var = current_omni_platform.device_control_env_var + visible_devices = visible_baseline if visible_baseline is not None else os.environ.get(env_var) + if visible_devices is not None: + visible_device_list = _parse_device_list(visible_devices) + device_list = _map_device_list(stage_id, device_list, visible_device_list) + return ",".join(device_list) + + def set_stage_devices( stage_id: int, devices: str | int | None, + *, + visible_baseline: str | None = None, ) -> str | None: """Configure per-stage device visibility and current device (CUDA or NPU). @@ -42,19 +72,18 @@ def set_stage_devices( The list of physical devices that were set for the given stage or None if we have no passed devices / are using cpu. """ - env_var = current_omni_platform.device_control_env_var - vis = os.environ.get(env_var) - if devices in (None, "cpu"): logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices) return None elif isinstance(devices, (int, str)): - device_list = _parse_device_list(devices) - if vis is not None: - visible_device_list = _parse_device_list(vis) - device_list = _map_device_list(stage_id, device_list, visible_device_list) - device_str = ",".join(device_list) + device_str = resolve_stage_physical_devices( + stage_id, + devices, + visible_baseline=visible_baseline, + ) + if device_str is None: + return None current_omni_platform.set_device_control_env_var(device_str) return device_str diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 23f386257d9..6d32f4ab117 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,3 +1,4 @@ +import json import os import types from collections import Counter @@ -11,7 +12,7 @@ from vllm.transformers_utils.config import get_config, get_hf_file_to_dict from vllm.transformers_utils.repo_utils import file_or_path_exists -from vllm_omni.config.stage_config import StageConfigFactory +from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.yaml_util import create_config, load_yaml_config, merge_configs from vllm_omni.diffusion.utils.hf_utils import _looks_like_dreamzero from vllm_omni.entrypoints.stage_utils import _to_dict @@ -313,56 +314,89 @@ def resolve_model_config_path(model: str) -> str: def load_stage_configs_from_model( model: str, + *, + trust_remote_code: bool, base_engine_args: dict | None = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, -) -> list: + strategy_config_path: str | None = None, +) -> tuple[list, str | None]: """Load stage configurations from model's default config file. For models registered in the pipeline registry (new path), uses - ``StageConfigFactory.create_from_model()`` which merges + ``StageConfigFactory.create_legacy_stage_configs_from_model()`` which merges PipelineConfig + DeployConfig + CLI overrides. For other models (legacy path), loads stage configs from YAML. Args: model: Model name or path (used to determine model_type) + trust_remote_code: Whether to trust remote code while resolving the model config. base_engine_args: Base engine args to merge as CLI overrides. deploy_config_path: Optional explicit deploy config path. stage_overrides: Per-stage overrides from --stage-overrides. + strategy_config_path: Optional path to a composable-parallel + ``strategy.yaml`` whose derived sizing is overlaid onto the + registry-merged stages (opt-in; ignored on the legacy YAML path). Returns: - List of stage configuration dictionaries + ``(stage_configs, omni_lb_policy)``: the list of stage configuration + dictionaries plus the strategy-derived pipeline-wide ``omni_lb_policy`` + (``None`` when no strategy set one). The policy is returned rather than + written into a caller-provided mutable dict. """ if base_engine_args is None: base_engine_args = {} cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args)) + cli_overrides["trust_remote_code"] = trust_remote_code if stage_overrides: for stage_id_str, overrides in stage_overrides.items(): for key, val in overrides.items(): cli_overrides[f"stage_{stage_id_str}_{key}"] = val - stages = StageConfigFactory.create_from_model( + # Current runtime initialization still consumes legacy OmegaConf stage + # configs. ``StageConfigFactory.create_from_model`` now produces the + # structured ``VllmOmniConfig`` object; future RFC #4021 changes will + # migrate the engine/runtime consumers and replace this legacy resolver. + strategy_specs = None + if strategy_config_path is not None: + from vllm_omni.config.composable_parallel.strategy_loader import load_strategy_specs + + strategy_specs = load_strategy_specs(strategy_config_path) + + stages, omni_lb_policy = StageConfigFactory.create_legacy_stage_configs_from_model( model, - trust_remote_code=cli_overrides.get("trust_remote_code", False), + trust_remote_code=trust_remote_code, cli_overrides=cli_overrides, deploy_config_path=deploy_config_path, + strategy_specs=strategy_specs, ) if stages is not None: # Convert StageConfig objects to OmegaConf for backward compat - return [stage.to_omegaconf() for stage in stages] + return [stage.to_omegaconf() for stage in stages], omni_lb_policy - # Legacy fallback: load from YAML + # Legacy fallback: load from YAML. A composable-parallel strategy cannot be + # applied here (it overlays onto registry-merged stages), so warn rather than + # silently dropping the operator's --strategy-config. + if strategy_config_path is not None: + logger.warning( + "--strategy-config (%s) was provided but model %r resolves via the " + "legacy stage_configs YAML path, which does not support " + "composable-parallel strategies; the strategy is ignored. Use a " + "registry-based model to apply it.", + strategy_config_path, + model, + ) stage_config_path = resolve_model_config_path(model) if stage_config_path is None: - return [] + return [], None stage_configs = load_stage_configs_from_yaml( config_path=stage_config_path, base_engine_args=base_engine_args, prefer_stage_engine_args=True, ) - return stage_configs + return stage_configs, None def load_stage_configs_from_yaml( @@ -501,29 +535,56 @@ def filter_stages( return stage_configs +def parse_stage_overrides(value: Any) -> dict[str, dict[str, Any]] | None: + """Parse the ``--stage-overrides`` value into a per-stage override dict. + + ``value`` may be a raw JSON string (as supplied on the CLI) or an + already-parsed mapping. Returns ``None`` when no overrides are given. + + Raises: + ValueError: when ``value`` is a string that is not valid JSON. + """ + if not value: + return None + if isinstance(value, str): + try: + return json.loads(value) + except json.JSONDecodeError as exc: + raise ValueError(f"--stage-overrides is not valid JSON: {exc}. Got: {value!r}") from exc + return value + + def load_and_resolve_stage_configs( model: str, stage_configs_path: str | None, kwargs: dict | None, + *, + trust_remote_code: bool, default_stage_cfg_factory: Any = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, -) -> tuple[str, list]: + strategy_config_path: str | None = None, +) -> tuple[str, list, str | None]: """Load stage configurations from model or YAML file with fallback to defaults. Args: model: Model name or path stage_configs_path: Optional path to legacy YAML (stage_args format) kwargs: Engine arguments to merge with stage configs + trust_remote_code: Whether to trust remote code while resolving the model config. default_stage_cfg_factory: Optional callable that takes no args and returns default stage config list when no configs are found deploy_config_path: Optional path to deploy YAML (new format). Mutually exclusive with ``stage_configs_path``. stage_overrides: Per-stage overrides from ``--stage-overrides`` JSON. Keys are stage_id strings, values are dicts of overrides. + strategy_config_path: Optional path to a composable-parallel + ``strategy.yaml`` overlaid onto the registry-merged stages. Returns: - Tuple of (config_path, stage_configs) + Tuple of ``(config_path, stage_configs, omni_lb_policy)`` — the last is + the strategy-derived pipeline-wide load-balancer policy (``None`` when no + strategy set one), returned for the engine to apply. """ if stage_configs_path is not None and deploy_config_path is not None: raise ValueError( @@ -551,13 +612,16 @@ def load_and_resolve_stage_configs( stage_configs_path, ) + omni_lb_policy: str | None = None if deploy_config_path is not None: config_path = deploy_config_path - stage_configs = load_stage_configs_from_model( + stage_configs, omni_lb_policy = load_stage_configs_from_model( model, + trust_remote_code=trust_remote_code, base_engine_args=kwargs, deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, ) if not stage_configs: if default_stage_cfg_factory is not None: @@ -567,10 +631,12 @@ def load_and_resolve_stage_configs( stage_configs = [] elif stage_configs_path is None: config_path = resolve_model_config_path(model) - stage_configs = load_stage_configs_from_model( + stage_configs, omni_lb_policy = load_stage_configs_from_model( model, + trust_remote_code=trust_remote_code, base_engine_args=kwargs, stage_overrides=stage_overrides, + strategy_config_path=strategy_config_path, ) if not stage_configs: if default_stage_cfg_factory is not None: @@ -585,7 +651,7 @@ def load_and_resolve_stage_configs( stage_configs = filter_stages(config_path, stage_configs, kwargs) logger.debug(f"stage_configs: {stage_configs}") - return config_path, stage_configs + return config_path, stage_configs, omni_lb_policy def get_final_stage_id_for_e2e( diff --git a/vllm_omni/experimental/__init__.py b/vllm_omni/experimental/__init__.py new file mode 100644 index 00000000000..c7d59e2e262 --- /dev/null +++ b/vllm_omni/experimental/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Experimental vLLM-Omni subsystems. + +Modules under ``vllm_omni.experimental`` are in active development and not yet +covered by the stability guarantees of the rest of the package — their APIs, +config surfaces, and defaults may change without notice. The AR-Diffusion +Engine (:mod:`vllm_omni.experimental.ar_diffusion`) lives here. +""" diff --git a/vllm_omni/experimental/ar_diffusion/__init__.py b/vllm_omni/experimental/ar_diffusion/__init__.py new file mode 100644 index 00000000000..b312660ec72 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +"""AR-Diffusion Engine (AR-Diffusion). + +The AR-Diffusion engine: a ``DiffusionEngine`` subclass that adds engine-level KV +cache management for autoregressive / chunked "world-model" diffusion models +(DreamZero and the AR-DiT family). Selected via +``OmniDiffusionConfig.engine_backend = "ar_diffusion"``. +""" + +from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine + +__all__ = ["ARDiffusionEngine"] diff --git a/vllm_omni/experimental/ar_diffusion/engine.py b/vllm_omni/experimental/ar_diffusion/engine.py new file mode 100644 index 00000000000..6fb6ee7c940 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/engine.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The AR-Diffusion Engine (AR-Diffusion).""" + +from __future__ import annotations + +from vllm.logger import init_logger + +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine + +logger = init_logger(__name__) + +#: Import path of the runner the AR-Diffusion engine routes its workers to. +AR_DIFFUSION_MODEL_RUNNER_CLS = "vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner" + + +class ARDiffusionEngine(DiffusionEngine): + """AR-Diffusion engine with engine-level KV cache management. + + AR-Diffusion serves autoregressive / chunked blockwise-causal diffusion models + (world models, AR-DiT) that materialize persistent attention KV. It reuses + vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / ``BlockTables``) + as a library, driven from the engine rather than hand-rolled inside each + model. See ``BDE_doc/diffusion_kv_cache_management.md`` for the design and + ``BDE_doc/dreamzero_kv_phase1_plan.md`` for the rollout. + + It is selected per model via ``OmniDiffusionConfig.engine_backend = "ar_diffusion"`` + (resolved by :meth:`DiffusionEngine.make_engine`), so models that do not opt + in keep using the base ``DiffusionEngine`` unchanged. + + Architecture note: in the multiproc setup the KV cache lives in the worker / + runner process (GPU side), co-located with the model and KV tensors — so the + actual KV *body* is :class:`~vllm_omni.experimental.ar_diffusion.kv_cache.manager.ARDiffusionKVCache`, owned + by :class:`~vllm_omni.experimental.ar_diffusion.runner.ARDiffusionModelRunner`. ``ARDiffusionEngine`` itself is the + thin selection / injection seam; it wires the AR-Diffusion executor → worker → runner + so DreamZero's rollout runs against the runner-owned KV cache. + """ + + #: Workers of this engine build the AR-Diffusion runner (unless the config + #: sets an explicit ``diffusion_model_runner_cls`` override). Declared on + #: the class so routing needs no od_config mutation at engine init. + default_diffusion_model_runner_cls = AR_DIFFUSION_MODEL_RUNNER_CLS diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py new file mode 100644 index 00000000000..ab29f6093ab --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/__init__.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: Apache-2.0 +"""AR-Diffusion engine-level KV cache helpers. + +Thin glue over vLLM's paged KV stack (``KVCacheManager`` / ``BlockPool`` / +``SlidingWindowManager``), used by the AR-Diffusion Engine to manage KV for +AR-diffusion models. See ``BDE_doc/dreamzero_kv_phase1_plan.md``. + +Layout: ``config`` (public knob) · ``paged`` (engine-generic paging mechanics + +chunk-window eviction spec) · ``manager`` (the ARDiffusionKVCache orchestrator + its +request adapter / pool builders) · ``state`` (the model-facing ARDiffusionKVState bridge). +""" + +from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig +from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ( + ARDiffusionKVCache, + ARDiffusionRequestAdapter, + build_kv_manager, + compute_num_blocks, +) +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( + ChunkWindowManager, + ChunkWindowSpec, + allocate_kv_pool_with_views, + chunk_slot_mapping, + compute_slot_mapping, + pool_write_chunk, + resident_block_ids, +) +from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( + ARDiffusionPagedForwardContext, + ARDiffusionPagedLayerContext, + ARDiffusionPagedLayerInputs, + ar_diffusion_paged_attention, + paged_write_attn, + set_current_paged_kv_cache, +) + +__all__ = [ + "ARDiffusionKVCache", + "ARDiffusionKVConfig", + "ARDiffusionPagedForwardContext", + "ARDiffusionPagedLayerContext", + "ARDiffusionPagedLayerInputs", + "ARDiffusionRequestAdapter", + "ChunkWindowManager", + "ChunkWindowSpec", + "allocate_kv_pool_with_views", + "ar_diffusion_paged_attention", + "build_kv_manager", + "chunk_slot_mapping", + "compute_num_blocks", + "compute_slot_mapping", + "paged_write_attn", + "pool_write_chunk", + "resident_block_ids", + "set_current_paged_kv_cache", +] diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/config.py b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py new file mode 100644 index 00000000000..4b3c737751d --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/config.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Configuration for AR-Diffusion engine-level KV cache management.""" + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass +class ARDiffusionKVConfig: + """Settings for the AR-Diffusion paged KV cache. + + Disabled by default: when ``enable`` is False the AR-Diffusion engine behaves exactly + like the base ``DiffusionEngine`` (no pool, no paged KV). + """ + + enable: bool = False + # Persistent KV tokens materialized per AR chunk. For DreamZero this is + # ``num_frame_per_block * frame_seqlen``. + chunk_size: int = 0 + # Resident window in chunks. ``None`` means full attention (no eviction). + window_chunks: int | None = None + # Protected leading chunks (attention sink); never evicted. + sink_chunks: int = 0 + # DreamZero-style window reset vs. VGGT-style sliding replace. + reset_at_boundary: bool = False + # Fraction of free device memory budgeted for the AR-Diffusion KV pool. + gpu_memory_fraction: float = 0.1 + # When CUDA graph / torch.compile is on (not enforce_eager), pre-capture the + # DiT graphs for every window-fill shape at load time via a synthetic rollout, + # so the serving run is fast from the first chunk. No effect when eager. + warmup_cudagraph: bool = True + # Also capture the post-window-boundary (reset-cycle) forward during warm-up. + warmup_capture_reset: bool = False + + @property + def sliding_window(self) -> int | None: + """Window size in tokens, or ``None`` for full attention.""" + if self.window_chunks is None: + return None + return self.window_chunks * self.chunk_size diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py new file mode 100644 index 00000000000..4b3b5d788d4 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/manager.py @@ -0,0 +1,497 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ARDiffusionKVCache — the engine-level KV cache orchestrator for one AR-Diffusion model. + +This is the *body* of AR-Diffusion's KV management: it owns a vLLM ``KVCacheManager`` (a +single chunk-window group) and the per-request adapter lifecycle, and exposes the +per-chunk operations a rollout needs — allocate, slot mapping, commit, window +lookup, free. It lives in the model runner (worker / GPU side), co-located with +the model and the KV tensors; the DreamZero pipeline calls these methods during a +rollout. The main-process ``ARDiffusionEngine`` only selects the engine and is otherwise +thin. +""" + +from __future__ import annotations + +import inspect +import os +from collections.abc import Sequence + +import torch +from vllm.logger import init_logger +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + KVCacheTensor, +) +from vllm.v1.request import RequestStatus + +from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import ( + ChunkWindowSpec, + allocate_kv_pool_with_views, + chunk_slot_mapping, + pool_write_chunk, + resident_block_ids, +) + +_log = init_logger(__name__) + + +class ARDiffusionRequestAdapter: + """Duck-types the subset of ``vllm.v1.request.Request`` that the + ``KVCacheManager`` reads (``allocate_slots`` / ``get_computed_blocks`` / + ``free`` and the coordinator they call into). + + It is intentionally NOT a full ``Request``. The conformance test exercises a + real ``KVCacheManager`` against this adapter so the surface cannot silently + drift across vLLM versions. + + A AR-Diffusion request advances one *chunk* at a time: ``allocate_slots`` is called + once per chunk and ``num_computed_tokens`` advances only when a chunk is + committed (:meth:`on_chunk_committed`), so the ``T`` denoise steps of a chunk + reuse the same slots. + """ + + def __init__( + self, + request_id: str, + *, + chunk_size: int, + prefill_prefix_tokens: int = 0, + ) -> None: + self.request_id = request_id + self._chunk_size = chunk_size + self._prefill = prefill_prefix_tokens + self._completed_chunks = 0 + # Filled only when cross-request prefix reuse is enabled (Phase 3). + self.block_hashes: list = [] + self.skip_reading_prefix_cache = True + self.num_preemptions = 0 + # vLLM watermark gate reads this; map the request lifecycle onto it. + self.status = RequestStatus.WAITING + + @property + def num_computed_tokens(self) -> int: + """Persistent KV already materialized (committed chunks + prefill).""" + return self._prefill + self._completed_chunks * self._chunk_size + + @property + def num_tokens(self) -> int: + """Total tokens once the in-flight chunk is committed.""" + return self._prefill + (self._completed_chunks + 1) * self._chunk_size + + @property + def num_prompt_tokens(self) -> int: + """The prefill prefix length (read by ``cache_blocks`` when caching).""" + return self._prefill + + @property + def completed_chunks(self) -> int: + return self._completed_chunks + + def on_chunk_committed(self) -> None: + """Advance by one chunk. Call once per chunk, not per denoise step.""" + self._completed_chunks += 1 + + +def compute_num_blocks( + available_bytes: int, + gpu_memory_fraction: float, + page_size_bytes: int, +) -> int: + """Number of KV blocks that fit in ``fraction`` of the memory budget.""" + if page_size_bytes <= 0: + raise ValueError(f"page_size_bytes must be positive, got {page_size_bytes}") + if not 0.0 < gpu_memory_fraction <= 1.0: + raise ValueError(f"gpu_memory_fraction must be in (0, 1], got {gpu_memory_fraction}") + budget = int(available_bytes * gpu_memory_fraction) + return max(0, budget // page_size_bytes) + + +def build_kv_manager( + spec: KVCacheSpec, + layer_names: Sequence[str], + num_blocks: int, + max_model_len: int, + *, + enable_caching: bool = False, +) -> KVCacheManager: + """Build a ``KVCacheManager`` with a single KV cache group for ``spec``. + + Args: + spec: The KV cache spec for the group (e.g. a ``ChunkWindowSpec``). + layer_names: Attention layers sharing this group's block table. + num_blocks: Total physical blocks in the pool. + max_model_len: Upper bound on a request's sequence length. + enable_caching: Cross-request prefix caching (Phase 3); off in Phase 1. + """ + layer_names = list(layer_names) + group = KVCacheGroupSpec(layer_names=layer_names, kv_cache_spec=spec) + tensors = [KVCacheTensor(size=spec.page_size_bytes * num_blocks, shared_by=layer_names)] + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=tensors, + kv_cache_groups=[group], + ) + kwargs = dict(max_model_len=max_model_len, hash_block_size=spec.block_size, enable_caching=enable_caching) + params = inspect.signature(KVCacheManager).parameters + if "scheduler_block_size" in params: + kwargs["scheduler_block_size"] = spec.block_size + if "max_num_batched_tokens" in params: + kwargs["max_num_batched_tokens"] = max_model_len + return KVCacheManager(config, **kwargs) + + +class ARDiffusionKVCache: + """Owns the paged KV pool + per-request lifecycle for a AR-Diffusion model. + + Build once per loaded model (dimensions known); then per request: + ``begin_request`` → per chunk (``allocate_chunk`` → ``chunk_write_slots`` → + [model writes K/V] → ``commit_chunk``) → ``end_request``. + """ + + def __init__( + self, + config: ARDiffusionKVConfig, + *, + num_layers: int, + num_kv_heads: int, + head_size: int, + dtype: torch.dtype, + block_size: int, + max_model_len: int, + available_bytes: int, + cross_attn_length: int = 0, + cross_attn_img_length: int = 0, + device: torch.device | None = None, + local_branches: int = 2, + num_frame_per_block: int = 1, + ) -> None: + if not config.enable: + raise ValueError("ARDiffusionKVCache built with a disabled ARDiffusionKVConfig") + if config.window_chunks is None: + raise ValueError("Phase 1 requires a bounded window (window_chunks)") + if config.chunk_size <= 0: + raise ValueError("ARDiffusionKVConfig.chunk_size must be set (> 0)") + if local_branches not in (1, 2): + raise ValueError(f"local_branches must be 1 or 2, got {local_branches}") + + self.config = config + # How many CFG branches allocate from THIS rank's pool. Under + # CFG-parallel x2 each rank executes exactly one branch (rank0 pos, + # rank1 neg) and the other branch's lazy contexts never allocate, so + # sizing for both would leave ~half the pool as idle capacity. A + # single-process run (TP1 / offline) executes both branches -> 2. + self.local_branches = local_branches + self.num_frame_per_block = max(1, int(num_frame_per_block)) + self.block_size = block_size + self.num_layers = num_layers + self.num_kv_heads = num_kv_heads + self.head_size = head_size + self.dtype = dtype + self.cross_attn_length = cross_attn_length + self.device = device or torch.device("cpu") + self._adapters: dict[str, ARDiffusionRequestAdapter] = {} + + # -- cross-attention pool (static, one-time-fill) ----------------------- + # Cross-attn KV is computed once from the text encoder and never changes. + # Allocate a per-layer contiguous tensor for each branch: (2, text_len, + # kv_heads, head_dim) where dim-0 = [pos, neg]. Allocated separately + # from the self-attn block pool (static, small — ~400 MiB for + # DreamZero I2V). Both pools draw from the same GPU free memory but + # the cross-attn pool is sized directly rather than via the block pool + # budget (no eviction/chunk lifecycle needed). + self.cross_attn_img_length = cross_attn_img_length + self._cross_k: list[torch.Tensor] = [] + self._cross_v: list[torch.Tensor] = [] + # I2V image-token cross-attn pool. Like the text k/v, the image-token + # k_img/v_img are session-invariant (the conditioning image doesn't change), + # so they are cached once and read every denoise step — see + # WanI2VCrossAttention (#4154 caches these model-side too). Empty for T2V. + self._cross_k_img: list[torch.Tensor] = [] + self._cross_v_img: list[torch.Tensor] = [] + + # Bytes the cross-attn pools consume directly (K+V, pos+neg, all layers). + # Deducted from the self-attn paged-pool budget below so the two + # allocations together stay within the GPU free-memory budget. + def _cross_pool_bytes(length: int) -> int: + return 2 * 2 * length * num_kv_heads * head_size * dtype.itemsize * num_layers + + cross_total_bytes = 0 + if device is not None and cross_attn_length > 0: + cross_shape = (2, cross_attn_length, num_kv_heads, head_size) + cross_bytes = _cross_pool_bytes(cross_attn_length) + cross_total_bytes += cross_bytes + for _ in range(num_layers): + self._cross_k.append(torch.empty(cross_shape, dtype=dtype, device=device)) + self._cross_v.append(torch.empty(cross_shape, dtype=dtype, device=device)) + _log.info( + "AR-Diffusion cross-attn pool: %d layers × (%d tok × %d heads × %d) = %.1f MiB", + num_layers, + cross_attn_length, + num_kv_heads, + head_size, + cross_bytes / (1024 * 1024), + ) + if cross_attn_img_length > 0: + img_shape = (2, cross_attn_img_length, num_kv_heads, head_size) + cross_total_bytes += _cross_pool_bytes(cross_attn_img_length) + for _ in range(num_layers): + self._cross_k_img.append(torch.empty(img_shape, dtype=dtype, device=device)) + self._cross_v_img.append(torch.empty(img_shape, dtype=dtype, device=device)) + _log.info( + "AR-Diffusion cross-attn IMG pool: %d layers × %d img-tok (I2V)", + num_layers, + cross_attn_img_length, + ) + + self.spec = ChunkWindowSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + sliding_window=config.window_chunks * config.chunk_size, + chunk_size=config.chunk_size, + window_chunks=config.window_chunks, + sink_chunks=config.sink_chunks, + reset_at_boundary=config.reset_at_boundary, + ) + # Each pool block spans all layers' K/V, so size against the per-layer + # page size times the layer count. + # Size the self-attn pool against the memory left after the cross-attn + # pools (allocated above), so cross + self-attn stays within the budget. + num_blocks = compute_num_blocks( + max(0, available_bytes - cross_total_bytes), + config.gpu_memory_fraction, + self.spec.page_size_bytes * num_layers, + ) + # Floor: one forward needs the resident window plus the in-flight chunk + # (num_frame_per_block frame-blocks) for every branch THIS rank runs, + # with a little eviction-transient headroom. The memory-fraction + # heuristic can under-size this once block_size grows — e.g. frame-granular + # paging at the true frame_seqlen makes each block larger and the pool + # fewer-blocks — so guarantee the minimum the rollout cannot run without, + # otherwise allocate_chunk hits an exhausted pool mid-forward. + min_blocks = self.local_branches * (config.window_chunks + self.num_frame_per_block) + 2 + if num_blocks < min_blocks: + _log.warning( + "AR-Diffusion KV pool: memory-fraction sizing gave %d blocks; raising to the %d-block " + "floor (%d local CFG branch(es) x (window_chunks=%d + num_frame_per_block=%d) + 2 headroom)", + num_blocks, + min_blocks, + self.local_branches, + config.window_chunks, + self.num_frame_per_block, + ) + num_blocks = min_blocks + layer_names = [f"ar_diffusion.layer.{i}" for i in range(num_layers)] + self.manager = build_kv_manager(self.spec, layer_names, num_blocks, max_model_len) + self.managed_num_blocks = num_blocks + self.num_blocks = num_blocks + # Scratch blocks are outside KVCacheManager ownership. They hold current + # denoise-step video KV when update_kv_cache=False and action/state KV in + # both modes; they are reused every forward and never committed. + # Default 4: measured high-water is num_frame_per_block video blocks (2) + # + 1 action block per branch per forward, +1 margin. + scratch_per_branch = int(os.environ.get("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH", "4")) + if scratch_per_branch <= 0: + raise ValueError("AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH must be positive") + self.scratch_blocks_per_branch = scratch_per_branch + self.scratch_num_blocks = self.local_branches * scratch_per_branch + self.num_blocks_total = self.managed_num_blocks + self.scratch_num_blocks + self.null_block_id = self.manager.block_pool.null_block.block_id + + # Allocate the per-layer paged K/V pools on the given device. + self._kv_pools: list[torch.Tensor] = [] + self._k_pools: list[torch.Tensor] = [] + self._v_pools: list[torch.Tensor] = [] + if device is not None: + self._kv_pools, self._k_pools, self._v_pools = allocate_kv_pool_with_views( + self.num_blocks_total, + block_size, + num_layers, + num_kv_heads, + head_size, + dtype, + device, + ) + + # -- cross-attention pool access ------------------------------------------- + # Cross-attn KV is static once populated — write once (from text encoder), + # read many (every denoising step). Not managed through the paged block pool. + + def write_cross_kv( + self, + layer_idx: int, + is_negative: bool, + k: torch.Tensor | None, + v: torch.Tensor | None, + k_img: torch.Tensor | None = None, + v_img: torch.Tensor | None = None, + ) -> None: + """Write one layer's cross-attn K/V into the pool. + + ``k`` / ``v``: ``(B, text_len, tp_num_heads, head_dim)``; only batch-0 + is copied (B=1 for inference). ``None`` skips the text half (window + restarts keep the still-valid text K/V and rewrite only the image half). + The ``is_negative`` flag selects the correct CFG branch slot. ``k_img`` / + ``v_img`` (I2V image tokens, ``(B, 257, ...)``) are written when the + image pool is allocated. + """ + branch = 1 if is_negative else 0 + if k is not None: + self._cross_k[layer_idx][branch].copy_(k[0]) + self._cross_v[layer_idx][branch].copy_(v[0]) + if k_img is not None and self._cross_k_img: + self._cross_k_img[layer_idx][branch].copy_(k_img[0]) + self._cross_v_img[layer_idx][branch].copy_(v_img[0]) + + def read_cross_kv(self, layer_idx: int, is_negative: bool) -> dict: + """Return a pool-backed cross-attn cache dict for one layer. + + The dict matches the ``{"is_init": True, "k": Tensor, "v": Tensor}`` + convention the cross-attention module expects — it reads from the pool + slice rather than from the lazy-initialised model-local dict. For I2V, + ``k_img`` / ``v_img`` are added so the image-token cache (added by #4154) + reads from the pool too. + """ + branch = 1 if is_negative else 0 + k = self._cross_k[layer_idx][branch].unsqueeze(0) # (1, L, heads, dim) + v = self._cross_v[layer_idx][branch].unsqueeze(0) + cache = {"is_init": True, "k": k, "v": v} + if self._cross_k_img: + cache["k_img"] = self._cross_k_img[layer_idx][branch].unsqueeze(0) + cache["v_img"] = self._cross_v_img[layer_idx][branch].unsqueeze(0) + return cache + + # -- request lifecycle --------------------------------------------------- + + def begin_request(self, request_id: str, *, prefill_prefix_tokens: int = 0) -> ARDiffusionRequestAdapter: + adapter = ARDiffusionRequestAdapter( + request_id, + chunk_size=self.spec.chunk_size, + prefill_prefix_tokens=prefill_prefix_tokens, + ) + self._adapters[request_id] = adapter + _log.debug("AR-Diffusion begin_request: req=%s prefill=%d", request_id, prefill_prefix_tokens) + return adapter + + def end_request(self, adapter: ARDiffusionRequestAdapter) -> None: + _log.debug( + "AR-Diffusion end_request: req=%s chunks=%d free=%d", + adapter.request_id, + adapter.completed_chunks, + self.manager.block_pool.get_num_free_blocks(), + ) + self.manager.free(adapter) + self._adapters.pop(adapter.request_id, None) + + # -- per-chunk operations ------------------------------------------------ + + def allocate_chunk(self, adapter: ARDiffusionRequestAdapter) -> list[int]: + """Allocate a chunk's blocks (evicting out-of-window blocks first). + + Returns the request's full block table (incl. null_block placeholders). + """ + blocks = self.manager.allocate_slots(adapter, num_new_tokens=self.spec.chunk_size) + if blocks is None: + raise RuntimeError("AR-Diffusion KV pool exhausted while allocating a chunk") + table = self.block_table(adapter) + resident = resident_block_ids(table, self.null_block_id) + _log.debug( + "AR-Diffusion allocate_chunk: req=%s chunk=%d table_len=%d resident=%d free=%d", + adapter.request_id, + adapter.completed_chunks, + len(table), + len(resident), + self.manager.block_pool.get_num_free_blocks(), + ) + return table + + def allocate_token_slots(self, adapter: ARDiffusionRequestAdapter, num_tokens: int) -> list[int]: + """Allocate managed blocks for an in-flight video span without committing it.""" + if num_tokens <= 0: + raise ValueError(f"num_tokens must be positive, got {num_tokens}") + blocks = self.manager.allocate_slots(adapter, num_new_tokens=num_tokens) + if blocks is None: + raise RuntimeError("AR-Diffusion KV pool exhausted while allocating paged attention slots") + return self.block_table(adapter) + + def block_table(self, adapter: ARDiffusionRequestAdapter) -> list[int]: + return list(self.manager.get_block_ids(adapter.request_id)[0]) + + def chunk_write_slots(self, adapter: ARDiffusionRequestAdapter) -> torch.Tensor: + """Slot mapping for the in-flight chunk — the K/V write target.""" + return chunk_slot_mapping( + self.block_table(adapter), + adapter.num_computed_tokens, + self.spec.chunk_size, + self.block_size, + ) + + def scratch_block_ids(self, is_negative: bool, start: int, count: int) -> list[int]: + """Return branch-local scratch block ids outside manager ownership.""" + if count < 0 or start < 0: + raise ValueError(f"scratch start/count must be non-negative, got start={start}, count={count}") + if start + count > self.scratch_blocks_per_branch: + raise RuntimeError( + "AR-Diffusion paged attention scratch blocks exhausted: " + f"need [{start}, {start + count}) of {self.scratch_blocks_per_branch}. " + "Increase AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH." + ) + # With one local branch (CFG-parallel x2) the rank's only branch maps to + # slot 0 whichever CFG side it is; the other branch never allocates here. + branch_offset = self.scratch_blocks_per_branch if (is_negative and self.local_branches == 2) else 0 + base = self.managed_num_blocks + branch_offset + start + return list(range(base, base + count)) + + def key_cache(self, layer_idx: int) -> torch.Tensor: + return self._kv_pools[layer_idx][0] + + def value_cache(self, layer_idx: int) -> torch.Tensor: + return self._kv_pools[layer_idx][1] + + def window_block_ids(self, adapter: ARDiffusionRequestAdapter) -> list[int]: + """Resident (non-null) managed blocks visible to paged attention.""" + return resident_block_ids(self.block_table(adapter), self.null_block_id) + + def commit_chunk(self, adapter: ARDiffusionRequestAdapter) -> None: + """Advance the adapter by one chunk after its K/V is written. + + This standalone primitive is used by low-level manager tests. In the + DreamZero paged-attention path, :meth:`ARDiffusionKVState.commit_paged_context` + advances the adapter only after the forward succeeds. Call once per + committed chunk, not per denoise step. + """ + _log.debug("AR-Diffusion commit: req=%s before=%d", adapter.request_id, adapter.completed_chunks) + adapter.on_chunk_committed() + _log.debug("AR-Diffusion commit: req=%s after=%d", adapter.request_id, adapter.completed_chunks) + + # -- pool-backed K/V access -------------------------------------------- + + def write_chunk_kv( + self, + layer_index: int, + new_k: torch.Tensor, + new_v: torch.Tensor, + adapter: ARDiffusionRequestAdapter, + ) -> None: + """Write one layer's committed-chunk K/V into the pool.""" + slots = self.chunk_write_slots(adapter) + _log.debug( + "AR-Diffusion write: req=%s layer=%d chunk=%d shapes=%s dev=%s", + adapter.request_id, + layer_index, + adapter.completed_chunks, + (tuple(new_k.shape), tuple(new_v.shape)), + slots.device, + ) + pool_write_chunk( + self._k_pools[layer_index], + self._v_pools[layer_index], + new_k, + new_v, + slots, + ) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py new file mode 100644 index 00000000000..a39a66b910f --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Generic paging mechanics + chunk-window eviction spec for the AR-Diffusion engine. + +Engine-generic, model-agnostic primitives — the layer a second model (e.g. the +Cosmos port) reuses unchanged. Three concerns live here: + +* **Slot mapping** — absolute token positions → physical KV-cache slots, the + standard PagedAttention layout ``slot(pos) = block_id(pos) * block_size + + (pos % block_size)``. +* **Pool I/O** — allocate per-layer FlashAttention-compatible paged K/V pools + plus flat compatibility views for direct slot writes. +* **Chunk-window eviction** — a ``SlidingWindowSpec`` subclass whose unit is a + *chunk* (``sliding_window = window_chunks * chunk_size``) plus a manager that + evicts at chunk boundaries. Memory policy / refcounting / ``null_block`` + replacement stay in vLLM's ``BlockPool``; only the token-skip math is here. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager +from vllm.v1.kv_cache_interface import SlidingWindowSpec +from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec + +# ── Slot mapping ──────────────────────────────────────────────────────────── + + +def compute_slot_mapping( + block_ids: Sequence[int], + positions: torch.Tensor | Sequence[int], + block_size: int, +) -> torch.Tensor: + """Map absolute token positions to physical KV-cache slots. + + Args: + block_ids: physical block id per block index (the request's block table). + positions: absolute token positions to map (1-D). + block_size: tokens per block. + + Returns: + ``LongTensor`` of physical slots, one per position. + """ + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + table = torch.as_tensor(block_ids, dtype=torch.long) + pos = torch.as_tensor(positions, dtype=torch.long) + block_index = torch.div(pos, block_size, rounding_mode="floor") + offset = pos % block_size + return table[block_index] * block_size + offset + + +def chunk_slot_mapping( + block_ids: Sequence[int], + num_computed_tokens: int, + chunk_size: int, + block_size: int, +) -> torch.Tensor: + """Slot mapping for the in-flight chunk's tokens (the commit write target). + + The chunk occupies absolute positions + ``[num_computed_tokens, num_computed_tokens + chunk_size)``. + """ + positions = torch.arange( + num_computed_tokens, + num_computed_tokens + chunk_size, + dtype=torch.long, + ) + return compute_slot_mapping(block_ids, positions, block_size) + + +def resident_block_ids(block_ids: Sequence[int], null_block_id: int) -> list[int]: + """Real (non-null) blocks currently resident, in table order. + + These are the blocks the read path gathers the attention window from; + out-of-window positions are the shared ``null_block`` and are excluded. + """ + return [int(b) for b in block_ids if int(b) != null_block_id] + + +# ── Pool I/O (allocate / write) ───────────────────────────────────────────── + + +def allocate_kv_pool_with_views( + num_blocks: int, + block_size: int, + num_layers: int, + num_kv_heads: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, +) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]]: + """Allocate vLLM-style paged KV pools plus flat slot-write views. + + The owning tensor follows FlashAttention's block-table cache layout: + ``(2, num_blocks, block_size, num_kv_heads, head_dim)``, where dim-0 is + ``[K, V]``. The flat K/V views keep ``slot = block_id * block_size + + offset`` writes simple without changing the kernel-facing cache layout. + """ + kv_pools: list[torch.Tensor] = [] + k_pools: list[torch.Tensor] = [] + v_pools: list[torch.Tensor] = [] + for _ in range(num_layers): + kv = torch.empty(2, num_blocks, block_size, num_kv_heads, head_dim, dtype=dtype, device=device) + kv_pools.append(kv) + k_pools.append(kv[0].reshape(num_blocks * block_size, num_kv_heads, head_dim)) + v_pools.append(kv[1].reshape(num_blocks * block_size, num_kv_heads, head_dim)) + return kv_pools, k_pools, v_pools + + +def pool_write_chunk( + k_pool: torch.Tensor, + v_pool: torch.Tensor, + new_k: torch.Tensor, + new_v: torch.Tensor, + slot_mapping: torch.Tensor, +) -> None: + """Write the committed chunk's K/V into pool slots (in place). + + ``new_k`` / ``new_v`` shape: ``(batch, chunk_size, num_kv_heads, head_dim)``. + ``slot_mapping`` shape: ``(chunk_size,)`` — one physical slot per token position. + The batch dim is written to the same slots (multi-batch writes the same + positions; the caller is responsible for per-sequence bookkeeping). + + This is the *commit* write — called once per chunk, after the last denoise step. + """ + k_pool[slot_mapping] = new_k[0] # batch index 0 + v_pool[slot_mapping] = new_v[0] + + +def _drop_batch(t: torch.Tensor) -> torch.Tensor: + """``(1, L, n_heads, head_dim)`` -> ``(L, n_heads, head_dim)``; pass 3-D through.""" + if t.dim() == 4 and t.shape[0] == 1: + return t[0] + return t + + +# ── Chunk-window eviction (spec + manager) ────────────────────────────────── + + +def chunk_window_skipped_tokens( + num_computed_tokens: int, + *, + chunk_size: int, + sliding_window: int, + sink_chunks: int, + reset_at_boundary: bool, +) -> int: + """Tokens outside the resident chunk window, snapped to a chunk boundary. + + Pure function so the eviction policy is unit-testable without constructing a + manager. Two strategies: + + - ``reset_at_boundary`` (DreamZero): at each chunk boundary everything past + the sink is dropped. + - otherwise (VGGT-style sliding replace): keep the last ``window`` tokens + (plus the sink); the skip count snaps down to a chunk boundary so a chunk + is never half-evicted. + """ + sink = sink_chunks * chunk_size + if reset_at_boundary: + completed = (num_computed_tokens // chunk_size) * chunk_size + return max(0, completed - sink) + skipped = max(0, num_computed_tokens - sliding_window - sink) + return (skipped // chunk_size) * chunk_size + + +class ChunkWindowManager(SlidingWindowManager): + """``SlidingWindowManager`` that evicts at chunk boundaries. + + ``self.sliding_window`` is set by the base ``__init__``; the chunk fields are + read from ``self.kv_cache_spec`` (a :class:`ChunkWindowSpec`). + """ + + def get_num_skipped_tokens(self, num_computed_tokens: int) -> int: + spec = self.kv_cache_spec + return chunk_window_skipped_tokens( + num_computed_tokens, + chunk_size=spec.chunk_size, + sliding_window=self.sliding_window, + sink_chunks=spec.sink_chunks, + reset_at_boundary=spec.reset_at_boundary, + ) + + +# Register so KVCacheManager resolves ChunkWindowSpec to ChunkWindowManager. +# Dispatch walks the spec's MRO, so without explicit registration the subclass +# would silently fall back to the parent SlidingWindowManager (override ignored). +# uniform_type_base_spec=None => its own KV cache group. +@register_kv_cache_spec(manager_class=ChunkWindowManager, uniform_type_base_spec=None) +@dataclass(frozen=True, kw_only=True) +class ChunkWindowSpec(SlidingWindowSpec): + # sliding_window (inherited) MUST equal window_chunks * chunk_size. + chunk_size: int + window_chunks: int + sink_chunks: int = 0 + reset_at_boundary: bool = False + + def __post_init__(self): + super().__post_init__() + if self.sliding_window != self.window_chunks * self.chunk_size: + raise ValueError( + "ChunkWindowSpec.sliding_window must equal " + f"window_chunks * chunk_size ({self.window_chunks} * " + f"{self.chunk_size}), got {self.sliding_window}" + ) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py new file mode 100644 index 00000000000..2038e496d31 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py @@ -0,0 +1,550 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Paged self-attention helpers for AR-Diffusion DreamZero KV reuse.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, ClassVar, NamedTuple + +import torch + +from vllm_omni.experimental.ar_diffusion.kv_cache.paged import compute_slot_mapping + +# Set by ARDiffusionPagedForwardContext.prepare() before each branch forward and +# read by the fused write+attend custom op below. The pools are process-lifetime +# allocations owned by one ARDiffusionKVCache per worker process (CFG-parallel / +# TP ranks are separate processes), mirroring vLLM's forward-context pattern for +# keeping multi-GiB cache tensors out of the compiled graph's inputs. +_CURRENT_PAGED_KV_CACHE: Any = None + + +def set_current_paged_kv_cache(kv_cache: Any) -> None: + global _CURRENT_PAGED_KV_CACHE + _CURRENT_PAGED_KV_CACHE = kv_cache + + +_LAYER_IDX_TENSORS: dict[int, torch.Tensor] = {} + + +def _layer_idx_tensor(layer_idx: int) -> torch.Tensor: + t = _LAYER_IDX_TENSORS.get(layer_idx) + if t is None: + t = torch.tensor(layer_idx, dtype=torch.int64) + _LAYER_IDX_TENSORS[layer_idx] = t + return t + + +class ARDiffusionPagedLayerInputs(NamedTuple): + """Compiled-region payload for one layer's paged self-attention. + + A NamedTuple of plain tensors + ints so ``torch.compile`` treats every field + as a pytree graph input (no object-attribute guards, no recompiles when only + tensor *values* change). All layers of one branch forward share the same + metadata tensor objects, built once by ``prepare()``. + + ``layer_idx`` is a 0-dim CPU tensor, NOT a python int: all 40 DiT blocks + share one compiled code object, and an int here becomes a per-layer dynamo + value guard (``layer_idx == k``) — 40 cache variants that blow the + recompile limit. A tensor input guards on shape/dtype only, so one graph + serves every layer. + """ + + layer_idx: torch.Tensor + seq_len: int + video_slots: torch.Tensor + action_slots: torch.Tensor + block_table: torch.Tensor + query_start_loc: torch.Tensor + seq_lens: torch.Tensor + max_query_len: int + max_seq_len: int + + +@dataclass +class ARDiffusionPagedForwardContext: + """Mutable branch-level state shared by all layer contexts in one forward.""" + + kv_cache: Any + adapter: Any + is_negative: bool + history_block_ids: list[int] + seq_len: int + commit_current: bool + max_video_tokens: int + current_video_block_ids: list[int] = field(default_factory=list) + current_video_slot_mapping: torch.Tensor | None = None + action_scratch_block_ids: list[int] = field(default_factory=list) + action_slot_mapping: torch.Tensor | None = None + query_len: int = 0 + kv_len: int = 0 + _allocated_video: bool = False + _committed: bool = False + _action_len: int = 0 + # Set once by prepare(); shared by all layers of the branch forward. + block_table: torch.Tensor | None = None + query_start_loc: torch.Tensor | None = None + seq_lens: torch.Tensor | None = None + max_query_len: int = 0 + max_seq_len: int = 0 + _prepared: bool = False + + @property + def block_size(self) -> int: + return int(self.kv_cache.block_size) + + @property + def num_current_video_blocks(self) -> int: + if self.seq_len % self.block_size != 0: + raise AssertionError( + "AR-Diffusion paged attention expects frame-aligned seq_len " + f"(multiple of block_size={self.block_size}), got {self.seq_len}" + ) + return self.seq_len // self.block_size + + def ensure_video_slots(self, device: torch.device) -> None: + """Allocate/write targets for the current video tokens, once per branch.""" + if self._allocated_video: + return + + n_blocks = self.num_current_video_blocks + if self.commit_current: + start = int(self.adapter.num_computed_tokens) + self.kv_cache.allocate_token_slots(self.adapter, self.seq_len) + table = self.kv_cache.block_table(self.adapter) + start_block = start // self.block_size + self.current_video_block_ids = [int(b) for b in table[start_block : start_block + n_blocks]] + positions = torch.arange(start, start + self.seq_len, dtype=torch.long) + self.current_video_slot_mapping = compute_slot_mapping(table, positions, self.block_size).to(device=device) + else: + self.current_video_block_ids = self.kv_cache.scratch_block_ids(self.is_negative, 0, n_blocks) + positions = torch.arange(self.seq_len, dtype=torch.long) + self.current_video_slot_mapping = compute_slot_mapping( + self.current_video_block_ids, + positions, + self.block_size, + ).to(device=device) + self._allocated_video = True + + def ensure_action_slots(self, action_len: int, device: torch.device) -> None: + """Reserve scratch slots for action/state K/V, if present.""" + if action_len <= 0: + self.action_scratch_block_ids = [] + self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device) + self._action_len = 0 + return + + self.ensure_video_slots(device) + if self.action_slot_mapping is not None and self._action_len == action_len: + return + + action_blocks = (action_len + self.block_size - 1) // self.block_size + scratch_offset = 0 if self.commit_current else len(self.current_video_block_ids) + self.action_scratch_block_ids = self.kv_cache.scratch_block_ids( + self.is_negative, + scratch_offset, + action_blocks, + ) + positions = torch.arange(action_len, dtype=torch.long) + self.action_slot_mapping = compute_slot_mapping( + self.action_scratch_block_ids, + positions, + self.block_size, + ).to(device=device) + self._action_len = action_len + + def video_block_table(self, device: torch.device) -> tuple[list[int], int]: + self.ensure_video_slots(device) + if self.max_video_tokens % self.block_size != 0: + raise AssertionError( + "AR-Diffusion paged attention requires max_video_tokens to be block-aligned, " + f"got max_video_tokens={self.max_video_tokens}, block_size={self.block_size}" + ) + all_video_blocks = self.history_block_ids + self.current_video_block_ids + max_video_blocks = self.max_video_tokens // self.block_size + visible_video_blocks = all_video_blocks[-max_video_blocks:] if max_video_blocks else [] + video_len = min(len(all_video_blocks) * self.block_size, self.max_video_tokens) + return visible_video_blocks, video_len + + def build_block_table( + self, + *, + action_len: int, + query_len: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]: + """Build FlashAttention block-table metadata for one self-attn call. + + The block table is tail-padded to a fixed width and ``max_seq_len`` is a + constant upper bound, so across window growth only tensor *values* change + — tensor shapes and the int consts stay stable for ``torch.compile`` + (and, later, CUDA-graph capture). The kernel only dereferences the first + ``ceil(seq_lens/block_size)`` entries, so padding is never read. + """ + video_blocks, video_len = self.video_block_table(device) + self.ensure_action_slots(action_len, device) + action_blocks = self.action_scratch_block_ids if action_len > 0 else [] + block_ids = video_blocks + action_blocks + if not block_ids: + raise RuntimeError("AR-Diffusion paged attention needs at least current video KV blocks") + + # Fixed capacity: full visible video window + one action-capacity block. + action_capacity_blocks = max(1, (action_len + self.block_size - 1) // self.block_size) + width = max(self.max_video_tokens // self.block_size + action_capacity_blocks, len(block_ids)) + padded = block_ids + [0] * (width - len(block_ids)) + + self.query_len = int(query_len) + self.kv_len = int(video_len + action_len) + max_seq_len = int(self.max_video_tokens + action_capacity_blocks * self.block_size) + block_table = torch.tensor([padded], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, self.query_len], dtype=torch.int32, device=device) + seq_lens = torch.tensor([self.kv_len], dtype=torch.int32, device=device) + return block_table, query_start_loc, seq_lens, self.query_len, max_seq_len + + def prepare(self, device: torch.device, action_len: int, query_len: int) -> None: + """Host-side, once-per-branch-forward setup (called OUTSIDE torch.compile). + + Allocates the current video/action slots (still lazy: only the branch a + CFG-parallel rank actually runs reaches its ``_forward_blocks``), builds + the padded block-table metadata ONCE for all layers, and publishes the + pool registry for the fused custom op. The compiled per-layer code then + only consumes prebuilt tensors via ``ARDiffusionPagedLayerInputs``. + """ + if getattr(self, "_prepared", False): + return + self.ensure_video_slots(device) + ( + self.block_table, + self.query_start_loc, + self.seq_lens, + self.max_query_len, + self.max_seq_len, + ) = self.build_block_table(action_len=action_len, query_len=query_len, device=device) + if self.action_slot_mapping is None: + self.action_slot_mapping = torch.empty(0, dtype=torch.long, device=device) + set_current_paged_kv_cache(self.kv_cache) + self._prepared = True + + def layer_inputs(self, layer_idx: int) -> ARDiffusionPagedLayerInputs: + if not getattr(self, "_prepared", False): + raise RuntimeError("ARDiffusionPagedForwardContext.layer_inputs() before prepare()") + return ARDiffusionPagedLayerInputs( + layer_idx=_layer_idx_tensor(layer_idx), + seq_len=int(self.seq_len), + video_slots=self.current_video_slot_mapping, + action_slots=self.action_slot_mapping, + block_table=self.block_table, + query_start_loc=self.query_start_loc, + seq_lens=self.seq_lens, + max_query_len=int(self.max_query_len), + max_seq_len=int(self.max_seq_len), + ) + + def mark_committed(self) -> None: + self._committed = True + + +@dataclass +class ARDiffusionPagedLayerContext: + """Layer-specific handle passed through DreamZero's existing ``kv_cache`` slot.""" + + is_ar_diffusion_paged_context: ClassVar[bool] = True + layer_idx: int + forward_ctx: ARDiffusionPagedForwardContext + + @property + def kv_cache(self): + return self.forward_ctx.kv_cache + + @property + def adapter(self): + return self.forward_ctx.adapter + + @property + def is_negative(self) -> bool: + return self.forward_ctx.is_negative + + @property + def history_block_ids(self) -> list[int]: + return self.forward_ctx.history_block_ids + + @property + def current_video_block_ids(self) -> list[int]: + return self.forward_ctx.current_video_block_ids + + @property + def current_video_slot_mapping(self) -> torch.Tensor | None: + return self.forward_ctx.current_video_slot_mapping + + @property + def action_scratch_block_ids(self) -> list[int]: + return self.forward_ctx.action_scratch_block_ids + + @property + def action_slot_mapping(self) -> torch.Tensor | None: + return self.forward_ctx.action_slot_mapping + + @property + def seq_len(self) -> int: + return self.forward_ctx.seq_len + + @property + def query_len(self) -> int: + return self.forward_ctx.query_len + + @property + def kv_len(self) -> int: + return self.forward_ctx.kv_len + + @property + def commit_current(self) -> bool: + return self.forward_ctx.commit_current + + def to_layer_inputs(self) -> ARDiffusionPagedLayerInputs: + """Compiled-region payload; requires ``forward_ctx.prepare()`` first.""" + return self.forward_ctx.layer_inputs(self.layer_idx) + + +def is_ar_diffusion_paged_context(value: object) -> bool: + return isinstance(value, ARDiffusionPagedLayerContext) + + +def _reference_paged_attention( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + softmax_scale: float, + *, + causal: bool, +) -> torch.Tensor: + if causal: + raise NotImplementedError("AR-Diffusion paged self-attention uses causal=False") + outs: list[torch.Tensor] = [] + block_size = key_cache.shape[1] + for i in range(seq_lens.shape[0]): + q_start = int(query_start_loc[i].item()) + q_end = int(query_start_loc[i + 1].item()) + kv_len = int(seq_lens[i].item()) + q = query[q_start:q_end] + positions = torch.arange(kv_len, device=query.device) + logical_blocks = torch.div(positions, block_size, rounding_mode="floor") + offsets = positions % block_size + physical_blocks = block_table[i, logical_blocks].long() + k = key_cache[physical_blocks, offsets] + v = value_cache[physical_blocks, offsets] + scores = torch.einsum("qhd,khd->hqk", q.float(), k.float()) * float(softmax_scale) + probs = torch.softmax(scores, dim=-1).to(v.dtype) + outs.append(torch.einsum("hqk,khd->qhd", probs, v)) + return torch.cat(outs, dim=0) + + +_FA_VERSION_BY_HEAD_SIZE: dict[int, int] = {} + + +def _resolve_fa_version(head_size: int) -> int: + # get_flash_attn_version -> current_platform.get_device_capability() is not + # dynamo-traceable, and the answer is fixed per head size for the process. + version = _FA_VERSION_BY_HEAD_SIZE.get(head_size) + if version is None: + try: + from vllm.v1.attention.backends.fa_utils import get_flash_attn_version + + version = int(get_flash_attn_version(requires_alibi=False, head_size=head_size) or 2) + except Exception: + version = 2 + _FA_VERSION_BY_HEAD_SIZE[head_size] = version + return version + + +def ar_diffusion_paged_attention( + query: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + max_query_len: int, + max_seq_len: int, + softmax_scale: float, + causal: bool = False, +) -> torch.Tensor: + """Run non-causal paged attention over a vLLM block table. + + ``query`` may be ``(B, L, H, D)`` or already flattened as ``(T, H, D)``. + ``key_cache`` / ``value_cache`` are ``(num_blocks, block_size, H, D)``. + """ + batched = query.dim() == 4 + if batched: + batch, q_len = query.shape[:2] + query_flat = query.reshape(batch * q_len, *query.shape[2:]) + else: + query_flat = query + + if not query_flat.is_cuda: + out = _reference_paged_attention( + query_flat, + key_cache, + value_cache, + block_table, + query_start_loc, + seq_lens, + softmax_scale, + causal=causal, + ) + else: + from vllm.vllm_flash_attn import flash_attn_varlen_func + + fa_version = _resolve_fa_version(query_flat.shape[-1]) + + out = torch.empty_like(query_flat) + flash_attn_varlen_func( + q=query_flat, + k=key_cache, + v=value_cache, + out=out, + cu_seqlens_q=query_start_loc, + max_seqlen_q=int(max_query_len), + seqused_k=seq_lens, + max_seqlen_k=int(max_seq_len), + softmax_scale=float(softmax_scale), + causal=causal, + block_table=block_table, + fa_version=fa_version, + ) + + if batched: + return out.reshape(query.shape) + return out + + +# ── Fused write+attend custom op (torch.compile-safe) ────────────────────── +# +# One opaque op per layer keeps the compiled DiT block fullgraph: dynamo treats +# it as a single graph node (no eager island, no graph breaks), and the K/V slot +# writes happen inside the op so write→read ordering with the block-table kernel +# is internal — no hidden-mutation ordering hazards against a separate reader op +# and no multi-GiB pool tensors as graph inputs (the flat and paged pool views +# alias the same storage, which AOTAutograd handles poorly as inputs). Pattern +# follows sage_attn3.py in-repo and vLLM's own unified attention ops. +def _paged_write_attn_impl( + query: torch.Tensor, + k_curr: torch.Tensor, + v_curr: torch.Tensor, + k_act: torch.Tensor | None, + v_act: torch.Tensor | None, + layer_idx: torch.Tensor, + video_slots: torch.Tensor, + action_slots: torch.Tensor, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + max_query_len: int, + max_seq_len: int, + softmax_scale: float, +) -> torch.Tensor: + kv = _CURRENT_PAGED_KV_CACHE + if kv is None: + raise RuntimeError("ar_diffusion_paged_write_attn called before prepare() set the KV pool registry") + layer_idx = int(layer_idx) + k_pool = kv._k_pools[layer_idx] + v_pool = kv._v_pools[layer_idx] + k_pool[video_slots] = k_curr.to(k_pool.dtype) + v_pool[video_slots] = v_curr.to(v_pool.dtype) + if k_act is not None and v_act is not None and k_act.shape[0] > 0: + k_pool[action_slots] = k_act.to(k_pool.dtype) + v_pool[action_slots] = v_act.to(v_pool.dtype) + return ar_diffusion_paged_attention( + query, + kv.key_cache(layer_idx), + kv.value_cache(layer_idx), + block_table=block_table, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + max_query_len=max_query_len, + max_seq_len=max_seq_len, + softmax_scale=softmax_scale, + causal=False, + ) + + +# hasattr guard keeps registration idempotent across test re-imports that pop +# the module from sys.modules (same as sage_attn3.py). +if not hasattr(torch.ops.vllm_omni, "ar_diffusion_paged_write_attn"): + + @torch.library.custom_op("vllm_omni::ar_diffusion_paged_write_attn", mutates_args=()) + def _paged_write_attn_op( + query: torch.Tensor, + k_curr: torch.Tensor, + v_curr: torch.Tensor, + k_act: torch.Tensor | None, + v_act: torch.Tensor | None, + layer_idx: torch.Tensor, + video_slots: torch.Tensor, + action_slots: torch.Tensor, + block_table: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + max_query_len: int, + max_seq_len: int, + softmax_scale: float, + ) -> torch.Tensor: + return _paged_write_attn_impl( + query, + k_curr, + v_curr, + k_act, + v_act, + layer_idx, + video_slots, + action_slots, + block_table, + query_start_loc, + seq_lens, + max_query_len, + max_seq_len, + softmax_scale, + ) + + @_paged_write_attn_op.register_fake + def _( + query, + k_curr, + v_curr, + k_act, + v_act, + layer_idx, + video_slots, + action_slots, + block_table, + query_start_loc, + seq_lens, + max_query_len, + max_seq_len, + softmax_scale, + ): + return torch.empty_like(query) + + +def paged_write_attn( + inputs: ARDiffusionPagedLayerInputs, query, k_curr, v_curr, k_act, v_act, softmax_scale: float +) -> torch.Tensor: + """Model-facing entry: routes through the custom op (traceable in fullgraph).""" + return torch.ops.vllm_omni.ar_diffusion_paged_write_attn( + query, + k_curr, + v_curr, + k_act, + v_act, + inputs.layer_idx, + inputs.video_slots, + inputs.action_slots, + inputs.block_table, + inputs.query_start_loc, + inputs.seq_lens, + inputs.max_query_len, + inputs.max_seq_len, + softmax_scale, + ) diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/state.py b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py new file mode 100644 index 00000000000..510bbe043a0 --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/state.py @@ -0,0 +1,183 @@ +# SPDX-License-Identifier: Apache-2.0 +"""DreamZero-facing AR-Diffusion KV bridge backed by paged self-attention.""" + +from __future__ import annotations + +from vllm.logger import init_logger + +from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( + ARDiffusionPagedForwardContext, + ARDiffusionPagedLayerContext, +) + +_log = init_logger(__name__) + + +class ARDiffusionKVState: + """Per-request bridge between a DreamZero rollout and the AR-Diffusion KV pool. + + The self-attention path is paged-only: ``get_kv_caches`` returns lightweight + per-layer contexts, and each attention layer writes current K/V into the + pool before calling the FlashAttention block-table kernel. No resident + window is gathered into a contiguous tensor. + """ + + def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None: + """One state per DreamZero session, holding BOTH CFG branches' adapters. + + ``pos``/``neg`` are the classifier-free-guidance branches: ``pos`` is + the conditional pass (real prompt), ``neg`` the unconditional pass + (negative prompt). Each owns an independent request adapter — and + therefore independent pool blocks — because the two branches produce + different K/V for the same frames. Every ``is_negative`` flag in this + module selects between them (``False`` -> ``pos``, ``True`` -> ``neg``). + """ + self.kv_cache = kv_cache + self.pos = pos_adapter + self.neg = neg_adapter + self.num_layers = num_layers + self._committed: dict[bool, int] = {False: 0, True: 0} + self._paged_pending: dict[bool, ARDiffusionPagedForwardContext | None] = {False: None, True: None} + # Cross-attn KV is populated once after text encoding (eager), then + # read every denoising step. Tracked per half: the text K/V depend only + # on the session prompt (survive window-boundary resets), while the I2V + # image-token K/V come from the current observation's CLIP features and + # must be re-projected on every window restart. + self._cross_text_populated: dict[bool, bool] = {False: False, True: False} + self._cross_img_populated: dict[bool, bool] = {False: False, True: False} + + def _adapter(self, is_negative: bool): + return self.neg if is_negative else self.pos + + def get_kv_caches( + self, + is_negative: bool, + seq_len: int | None = None, + commit_current: bool = False, + ) -> list[ARDiffusionPagedLayerContext]: + if seq_len is None: + raise ValueError("AR-Diffusion paged self-attention requires seq_len in get_kv_caches()") + return self.prepare_paged_context(is_negative, seq_len, commit_current) + + def prepare_paged_context( + self, + is_negative: bool, + seq_len: int, + commit_current: bool, + ) -> list[ARDiffusionPagedLayerContext]: + """Return per-layer paged attention contexts for one branch forward. + + Allocation is lazy: constructing both CFG branches' kwargs must not + allocate manager blocks for a branch that a CFG-parallel rank will not run. + The first self-attn layer that consumes the context allocates the current + video slots, then all layers share that branch-level state. + """ + cs = self.kv_cache.spec.chunk_size + if int(seq_len) % cs != 0: + raise AssertionError( + f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {seq_len}" + ) + + pending = self._paged_pending.get(is_negative) + if pending is not None and pending.commit_current and pending._allocated_video and not pending._committed: + raise RuntimeError("AR-Diffusion paged context replaced before its managed current chunk was committed") + + adapter = self._adapter(is_negative) + branch = "neg" if is_negative else "pos" + forward_ctx = ARDiffusionPagedForwardContext( + kv_cache=self.kv_cache, + adapter=adapter, + is_negative=is_negative, + history_block_ids=self.kv_cache.window_block_ids(adapter), + seq_len=int(seq_len), + commit_current=bool(commit_current), + max_video_tokens=int(self.kv_cache.spec.sliding_window), + ) + self._paged_pending[is_negative] = forward_ctx + _log.debug( + "AR-Diffusion GET [%s] source=paged-attn layers=%d history_blocks=%d seq_len=%d commit_current=%s", + branch, + self.num_layers, + len(forward_ctx.history_block_ids), + int(seq_len), + bool(commit_current), + ) + return [ARDiffusionPagedLayerContext(layer_idx=i, forward_ctx=forward_ctx) for i in range(self.num_layers)] + + def commit_paged_context(self, is_negative: bool) -> None: + """Commit the managed current video blocks after a successful forward.""" + ctx = self._paged_pending.get(is_negative) + if ctx is None: + return + branch = "neg" if is_negative else "pos" + if ctx.commit_current and ctx._allocated_video: + n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size + for _ in range(n_chunks): + ctx.adapter.on_chunk_committed() + self._committed[is_negative] += ctx.seq_len + _log.debug( + "AR-Diffusion COMMIT [%s] paged-attn new_tokens=%d chunks=%d resident=%d/%d", + branch, + ctx.seq_len, + n_chunks, + len(self.kv_cache.window_block_ids(ctx.adapter)), + self.kv_cache.spec.window_chunks, + ) + ctx.mark_committed() + self._paged_pending[is_negative] = None + + def get_cross_kv_caches(self, is_negative: bool) -> list[dict]: + """Return pool-backed cross-attn cache dicts. + + Under AR-Diffusion the cross-attn pool is always populated by ``_kv_populate_cross`` + before the first read, so this must never be reached unpopulated — if it + were, the model would lazily project and write cross KV itself, violating + engine ownership. Fail loud rather than fall back to the model. + """ + img_ok = not self.kv_cache._cross_k_img or self._cross_img_populated[is_negative] + if not (self._cross_text_populated[is_negative] and img_ok and self.kv_cache.cross_attn_length > 0): + raise RuntimeError( + f"AR-Diffusion cross-attn read before _kv_populate_cross (neg={is_negative}, " + f"cross_attn_length={self.kv_cache.cross_attn_length}, " + f"text={self._cross_text_populated[is_negative]} " + f"img={self._cross_img_populated[is_negative]}) — the engine must own all cross KV" + ) + return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)] + + def close(self) -> None: + """Final teardown: free both branches' resident pool blocks.""" + self.kv_cache.end_request(self.pos) + self.kv_cache.end_request(self.neg) + self._paged_pending = {False: None, True: None} + + def reset(self, *, keep_cross_text: bool = False) -> None: + """Drop this session's KV — mirrors the model-local ``state.reset()``. + + DreamZero resets at the attention-window boundary (``should_reset``); the + model starts a fresh sliding window, so AR-Diffusion must free the resident pool + blocks and start a fresh adapter for each branch. The ``ARDiffusionKVState`` + object (and thus ``pipeline._ar_diffusion_kv_state``) is preserved across the reset + so the runner's session mapping stays valid — only the pool-backed state + is recycled. + + Args: + keep_cross_text: On window ("inference") resets the session prompt is + unchanged, so the text cross-attn K/V in the static pool buffers are + still valid — keep them and only force the image half to repopulate. + Session resets (prompt change / explicit reset) pass ``False``. + """ + pos_id, neg_id = self.pos.request_id, self.neg.request_id + self.close() + self.pos = self.kv_cache.begin_request(pos_id) + self.neg = self.kv_cache.begin_request(neg_id) + self._committed = {False: 0, True: 0} + self._paged_pending = {False: None, True: None} + if not keep_cross_text: + self._cross_text_populated = {False: False, True: False} + self._cross_img_populated = {False: False, True: False} + _log.info( + "AR-Diffusion RESET [%s/%s] session KV cleared (window boundary; cross text %s)", + pos_id, + neg_id, + "kept" if keep_cross_text else "cleared", + ) diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py new file mode 100644 index 00000000000..862b3d9e8bc --- /dev/null +++ b/vllm_omni/experimental/ar_diffusion/runner.py @@ -0,0 +1,409 @@ +# SPDX-License-Identifier: Apache-2.0 +"""ARDiffusionModelRunner — the diffusion model runner for the AR-Diffusion engine. + +Subclasses ``DiffusionModelRunner`` and owns the engine-level KV cache +(:class:`ARDiffusionKVCache`). It brackets a request's rollout with the KV lifecycle +(``begin_request`` / ``end_request``) and exposes the live ``ARDiffusionKVCache`` so the +model pipeline (DreamZero) can drive the per-chunk allocate / slot-mapping / +gather / commit operations during ``pipeline.forward``. + +Selecting the AR-Diffusion engine enables KV management (no env gate); optional KV +parameters come from the deploy yaml's ``model_config["ar_diffusion_kv_config"]``. +""" + +from __future__ import annotations + +import dataclasses +import math +import time +from collections import OrderedDict + +import numpy as np +import torch +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import MAX_DREAMZERO_SESSIONS +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner +from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVConfig +from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache +from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + +logger = init_logger(__name__) + + +def resolve_ar_diffusion_kv_config(od_config: OmniDiffusionConfig) -> ARDiffusionKVConfig: + """Resolve the AR-Diffusion KV config from the deploy/command arg. + + Selecting the AR-Diffusion engine implies KV management, so the result is always + enabled. Optional overrides (``window_chunks`` / ``gpu_memory_fraction`` / + ``sink_chunks`` / ``reset_at_boundary``) may be supplied via the deploy yaml as + ``model_config["ar_diffusion_kv_config"]``. ``chunk_size`` / ``window_chunks`` are + finalized from the model geometry at load (see ``_preallocate_kv_cache``). + """ + raw = getattr(od_config, "ar_diffusion_kv_config", None) + if raw is None: + model_config = getattr(od_config, "model_config", None) + if isinstance(model_config, dict): + raw = model_config.get("ar_diffusion_kv_config") + if isinstance(raw, ARDiffusionKVConfig): + return dataclasses.replace(raw, enable=True) + if isinstance(raw, dict): + return ARDiffusionKVConfig(**{**raw, "enable": True}) + return ARDiffusionKVConfig(enable=True) + + +class ARDiffusionModelRunner(DiffusionModelRunner): + def __init__(self, vllm_config, od_config: OmniDiffusionConfig, device) -> None: + super().__init__(vllm_config, od_config, device) + self.ar_diffusion_kv_config = resolve_ar_diffusion_kv_config(od_config) + # Built after the model is loaded (dimensions known); stays None while KV + # management is disabled. + self.kv_cache: ARDiffusionKVCache | None = None + # DreamZero KV is session-scoped (persists across a session's forwards), + # so AR-Diffusion KV state is keyed by session_id and reused, not created per request. + # Bounded by the same LRU cap as DreamZeroPipeline._states: an evicted + # session's ARDiffusionKVState owns pool blocks (two CFG adapters), so without a + # bound, session-id churn in a long-running server would leak pool ownership + # until the KV pool is exhausted. Evicting frees those blocks (state.close()). + self._ar_diffusion_states: OrderedDict[str, ARDiffusionKVState] = OrderedDict() + self._max_ar_diffusion_states = MAX_DREAMZERO_SESSIONS + # Per-request server-side forward E2E times (seconds), recorded in + # execute_model. This is the worker-side compute time — the analog of the + # upstream server's per-request E2E — and excludes the engine<->worker IPC + # that the client's omni.generate wall time includes. Read (and cleared) + # by the ar_diffusion_perf_stats worker RPC for the perf-compare summary. + self._perf_e2e_times: list[float] = [] + + def build_kv_cache( + self, + *, + num_layers: int, + num_kv_heads: int, + head_size: int, + dtype, + block_size: int, + max_model_len: int, + available_bytes: int, + cross_attn_length: int = 0, + cross_attn_img_length: int = 0, + local_branches: int = 2, + num_frame_per_block: int = 1, + ) -> None: + """Construct the ARDiffusionKVCache (preallocating GPU pools on ``self.device``). + + No-op when KV management is disabled. + """ + if not self.ar_diffusion_kv_config.enable: + return + self.kv_cache = ARDiffusionKVCache( + self.ar_diffusion_kv_config, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=dtype, + block_size=block_size, + max_model_len=max_model_len, + available_bytes=available_bytes, + cross_attn_length=cross_attn_length, + cross_attn_img_length=cross_attn_img_length, + local_branches=local_branches, + num_frame_per_block=num_frame_per_block, + device=self.device, + ) + logger.info( + "AR-Diffusion KV cache enabled: %d blocks, chunk_size=%d, window_chunks=%s, " + "layers=%d kv_heads=%d head_dim=%d block_size=%d cross_attn_len=%d device=%s", + self.kv_cache.num_blocks, + self.ar_diffusion_kv_config.chunk_size, + self.ar_diffusion_kv_config.window_chunks, + num_layers, + num_kv_heads, + head_size, + block_size, + cross_attn_length, + self.device, + ) + + # -- preallocation at load ------------------------------------------------- + + def load_model(self, *args, **kwargs): + super().load_model(*args, **kwargs) + if self.ar_diffusion_kv_config.enable and self.pipeline is not None: + self._preallocate_kv_cache() + # Pre-capture the CUDA graphs for every window-fill shape at load time + # (only when compiling/cuda-graph is on) so serving is fast from chunk 0. + if not self.od_config.enforce_eager and self.ar_diffusion_kv_config.warmup_cudagraph: + self._warmup_ar_rollout() + + def _infer_frame_seqlen(self) -> int: + """frame_seqlen = (H//8)*(W//8)//4 from the configured image_resolution.""" + mc = getattr(self.od_config, "model_config", None) or {} + psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {} + res = psc.get("image_resolution", [180, 320]) + h, w = int(res[0]), int(res[1]) + return (h // 8) * (w // 8) // 4 + + def _preallocate_kv_cache(self) -> None: + """Build the KV pool once, from the loaded DreamZero transformer geometry.""" + # The paged self-attention path carries one rollout per request: slot + # mappings/adapters are single-sequence and K/V writes address batch + # index 0 only. Reject a multi-sequence config at init instead of + # failing (or silently dropping sequences) mid-forward. + max_num_seqs = int(getattr(self.od_config, "max_num_seqs", 1) or 1) + if max_num_seqs > 1: + raise ValueError( + "AR-Diffusion paged KV supports max_num_seqs=1 (single-sequence " + f"rollouts); got max_num_seqs={max_num_seqs}. Per-sequence " + "adapters/slot maps are not implemented." + ) + t = self.pipeline.transformer + num_layers = int(t.num_layers) + num_kv_heads = int(getattr(t.blocks[0].self_attn, "tp_num_heads", t.num_heads)) + head_size = int(t.dim // t.num_heads) + num_frame_per_block = int(t.num_frame_per_block) + local_attn_size = int(t.local_attn_size) + # The loaded transformer is the source of truth for frame_seqlen (it derives + # max_attention_size from it). The deploy-config inference is only a + # cross-check — warn on drift rather than silently sizing chunks wrong. + frame_seqlen = int(getattr(t, "frame_seqlen", 0)) or self._infer_frame_seqlen() + inferred = self._infer_frame_seqlen() + if frame_seqlen != inferred: + logger.warning( + "AR-Diffusion frame_seqlen mismatch: transformer=%d deploy-config=%d; using transformer", + frame_seqlen, + inferred, + ) + # The model's own attention window, in tokens (= the [-max_attention_size:] + # slice it applies). Read it directly rather than recomputing. + max_attention_size = int(t.blocks[0].self_attn.max_attention_size) + # Cross-attn text length — static pool for the text-encoder output. + # The cached k/v spans the full text sequence (the image tokens + # prepended in _forward_blocks are stripped inside the cross-attn + # forward but are not part of the cached k/v). + cross_attn_length = int(getattr(t, "text_len", 0)) + # I2V image-token cross-attn length: the model splits context[:, :257] as the + # image tokens (the img_emb(clip_feature) output), caching k_img/v_img. T2V: 0. + cross_attn_img_length = 257 if getattr(t, "model_type", "t2v") == "i2v" else 0 + + # Frame-granular paging: 1 block = 1 frame = frame_seqlen tokens, so the + # resident window matches max_attention_size exactly (it need not be a whole + # number of num_frame_per_block causal blocks). + chunk_size = frame_seqlen + window_chunks = self.ar_diffusion_kv_config.window_chunks or (max_attention_size // frame_seqlen) + + self.ar_diffusion_kv_config = dataclasses.replace( + self.ar_diffusion_kv_config, chunk_size=chunk_size, window_chunks=window_chunks + ) + free_bytes = torch.cuda.mem_get_info(self.device)[0] + # Under CFG-parallel each rank executes exactly ONE branch (rank0 pos, + # rank1 neg; the other branch's lazy contexts never allocate on this + # rank), so its pool only needs one branch's capacity. A single-process + # run (cfg world 1) executes both branches from one pool. + try: + from vllm_omni.diffusion.distributed.parallel_state import ( + get_classifier_free_guidance_world_size, + ) + + cfg_world = int(get_classifier_free_guidance_world_size()) + except Exception: + cfg_world = 1 + local_branches = 1 if cfg_world >= 2 else 2 + logger.info( + "AR-Diffusion preallocating (paged): frame_seqlen=%d num_frame_per_block=%d " + "local_attn_size=%d -> chunk_size=%d window_chunks=%d (window=%d tokens)", + frame_seqlen, + num_frame_per_block, + local_attn_size, + chunk_size, + window_chunks, + window_chunks * chunk_size, + ) + self.build_kv_cache( + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_size=head_size, + dtype=self.od_config.dtype, + block_size=chunk_size, # one pool block per chunk + max_model_len=1 << 20, + available_bytes=free_bytes, + cross_attn_length=cross_attn_length, + cross_attn_img_length=cross_attn_img_length, + local_branches=local_branches, + num_frame_per_block=num_frame_per_block, + ) + + def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput: + # KV disabled -> base behavior, unchanged. + if self.kv_cache is None: + return super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs) + + kv = self.kv_cache + # DreamZero KV is session-scoped (the model state persists across a + # session's forwards), so the AR-Diffusion KV state is keyed by session_id and + # reused — matching how pipeline.forward resolves the model-local state. + extra_args = req.sampling_params.extra_args or {} + session_id = str(extra_args.get("session_id") or "default") + state = self._ar_diffusion_states.get(session_id) + if state is None: + pos = kv.begin_request(f"bde__{session_id}") + neg = kv.begin_request(f"bde__{session_id}__neg") + state = ARDiffusionKVState(kv, pos, neg, num_layers=kv.num_layers) + self._ar_diffusion_states[session_id] = state + # Evict the least-recently-used session(s) past the cap, freeing their + # pool blocks — mirrors DreamZeroPipeline._states so the AR-Diffusion pool can't + # outlive the model-local state map it shadows. + while len(self._ar_diffusion_states) > self._max_ar_diffusion_states: + old_id, old_state = self._ar_diffusion_states.popitem(last=False) + old_state.close() + logger.debug("AR-Diffusion evicted session=%s (LRU); freed pool blocks", old_id) + # Track recency on every access (hit or miss) so the LRU order is correct. + self._ar_diffusion_states.move_to_end(session_id) + logger.debug( + "AR-Diffusion execute_model: req=%s session=%s chunk_size=%d window_chunks=%s num_blocks=%d", + req.request_id, + session_id, + kv.spec.chunk_size, + kv.config.window_chunks, + kv.num_blocks, + ) + self.pipeline._ar_diffusion_kv_state = state + # Time the worker-side forward (true per-request compute E2E). The CUDA + # sync below makes the stop reflect completed GPU work, not just dispatch. + _e2e_t0 = time.perf_counter() + try: + out = super().execute_model(req, kv_prefetch_jobs=kv_prefetch_jobs) + except Exception: + # Transactional containment: a forward that died partway may have + # written some layers' K/V into allocated-but-uncommitted blocks + # and advanced model-local state. Tear the whole session down + # (engine pool blocks freed via close(); model-local state dropped) + # so the next request with this session id starts clean instead of + # tripping the pending-commit guard or reading half-written KV. + broken = self._ar_diffusion_states.pop(session_id, None) + if broken is not None: + broken.close() + states = getattr(self.pipeline, "_states", None) + if states is not None: + states.pop(session_id, None) + logger.warning( + "AR-Diffusion execute_model failed for session=%s; session state torn down " + "(pool blocks freed) — the next request starts a fresh session", + session_id, + ) + raise + finally: + self.pipeline._ar_diffusion_kv_state = None + if self.device is not None and torch.cuda.is_available(): + torch.accelerator.synchronize(self.device) + self._perf_e2e_times.append(time.perf_counter() - _e2e_t0) + return out + + # -- cuda-graph warm-up ---------------------------------------------------- + + #: Dedicated session id for the warm-up rollout; cleaned from both state maps. + _WARMUP_SID = "__ardiffusion_warmup__" + + def _synth_robot_obs(self, h: int, w: int, n_frames: int) -> dict: + """Zero-filled roboarena observation matching the client's obs layout. + + Mirrors ``examples/.../client_schedule.make_obs_from_video`` (2 ext + wrist + cameras, 7/6/1 state vectors) so it survives ``_transform_robot_obs`` exactly + like a real request. One frame -> ``(H,W,3)``; chunk -> ``(n_frames,H,W,3)``. + """ + img = np.zeros((h, w, 3), dtype=np.uint8) if n_frames == 1 else np.zeros((n_frames, h, w, 3), dtype=np.uint8) + return { + "observation/exterior_image_0_left": img, + "observation/exterior_image_1_left": img, + "observation/wrist_image_left": img, + "observation/joint_position": np.zeros(7, dtype=np.float32), + "observation/cartesian_position": np.zeros(6, dtype=np.float32), + "observation/gripper_position": np.zeros(1, dtype=np.float32), + "prompt": "warmup", + "session_id": self._WARMUP_SID, + } + + def _warmup_ar_rollout(self) -> None: + """Drive a synthetic rollout through one window-fill so every DiT graph is + captured at load (off the serving hot path). Frees the synthetic session and + leaves the KV pool exactly as found. Never raises — falls back to lazy capture. + """ + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + kv = self.kv_cache + if kv is None: + return + sid = self._WARMUP_SID + t = self.pipeline.transformer + nfpb = int(t.num_frame_per_block) + window_chunks = int(kv.config.window_chunks or 0) + if window_chunks <= 0: + return + # Distinct DiT window shapes are walked by one fill: prefill + ceil((W-1)/nfpb) + # chunk forwards. +1 optionally to also capture the reset-cycle forward. + n_forwards = 1 + math.ceil(max(0, window_chunks - 1) / nfpb) + if self.ar_diffusion_kv_config.warmup_capture_reset: + n_forwards += 1 + res = self._image_resolution() + h, w = int(res[0]), int(res[1]) + + free_before = kv.manager.block_pool.get_num_free_blocks() + # The step cache is left active: the DiT input shape is per-chunk (the resident + # window length), and the first denoise step of every chunk always computes, so + # every window shape is captured regardless of how many later steps are skipped. + logger.info( + "AR-Diffusion cudagraph warm-up: up to %d synthetic forwards (window_chunks=%d, nfpb=%d)", + n_forwards, + window_chunks, + nfpb, + ) + try: + for i in range(n_forwards): + n_frames = 1 if i == 0 else 4 # client convention: 1-frame prefill, 4-frame chunks + obs = self._synth_robot_obs(h, w, n_frames) + sp = OmniDiffusionSamplingParams(extra_args={"reset": i == 0, "session_id": sid, "robot_obs": obs}) + req = OmniDiffusionRequest(prompts=["warmup"], sampling_params=sp, request_id=f"ardiffusion-warmup-{i}") + self.execute_model(req) + # Stop once the resident window is full — the remaining shapes are + # already captured (the window caps/resets through the same set). + ar_state = self._ar_diffusion_states.get(sid) + if ( + not self.ar_diffusion_kv_config.warmup_capture_reset + and ar_state is not None + and len(kv.window_block_ids(ar_state.pos)) >= window_chunks + ): + break + except Exception as e: # noqa: BLE001 — warm-up must never break model load + logger.warning("AR-Diffusion cudagraph warm-up failed (%s); using lazy capture.", e) + finally: + # Free the synthetic session from BOTH maps (KV pool + pipeline state). + ar_state = self._ar_diffusion_states.pop(sid, None) + if ar_state is not None: + try: + ar_state.close() # return KV blocks to the BlockPool + except Exception as e: # noqa: BLE001 + logger.warning("AR-Diffusion warm-up: state.close() failed (%s)", e) + try: + self.pipeline._states.pop(sid, None) + except Exception: # noqa: BLE001 + pass + # Drop any per-forward timings the warm-up recorded so serving stats start clean. + self._perf_e2e_times.clear() + free_after = kv.manager.block_pool.get_num_free_blocks() + if free_after != free_before: + logger.warning( + "AR-Diffusion warm-up: KV pool not restored (free %d -> %d) — possible leak", + free_before, + free_after, + ) + else: + logger.info( + "AR-Diffusion cudagraph warm-up complete; KV pool restored (%d free blocks)", + free_after, + ) + + def _image_resolution(self) -> list[int]: + mc = getattr(self.od_config, "model_config", None) or {} + psc = (mc.get("policy_server_config") if isinstance(mc, dict) else None) or {} + return psc.get("image_resolution", [180, 320]) diff --git a/vllm_omni/experimental/fullduplex/README.md b/vllm_omni/experimental/fullduplex/README.md new file mode 100644 index 00000000000..7f9e22fa61a --- /dev/null +++ b/vllm_omni/experimental/fullduplex/README.md @@ -0,0 +1,80 @@ +# Full-duplex interaction framework + +A model-agnostic framework for real-time, full-duplex (streaming-in / streaming-out) +model interaction, plus the JoyVL implementation built on it. + +To **run the JoyVL model**, see the recipe: +[`recipes/JD/JoyAI-VL-Interaction.md`](../../../recipes/JD/JoyAI-VL-Interaction.md). +This README covers the framework itself and how to add a new model. + +## Layout + +``` +vllm_omni/experimental/fullduplex/ + core/ generic full-duplex framework (model-agnostic): DuplexRuntime (event + loop + epoch barge-in), DuplexSession, DuplexAdapter (ABC), protocol + joyvl/ JoyVL implementation (model-specific): + adapter.py JoyVLDuplexAdapter (implements core.DuplexAdapter) + decision/ policy + output_parser + prompts (speak/silence/delegate) + memory/ InteractionBrain — 3-tier summary memory (async) + serving/ OpenAI-compatible HTTP orchestrator + bridges/ model backend + delegation +``` + +`core/` is the only part shared across models; data planes differ by model and are +intentionally not shared. + +## Scope + +The runnable serving path is `joyvl/serving/` driving `decision/` + `memory/` directly. +`joyvl/adapter.py` + `core/` are a **demonstration** of how a model plugs into the +generic full-duplex framework (exercised by `tests/fullduplex/`), not the serving path +the HTTP orchestrator currently uses — a fused-audio model (e.g. MiniCPM-o) is the case +`core/` is built for. + +## Adding a full-duplex model + +The seam is `core.DuplexAdapter`. `core/` owns the session lifecycle, epoch-based +barge-in, playback cursor, and the event protocol — you implement only model policy. + +1. Create a sibling package `vllm_omni/experimental/fullduplex/<model>/` next to `joyvl/`; keep + model-specific code there and do not touch `core/`. +2. Implement one `DuplexAdapter` (three required methods; the rest have defaults): + + ```python + from collections.abc import AsyncIterator + from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk + from vllm_omni.experimental.fullduplex.core.session import DuplexSession + + class MyModelAdapter(DuplexAdapter): + def capabilities(self) -> DuplexCapability: + return DuplexCapability( + input_modalities=frozenset({"audio", "text"}), + output_modalities=frozenset({"audio", "text"}), + proactive=True, # speak without being asked? + ) + + async def on_input(self, session: DuplexSession, modality: str, data) -> None: + ... # buffer/route an incoming chunk + + async def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]: + async for piece in self._model_stream(session): + yield OutputChunk(modality="audio", data=piece) # runtime drops stale + # chunks after a barge-in + + # optional: should_respond / on_barge_in / on_playback_ack + ``` + +3. Run it through the shared runtime — no new control-plane code: + + ```python + from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime + from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig + + rt = DuplexRuntime(DuplexSession("sid", DuplexSessionConfig()), MyModelAdapter()) + await rt.run(input_events, emit) + ``` + +`joyvl/adapter.py` is the worked demonstration (currently exercised by tests, not the +HTTP serving path). Promote a helper from a model package up into `core/` only once a +second model actually needs it. diff --git a/vllm_omni/experimental/fullduplex/__init__.py b/vllm_omni/experimental/fullduplex/__init__.py new file mode 100644 index 00000000000..0615b7681fa --- /dev/null +++ b/vllm_omni/experimental/fullduplex/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk +from vllm_omni.experimental.fullduplex.core.runtime import DuplexRuntime +from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexSessionConfig, DuplexState + +__all__ = [ + "DuplexAdapter", + "DuplexCapability", + "DuplexRuntime", + "DuplexSession", + "DuplexSessionConfig", + "DuplexState", + "OutputChunk", +] diff --git a/vllm_omni/experimental/fullduplex/core/__init__.py b/vllm_omni/experimental/fullduplex/core/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/core/__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/experimental/fullduplex/core/adapter.py b/vllm_omni/experimental/fullduplex/core/adapter.py new file mode 100644 index 00000000000..3fe400d4986 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/core/adapter.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any + +from vllm_omni.experimental.fullduplex.core.session import DuplexSession + + +@dataclass +class DuplexCapability: + input_modalities: frozenset[str] + output_modalities: frozenset[str] + + proactive: bool = False + + +@dataclass +class OutputChunk: + modality: str + data: Any + + +class DuplexAdapter(ABC): + @abstractmethod + def capabilities(self) -> DuplexCapability: ... + + @abstractmethod + async def on_input(self, session: DuplexSession, modality: str, data: Any) -> None: ... + + @abstractmethod + def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]: ... + + def should_respond(self, session: DuplexSession) -> bool: + return True + + async def on_barge_in(self, session: DuplexSession) -> None: ... + + async def on_playback_ack(self, session: DuplexSession, cursor: int) -> None: ... diff --git a/vllm_omni/experimental/fullduplex/core/protocol.py b/vllm_omni/experimental/fullduplex/core/protocol.py new file mode 100644 index 00000000000..0056e1abafc --- /dev/null +++ b/vllm_omni/experimental/fullduplex/core/protocol.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +INPUT_APPEND = "input.append" +INPUT_COMMIT = "input.commit" +RESPONSE_CREATE = "response.create" +RESPONSE_CANCEL = "response.cancel" +PLAYBACK_ACK = "playback.ack" +CLOSE = "close" + + +RESPONSE_CREATED = "response.created" +RESPONSE_DELTA = "response.delta" +RESPONSE_DONE = "response.done" +RESPONSE_CANCELLED = "response.cancelled" +ERROR = "error" + + +def created(response_index: int) -> dict[str, Any]: + return {"type": RESPONSE_CREATED, "response_index": response_index} + + +def delta(response_index: int, modality: str, data: Any) -> dict[str, Any]: + return {"type": RESPONSE_DELTA, "response_index": response_index, "modality": modality, "data": data} + + +def done(response_index: int) -> dict[str, Any]: + return {"type": RESPONSE_DONE, "response_index": response_index} + + +def cancelled(response_index: int) -> dict[str, Any]: + return {"type": RESPONSE_CANCELLED, "response_index": response_index} + + +def error(message: str) -> dict[str, Any]: + return {"type": ERROR, "message": message} diff --git a/vllm_omni/experimental/fullduplex/core/runtime.py b/vllm_omni/experimental/fullduplex/core/runtime.py new file mode 100644 index 00000000000..86ca011ef82 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/core/runtime.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import AsyncIterator, Awaitable, Callable + +from vllm_omni.experimental.fullduplex.core import protocol as ev +from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter +from vllm_omni.experimental.fullduplex.core.session import DuplexSession, DuplexState + +Emit = Callable[[dict], Awaitable[None]] + + +class DuplexRuntime: + def __init__(self, session: DuplexSession, adapter: DuplexAdapter) -> None: + self.session = session + self.adapter = adapter + self._capabilities = adapter.capabilities() + + async def run(self, inputs: AsyncIterator[dict], emit: Emit) -> None: + task: asyncio.Task | None = None + async for event in inputs: + etype = event.get("type") + if etype == ev.INPUT_APPEND: + modality = event.get("modality", "") + if modality not in self._capabilities.input_modalities: + await emit(ev.error(f"unsupported input modality: {modality}")) + continue + await self.adapter.on_input(self.session, modality, event.get("data")) + self.session.state = DuplexState.LISTENING + if self.session.config.proactive and self.adapter.should_respond(self.session): + task = await self._start_response(task, emit) + elif etype in (ev.INPUT_COMMIT, ev.RESPONSE_CREATE): + if self.adapter.should_respond(self.session): + task = await self._start_response(task, emit) + elif etype == ev.RESPONSE_CANCEL: + await self._barge_in() + await self._cancel(task) + task = None + await emit(ev.cancelled(self.session.response_index)) + elif etype == ev.PLAYBACK_ACK: + await self.adapter.on_playback_ack(self.session, int(event.get("cursor", 0))) + elif etype == ev.CLOSE: + break + if task is not None and not task.done(): + await task + self.session.state = DuplexState.CLOSED + + async def _start_response(self, prev: asyncio.Task | None, emit: Emit) -> asyncio.Task: + # A new response supersedes any in-flight one. Mark the old response stale and + # cancel it (barge-in) instead of awaiting its full stream — otherwise the input + # loop would block until the old response finishes and could not service a + # cancel / barge-in event in time, breaking the full-duplex contract. + if prev is not None and not prev.done(): + await self._barge_in() + await self._cancel(prev) + return asyncio.create_task(self._respond(emit)) + + @staticmethod + async def _cancel(task: asyncio.Task | None) -> None: + if task is not None and not task.done(): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task + + async def _respond(self, emit: Emit) -> None: + response_index, epoch = self.session.begin_response() + await emit(ev.created(response_index)) + try: + async for chunk in self.adapter.respond(self.session): + if self.session.is_stale(epoch): + return + await emit(ev.delta(response_index, chunk.modality, chunk.data)) + except Exception as err: + await emit(ev.error(f"response failed: {err}")) + finally: + if not self.session.is_stale(epoch): + await emit(ev.done(response_index)) + self.session.end_response() + + async def _barge_in(self) -> None: + self.session.barge_in() + await self.adapter.on_barge_in(self.session) diff --git a/vllm_omni/experimental/fullduplex/core/session.py b/vllm_omni/experimental/fullduplex/core/session.py new file mode 100644 index 00000000000..844f5a4b6f3 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/core/session.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import enum +from dataclasses import dataclass, field + + +class DuplexState(enum.Enum): + IDLE = "idle" + LISTENING = "listening" + RESPONDING = "responding" + CLOSED = "closed" + + +@dataclass +class DuplexSessionConfig: + input_modalities: tuple[str, ...] = ("audio",) + output_modalities: tuple[str, ...] = ("audio",) + sample_rate: int = 24000 + + proactive: bool = False + + +@dataclass +class DuplexSession: + session_id: str + config: DuplexSessionConfig = field(default_factory=DuplexSessionConfig) + state: DuplexState = DuplexState.IDLE + + epoch: int = 0 + response_index: int = 0 + + playback_cursor: int = 0 + + def begin_response(self) -> tuple[int, int]: + self.response_index += 1 + self.state = DuplexState.RESPONDING + return self.response_index, self.epoch + + def end_response(self) -> None: + if self.state is DuplexState.RESPONDING: + self.state = DuplexState.LISTENING + + def barge_in(self) -> int: + self.epoch += 1 + self.state = DuplexState.LISTENING + return self.epoch + + def is_stale(self, epoch: int) -> bool: + return epoch != self.epoch diff --git a/vllm_omni/experimental/fullduplex/joyvl/__init__.py b/vllm_omni/experimental/fullduplex/joyvl/__init__.py new file mode 100644 index 00000000000..d3429ef820d --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import ( + Action, + ParsedAction, + parse_action, +) +from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig +from vllm_omni.experimental.fullduplex.joyvl.serving.session import InteractionSession, StepResult + +__all__ = [ + "Action", + "InteractionConfig", + "InteractionSession", + "ParsedAction", + "StepResult", + "parse_action", +] diff --git a/vllm_omni/experimental/fullduplex/joyvl/adapter.py b/vllm_omni/experimental/fullduplex/joyvl/adapter.py new file mode 100644 index 00000000000..16e220c29ed --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/adapter.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import AsyncIterator, Awaitable, Callable +from typing import Any + +from vllm_omni.experimental.fullduplex.core.adapter import DuplexAdapter, DuplexCapability, OutputChunk +from vllm_omni.experimental.fullduplex.core.session import DuplexSession +from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationBridge, StubDelegationBridge +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import Action +from vllm_omni.experimental.fullduplex.joyvl.decision.policy import JoyVLPolicy, sample_frames + +GenerateFn = Callable[[list[dict[str, Any]]], Awaitable[str]] + + +class JoyVLDuplexAdapter(DuplexAdapter): + def __init__( + self, + generate: GenerateFn, + *, + persona: str = "default", + num_frames: int = 4, + chunk_frames: int = 200, + frame_seconds: float = 1.0, + delegation: DelegationBridge | None = None, + ) -> None: + self._generate = generate + self._policy = JoyVLPolicy( + persona=persona, + num_frames=num_frames, + chunk_frames=chunk_frames, + frame_seconds=frame_seconds, + delegation=delegation or StubDelegationBridge(), + ) + self._frames: list[str] = [] + self._pending_query: str | None = None + + def capabilities(self) -> DuplexCapability: + return DuplexCapability(frozenset({"video", "text"}), frozenset({"text"}), proactive=True) + + async def on_input(self, session: DuplexSession, modality: str, data: Any) -> None: + if modality == "video": + self._policy.tick() + self._frames.append(data) + elif modality == "text": + self._pending_query = data + + def should_respond(self, session: DuplexSession) -> bool: + return bool(self._frames) + + async def respond(self, session: DuplexSession) -> AsyncIterator[OutputChunk]: + policy = self._policy + await policy.fold_delegations() + policy.set_query(self._pending_query) + self._pending_query = None + + parts = [{"type": "image_url", "image_url": {"url": f}} for f in sample_frames(self._frames, policy.num_frames)] + messages, _ = policy.build_messages(parts) + action = policy.commit(await self._generate(messages)) + + frames = [(str(i), f) for i, f in enumerate(self._frames)] + await policy.submit_if_delegate(action, frames) + if action.action is not Action.SILENCE and action.text: + yield OutputChunk("text", action.text) + if policy.needs_flush(): + await policy.flush(frames) + self._frames.clear() + + async def on_barge_in(self, session: DuplexSession) -> None: + self._pending_query = None diff --git a/vllm_omni/experimental/fullduplex/joyvl/bridges/__init__.py b/vllm_omni/experimental/fullduplex/joyvl/bridges/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/bridges/__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/experimental/fullduplex/joyvl/bridges/backend.py b/vllm_omni/experimental/fullduplex/joyvl/bridges/backend.py new file mode 100644 index 00000000000..1aa2d57d2d8 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/bridges/backend.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any, Protocol + +from openai import AsyncOpenAI + + +class ModelBackend(Protocol): + async def generate( + self, + messages: list[dict[str, Any]], + *, + max_tokens: int, + temperature: float, + top_p: float, + extra_body: dict[str, Any] | None = None, + ) -> tuple[str, dict[str, Any] | None]: ... + + async def aclose(self) -> None: ... + + +class OpenAIBackend: + def __init__( + self, + base_url: str, + model: str, + api_key: str = "EMPTY", + timeout: float = 300.0, + ) -> None: + self.model = model + self._client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=timeout) + + async def generate( + self, + messages: list[dict[str, Any]], + *, + max_tokens: int, + temperature: float, + top_p: float, + extra_body: dict[str, Any] | None = None, + ) -> tuple[str, dict[str, Any] | None]: + response = await self._client.chat.completions.create( + model=self.model, + messages=messages, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + extra_body=extra_body or {}, + ) + if not response.choices: + return "", None + text = response.choices[0].message.content or "" + usage = response.usage.model_dump() if getattr(response, "usage", None) else None + return text, usage + + async def aclose(self) -> None: + await self._client.close() diff --git a/vllm_omni/experimental/fullduplex/joyvl/bridges/delegation.py b/vllm_omni/experimental/fullduplex/joyvl/bridges/delegation.py new file mode 100644 index 00000000000..c61ca474b8a --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/bridges/delegation.py @@ -0,0 +1,389 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from typing import Any, Protocol + +import httpx +from openai import AsyncOpenAI + +from vllm_omni.experimental.fullduplex.joyvl.decision.prompts import DELEGATION_SOLVER_PROMPT + + +def _cancel_task(tasks: dict[str, asyncio.Task], task_id: str) -> None: + task = tasks.pop(task_id, None) + if task is not None and not task.done(): + task.cancel() + + +async def _aclose_tasks(tasks: dict[str, asyncio.Task]) -> None: + pending = list(tasks.values()) + for task in pending: + if not task.done(): + task.cancel() + if pending: + await asyncio.gather(*pending, return_exceptions=True) + tasks.clear() + + +@dataclass +class DelegationResult: + task_id: str + status: str + digest: str = "" + media: str = "" + + @property + def is_ready(self) -> bool: + return self.status == "ready" + + +class DelegationBridge(Protocol): + async def submit(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: ... + + async def poll(self, task_id: str) -> DelegationResult: ... + + def cancel(self, task_id: str) -> None: + """Drop a task and cancel its background work if still running. + + Called when a session is reset/evicted so finished-but-unpolled tasks do not + accumulate on a long-lived shared bridge. + """ + ... + + async def aclose(self) -> None: + """Cancel+await all pending tasks and close any held client (server shutdown).""" + ... + + +class StubDelegationBridge: + def __init__(self, ready_after_ticks: int = 2) -> None: + self._ready_after = max(1, ready_after_ticks) + self._tasks: dict[str, dict[str, Any]] = {} + self._counter = 0 + + async def submit(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: + self._counter += 1 + task_id = f"deleg-{self._counter}" + self._tasks[task_id] = {"question": question, "polls": 0} + return task_id + + async def poll(self, task_id: str) -> DelegationResult: + task = self._tasks.get(task_id) + if task is None: + return DelegationResult(task_id, "error", "unknown task") + task["polls"] += 1 + if task["polls"] < self._ready_after: + return DelegationResult(task_id, "pending") + question = task["question"] + return DelegationResult( + task_id, + "ready", + digest=f"(background result for: {question}) — stub digest; wire a real brain here.", + ) + + def cancel(self, task_id: str) -> None: + self._tasks.pop(task_id, None) + + async def aclose(self) -> None: + self._tasks.clear() + + +class OpenAIDelegationBridge: + """Hand a delegated question to a stronger background model over an OpenAI-compatible API. + + Non-blocking by design: ``submit`` launches the request as a background task and returns + immediately; ``poll`` only checks whether that task has finished, so the per-tick interaction + loop never stalls waiting on the slower background brain. + """ + + def __init__( + self, + base_url: str, + model: str, + api_key: str = "EMPTY", + *, + max_frames: int = 8, + max_tokens: int = 512, + temperature: float = 0.2, + timeout: float = 120.0, + system_prompt: str = DELEGATION_SOLVER_PROMPT, + ) -> None: + self.model = model + self._client = AsyncOpenAI(base_url=base_url, api_key=api_key, timeout=timeout) + self._max_frames = max_frames + self._max_tokens = max_tokens + self._temperature = temperature + self._system_prompt = system_prompt + self._tasks: dict[str, asyncio.Task] = {} + self._counter = 0 + + async def submit(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: + self._counter += 1 + task_id = f"deleg-{self._counter}" + self._tasks[task_id] = asyncio.create_task(self._solve(question, note, frames)) + return task_id + + async def poll(self, task_id: str) -> DelegationResult: + task = self._tasks.get(task_id) + if task is None: + return DelegationResult(task_id, "error", "unknown task") + if not task.done(): + return DelegationResult(task_id, "pending") + self._tasks.pop(task_id, None) + try: + digest = task.result() + except Exception as exc: # surface any background failure as an error result, never crash the loop + return DelegationResult(task_id, "error", str(exc)) + return DelegationResult(task_id, "ready", digest=digest) + + async def _solve(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: + content: list[dict[str, Any]] = [] + for time_range, data_url in frames[-self._max_frames :]: + content.append({"type": "text", "text": f"<{time_range}>"}) + content.append({"type": "image_url", "image_url": {"url": data_url}}) + content.append({"type": "text", "text": question or note or "Answer the user's question."}) + messages = [ + {"role": "system", "content": self._system_prompt}, + {"role": "user", "content": content}, + ] + return await self._complete(messages) + + async def _complete(self, messages: list[dict[str, Any]]) -> str: + response = await self._client.chat.completions.create( + model=self.model, + messages=messages, + max_tokens=self._max_tokens, + temperature=self._temperature, + ) + if not response.choices: + return "" + return (response.choices[0].message.content or "").strip() + + def cancel(self, task_id: str) -> None: + _cancel_task(self._tasks, task_id) + + async def aclose(self) -> None: + await _aclose_tasks(self._tasks) + await self._client.close() + + +class ImageGenDelegationBridge: + """Delegate an image-generation request to a text-to-image model over the + OpenAI-compatible ``/v1/images/generations`` endpoint (e.g. vLLM-Omni Qwen-Image). + + Non-blocking: ``submit`` launches generation as a background task and returns at once; + ``poll`` only checks completion. The generated image is returned as a ``data:`` URL in + ``DelegationResult.media`` (surfaced to the client), while ``digest`` stays a short text + placeholder so the base64 image never enters the model's context. + """ + + def __init__( + self, + base_url: str, + *, + size: str = "1024x1024", + timeout: float = 300.0, + ) -> None: + self._base_url = base_url.rstrip("/") + self._size = size + self._timeout = timeout + self._tasks: dict[str, asyncio.Task] = {} + self._counter = 0 + + async def submit(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: + self._counter += 1 + task_id = f"deleg-{self._counter}" + self._tasks[task_id] = asyncio.create_task(self._generate(question or note)) + return task_id + + async def poll(self, task_id: str) -> DelegationResult: + task = self._tasks.get(task_id) + if task is None: + return DelegationResult(task_id, "error", "unknown task") + if not task.done(): + return DelegationResult(task_id, "pending") + self._tasks.pop(task_id, None) + try: + prompt, data_url = task.result() + except Exception as exc: # surface any background failure as an error result, never crash the loop + return DelegationResult(task_id, "error", str(exc)) + return DelegationResult(task_id, "ready", digest=f"(generated an image for: {prompt})", media=data_url) + + async def _generate(self, prompt: str) -> tuple[str, str]: + data = await self._call_images({"prompt": prompt, "size": self._size}) + items = data.get("data") or [] + b64 = items[0].get("b64_json") if items and isinstance(items[0], dict) else None + if not b64: + raise RuntimeError("image endpoint returned no image data") + return prompt, f"data:image/png;base64,{b64}" + + async def _call_images(self, payload: dict[str, Any]) -> dict[str, Any]: + async with httpx.AsyncClient(timeout=self._timeout) as client: + resp = await client.post(f"{self._base_url}/images/generations", json=payload) + resp.raise_for_status() + return resp.json() + + def cancel(self, task_id: str) -> None: + _cancel_task(self._tasks, task_id) + + async def aclose(self) -> None: + await _aclose_tasks(self._tasks) + + +def _extract_image_url(data: dict[str, Any]) -> str | None: + choices = data.get("choices") or [] + if not choices: + return None + content = (choices[0].get("message") or {}).get("content") + if isinstance(content, list): + for part in content: + if isinstance(part, dict) and part.get("image_url"): + return (part["image_url"] or {}).get("url") + if isinstance(content, str) and content.startswith("data:image/"): + return content + return None + + +class ImageEditDelegationBridge: + """Delegate "restyle / cartoonify what you currently see" to an image-editing model + (e.g. vLLM-Omni Qwen-Image-Edit) over the OpenAI-compatible ``/v1/chat/completions`` + endpoint, which takes an input image + instruction and returns the edited image. + + Conditions on the latest frame passed to ``submit`` — so it preserves the real + composition instead of re-imagining it. Non-blocking submit/poll; the edited image + rides in ``DelegationResult.media`` while ``digest`` stays a short text placeholder. + """ + + _DEFAULT_INSTRUCTION = ( + "Convert this image into a clean cartoon / anime illustration style. " + "Preserve the composition, subjects, poses, and overall layout." + ) + + def __init__(self, base_url: str, model: str, *, timeout: float = 300.0, instruction: str | None = None) -> None: + self._base_url = base_url.rstrip("/") + self._model = model + self._timeout = timeout + self._instruction = instruction or self._DEFAULT_INSTRUCTION + self._tasks: dict[str, asyncio.Task] = {} + self._counter = 0 + + async def submit(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: + self._counter += 1 + task_id = f"deleg-{self._counter}" + frame = frames[-1][1] if frames else None + self._tasks[task_id] = asyncio.create_task(self._edit(question or note, frame)) + return task_id + + async def poll(self, task_id: str) -> DelegationResult: + task = self._tasks.get(task_id) + if task is None: + return DelegationResult(task_id, "error", "unknown task") + if not task.done(): + return DelegationResult(task_id, "pending") + self._tasks.pop(task_id, None) + try: + label, data_url = task.result() + except Exception as exc: # surface any background failure as an error result, never crash the loop + return DelegationResult(task_id, "error", str(exc)) + return DelegationResult(task_id, "ready", digest=f"(restyled the current view: {label})", media=data_url) + + async def _edit(self, request_text: str, frame: str | None) -> tuple[str, str]: + if not frame: + raise RuntimeError("no frame available to restyle") + instruction = request_text or self._instruction + data = await self._call_chat(frame, instruction) + url = _extract_image_url(data) + if not url: + raise RuntimeError("image-edit endpoint returned no image") + return (request_text or "cartoon"), url + + async def _call_chat(self, frame: str, instruction: str) -> dict[str, Any]: + body = { + "model": self._model, + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": frame}}, + {"type": "text", "text": instruction}, + ], + } + ], + } + async with httpx.AsyncClient(timeout=self._timeout) as client: + resp = await client.post(f"{self._base_url}/chat/completions", json=body) + resp.raise_for_status() + return resp.json() + + def cancel(self, task_id: str) -> None: + _cancel_task(self._tasks, task_id) + + async def aclose(self) -> None: + await _aclose_tasks(self._tasks) + + +_EDIT_STYLE_KW = ("卡通", "漫画", "动漫", "cartoon", "anime", "风格", "restyle", "变成", "画成") +_EDIT_SCENE_KW = ("看到", "画面", "当前", "现在", "this", "what you see", "场景", "镜头") +_IMAGE_KW = ("画", "生成", "draw", "generate", "picture", "image", "图", "render") + + +class RoutingDelegationBridge: + """Dispatch each delegation to the right backend by inspecting the request text. + + - "restyle / cartoonify what you currently see" -> image-edit (conditions on the frame) + - "draw / generate an image of X" -> text-to-image + - everything else (hard questions) -> the chat / VL background brain + + Wraps sub-bridges that each implement the DelegationBridge protocol; any may be None. + Routing is heuristic and demo-grade, not a trained classifier. + """ + + def __init__(self, *, chat: Any = None, image: Any = None, edit: Any = None) -> None: + self._chat = chat + self._image = image + self._edit = edit + self._route: dict[str, tuple[Any, str]] = {} + self._counter = 0 + + def _pick(self, question: str, note: str) -> Any: + text = f"{question} {note}".lower() + if self._edit and any(k in text for k in _EDIT_STYLE_KW) and any(k in text for k in _EDIT_SCENE_KW): + return self._edit + if self._image and any(k in text for k in _IMAGE_KW): + return self._image + return self._chat or self._image or self._edit + + async def submit(self, question: str, note: str, frames: list[tuple[str, str]]) -> str: + bridge = self._pick(question, note) + if bridge is None: + raise RuntimeError("no delegation backend configured") + self._counter += 1 + key = f"route-{self._counter}" + self._route[key] = (bridge, await bridge.submit(question, note, frames)) + return key + + async def poll(self, task_id: str) -> DelegationResult: + entry = self._route.get(task_id) + if entry is None: + return DelegationResult(task_id, "error", "unknown task") + bridge, inner_id = entry + res = await bridge.poll(inner_id) + if res.status in ("ready", "error"): + self._route.pop(task_id, None) + return DelegationResult(task_id, res.status, res.digest, res.media) + + def cancel(self, task_id: str) -> None: + entry = self._route.pop(task_id, None) + if entry is not None: + bridge, inner_id = entry + bridge.cancel(inner_id) + + async def aclose(self) -> None: + self._route.clear() + for bridge in (self._chat, self._image, self._edit): + if bridge is not None: + await bridge.aclose() diff --git a/vllm_omni/experimental/fullduplex/joyvl/decision/__init__.py b/vllm_omni/experimental/fullduplex/joyvl/decision/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/decision/__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/experimental/fullduplex/joyvl/decision/output_parser.py b/vllm_omni/experimental/fullduplex/joyvl/decision/output_parser.py new file mode 100644 index 00000000000..cea17fd1c72 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/decision/output_parser.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import enum +from dataclasses import dataclass + +SILENCE_TOKEN = "</silence>" +RESPONSE_TOKEN = "</response>" +DELEGATION_TAG = "</delegation>" +_DELEGATION_MARKERS = ("</delegation>", "<delegation>") + + +class Action(enum.Enum): + SILENCE = "silence" + RESPONSE = "response" + DELEGATE = "delegate" + + +@dataclass +class ParsedAction: + action: Action + text: str = "" + delegated_question: str | None = None + raw: str = "" + + @property + def spoke(self) -> bool: + return self.action is not Action.SILENCE + + +def parse_action(raw: str) -> ParsedAction: + text = (raw or "").replace("\r\n", "\n").replace("\r", "\n").strip() + if not text: + return ParsedAction(Action.SILENCE, raw=raw or "") + + marker = _first_marker(text) + if marker != -1 and text[marker:].startswith(SILENCE_TOKEN): + return ParsedAction(Action.SILENCE, raw=raw) + + if marker == -1: + payload = _collapse(text.splitlines()[0]) + else: + after = text[marker + len(RESPONSE_TOKEN) :].strip() + payload = _collapse(after.splitlines()[0]) if after else "" + + for tag in _DELEGATION_MARKERS: + idx = payload.find(tag) + if idx != -1: + return ParsedAction( + Action.DELEGATE, + text=_collapse(payload[:idx]), + delegated_question=_collapse(payload[idx + len(tag) :]) or None, + raw=raw, + ) + + return ParsedAction(Action.RESPONSE, text=payload, raw=raw) + + +def to_token_form(action: ParsedAction) -> str: + if action.action is Action.SILENCE: + return SILENCE_TOKEN + if action.action is Action.DELEGATE: + return f"{RESPONSE_TOKEN} {action.text} {DELEGATION_TAG} {action.delegated_question or ''}".strip() + return f"{RESPONSE_TOKEN} {action.text}".strip() + + +def _first_marker(text: str) -> int: + present = [i for i in (text.find(SILENCE_TOKEN), text.find(RESPONSE_TOKEN)) if i != -1] + return min(present) if present else -1 + + +def _collapse(text: str) -> str: + return " ".join(text.split()) diff --git a/vllm_omni/experimental/fullduplex/joyvl/decision/policy.py b/vllm_omni/experimental/fullduplex/joyvl/decision/policy.py new file mode 100644 index 00000000000..118fcdffe3f --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/decision/policy.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from difflib import SequenceMatcher +from typing import Any + +from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationBridge +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import Action, ParsedAction, parse_action +from vllm_omni.experimental.fullduplex.joyvl.decision.prompts import SYSTEM_PROMPTS, USER_QUERY_HEADER +from vllm_omni.experimental.fullduplex.joyvl.memory.brain import InteractionBrain +from vllm_omni.experimental.fullduplex.joyvl.memory.memory import Summarizer + + +def sample_frames(frames: list[str], num_frames: int) -> list[str]: + n = len(frames) + if num_frames <= 0 or n <= num_frames: + return list(frames) + stride = max(1, n // num_frames) + idx = [i * stride for i in range(num_frames - 1)] + [n - 1] + return [frames[i] for i in idx] + + +class JoyVLPolicy: + def __init__( + self, + *, + persona: str = "default", + system_prompt: str | None = None, + num_frames: int = 4, + summarizer: Summarizer | None = None, + delegation: DelegationBridge | None = None, + chunk_frames: int = 16, + long_term_every_n_chunks: int = 5, + long_term_window: int = 15, + keep_qa_history: bool = True, + frame_seconds: float = 1.0, + enable_delegation: bool = True, + response_dedup_threshold: float = 1.0, + ) -> None: + self.response_dedup_threshold = response_dedup_threshold + self.brain = InteractionBrain( + summarizer=summarizer, + delegation=delegation, + chunk_frames=chunk_frames, + long_term_every_n_chunks=long_term_every_n_chunks, + long_term_window=long_term_window, + keep_qa_history=keep_qa_history, + frame_seconds=frame_seconds, + enable_delegation=enable_delegation, + ) + self.system_prompt = system_prompt or SYSTEM_PROMPTS.get(persona, SYSTEM_PROMPTS["default"]) + self.num_frames = num_frames + + def tick(self, n: int = 1) -> None: + self.brain.tick(n) + + def observe(self, time_range: str, data_url: str) -> None: + self.brain.observe(time_range, data_url) + + @property + def working_frames(self) -> list[tuple[str, str]]: + return self.brain.working_frames + + def take_working_frames(self) -> list[tuple[str, str]]: + return self.brain.take_working_frames() + + def set_query(self, query: str | None) -> bool: + return self.brain.update_query(query) + + def should_respond(self) -> bool: + return True + + def user_content(self, frame_parts: list[dict[str, Any]]) -> list[dict[str, Any]]: + content: list[dict[str, Any]] = [] + prefix = self.brain.build_prefix() + if prefix: + content.append({"type": "text", "text": prefix}) + content.extend(frame_parts) + if self.brain.current_query and self.brain.query_in_current_chunk: + content.append({"type": "text", "text": f"{USER_QUERY_HEADER}\n{self.brain.current_query}"}) + return content + + def build_messages(self, frame_parts: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, Any]]: + user = {"role": "user", "content": self.user_content(frame_parts)} + return [{"role": "system", "content": self.system_prompt}, user], user + + def _is_repeat(self, text: str) -> bool: + last = self.brain.last_response_text + if not last: + return False + if text == last: + return True + if self.response_dedup_threshold >= 1.0: + return False + similarity = SequenceMatcher(None, text.strip().lower(), last.strip().lower()).ratio() + return similarity >= self.response_dedup_threshold + + def commit(self, response_text: str) -> ParsedAction: + action = parse_action(response_text) + if action.action is Action.SILENCE: + self.brain.last_response_text = "" + return action + if action.action is Action.RESPONSE and action.text and self._is_repeat(action.text): + return ParsedAction(Action.SILENCE, raw="</silence>") + self.brain.last_response_text = action.text + if self.brain.current_query: + self.brain.record_response(action.text) + return action + + async def fold_delegations(self) -> dict[str, Any] | None: + return await self.brain.fold_delegations() + + async def submit_if_delegate( + self, action: ParsedAction, frames: list[tuple[str, str]] | None = None + ) -> dict | None: + return await self.brain.submit_delegation(action, frames or []) + + def needs_flush(self) -> bool: + return self.brain.should_flush() + + def close_chunk(self) -> int: + return self.brain.close_chunk() + + async def consolidate(self, chunk_index: int, frames: list[tuple[str, str]] | None = None) -> None: + await self.brain.consolidate(chunk_index, frames or []) + + async def flush(self, frames: list[tuple[str, str]] | None = None) -> None: + await self.brain.flush(frames or []) diff --git a/vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py b/vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py new file mode 100644 index 00000000000..b7fa95f0a92 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/decision/prompts.py @@ -0,0 +1,182 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +DEFAULT_SYSTEM_PROMPT = """You are a real-time video streaming assistant observing a continuous camera feed frame by frame. The last frame represents the current moment. +## Action Format +At every inference step you MUST choose exactly one of the following three actions: +**Stay silent** — output ONLY: +</silence> +Choose this when nothing noteworthy has changed in the scene, no user query is pending, or there is nothing useful to say. +**Speak** — output the token followed by a concise reply: +</response> Your reply here. +Choose this when you observe something worth reporting or a significant state change, or when you can answer a user question based on available evidence. + +**Delegate** — when a question is too hard or error-prone to answer reliably yourself, speak a brief note that you're delegating, then hand the question to the background solver: +</response> Brief note that you're delegating. </delegation> <the question>""".strip() + +SILENT_SYSTEM_PROMPT = """You are a real-time video streaming assistant observing a continuous camera feed frame by frame. The last frame represents the current moment. +## Action Format +At every inference step you MUST choose exactly one of the following three actions: +**Stay silent** — output ONLY: +</silence> +Choose this when nothing noteworthy has changed in the scene, no user query is pending, or there is nothing useful to say. +**Speak** — output the token followed by a concise reply: +</response> Your reply here. +Choose this when you observe something worth reporting or a significant state change, or when you can answer a user question based on available evidence. + +Do NOT delegate or defer questions to other models. Answer user queries directly based on everything you have seen.""".strip() + +TALKATIVE_SYSTEM_PROMPT = """You are a real-time video streaming assistant observing a continuous camera feed frame by frame. The last frame represents the current moment. +## Action Format +At every inference step you MUST choose exactly one of the following three actions: +**Stay silent** — output ONLY: +</silence> +Choose this when nothing noteworthy has changed in the scene, no user query is pending, or there is nothing useful to say. +**Speak** — output the token followed by a concise reply: +</response> Your reply here. +Choose this when you observe something worth reporting or a significant state change, or when you can answer a user question based on available evidence. + +**Delegate** — when a question is too hard or error-prone to answer reliably yourself, speak a brief note that you're delegating, then hand the question to the background solver: +</response> Brief note that you're delegating. <delegation> <the question> + +## Style +Proactively speak whenever you observe a meaningful event, state change, or anomaly — do not wait for the user to ask. However, avoid repeating information you have already reported. When the scene shows no obvious change, choose silence.""".strip() + +SYSTEM_PROMPTS = { + "default": DEFAULT_SYSTEM_PROMPT, + "silent": SILENT_SYSTEM_PROMPT, + "talkative": TALKATIVE_SYSTEM_PROMPT, +} + + +DELEGATION_SOLVER_PROMPT = ( + "You are the background solver for a real-time video agent. The foreground agent delegated a " + "question it could not answer quickly and reliably. Use the attached recent frames to answer it. " + "Reply with a single concise, factual digest the foreground agent can relay to the user — no " + "preamble, no restating the question. If the frames are insufficient, say so plainly." +) + + +USER_QUERY_HEADER = "[User Query (IMPORTANT — follow this instruction)]" + +VIDEO_HISTORY_HEADER = ( + "[Video History]\n" + "The following are summaries of earlier video segments you can no longer see. " + "Use them as background context, but always prioritize the current visual frames " + "and the User Query below when making decisions.\n" + "IMPORTANT: These summaries are written by an external system in a descriptive style. " + "Do NOT imitate their writing style in your responses.\n" +) + +QA_HISTORY_HEADER = "[Q&A History]\nThe following are previous queries and the system's responses.\n\n" + +QA_QUERY_LABEL = "Query" +QA_RESPONSE_LABEL = "Response" + + +EMPTY_CHUNK_SUMMARY = "在帧 {frame_range} 中未观察到明显的视觉变化。" + +MID_TERM_SUMMARY_PROMPT = """\ +你正在为一个长时间运行的视频智能体编写中期记忆。用户消息包含第 {chunk_index} 个片段的带时间戳关键帧，覆盖时间范围 {frame_range}；每张图片前标注其采样时间段。这些帧在本次之后不再可用，因此你的段落需要保留下游模型在没有视觉信息时进行回忆、推理和回答后续问题所需的关键证据。 + +【输出格式】 +- 撰写一个事实性段落{length_instruction}。信息密集时充分利用空间；接近静止或几乎无新信息时（如长时间静止画面、近乎重复的帧），用 1-2 句简短记录即可，不要凑长度。 +- 身份、文字、数值不确定时不要猜测，用"疑似 X"、"约 X"或"不可读"等方式标注；不可读时仍保留可观察到的类别、外观、动作或位置特征。 +- 段落主体使用中文叙述；画面里能看到的英文（术语、人名、作品名、短语、年代标注等）每次出现都写成"中文译名（English 原文）"，不许只留一种，不许半中半英，画面没出现的英文不要补。 + +【需保留的信息（按优先级从高到低）】 + +最高优先级 — 任务连续性必需，必须保留： +- 片段结束时的交接状态：屏幕上剩余什么、任务推进到哪一步、哪些物体仍然可用、哪些动作未完成或仍在进行。 +- 不可逆事件与状态变化：物体的出现、消失、移动、开合、组合、分离、转移；重要操作的来源、目的地与结果状态。 +- 中间结果、任务进度、部分完成或被中断的步骤。 +- 因果链：当一个动作引发另一个动作时，保持前后关系，不要列成孤立片段。 + +高优先级 — 信息价值高且难以重建： +- 可读的文字、标签、标题、数字、测量值、计数、日期、标识符、规格。多值同时出现时保留"项目→值"的映射，不要列出脱离上下文的数字。 +- 在本片段内首次出现或明显异常的事件、实体或行为（仅按本片段判断，不需考虑跨片段历史）。 +- 重要物体的最终位置与关键空间关系。 + +中优先级 — 背景与陪衬： +- 出现的人物、身体部位、工具、容器、食材、设备等实体清单。 +- 场景布局与持续存在的背景：首次出现时记录一次，发生有意义变化时更新，未变化则不重复。 +- 重要事件与转换的稀疏时间锚点。 + +【撰写规则】 + +时间引用： +- 整段最多使用 3-5 个显式时间锚点，且仅在能改善事件区分或转换识别时使用。优先采用约 {preferred_time_span} 的合并范围，一个连续动作阶段最多使用一个显式时间范围。例如写"从 349s 到 364s，吉他手持续弹奏、抬起左手并鞠躬"，而不是 349s-350s、351s-352s、353s-355s 分别写。 +- 图片前的时间标注是证据锚点，不要逐个复制到输出中；相邻帧中重复或连续的动作合并为一个范围，并在其中描述可见变化。 + +详略与抽象层级： +- 详写最高与高优先级信息；略写稳定背景与重复动作。 +- 保留具体步骤，不要用"准备继续"、"做相关操作"等空话替代可见动作。 +- 镜头切换、机位移动或视角变化本身无需逐次提及；仅当它们揭示新元素、改变可读信息或影响动作连续性时才记录。 +- 当场景包含结构化信息（表格、菜单、表单、HUD、字幕等）时，保留字段-值对应关系，而不仅概述主题。 +- 这是记忆而非字幕：均衡保留状态、实体、动作、文字与数值，而不是只偏重叙事、新颖性或运动。 + +事实性约束： +- 仅使用所提供关键帧中直接可见的细节；不要补全画面外的内容或推测后续发展。 +- 按关键帧的时间顺序描述，但不要逐一叙述每个采样间隔。 +- 避免泛泛的总结性短语、元评论（如"本片段展示了……"）和重复句式。 + +仅输出该段落。""" + +LONG_TERM_COMPRESS_PROMPT = """\ +你正在为一个长时间运行的视频智能体将多段中期记忆压缩为长期记忆。时间段 {merged_range} 的原始帧将不再可用，因此输出需要保留下游模型在此之后进行回忆、推理和回答后续问题所需的关键证据。 + +需要合并的中期摘要： +{summaries_text} + +【任务定位】 +这是压缩任务，不是简单合并。需要按优先级取舍以降低信息密度，而不是把所有内容堆在一起。输入由 N 段中期摘要组成，每段都有自己的期末交接状态；合并后只有**最后一段的期末状态**对下游有意义，前面各段的期末状态若已被后续事件覆盖（如"拿起又放下的杯子"、"打开又关闭的菜单"），应作为过程信息处理或省略。 + +【输出格式】 +- 撰写一个统一的事实性段落（不是按片段分开的摘要），在目标长度内{length_instruction}按优先级取舍。 +- 中期摘要中的不确定性标注（"疑似 X"、"约 X"、"不可读"）必须传递到长期记忆，不要在合并时被"洗"成确定表述。 + +【保留优先级】 + +最高优先级 — 必须保留： +- 合并期**结束时**的最终交接状态：屏幕上剩余什么、任务推进到哪一步、哪些物体仍然可用、哪些动作未完成或仍在进行。 +- 跨片段的不可逆事件与状态变化：物体的最终出现/消失、关键移动、组合、分离、转移；重要操作的来源、目的地与结果状态。 +- 可读的文字、标签、数字、测量值、计数、日期、标识符、规格，以及"项目→值"的映射关系。 +- 因果链：当一个动作引发后续动作时，保持依赖关系，不要扁平化为孤立事实。 + +高优先级 — 应当保留： +- 跨多个片段持续存在的实体（人物、关键工具、关键物体）及其演变。 +- 任务的整体进度结构与阶段性里程碑。 +- 在合并范围内首次出现或明显异常的事件。 +- 重要物体在合并期末的位置与关键空间关系。 + +可压缩 — 主动降低密度或丢弃： +- 已被后续事件覆写的中间状态——仅在该过程本身有信息价值时才保留。 +- 跨片段未变化的稳定背景，只在首次描述时记录一次。 +- 同一动作的多次相似重复，按下方"重复处理"规则归并。 +- 过程性的细微步骤，在更高层动作已被记录时可以省略。 + +【合并规则】 + +跨片段一致性： +- 同一实体在不同片段中名称不一致时（如"红色杯子" → "马克杯" → "容器"），统一为信息量最高、最具体的那个，并在首次出现时简要标注其他称呼。 +- 多个片段对同一事件描述冲突时，优先采信时间靠后或信息更具体的版本；如果无法判定，用"X 或 Y"并列保留，不要单方面选边。 +- 同一事件在相邻片段边界处被重复描述时，合并为一次记录。 + +重复处理： +- 当重复次数本身携带信息时（如"敲门 5 次"、"刷新页面 3 次"），保留次数。 +- 当重复是过程性持续时（如"持续搅拌"、"反复检查"），用一个时间范围加动作描述合并，例如"从 349s 到 364s 持续弹奏并多次抬手致意"。 +- 不要将多个不同的动作或状态变化压缩为模糊总结。 + +【时间引用】 +- 整段使用 3-7 个显式时间锚点（按 {merged_range} 跨度调整），仅在能改善事件区分或阶段切换识别时使用。 +- 不要保留源摘要的 <Xs-Ys> 标题，也不要复制源摘要中"在 <X>s..."的连续时间戳节奏。 +- 连续或重复的相邻动作合并为一个范围，并在范围内描述可见变化。 + +【事实性约束】 +- 仅使用中期摘要中已有的内容；不要补全画面外信息或推测后续发展。 +- 保留具体步骤与可读数值，不要用"准备继续"、"做相关操作"等空话替代。 +- 避免泛泛的总结性短语、元评论（如"这段记忆包含……"）和重复句式。 + +仅输出统一的叙述文本，不要输出其他内容。""" diff --git a/vllm_omni/experimental/fullduplex/joyvl/memory/__init__.py b/vllm_omni/experimental/fullduplex/joyvl/memory/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/memory/__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/experimental/fullduplex/joyvl/memory/brain.py b/vllm_omni/experimental/fullduplex/joyvl/memory/brain.py new file mode 100644 index 00000000000..0cb7f450d03 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/memory/brain.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +from typing import Any + +from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationBridge +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import Action, ParsedAction +from vllm_omni.experimental.fullduplex.joyvl.memory.memory import ( + MidTermSummary, + QAEntry, + SessionMemory, + Summarizer, + build_memory_prefix, +) + + +class InteractionBrain: + def __init__( + self, + *, + summarizer: Summarizer | None = None, + delegation: DelegationBridge | None = None, + chunk_frames: int = 16, + long_term_every_n_chunks: int = 5, + long_term_window: int = 15, + keep_qa_history: bool = True, + frame_seconds: float = 1.0, + enable_delegation: bool = True, + ) -> None: + self._summarizer = summarizer + self._delegation = delegation + self._chunk_frames = chunk_frames + self._long_term_every_n = long_term_every_n_chunks + self._long_term_window = long_term_window + self._keep_qa_history = keep_qa_history + self._frame_seconds = frame_seconds + self._enable_delegation = enable_delegation + self.reset() + + def reset(self) -> None: + # Cancel any in-flight/finished-but-unpolled delegations so the shared bridge + # does not accumulate abandoned tasks when a session is reset or evicted. + if self._delegation is not None: + for task in getattr(self, "_pending_delegations", []): + self._delegation.cancel(task["task_id"]) + self.memory = SessionMemory() + self.current_query: str | None = None + self.query_time: str | None = None + self.query_in_current_chunk = False + self.frame_index = 0 + self.chunk_index = 1 + self._chunk_frame_count = 0 + self.working_frames: list[tuple[str, str]] = [] + self.last_response_text = "" + self.response_records: list[tuple[str, str]] = [] + self._pending_delegations: list[dict[str, str]] = [] + self._consolidate_lock = asyncio.Lock() + + def now(self) -> str: + return f"{self.frame_index * self._frame_seconds:.1f} seconds" + + def last_frame_time(self) -> str: + return f"{max(0, self.frame_index - 1) * self._frame_seconds:.1f} seconds" + + def tick(self, n: int = 1) -> None: + self.frame_index += n + self._chunk_frame_count += n + + def observe(self, time_range: str, data_url: str) -> None: + self.tick() + self.working_frames.append((time_range, data_url)) + + def take_working_frames(self) -> list[tuple[str, str]]: + frames = self.working_frames + self.working_frames = [] + return frames + + def should_flush(self) -> bool: + return self._chunk_frames > 0 and self._chunk_frame_count >= self._chunk_frames + + def update_query(self, query: str | None) -> bool: + q = (query or "").strip() + if not q or q == self.current_query: + return False + if self.current_query is not None: + self.archive_query() + self.current_query = q + self.query_time = self.now() + self.query_in_current_chunk = True + return True + + def archive_query(self) -> None: + if self.current_query and self.response_records: + self.memory.qa_history.append( + QAEntry( + self.current_query, + self.query_time or "", + list(self.response_records), + self.chunk_index, + ) + ) + self.response_records = [] + + def record_response(self, text: str) -> None: + if text and self.current_query: + self.response_records.append((self.last_frame_time(), text)) + + def build_prefix(self) -> str: + return build_memory_prefix( + self.memory, + current_query=self.current_query, + query_in_current_chunk=self.query_in_current_chunk, + keep_qa_history=self._keep_qa_history, + current_chunk_index=self.chunk_index, + ) + + def close_chunk(self) -> int: + self.archive_query() + closed = self.chunk_index + self.chunk_index += 1 + self._chunk_frame_count = 0 + self.query_in_current_chunk = False + return closed + + async def consolidate(self, chunk_index: int, frames: list[tuple[str, str]]) -> None: + if self._summarizer is None or not frames: + return + frame_range = f"{frames[0][0]}-{frames[-1][0]}" + async with self._consolidate_lock: + summary = await self._summarizer.summarize_chunk(chunk_index, frame_range, frames) + self.memory.mid_term_summaries.append(MidTermSummary(chunk_index, frame_range, summary)) + if len(self.memory.mid_term_summaries) >= self._long_term_every_n: + block = await self._summarizer.compress_to_long_term(self.memory.mid_term_summaries) + if block: + self.memory.long_term_blocks.append(block) + window = self._long_term_window + if window > 0 and len(self.memory.long_term_blocks) > window: + # drop oldest blocks so long-term memory stays bounded over hours + del self.memory.long_term_blocks[: len(self.memory.long_term_blocks) - window] + self.memory.long_term_memory = "\n\n".join(self.memory.long_term_blocks) + self.memory.mid_term_summaries.clear() + + async def flush(self, frames: list[tuple[str, str]]) -> None: + await self.consolidate(self.close_chunk(), frames) + + async def submit_delegation(self, action: ParsedAction, frames: list[tuple[str, str]]) -> dict[str, Any] | None: + if action.action is not Action.DELEGATE or self._delegation is None or not self._enable_delegation: + return None + task_id = await self._delegation.submit(action.delegated_question or "", action.text, frames) + self._pending_delegations.append({"task_id": task_id, "question": action.delegated_question or ""}) + return {"task_id": task_id, "status": "submitted", "question": action.delegated_question} + + async def fold_delegations(self) -> dict[str, Any] | None: + if not self._pending_delegations or self._delegation is None: + return None + folded: dict[str, Any] | None = None + still_pending: list[dict[str, str]] = [] + for task in self._pending_delegations: + result = await self._delegation.poll(task["task_id"]) + if result.is_ready: + self.memory.qa_history.append( + QAEntry( + f"[delegated] {task['question']}", + self.query_time or "", + [("", result.digest)], + self.chunk_index, + ) + ) + folded = { + "task_id": result.task_id, + "status": "ready", + "question": task["question"], + "digest": result.digest, + } + if result.media: + folded["media"] = result.media + elif result.status == "error": + folded = {"task_id": result.task_id, "status": "error", "question": task["question"]} + else: + still_pending.append(task) + self._pending_delegations = still_pending + return folded diff --git a/vllm_omni/experimental/fullduplex/joyvl/memory/memory.py b/vllm_omni/experimental/fullduplex/joyvl/memory/memory.py new file mode 100644 index 00000000000..533982a1238 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/memory/memory.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +from vllm_omni.experimental.fullduplex.joyvl.bridges.backend import ModelBackend +from vllm_omni.experimental.fullduplex.joyvl.decision import prompts + + +@dataclass +class WorkingChunk: + messages: list[dict[str, Any]] = field(default_factory=list) + + +@dataclass +class MidTermSummary: + chunk_index: int + frame_range: str + summary_text: str + + +@dataclass +class QAEntry: + query: str + query_time: str + responses: list[tuple[str, str]] = field(default_factory=list) + archived_in_chunk: int = 0 + + +@dataclass +class SessionMemory: + long_term_memory: str = "" + long_term_blocks: list[str] = field(default_factory=list) + mid_term_summaries: list[MidTermSummary] = field(default_factory=list) + qa_history: list[QAEntry] = field(default_factory=list) + + +def build_memory_prefix( + memory: SessionMemory, + *, + current_query: str | None, + query_in_current_chunk: bool, + keep_qa_history: bool, + current_chunk_index: int, +) -> str: + has_query = bool(current_query) + sections: list[str] = [] + + history_parts: list[str] = [] + if has_query and memory.long_term_memory: + history_parts.append(memory.long_term_memory) + for entry in memory.mid_term_summaries: + history_parts.append(f"<{entry.frame_range}>\n{entry.summary_text}") + if history_parts: + sections.append(prompts.VIDEO_HISTORY_HEADER + "\n\n".join(history_parts)) + + if keep_qa_history and has_query: + qa_entries = [e for e in memory.qa_history if e.archived_in_chunk < current_chunk_index] + if qa_entries: + lines = [] + for idx, entry in enumerate(qa_entries, 1): + parts = [f"#{idx} [{prompts.QA_QUERY_LABEL}@{entry.query_time or 'N/A'}] {entry.query}"] + for resp_time, payload in entry.responses: + parts.append(f"[{prompts.QA_RESPONSE_LABEL}@{resp_time}] {payload}") + lines.append("\n".join(parts)) + sections.append(prompts.QA_HISTORY_HEADER + "\n".join(lines)) + + if has_query and not query_in_current_chunk: + sections.append(prompts.USER_QUERY_HEADER + "\n" + current_query.strip()) + + return "\n\n".join(sections) + + +def _sample_indices(n: int, budget: int) -> list[int]: + if n == 0: + return [] + if budget <= 0 or n <= budget: + return list(range(n)) + if budget == 1: + return [n // 2] + selected: list[int] = [] + last_pos = -1 + for slot in range(budget): + target = round(slot * (n - 1) / (budget - 1)) + target = max(target, last_pos + 1) + target = min(target, n - (budget - slot)) + selected.append(target) + last_pos = target + return selected + + +def _resize_data_url(data_url: str, max_pixels: int) -> str: + """Downscale a base64 image data URL so width*height <= max_pixels (preserving + aspect ratio), mirroring the reference summarizer. No-op if already within budget, + not a data URL, or PIL is unavailable.""" + if max_pixels <= 0 or not data_url.startswith("data:"): + return data_url + import base64 + import io + import re + + match = re.match(r"data:image/\w+;base64,(.+)", data_url) + if not match: + return data_url + try: + from PIL import Image + + img = Image.open(io.BytesIO(base64.b64decode(match.group(1)))) + w, h = img.size + if w * h <= max_pixels: + return data_url + scale = (max_pixels / (w * h)) ** 0.5 + img = img.resize((max(1, int(w * scale)), max(1, int(h * scale))), Image.LANCZOS) + buf = io.BytesIO() + img.convert("RGB").save(buf, format="JPEG") + return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode() + except Exception: + return data_url + + +class Summarizer: + def __init__( + self, + backend: ModelBackend, + *, + key_frames_per_chunk: int = 0, + mid_term_max_tokens: int = 4000, + long_term_max_tokens: int = 4000, + preferred_time_span: float = 10.0, + max_pixels: int = 262144, + ) -> None: + self._backend = backend + self._key_frames = key_frames_per_chunk + self._mid_max_tokens = mid_term_max_tokens + self._long_max_tokens = long_term_max_tokens + self._preferred_time_span = preferred_time_span + self._max_pixels = max_pixels + + async def summarize_chunk( + self, + chunk_index: int, + frame_range: str, + frames: list[tuple[str, str]], + ) -> str: + if not frames: + return prompts.EMPTY_CHUNK_SUMMARY.format(frame_range=frame_range) + + picked = [frames[i] for i in _sample_indices(len(frames), self._key_frames)] + prompt = prompts.MID_TERM_SUMMARY_PROMPT.format( + chunk_index=chunk_index, + frame_range=frame_range, + length_instruction="", + preferred_time_span=f"{self._preferred_time_span:g} seconds", + ) + content: list[dict[str, Any]] = [{"type": "text", "text": prompt}] + for time_range, data_url in picked: + content.append({"type": "text", "text": f"<{time_range}>"}) + content.append({"type": "image_url", "image_url": {"url": _resize_data_url(data_url, self._max_pixels)}}) + + text, _ = await self._backend.generate( + [{"role": "user", "content": content}], + max_tokens=self._mid_max_tokens, + temperature=0.8, + top_p=0.9, + extra_body={"top_k": 40, "repetition_penalty": 1.1, "presence_penalty": 1.0, "greedy": False}, + ) + return text.strip() + + async def compress_to_long_term(self, mid_terms: list[MidTermSummary]) -> str: + """Compress a batch of mid-term summaries into a single long-term block. + + Returns one block; the caller (InteractionBrain) accumulates blocks and applies + the sliding window, matching the reference long-term memory behavior. + """ + if not mid_terms: + return "" + + merged_range = f"{_range_start(mid_terms[0].frame_range)}-{_range_end(mid_terms[-1].frame_range)}" + summaries_text = "\n\n".join(f"<{m.frame_range}>\n{m.summary_text}" for m in mid_terms) + prompt = prompts.LONG_TERM_COMPRESS_PROMPT.format( + merged_range=merged_range, + summaries_text=summaries_text, + length_instruction="", + ) + compressed, _ = await self._backend.generate( + [{"role": "user", "content": prompt}], + max_tokens=self._long_max_tokens, + temperature=0.3, + top_p=0.7, + extra_body={"top_k": 30, "repetition_penalty": 1.1, "presence_penalty": 0.5, "greedy": False}, + ) + return f"<{merged_range}>\n{compressed.strip()}" + + async def aclose(self) -> None: + await self._backend.aclose() + + +def _range_start(frame_range: str) -> str: + for sep in (" ~ ", "-"): + if sep in frame_range: + return frame_range.split(sep, 1)[0].strip() + return frame_range.strip() + + +def _range_end(frame_range: str) -> str: + for sep in (" ~ ", "-"): + if sep in frame_range: + return frame_range.split(sep, 1)[-1].strip() + return frame_range.strip() diff --git a/vllm_omni/experimental/fullduplex/joyvl/serving/__init__.py b/vllm_omni/experimental/fullduplex/joyvl/serving/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/serving/__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/experimental/fullduplex/joyvl/serving/config.py b/vllm_omni/experimental/fullduplex/joyvl/serving/config.py new file mode 100644 index 00000000000..fe2a4c19ea9 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/serving/config.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass, field + +from vllm_omni.experimental.fullduplex.joyvl.decision.prompts import SYSTEM_PROMPTS + + +@dataclass +class SamplingConfig: + max_tokens: int = 256 + temperature: float = 0.8 + top_p: float = 0.9 + top_k: int = 40 + repetition_penalty: float = 1.1 + presence_penalty: float = 1.5 + + +@dataclass +class InteractionConfig: + main_backend_url: str = "http://127.0.0.1:8061/v1" + main_model: str = "JoyAI-VL-Interaction-Preview" + api_key: str = "EMPTY" + + persona: str = "default" + frame_seconds: float = 1.0 + max_pixels: int = 262144 + sampling: SamplingConfig = field(default_factory=SamplingConfig) + + force_silence_before_query: bool = True + + response_dedup_threshold: float = 1.0 + + enable_memory: bool = True + + summarizer_backend_url: str | None = None + summarizer_model: str | None = None + + chunk_frames: int = 100 + + mid_term_key_frames: int = 0 + + long_term_every_n_chunks: int = 5 + long_term_memory_window: int = 15 + mid_term_max_tokens: int = 4000 + long_term_max_tokens: int = 4000 + keep_qa_history: bool = True + + enable_delegation: bool = True + delegation_backend_url: str | None = None + delegation_model: str | None = None + delegation_api_key: str | None = None + delegation_max_tokens: int = 512 + delegation_kind: str = "chat" + delegation_image_url: str | None = None + delegation_edit_url: str | None = None + delegation_edit_model: str | None = None + + session_timeout_seconds: float = 3600.0 + request_timeout_seconds: float = 300.0 + + @property + def system_prompt(self) -> str: + return SYSTEM_PROMPTS.get(self.persona, SYSTEM_PROMPTS["default"]) + + @property + def resolved_summarizer_url(self) -> str: + return self.summarizer_backend_url or self.main_backend_url + + @property + def resolved_summarizer_model(self) -> str: + return self.summarizer_model or self.main_model + + @property + def resolved_delegation_model(self) -> str: + return self.delegation_model or self.main_model + + @property + def resolved_delegation_api_key(self) -> str: + return self.delegation_api_key or self.api_key diff --git a/vllm_omni/experimental/fullduplex/joyvl/serving/server.py b/vllm_omni/experimental/fullduplex/joyvl/serving/server.py new file mode 100644 index 00000000000..259d8eb632e --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/serving/server.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import argparse +import asyncio +import contextlib +import time +import uuid +from collections.abc import AsyncIterator +from typing import Any + +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from vllm_omni.experimental.fullduplex.joyvl.bridges.backend import OpenAIBackend +from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import ( + DelegationBridge, + ImageEditDelegationBridge, + ImageGenDelegationBridge, + OpenAIDelegationBridge, + RoutingDelegationBridge, + StubDelegationBridge, +) +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import to_token_form +from vllm_omni.experimental.fullduplex.joyvl.memory.memory import Summarizer +from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig +from vllm_omni.experimental.fullduplex.joyvl.serving.session import InteractionSession, StepResult + + +class SessionManager: + def __init__(self, config: InteractionConfig) -> None: + self.config = config + self._backend = OpenAIBackend( + config.main_backend_url, config.main_model, config.api_key, config.request_timeout_seconds + ) + self._summarizer: Summarizer | None = None + if config.enable_memory: + summarizer_backend = OpenAIBackend( + config.resolved_summarizer_url, + config.resolved_summarizer_model, + config.api_key, + config.request_timeout_seconds, + ) + self._summarizer = Summarizer( + summarizer_backend, + key_frames_per_chunk=config.mid_term_key_frames, + mid_term_max_tokens=config.mid_term_max_tokens, + long_term_max_tokens=config.long_term_max_tokens, + max_pixels=config.max_pixels, + ) + self._delegation = self._build_delegation(config) if config.enable_delegation else None + self._sessions: dict[str, InteractionSession] = {} + self._locks: dict[str, asyncio.Lock] = {} + + @staticmethod + def _build_delegation(config: InteractionConfig) -> DelegationBridge | None: + # Returns None (delegation off) unless a real backend is configured. The stub is + # only returned on explicit opt-in (--delegation-kind stub) for tests/demos — it + # would otherwise fold fake answers into session memory. + kind = config.delegation_kind + url = config.delegation_backend_url + + def chat_bridge(backend_url: str) -> OpenAIDelegationBridge: + return OpenAIDelegationBridge( + backend_url, + config.resolved_delegation_model, + config.resolved_delegation_api_key, + max_tokens=config.delegation_max_tokens, + timeout=config.request_timeout_seconds, + ) + + if kind == "stub": + return StubDelegationBridge() + if kind == "router": + if not (url or config.delegation_image_url or config.delegation_edit_url): + return None + return RoutingDelegationBridge( + chat=chat_bridge(url) if url else None, + image=ImageGenDelegationBridge(config.delegation_image_url, timeout=config.request_timeout_seconds) + if config.delegation_image_url + else None, + edit=ImageEditDelegationBridge( + config.delegation_edit_url, + config.delegation_edit_model or config.main_model, + timeout=config.request_timeout_seconds, + ) + if config.delegation_edit_url + else None, + ) + if url and kind == "image": + return ImageGenDelegationBridge(url, timeout=config.request_timeout_seconds) + if url and kind == "edit": + return ImageEditDelegationBridge( + url, config.resolved_delegation_model, timeout=config.request_timeout_seconds + ) + if url: + return chat_bridge(url) + return None + + def _get(self, session_id: str) -> InteractionSession: + session = self._sessions.get(session_id) + if session is None: + session = InteractionSession( + session_id, + self.config, + self._backend, + summarizer=self._summarizer, + delegation=self._delegation, + ) + self._sessions[session_id] = session + self._locks[session_id] = asyncio.Lock() + return session + + async def step(self, session_id: str, frames: list[str], query: str | None) -> StepResult: + await self._evict_expired() + # capture session + lock together (no await between) so a concurrent reset cannot + # swap them out from under us; the lock then serializes step against reset. + session = self._get(session_id) + lock = self._locks[session_id] + async with lock: + return await session.step(frames, query) + + async def reset(self, session_id: str) -> None: + # Acquire the per-session lock so reset waits for any in-flight step() to finish + # instead of mutating/dropping a session mid-request. + lock = self._locks.get(session_id) + if lock is None: + return + async with lock: + session = self._sessions.pop(session_id, None) + self._locks.pop(session_id, None) + if session is not None: + await session.reset() + + async def set_persona(self, session_id: str, persona: str) -> bool: + session = self._get(session_id) + lock = self._locks[session_id] + async with lock: + return session.set_persona(persona) + + async def _evict_expired(self) -> None: + ttl = self.config.session_timeout_seconds + if ttl <= 0: + return + now = time.monotonic() + expired = [ + sid + for sid, sess in self._sessions.items() + if now - sess.last_access > ttl and not self._locks[sid].locked() + ] + for sid in expired: + await self.reset(sid) + + async def aclose(self) -> None: + for sid in list(self._sessions): + lock = self._locks.get(sid) + if lock is not None: + async with lock: + session = self._sessions.pop(sid, None) + self._locks.pop(sid, None) + if session is not None: + await session.aclose() + await self._backend.aclose() + if self._summarizer is not None: + await self._summarizer.aclose() + if self._delegation is not None: + await self._delegation.aclose() + + +def _extract_frames_and_query(payload: dict[str, Any]) -> tuple[list[str], str | None]: + messages = payload.get("messages") or [] + frames: list[str] = [] + texts: list[str] = [] + for message in messages: + if message.get("role") != "user": + continue + content = message.get("content") + if isinstance(content, str): + texts.append(content) + continue + for part in content or []: + ptype = part.get("type") + if ptype == "image_url": + url = (part.get("image_url") or {}).get("url") + if url: + frames.append(url) + elif ptype == "text" and part.get("text"): + texts.append(part["text"]) + query = "\n".join(t.strip() for t in texts if t.strip()) or None + return frames, query + + +def _session_id(request: Request, payload: dict[str, Any]) -> str: + return ( + request.headers.get("x-streaming-session") + or request.headers.get("x-session-id") + or payload.get("session_id") + or payload.get("user") + or "default" + ) + + +def _completion_response(model: str, result: StepResult) -> dict[str, Any]: + action = result.action + memory = {"long_term_memory": result.long_term_memory, "mid_term_summaries": result.mid_term_summaries} + return { + "id": f"intchat-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": to_token_form(action)}, + "finish_reason": "stop", + } + ], + "usage": None, + "interaction": { + "action": action.action.value, + "spoke": action.spoke, + "text": action.text, + "delegated_question": action.delegated_question, + "delegation": result.delegation, + "chunk_index": result.chunk_index, + "frame_index": result.frame_index, + "inference_skipped": result.inference_skipped, + "latency_ms": result.latency_ms, + "memory": memory, + }, + } + + +def create_app(config: InteractionConfig) -> FastAPI: + manager = SessionManager(config) + + @contextlib.asynccontextmanager + async def lifespan(_app: FastAPI) -> AsyncIterator[None]: + try: + yield + finally: + # close httpx/AsyncOpenAI clients and cancel+await pending tasks on shutdown/reload + await manager.aclose() + + app = FastAPI(title="vLLM-Omni Interaction Server", lifespan=lifespan) + + @app.get("/health") + async def health() -> dict[str, str]: + return {"status": "ok"} + + @app.get("/v1/models") + async def models() -> dict[str, Any]: + return {"object": "list", "data": [{"id": config.main_model, "object": "model", "owned_by": "vllm-omni"}]} + + @app.post("/v1/chat/completions") + async def chat_completions(request: Request) -> JSONResponse: + payload = await request.json() + frames, query = _extract_frames_and_query(payload) + if not frames: + return JSONResponse({"error": "interaction server requires at least one image_url frame"}, status_code=400) + result = await manager.step(_session_id(request, payload), frames, query) + return JSONResponse(_completion_response(config.main_model, result)) + + @app.post("/reset") + @app.post("/v1/streaming/reset") + async def reset(request: Request) -> dict[str, str]: + payload = await request.json() if await request.body() else {} + await manager.reset(_session_id(request, payload)) + return {"status": "reset"} + + @app.post("/v1/streaming/persona") + async def persona(request: Request) -> dict[str, Any]: + payload = await request.json() if await request.body() else {} + ok = await manager.set_persona(_session_id(request, payload), payload.get("persona", "default")) + return {"status": "ok" if ok else "unknown_persona"} + + return app + + +def _build_config(args: argparse.Namespace) -> InteractionConfig: + config = InteractionConfig( + main_backend_url=args.main_backend_url, + main_model=args.main_model, + persona=args.persona, + enable_memory=not args.no_memory, + summarizer_backend_url=args.summarizer_backend_url, + summarizer_model=args.summarizer_model, + enable_delegation=not args.no_delegation, + delegation_backend_url=args.delegation_backend_url, + delegation_model=args.delegation_model, + delegation_api_key=args.delegation_api_key, + delegation_kind=args.delegation_kind, + delegation_image_url=args.delegation_image_url, + delegation_edit_url=args.delegation_edit_url, + delegation_edit_model=args.delegation_edit_model, + force_silence_before_query=not args.no_force_silence, + ) + if args.chunk_frames is not None: + config.chunk_frames = args.chunk_frames + config.response_dedup_threshold = args.response_dedup_threshold + config.sampling.max_tokens = args.max_tokens + config.sampling.temperature = args.temperature + return config + + +def main() -> None: + parser = argparse.ArgumentParser(description="vLLM-Omni streaming interaction server") + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, default=8070) + parser.add_argument("--main-backend-url", default="http://127.0.0.1:8061/v1") + parser.add_argument("--main-model", default="JoyAI-VL-Interaction-Preview") + parser.add_argument("--persona", default="default", choices=["default", "silent", "talkative"]) + parser.add_argument("--summarizer-backend-url", default=None) + parser.add_argument("--summarizer-model", default=None) + parser.add_argument("--chunk-frames", type=int, default=None) + parser.add_argument( + "--response-dedup-threshold", + type=float, + default=1.0, + help="1.0 drops only exact repeats (reference); < 1.0 also drops near-duplicate narration", + ) + parser.add_argument("--max-tokens", type=int, default=128) + parser.add_argument("--temperature", type=float, default=0.8) + parser.add_argument("--no-memory", action="store_true", help="disable mid/long-term summarizer memory") + parser.add_argument("--no-delegation", action="store_true", help="disable the delegation bridge") + parser.add_argument( + "--delegation-backend-url", + default=None, + help="OpenAI-compatible endpoint for the background brain; unset falls back to the stub", + ) + parser.add_argument("--delegation-model", default=None, help="model name for the background brain") + parser.add_argument( + "--delegation-api-key", + default=None, + help="API key for the background brain endpoint (e.g. an Anthropic key for a claude-* model)", + ) + parser.add_argument( + "--delegation-kind", + default="chat", + choices=["chat", "image", "edit", "router", "stub"], + help="chat = text/VL brain; image = text-to-image; edit = restyle the frame; " + "router = dispatch by request; stub = canned demo/test answers. chat/image/edit/router " + "need a backend URL — without one, delegation stays off.", + ) + parser.add_argument("--delegation-image-url", default=None, help="router mode: text-to-image endpoint") + parser.add_argument("--delegation-edit-url", default=None, help="router mode: image-edit endpoint") + parser.add_argument("--delegation-edit-model", default=None, help="router mode: image-edit model name") + parser.add_argument("--no-force-silence", action="store_true", help="run the model before any user query") + args = parser.parse_args() + + app = create_app(_build_config(args)) + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/experimental/fullduplex/joyvl/serving/session.py b/vllm_omni/experimental/fullduplex/joyvl/serving/session.py new file mode 100644 index 00000000000..8ede35f47e0 --- /dev/null +++ b/vllm_omni/experimental/fullduplex/joyvl/serving/session.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import asyncio +import logging +import time +from dataclasses import dataclass, field +from typing import Any + +from vllm_omni.experimental.fullduplex.joyvl.bridges.backend import ModelBackend +from vllm_omni.experimental.fullduplex.joyvl.bridges.delegation import DelegationBridge +from vllm_omni.experimental.fullduplex.joyvl.decision.output_parser import Action, ParsedAction +from vllm_omni.experimental.fullduplex.joyvl.decision.policy import JoyVLPolicy +from vllm_omni.experimental.fullduplex.joyvl.decision.prompts import SYSTEM_PROMPTS, USER_QUERY_HEADER +from vllm_omni.experimental.fullduplex.joyvl.memory.memory import Summarizer, WorkingChunk +from vllm_omni.experimental.fullduplex.joyvl.serving.config import InteractionConfig + +logger = logging.getLogger("vllm_omni.experimental.fullduplex.joyvl") + + +@dataclass +class StepResult: + action: ParsedAction + chunk_index: int + frame_index: int + inference_skipped: bool + latency_ms: float + long_term_memory: str + mid_term_summaries: list[dict[str, Any]] = field(default_factory=list) + delegation: dict[str, Any] | None = None + + +class InteractionSession: + def __init__( + self, + session_id: str, + config: InteractionConfig, + backend: ModelBackend, + summarizer: Summarizer | None = None, + delegation: DelegationBridge | None = None, + ) -> None: + self.session_id = session_id + self.config = config + self._backend = backend + self._policy = JoyVLPolicy( + persona=config.persona, + summarizer=summarizer, + delegation=delegation, + chunk_frames=config.chunk_frames, + long_term_every_n_chunks=config.long_term_every_n_chunks, + long_term_window=config.long_term_memory_window, + keep_qa_history=config.keep_qa_history, + frame_seconds=config.frame_seconds, + enable_delegation=config.enable_delegation, + response_dedup_threshold=config.response_dedup_threshold, + ) + self.chunk = WorkingChunk() + self._system_prompt = config.system_prompt + self.last_access = time.monotonic() + self._consolidating: set[asyncio.Task] = set() + + def set_persona(self, persona: str) -> bool: + prompt = SYSTEM_PROMPTS.get(persona) + if prompt is None: + return False + self._system_prompt = prompt + return True + + async def step(self, frames: list[str], query: str | None = None, t: float | None = None) -> StepResult: + self.last_access = time.monotonic() + started = time.perf_counter() + policy = self._policy + brain = policy.brain + + base = t if t is not None else brain.frame_index * self.config.frame_seconds + time_ranges = [f"{base + i * self.config.frame_seconds:.1f} seconds" for i in range(len(frames))] + + delegation_info = await policy.fold_delegations() + + if policy.needs_flush(): + self._spawn_consolidation(policy.close_chunk(), policy.take_working_frames()) + self.chunk = WorkingChunk() + + query_is_fresh = policy.set_query(query) + + for tr, url in zip(time_ranges, frames): + policy.observe(tr, url) + self.chunk.messages.append(self._frame_message(time_ranges, frames, include_query=query_is_fresh)) + + if self.config.force_silence_before_query and not brain.current_query: + action = ParsedAction(Action.SILENCE, raw="</silence>") + skipped = True + else: + action = policy.commit(await self._infer()) + skipped = False + self.chunk.messages.append({"role": "assistant", "content": action.raw or "</silence>"}) + + submitted = await policy.submit_if_delegate(action, list(policy.working_frames)) + if submitted: + delegation_info = submitted + + return StepResult( + action=action, + chunk_index=brain.chunk_index, + frame_index=brain.frame_index, + inference_skipped=skipped, + latency_ms=round((time.perf_counter() - started) * 1000, 1), + long_term_memory=brain.memory.long_term_memory, + mid_term_summaries=[ + {"chunk_index": m.chunk_index, "frame_range": m.frame_range, "summary_text": m.summary_text} + for m in brain.memory.mid_term_summaries + ], + delegation=delegation_info, + ) + + async def reset(self) -> None: + await self._cancel_consolidation() + self._policy.brain.reset() + self.chunk = WorkingChunk() + + async def aclose(self) -> None: + await self._cancel_consolidation() + + async def _cancel_consolidation(self) -> None: + tasks = list(self._consolidating) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + self._consolidating.clear() + + def _spawn_consolidation(self, chunk_index: int, frames: list[tuple[str, str]]) -> None: + task = asyncio.create_task(self._policy.consolidate(chunk_index, frames)) + self._consolidating.add(task) + task.add_done_callback(self._on_consolidation_done) + + def _on_consolidation_done(self, task: asyncio.Task) -> None: + self._consolidating.discard(task) + if not task.cancelled() and task.exception() is not None: + logger.warning("memory consolidation failed for session %s: %s", self.session_id, task.exception()) + + async def _infer(self) -> str: + s = self.config.sampling + extra_body = { + "skip_special_tokens": False, + "greedy": False, + "top_k": s.top_k, + "repetition_penalty": s.repetition_penalty, + "presence_penalty": s.presence_penalty, + } + raw, _ = await self._backend.generate( + self._build_api_messages(), + max_tokens=s.max_tokens, + temperature=s.temperature, + top_p=s.top_p, + extra_body=extra_body, + ) + return raw + + def _build_api_messages(self) -> list[dict[str, Any]]: + messages: list[dict[str, Any]] = [{"role": "system", "content": self._system_prompt}] + chunk_messages = [dict(m) for m in self.chunk.messages] + prefix = self._policy.brain.build_prefix() + if prefix and chunk_messages: + head = chunk_messages[0] + head["content"] = [{"type": "text", "text": prefix}] + list(head["content"]) + messages.extend(chunk_messages) + return messages + + def _frame_message(self, time_ranges: list[str], frames: list[str], include_query: bool) -> dict[str, Any]: + content: list[dict[str, Any]] = [] + query = self._policy.brain.current_query + if include_query and query: + content.append({"type": "text", "text": f"{USER_QUERY_HEADER}\n{query.strip()}"}) + for tr, url in zip(time_ranges, frames): + content.append({"type": "text", "text": f"<{tr}>"}) + content.append({"type": "image_url", "image_url": {"url": url}, "max_pixels": self.config.max_pixels}) + return {"role": "user", "content": content} diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index a2b6bf722d5..73ac8064280 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -326,6 +326,15 @@ def batch_size(self): @property def resolved_frame_rate(self) -> float | None: + """Frame rate the model should run at, or ``None`` if the user did not provide one. + + Precedence is ``frame_rate`` (the model-internal rate) over ``fps`` (the + request/output-encoding alias). Returns ``None`` when neither was set, which is + the model-agnostic "fps not provided" signal the video serving layer emits when + the request omits fps. Consumers MUST treat ``None`` as "use your own default" and + guard the read (e.g. ``req.sampling_params.resolved_frame_rate or 24.0``); never + feed it into arithmetic unguarded. + """ if self.frame_rate is not None: return float(self.frame_rate) diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 6c1288edf81..9660e4f343c 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -11,11 +11,14 @@ ``_total`` (auto-suffixed by the prometheus client), sizes use ``_bytes``. """ +import logging +import os from typing import Any -# vllm:omni_ avoids upstream's unregister_vllm_metrics() stripping, which -# removes every collector whose ``_name`` does not start with ``vllm``. -METRIC_PREFIX = "vllm:omni_" +# vllm_omni: namespace for omni-specific Prometheus families, distinct from +# the upstream vllm:* families. +METRIC_PREFIX = "vllm_omni:" +logger = logging.getLogger(__name__) # ============================================================================ @@ -76,10 +79,8 @@ VLLM_TPOTS = "vllm_tpots" VLLM_ITLS = "vllm_itls" AUDIO_TTFPS = "audio_ttfps" -AUDIO_RTFS = "audio_rtfs" AUDIO_DURATIONS = "audio_durations" MISSING_AUDIO_DURATION_COUNT = "missing_audio_duration_count" -MISSING_AUDIO_RTF_COUNT = "missing_audio_rtf_count" STAGE_GEN_TIME = "stage_gen_time" STAGE_GEN_TIME_MS = f"{STAGE_GEN_TIME}_ms" STAGE_GEN_TIMES_MS = f"{STAGE_GEN_TIME}s_ms" @@ -114,6 +115,10 @@ # ``_total`` at exposition time. REQUESTS_SUCCESS = METRIC_PREFIX + "requests_success" +# Token counters — aggregated across all pipeline stages per request. +PROMPT_TOKENS = METRIC_PREFIX + "prompt_tokens" +GENERATION_TOKENS = METRIC_PREFIX + "generation_tokens" + # ============================================================================ # Audio family (per-stage + per-replica audio path metrics) @@ -278,6 +283,9 @@ def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int # ming_flash 16000 — these models populate multimodal_output["audio_sample_rate"] # at runtime so this default only kicks in when the field is missing. DEFAULT_AUDIO_SAMPLE_RATE = 24000 +DEFAULT_AUDIO_CHANNELS = 1 +AUDIO_SAMPLE_RATE_ENV = "VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE" +AUDIO_CHANNELS_ENV = "VLLM_OMNI_BENCH_AUDIO_CHANNELS" _SAMPLE_RATE_KEYS = ("output_sample_rate", "audio_sample_rate", "sample_rate", "sampling_rate", "sr") @@ -304,3 +312,26 @@ def resolve_audio_sample_rate(source: dict[str, Any] | Any | None) -> int: if value > 0: return value return DEFAULT_AUDIO_SAMPLE_RATE + + +def stream_pcm_format_from_env( + *, + default_sample_rate: int = DEFAULT_AUDIO_SAMPLE_RATE, + default_channels: int = DEFAULT_AUDIO_CHANNELS, +) -> tuple[int, int]: + """Return the sample rate and channel count for streamed raw PCM.""" + sample_rate = default_sample_rate + channels = default_channels + raw_sr = os.environ.get(AUDIO_SAMPLE_RATE_ENV) + if raw_sr: + try: + sample_rate = int(raw_sr) + except ValueError: + logger.warning("Invalid %s=%r; using default %d", AUDIO_SAMPLE_RATE_ENV, raw_sr, sample_rate) + raw_ch = os.environ.get(AUDIO_CHANNELS_ENV) + if raw_ch: + try: + channels = int(raw_ch) + except ValueError: + logger.warning("Invalid %s=%r; using default %d", AUDIO_CHANNELS_ENV, raw_ch, channels) + return max(sample_rate, 1), max(channels, 1) diff --git a/vllm_omni/metrics/prometheus.py b/vllm_omni/metrics/prometheus.py index fc971df03da..d372c3d6bc1 100644 --- a/vllm_omni/metrics/prometheus.py +++ b/vllm_omni/metrics/prometheus.py @@ -27,14 +27,23 @@ labelnames=_labelnames, buckets=defs.SECONDS_BUCKETS, ) +_prompt_tokens_family = Counter( + defs.PROMPT_TOKENS, + "Total prompt (input) tokens processed across all pipeline stages.", + labelnames=_labelnames, +) +_generation_tokens_family = Counter( + defs.GENERATION_TOKENS, + "Total generation (output) tokens produced across all pipeline stages.", + labelnames=_labelnames, +) class OmniPrometheusMetrics: """Label-bound wrapper around the raw Prometheus metrics. - Metric collectors use the ``vllm:omni_`` prefix to avoid being - removed by upstream vLLM's ``unregister_vllm_metrics()``, which - strips every collector whose ``_name`` contains ``"vllm"``. + Metric collectors use the ``vllm_omni:`` prefix, distinct from the + upstream ``vllm:*`` families. """ def __init__(self, model_name: str, log_stats: bool = True) -> None: @@ -43,6 +52,8 @@ def __init__(self, model_name: str, log_stats: bool = True) -> None: self._running = _running_family.labels(model_name=model_name) self._waiting = _waiting_family.labels(model_name=model_name) self._e2e_latency = _e2e_latency_family.labels(model_name=model_name) + self._prompt_tokens = _prompt_tokens_family.labels(model_name=model_name) + self._generation_tokens = _generation_tokens_family.labels(model_name=model_name) def set_running(self, n: int) -> None: if not self._log_stats: @@ -54,6 +65,14 @@ def set_waiting(self, n: int) -> None: return self._waiting.set(n) + def observe_tokens(self, prompt_tokens: int, generation_tokens: int) -> None: + if not self._log_stats: + return + if prompt_tokens > 0: + self._prompt_tokens.inc(prompt_tokens) + if generation_tokens > 0: + self._generation_tokens.inc(generation_tokens) + def request_succeeded( self, e2e_seconds: float, diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 3db2fda14a2..a282d7d50d9 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -48,7 +48,6 @@ class StageRequestStats: audio_generated_frames: int = 0 audio_sample_rate: int = 0 audio_duration_s: float = 0.0 - audio_rtf: float = 0.0 image_pixels: int = 0 denoise_step_latency_ms: float = 0.0 pipeline_timings: dict[str, float] | None = None @@ -471,7 +470,6 @@ def _merge_stage_metric_event( defs.AUDIO_FRAMES: int(evt.audio_generated_frames), defs.AUDIO_SAMPLE_RATE: int(evt.audio_sample_rate), f"{defs.AUDIO_DURATION}_s": float(evt.audio_duration_s), - defs.AUDIO_RTF: float(evt.audio_rtf), defs.IMAGE_PIXELS: int(evt.image_pixels), defs.DENOISE_STEP_LATENCY_MS: float(evt.denoise_step_latency_ms), "output_unit_type": evt.output_unit_type, @@ -541,10 +539,6 @@ def _merge_stage_metric_event( 0.0, ) current[defs.TIME_PER_OUTPUT_UNIT_MS] = remaining_ms / float(output_count - 1) if output_count > 1 else 0.0 - duration_s = float(current.get(f"{defs.AUDIO_DURATION}_s", 0.0)) - current[defs.AUDIO_RTF] = defs.compute_audio_rtf( - float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) / 1000.0, duration_s - ) return current def _build_stage_metrics_snapshot(self, req_id: str) -> dict[str, dict[str, Any]]: diff --git a/vllm_omni/model_executor/layers/rotary_embedding/mrope.py b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py index 3d3a88d8772..9c65a3c99e3 100644 --- a/vllm_omni/model_executor/layers/rotary_embedding/mrope.py +++ b/vllm_omni/model_executor/layers/rotary_embedding/mrope.py @@ -16,12 +16,18 @@ import torch from transformers import PretrainedConfig from vllm.logger import init_logger -from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding + +from vllm_omni.platforms import current_omni_platform + +if current_omni_platform.is_npu(): + from vllm_ascend.ops.rotary_embedding import AscendMRotaryEmbedding as _BaseMRotaryEmbedding +else: + from vllm.model_executor.layers.rotary_embedding.mrope import MRotaryEmbedding as _BaseMRotaryEmbedding logger = init_logger(__name__) -class OmniMRotaryEmbedding(MRotaryEmbedding): +class OmniMRotaryEmbedding(_BaseMRotaryEmbedding): """Omni-extended MRotaryEmbedding with multimodal position computation. Extends the upstream MRotaryEmbedding with additional class methods for diff --git a/vllm_omni/model_executor/models/aura_omni/__init__.py b/vllm_omni/model_executor/models/aura_omni/__init__.py new file mode 100644 index 00000000000..fb6738a2967 --- /dev/null +++ b/vllm_omni/model_executor/models/aura_omni/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/aura_omni/pipeline.py b/vllm_omni/model_executor/models/aura_omni/pipeline.py new file mode 100644 index 00000000000..14aa5ff3004 --- /dev/null +++ b/vllm_omni/model_executor/models/aura_omni/pipeline.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AURA Omni pipeline topology. + +Semantic modules: + 1. Qwen3-ASR: microphone audio -> transcript text + 2. AURA/Qwen3-VL: transcript + video -> text or <|silent|> + 3. Qwen3-TTS: response text -> audio + +The Qwen3-TTS module is represented as two native engine stages, Talker and +Code2Wav, to reuse the existing vLLM-Omni implementation without wrapping or +duplicating codec streaming behavior. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_AURA_PROC = "vllm_omni.model_executor.stage_input_processors.aura_omni" +_QWEN3_TTS_PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_tts" + + +AURA_OMNI_PIPELINE = PipelineConfig( + model_type="aura_omni", + default_deploy_config_name="aura_omni.yaml", + model_arch="Qwen3ASRForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="asr", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + model_arch="Qwen3ASRForConditionalGeneration", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="aura", + execution_type=StageExecutionType.LLM_AR, + input_sources=(0,), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + model_arch="AuraQwen3VLForConditionalGeneration", + custom_process_input_func=f"{_AURA_PROC}.asr2aura", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=2, + model_stage="qwen3_tts", + execution_type=StageExecutionType.LLM_AR, + input_sources=(1,), + owns_tokenizer=True, + engine_output_type="latent", + model_arch="Qwen3TTSTalkerForConditionalGeneration", + custom_process_input_func=f"{_AURA_PROC}.aura2tts", + custom_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_full_payload", + async_chunk_process_next_stage_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_async_chunk", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [2150], + }, + ), + StagePipelineConfig( + stage_id=3, + model_stage="code2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(2,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="Qwen3TTSCode2Wav", + sync_process_input_func=f"{_QWEN3_TTS_PROC}.talker2code2wav_token_only", + sampling_constraints={"detokenize": True}, + extras={"tts_args": {"max_instructions_length": 500}}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/aura_omni/qwen3_vl.py b/vllm_omni/model_executor/models/aura_omni/qwen3_vl.py new file mode 100644 index 00000000000..65243893ba7 --- /dev/null +++ b/vllm_omni/model_executor/models/aura_omni/qwen3_vl.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""AURA-compatible Qwen3-VL wrapper. + +AURA ships Qwen3-VL weights with a repository-local ``configuration_qwen3_vl`` +module in ``auto_map``. With ``trust_remote_code=True``, Transformers therefore +returns a remote ``Qwen3VLConfig`` class. The model structure is still Qwen3-VL, +but vLLM's stock multimodal processor checks for the upstream Transformers +``Qwen3VLConfig`` type exactly. This wrapper keeps the stock model and processor +logic, while accepting structurally-compatible remote Qwen3-VL configs. +""" + +from __future__ import annotations + +from typing import Any + +from vllm.model_executor.models.qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLDummyInputsBuilder, + Qwen3VLForConditionalGeneration, + Qwen3VLMultiModalProcessor, + Qwen3VLProcessingInfo, + Qwen3VLProcessor, +) +from vllm.multimodal import MULTIMODAL_REGISTRY + + +class AuraQwen3VLProcessingInfo(Qwen3VLProcessingInfo): + def get_hf_config(self) -> Any: + try: + return self.ctx.get_hf_config(Qwen3VLConfig) + except TypeError: + config = self.ctx.get_hf_config() + if ( + getattr(config, "model_type", None) == "qwen3_vl" + and hasattr(config, "vision_config") + and hasattr(config, "text_config") + ): + return config + raise + + def get_hf_processor(self, **kwargs: object) -> Any: + # AURA's config advertises `processing_qwen3_vl.py` in auto_map, but + # the checkpoint snapshot does not include that file. Force the + # upstream processor class while still allowing AURA's remote config. + return self.ctx.get_hf_processor( + Qwen3VLProcessor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=AuraQwen3VLProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class AuraQwen3VLForConditionalGeneration(Qwen3VLForConditionalGeneration): + pass diff --git a/vllm_omni/model_executor/models/bagel/__init__.py b/vllm_omni/model_executor/models/bagel/__init__.py index 11e4c666710..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/bagel/__init__.py +++ b/vllm_omni/model_executor/models/bagel/__init__.py @@ -1,3 +1,8 @@ -from .bagel import OmniBagelForConditionalGeneration +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = ["OmniBagelForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 282abed23de..a3dfd7bf151 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -587,9 +587,9 @@ def prepare_runner_inputs( input_ids: torch.Tensor | None, positions: torch.Tensor | None, inputs_embeds: torch.Tensor | None, - req_ids: list[str], - num_computed_tokens: list[int], - num_scheduled_tokens: list[int], + req_ids: Sequence[str], + num_computed_tokens: Sequence[int], + num_scheduled_tokens: Sequence[int], input_ids_buffer: torch.Tensor | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: """Restore input_ids so _adjust_positions_for_img2img can locate @@ -599,7 +599,7 @@ def prepare_runner_inputs( input_ids = input_ids_buffer return input_ids, positions - def flush_pending_metadata(self, req_ids: list[str]) -> None: + def flush_pending_metadata(self, req_ids: Sequence[str]) -> None: """Map pending metadata (batch order) to req_ids after forward(). Guard: if a request already has metadata with ``image_shape`` diff --git a/vllm_omni/model_executor/models/bagel/pipeline.py b/vllm_omni/model_executor/models/bagel/pipeline.py index c68a531c294..b0544222fbd 100644 --- a/vllm_omni/model_executor/models/bagel/pipeline.py +++ b/vllm_omni/model_executor/models/bagel/pipeline.py @@ -27,6 +27,7 @@ BAGEL_PIPELINE = PipelineConfig( model_type="bagel", + default_deploy_config_name="bagel.yaml", model_arch="OmniBagelForConditionalGeneration", hf_architectures=("BagelForConditionalGeneration",), stages=( @@ -63,6 +64,7 @@ BAGEL_THINK_PIPELINE = PipelineConfig( model_type="bagel_think", + default_deploy_config_name="bagel_think.yaml", model_arch="OmniBagelForConditionalGeneration", hf_architectures=(), stages=( @@ -96,6 +98,7 @@ BAGEL_SINGLE_STAGE_PIPELINE = PipelineConfig( model_type="bagel_single_stage", + default_deploy_config_name="bagel_single_stage.yaml", model_arch="BagelForConditionalGeneration", hf_architectures=(), stages=( diff --git a/vllm_omni/model_executor/models/common/ming/audio_vae.py b/vllm_omni/model_executor/models/common/ming/audio_vae.py index 6924ae54846..aecc35da997 100644 --- a/vllm_omni/model_executor/models/common/ming/audio_vae.py +++ b/vllm_omni/model_executor/models/common/ming/audio_vae.py @@ -138,8 +138,8 @@ def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, pat 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) + 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) @@ -346,11 +346,9 @@ def encode_latent(self, waveform, waveform_length): h, _ = self.encoder(waveform) h = h.transpose(1, 2) - 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) + mean, scale = torch.chunk(h, 2, dim=1) + std = F.softplus(scale) + 1e-4 + latent = (mean + std * torch.randn_like(mean)).transpose(1, 2) return latent, frame_num @torch.inference_mode() diff --git a/vllm_omni/model_executor/models/common/ming/dit.py b/vllm_omni/model_executor/models/common/ming/dit.py index 3b4ab112230..56e65f47fae 100644 --- a/vllm_omni/model_executor/models/common/ming/dit.py +++ b/vllm_omni/model_executor/models/common/ming/dit.py @@ -6,6 +6,19 @@ from x_transformers.x_transformers import apply_rotary_pos_emb +def _apply_rope_cached(t, cos, sin): + """Bit-exact RoPE for ``scale==1`` / interleaved-half. + Uses precomputed ``cos``/``sin`` per seq_len to skip per-step trig and cat.""" + rot_dim = cos.shape[-1] + tr = t[..., :rot_dim] + x = tr.reshape(*tr.shape[:-1], -1, 2) + rot_half = torch.stack((-x[..., 1], x[..., 0]), dim=-1).reshape_as(tr) + out = tr * cos + rot_half * sin + if rot_dim < t.shape[-1]: # partial rotary (unused by ming heads); keep general + out = torch.cat((out, t[..., rot_dim:]), dim=-1) + return out.type(t.dtype) + + class RMSNorm(nn.Module): def __init__(self, dim, eps=1e-6): super().__init__() @@ -66,12 +79,33 @@ def __init__( self.to_out.append(nn.Dropout(dropout)) self.pe_attn_head = pe_attn_head self.attn_mask_enabled = attn_mask_enabled + # cos/sin cache keyed by seq_len (freqs are deterministic in seq_len). + self._rope_cache: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + # TODO: materialize fused weights eagerly outside any torch.compile + # scope (e.g. post-weight-load); + self._qkv_w: torch.Tensor | None = None + self._qkv_b: torch.Tensor | None = None + + def _qkv(self, x): + if self._qkv_w is None: + self._qkv_w = torch.cat([self.to_q.weight, self.to_k.weight, self.to_v.weight], dim=0) + if self.to_q.bias is not None: + self._qkv_b = torch.cat([self.to_q.bias, self.to_k.bias, self.to_v.bias], dim=0) + qkv = F.linear(x, self._qkv_w, self._qkv_b) + return qkv.split(self.inner_dim, dim=-1) + + def _rope_cos_sin(self, freqs): + seq_len = freqs.shape[-2] + cached = self._rope_cache.get(seq_len) + if cached is None: + f = freqs.unsqueeze(1) if freqs.ndim == 3 else freqs # b n d -> b 1 n d for 4D q/k + cached = (f.cos(), f.sin()) + self._rope_cache[seq_len] = cached + return cached 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) + query, key, value = self._qkv(x) inner_dim = key.shape[-1] head_dim = inner_dim // self.heads @@ -86,14 +120,21 @@ def forward(self, x, mask=None, rope=None): 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) + scale_is_one = xpos_scale is None or (not torch.is_tensor(xpos_scale) and xpos_scale == 1.0) + if scale_is_one and self.pe_attn_head is None: + # Fast path: cached, cat-free RoPE (bit-exact for scale==1, full rotary). + cos, sin = self._rope_cos_sin(freqs) + query = _apply_rope_cached(query, cos, sin) + key = _apply_rope_cached(key, cos, sin) else: - query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) - key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + 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) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index eca5addba79..7ac2ba123d8 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -3,7 +3,7 @@ Shared by Qwen3-Omni and Qwen3-TTS talker models. * SDPA attention (F.scaled_dot_product_attention) with native GQA support -* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers) +* HF-compatible CPU/CUDA numerics with NPU-only fused norm/RoPE fast paths * Per-call embedding buffer to avoid cross-request aliasing * Pre-allocated position_ids (read-only, safe to persist) * torch.compile (epilogue_fusion=False) on inner transformer by default @@ -14,7 +14,7 @@ from __future__ import annotations import dataclasses -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import torch import torch.nn as nn @@ -24,37 +24,56 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.layers.custom_op import CustomOp from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) +_GeneratorLike = torch.Generator | Sequence[torch.Generator | None] | None +_UNIFORM_EPS = 1e-20 + +if current_omni_platform.is_npu(): + import torch_npu + # =================================================================== -# HF-numerics-compatible layers for code predictor +# Portable layers for code predictor # =================================================================== # # These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32, # rotate_half RoPE) to produce outputs numerically identical to the -# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel, -# get_rope) introduce small precision differences that compound across -# the autoregressive steps of the code predictor, causing severe -# audio quality degradation. +# HuggingFace reference on CPU/CUDA. vLLM's fused kernels (RMSNorm, +# QKVParallel, get_rope) introduce small precision differences that compound +# across the autoregressive steps of the code predictor, causing severe audio +# quality degradation. The Ascend fused norm/RoPE kernels below are dispatched +# by the current device platform. # # See: https://github.com/vllm-project/vllm-omni/issues/2274 -class _RMSNorm(nn.Module): - """RMSNorm matching HuggingFace's implementation exactly. - - Computes variance in float32 to avoid bfloat16 precision loss. - """ +class _RMSNorm(CustomOp): + """RMSNorm with HuggingFace-compatible CPU/CUDA math and an NPU fast path.""" def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(hidden_size)) self.variance_epsilon = eps - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states, _ = torch_npu.npu_rms_norm( + hidden_states, + self.weight, + self.variance_epsilon, + ) + return hidden_states + + def forward_cuda(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.forward_native(hidden_states) + + def forward_xpu(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.forward_native(hidden_states) + + def forward_native(self, hidden_states: torch.Tensor) -> torch.Tensor: input_dtype = hidden_states.dtype hidden_states = hidden_states.to(torch.float32) variance = hidden_states.pow(2).mean(-1, keepdim=True) @@ -69,11 +88,8 @@ def _rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) -class _RotaryEmbedding(nn.Module): - """RoPE matching HuggingFace's implementation exactly. - - Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False). - """ +class _RotaryEmbedding(CustomOp): + """RoPE with HuggingFace-compatible CPU/CUDA math and cached NPU tables.""" def __init__(self, config) -> None: super().__init__() @@ -85,8 +101,24 @@ def __init__(self, config) -> None: rope_theta = getattr(config, "rope_theta", 10000.0) inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) + if current_omni_platform.is_npu(): + max_seq = int(getattr(config, "num_code_groups", 0) or 0) + 1 + positions = torch.arange(max_seq, dtype=torch.float32) + freqs = torch.outer(positions, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos(), persistent=False) + self.register_buffer("sin_cached", emb.sin(), persistent=False) + + def forward_npu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.cos_cached[position_ids].to(dtype=x.dtype), self.sin_cached[position_ids].to(dtype=x.dtype) + + def forward_cuda(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.forward_native(x, position_ids) + + def forward_xpu(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return self.forward_native(x, position_ids) - def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + def forward_native(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: # position_ids: [batch, seq_len] inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) position_ids_expanded = position_ids[:, None, :].float() @@ -142,7 +174,6 @@ def __init__(self, config, *, prefix: str = "") -> None: self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) - if current_omni_platform.is_npu(): if self.max_seq > 2048: raise ValueError( @@ -165,8 +196,6 @@ def _forward_npu_attention( bsz: int, seq_len: int, ) -> torch.Tensor: - import torch_npu - q_f, k_f, v_f = q, k, v if self.is_gqa: k_f = ( @@ -227,10 +256,13 @@ def forward( # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads cos = cos.unsqueeze(1) # [batch, 1, seq_len, head_dim] sin = sin.unsqueeze(1) - q = (q * cos) + (_rotate_half(q) * sin) - k = (k * cos) + (_rotate_half(k) * sin) - - if not current_omni_platform.is_npu(): + if current_omni_platform.is_npu(): + q = torch_npu.npu_rotary_mul(q, cos, sin) + k = torch_npu.npu_rotary_mul(k, cos, sin) + attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len) + else: + q = (q * cos) + (_rotate_half(q) * sin) + k = (k * cos) + (_rotate_half(k) * sin) attn_out = F.scaled_dot_product_attention( q, k, @@ -239,8 +271,6 @@ def forward( is_causal=True, enable_gqa=self.is_gqa, ) - else: - attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len) attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(attn_out) @@ -287,10 +317,17 @@ def forward( residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, position_embeddings) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + if current_omni_platform.is_npu(): + hidden_states, _, residual = torch_npu.npu_add_rms_norm( + hidden_states, + residual, + self.post_attention_layernorm.weight, + self.post_attention_layernorm.variance_epsilon, + ) + else: + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states @@ -605,6 +642,30 @@ def _parse_positive_int_set(value: object) -> set[int]: values.add(parsed) return values + @staticmethod + def _normalize_generators( + generator: _GeneratorLike, batch_size: int + ) -> torch.Generator | list[torch.Generator | None] | None: + if generator is None or isinstance(generator, torch.Generator): + return generator + + row_generators = list(generator) + if len(row_generators) != batch_size: + raise ValueError(f"Expected {batch_size} per-row generators, but got {len(row_generators)}.") + return row_generators + + @classmethod + def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor: + """Sample ``logits`` via Gumbel-max with optional per-row generators.""" + row_generators = cls._normalize_generators(generator, int(logits.shape[0])) + u = torch.empty_like(logits, dtype=torch.float32) + if isinstance(row_generators, list): + for row, row_generator in enumerate(row_generators): + u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator) + else: + u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators) + return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True) + def _prefix_seq_lens(self, max_seq: int) -> list[int]: all_seq_lens = list(range(2, max_seq)) if not self._prefix_graph_seq_lens: @@ -742,6 +803,27 @@ def _capture_npu_graphs(self) -> None: # Forward -- re-prefill + inline sampling # ------------------------------------------------------------------ + @staticmethod + def _multinomial( + probs: torch.Tensor, + generator: torch.Generator | None, + generators: Sequence[torch.Generator | None] | None, + ) -> torch.Tensor: + """Sample one code per row, optionally with per-row generators. + + Per-row generators keep explicitly-seeded requests deterministic in a + multi-row batch: each row consumes draws only from its own generator, + so the transformer forward can stay batched (#4883). + """ + if generators is None: + return torch.multinomial(probs, num_samples=1, generator=generator) + return torch.cat( + [ + torch.multinomial(probs[row : row + 1], num_samples=1, generator=row_generator) + for row, row_generator in enumerate(generators) + ] + ) + @torch.inference_mode() def forward( self, @@ -753,9 +835,13 @@ def forward( top_k: int = 50, top_p: float = 1.0, generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Predict residual codebooks 1..G-1 autoregressively via re-prefill.""" bsz = int(layer0_code.shape[0]) + if generators is not None and len(generators) != bsz: + raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") + sample_generator: _GeneratorLike = generators if generators is not None else generator num_groups = self._num_groups device = layer0_code.device @@ -823,7 +909,6 @@ def forward( # Use captured device graph if available, otherwise call compiled fn. device_graph_entry = self._device_graphs.get(graph_key) - # Run transformer (device graph replay or compiled forward) if device_graph_entry is not None: device_graph_entry[0].replay() hidden_out = device_graph_entry[1] @@ -832,9 +917,20 @@ def forward( logits = lm_heads[step - 1](hidden_out[:bsz, step, :]) - # Sample next code + # Sample next code via Gumbel-max. + # + # ``argmax_i(logits_i + Gumbel_i)`` with + # ``Gumbel_i = -log(-log(u_i)), u_i ~ Uniform(0, 1)`` is + # distributionally identical to sampling from ``softmax(logits)``. + # In this file the motivations are practical rather than graph + # related: it is measurably cheaper than ``softmax + multinomial`` + # on the B x 2048 shapes used here, it stays well-defined for + # degenerate masked rows with a surviving finite entry (and is more + # defensive than ``multinomial`` around fully-masked/NaN inputs), + # and the helper below can honor either one batch generator or one + # generator per seeded row. if stored_mode: - # "stored" mode: top-k -> top-p -> softmax -> multinomial + # "stored" mode: top-k -> top-p -> Gumbel-max if s_top_k > 0: topk_vals, _ = logits.topk(s_top_k, dim=-1) logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) @@ -845,17 +941,15 @@ def forward( remove_mask = (cumulative_probs - sorted_probs) >= s_top_p sorted_logits[remove_mask] = float("-inf") logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) - probs = F.softmax(logits, dim=-1, dtype=torch.float32) - code = torch.multinomial(probs, num_samples=1, generator=generator) + code = self._sample_codes_gumbel(logits, generator=sample_generator) else: - # "per_call" mode: temperature-scaled + top-k + # "per_call" mode: temperature-scaled + top-k -> Gumbel-max if use_sampling: scaled = logits * inv_temperature if top_k > 0: topk_vals, _ = scaled.topk(top_k, dim=-1) scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) - probs = F.softmax(scaled, dim=-1, dtype=torch.float32) - code = torch.multinomial(probs, num_samples=1, generator=generator) + code = self._sample_codes_gumbel(scaled, generator=sample_generator) else: code = logits.argmax(dim=-1, keepdim=True) @@ -878,6 +972,18 @@ def forward( # Weight loading # ------------------------------------------------------------------ + def _prepare_npu_weights(self) -> None: + from vllm_ascend.utils import maybe_trans_nz + + linear_count = 0 + with torch.no_grad(): + # Pack linear weights once for NPU matmul. + for module in self.modules(): + if isinstance(module, nn.Linear): + module.weight.data = maybe_trans_nz(module.weight.data) + linear_count += 1 + logger.info("Prepared NPU code predictor weights: linear=%d", linear_count) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights directly (no fused projection remapping needed).""" loaded: set[str] = set() @@ -904,4 +1010,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: weight_loader(param, w) loaded.add(name) + if current_omni_platform.is_npu(): + self._prepare_npu_weights() + return loaded diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 1a387cdcc39..b3e17856526 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -49,6 +49,7 @@ unpad_prompt_conditioning, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -198,7 +199,7 @@ def _ensure_s3_model(cls): ) from e model = _s3.load_model("speech_tokenizer_v3_25hz") - device = "cuda" if torch.cuda.is_available() else "cpu" + device = current_omni_platform.get_torch_device() model = model.to(device).eval() cls._s3_model = (model, _s3, device) return cls._s3_model @@ -479,14 +480,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Keep additional information synchronized for async_chunk updates. self.enable_update_additional_information = True - # Expose streaming parameters - self.token_overlap_len = self.code2wav.token_overlap_len - self.mel_overlap_len = self.code2wav.mel_overlap_len - self.mel_window = self.code2wav.mel_window - self.mel_cache_len = self.code2wav.mel_cache_len - self.source_cache_len = self.code2wav.source_cache_len - self.speech_window = self.code2wav.speech_window - self._stream_audio_cache_by_req: dict[str, torch.Tensor] = {} self._stream_audio_cache_lock = Lock() self._stream_vocoder_cache_by_req: dict[str, dict[str, torch.Tensor]] = {} else: @@ -511,27 +504,6 @@ def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig: # Use parent's cache config - critical for PagedAttention to work correctly return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"]) - @staticmethod - def _cross_fade_audio(audio: torch.Tensor, prev_tail: torch.Tensor) -> torch.Tensor: - """Blend previous chunk tail into current chunk head using a Hamming window. - - This mirrors upstream CosyVoice's `fade_in_out(...)` semantics: - update the current head in-place using a 2*overlap window, then - concatenate the unchanged remainder. - """ - if audio.numel() == 0 or prev_tail.numel() == 0: - return audio - overlap = min(int(audio.numel()), int(prev_tail.numel())) - if overlap <= 0: - return audio - window = torch.hamming_window(2 * overlap, periodic=False, dtype=audio.dtype, device=audio.device) - fade_in = window[:overlap] - fade_out = window[overlap:] - blended = audio[:overlap] * fade_in + prev_tail[-overlap:].to(device=audio.device, dtype=audio.dtype) * fade_out - if overlap == int(audio.numel()): - return blended - return torch.cat([blended, audio[overlap:]], dim=0) - def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor: """Pass-through stitching for async_chunk. @@ -539,11 +511,9 @@ def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_f Applying an additional waveform-domain fade/cache step introduces either duplicated overlap (if no tail trim) or duration shrink (if tail trim). """ - if req_id is not None and stream_finished and hasattr(self, "_stream_audio_cache_by_req"): + if req_id is not None and stream_finished and hasattr(self, "_stream_vocoder_cache_by_req"): with self._stream_audio_cache_lock: - self._stream_audio_cache_by_req.pop(req_id, None) - if hasattr(self, "_stream_vocoder_cache_by_req"): - self._stream_vocoder_cache_by_req.pop(req_id, None) + self._stream_vocoder_cache_by_req.pop(req_id, None) return audio @staticmethod diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index cb3228c13a7..49e8e515b25 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -11,8 +11,6 @@ from __future__ import annotations -from contextlib import nullcontext - import numpy as np import torch import torch.nn as nn @@ -175,19 +173,18 @@ def _forward_mel( prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32) prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32) - with nullcontext(): - feat, _ = self.flow_model.inference( - token=token, - token_len=token_len, - prompt_token=prompt_token, - prompt_token_len=prompt_token_len, - prompt_feat=prompt_feat, - prompt_feat_len=prompt_feat_len, - embedding=embedding, - streaming=streaming, - finalize=finalize, - n_timesteps=n_timesteps, - ) + feat, _ = self.flow_model.inference( + token=token, + token_len=token_len, + prompt_token=prompt_token, + prompt_token_len=prompt_token_len, + prompt_feat=prompt_feat, + prompt_feat_len=prompt_feat_len, + embedding=embedding, + streaming=streaming, + finalize=finalize, + n_timesteps=n_timesteps, + ) trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio) if trim_mel > 0: @@ -195,28 +192,6 @@ def _forward_mel( return feat - @staticmethod - def _fade_speech( - speech: torch.Tensor, - prev_speech: torch.Tensor, - ) -> torch.Tensor: - """Blend previous speech tail into current speech head.""" - if speech.numel() == 0 or prev_speech.numel() == 0: - return speech - overlap = min(int(speech.shape[-1]), int(prev_speech.shape[-1])) - if overlap <= 0: - return speech - window = torch.hamming_window(2 * overlap, periodic=False, dtype=speech.dtype, device=speech.device) - fade_in = window[:overlap].view(1, -1) - fade_out = window[overlap:].view(1, -1) - blended_head = ( - speech[:, :overlap] * fade_in - + prev_speech[:, -overlap:].to(device=speech.device, dtype=speech.dtype) * fade_out - ) - if overlap == int(speech.shape[-1]): - return blended_head - return torch.cat([blended_head, speech[:, overlap:]], dim=-1) - @torch.inference_mode() def forward_streaming( self, @@ -238,17 +213,16 @@ def forward_streaming( suffix. That preserves causal look-right handling without double trimming or duplicated overlap at chunk boundaries. """ - with nullcontext(): - feat = self._forward_mel( - token=token, - prompt_token=prompt_token, - prompt_feat=prompt_feat, - embedding=embedding, - n_timesteps=n_timesteps, - token_offset_tokens=token_offset_tokens, - streaming=True, - finalize=finalize, - ) + feat = self._forward_mel( + token=token, + prompt_token=prompt_token, + prompt_feat=prompt_feat, + embedding=embedding, + n_timesteps=n_timesteps, + token_offset_tokens=token_offset_tokens, + streaming=True, + finalize=finalize, + ) hift_weight = self.hift.m_source.l_linear.weight chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) @@ -268,8 +242,7 @@ def forward_streaming( if tts_mel.shape[-1] == 0: tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype) else: - with nullcontext(): - tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize) + tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize) tts_speech = tts_speech.reshape(tts_speech.shape[0], -1) speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1]))) diff --git a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py index ed35c93bd13..413dc1dc33e 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/pipeline.py +++ b/vllm_omni/model_executor/models/cosyvoice3/pipeline.py @@ -4,8 +4,9 @@ Stage 0: Talker — text prompt → speech tokens (LLM autoregressive). Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform. - * ``sync_process_input_func`` runs when ``deploy.async_chunk=false``: - stage 1 builds full-sequence flow input via ``text2flow``. + * ``sync_process_input_func`` (``text2flow_token_only``) runs when + ``deploy.async_chunk=false``: stage 1 allocates placeholder slots while + bulk tensors arrive via ``text2flow_full_payload`` on the connector. * ``async_chunk_process_next_stage_input_func`` runs when ``deploy.async_chunk=true``: stage 0 streams codec chunks to stage 1 through the shared-memory connector. @@ -21,6 +22,7 @@ COSYVOICE3_PIPELINE = PipelineConfig( model_type="cosyvoice3", + default_deploy_config_name="cosyvoice3.yaml", model_arch="CosyVoice3Model", stages=( StagePipelineConfig( @@ -45,7 +47,6 @@ final_output=True, final_output_type="audio", engine_output_type="latent", - custom_process_input_func=f"{_PROC}.text2flow", sync_process_input_func=f"{_PROC}.text2flow_token_only", ), ), diff --git a/vllm_omni/model_executor/models/cosyvoice3/utils.py b/vllm_omni/model_executor/models/cosyvoice3/utils.py index 5827e6bbbc2..d2c3c686574 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/utils.py +++ b/vllm_omni/model_executor/models/cosyvoice3/utils.py @@ -13,8 +13,6 @@ logger = logging.getLogger(__name__) -IGNORE_ID = -1 - def dynamic_range_compression_torch(x, c=1, clip_val=1e-5): return torch.log(torch.clamp(x, min=clip_val) * c) @@ -182,26 +180,6 @@ def log_mel_spectrogram( return log_spec -def extract_speech_token(prompt_wav, speech_tokenizer_session, device): - speech = load_wav(prompt_wav, 16000) - assert speech.shape[1] / 16000 <= 30, "do not support extract speech token for audio longer than 30s" - feat = log_mel_spectrogram(speech, n_mels=128) - speech_token = ( - speech_tokenizer_session.run( - None, - { - speech_tokenizer_session.get_inputs()[0].name: feat.detach().cpu().numpy(), - speech_tokenizer_session.get_inputs()[1].name: np.array([feat.shape[2]], dtype=np.int32), - }, - )[0] - .flatten() - .tolist() - ) - speech_token = torch.tensor([speech_token], dtype=torch.int32).to(device) - speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32).to(device) - return speech_token, speech_token_len - - def extract_spk_embedding(prompt_wav, campplus_session, device): speech = load_wav(prompt_wav, 16000) feat = kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000) diff --git a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py index 8153c2f385e..b9b06857d07 100644 --- a/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py +++ b/vllm_omni/model_executor/models/covo_audio/config_covo_audio.py @@ -67,6 +67,5 @@ def __init__(self, **kwargs): self.inference = { "s_steps": 10, "cfg_alpha": 1.0, - "dynamic_the": False, "rescale_logits": False, } diff --git a/vllm_omni/model_executor/models/covo_audio/pipeline.py b/vllm_omni/model_executor/models/covo_audio/pipeline.py index 97053e3286f..c2e1911efaf 100644 --- a/vllm_omni/model_executor/models/covo_audio/pipeline.py +++ b/vllm_omni/model_executor/models/covo_audio/pipeline.py @@ -17,6 +17,7 @@ COVO_AUDIO_PIPELINE = PipelineConfig( model_type="covo_audio", + default_deploy_config_name="covo_audio.yaml", model_arch="CovoAudioForConditionalGeneration", stages=( StagePipelineConfig( @@ -44,7 +45,6 @@ final_output=True, final_output_type="audio", engine_output_type="audio", - custom_process_input_func=f"{_PROC}.llm2code2wav", sync_process_input_func=f"{_PROC}.llm2code2wav_token_only", sampling_constraints={"detokenize": False}, ), diff --git a/vllm_omni/model_executor/models/covo_audio/token2wav.py b/vllm_omni/model_executor/models/covo_audio/token2wav.py index 4a6ce4b7e3a..8a96e1f398b 100644 --- a/vllm_omni/model_executor/models/covo_audio/token2wav.py +++ b/vllm_omni/model_executor/models/covo_audio/token2wav.py @@ -598,7 +598,7 @@ def inference( latent_cond = torch.zeros(b, tgt_lens, self.target_dim).to(device) if prefix_target is not None: latent_cond[:, : prefix_target.size(1), :] = prefix_target - sample, trajectory = self.sample( + sample, _ = self.sample( tokens=token, audio=latent_cond, steps=s_steps, alpha=cfg_alpha, g_cond=ge, rescale_logits=rescale_logits ) return sample @@ -612,9 +612,8 @@ def solver(t, z): output = torch.cat([tokens, audio, z], dim=-1) return self.vectorfield_forward(inputs=output, times=t.unsqueeze(0), self_attn_mask=None, g_cond=g_cond) tokens_empty = torch.zeros(*audio.shape[:2], self.model_dim, device=tokens.device, dtype=tokens.dtype) - audio_empty = audio tokens_t = torch.cat([tokens_empty, tokens], dim=0) - audio_t = torch.cat([audio_empty, audio], dim=0) + audio_t = torch.cat([audio, audio], dim=0) audio_noizy_t = torch.cat([z, z], dim=0) t_t = torch.stack([t, t], dim=0) c = g_cond diff --git a/vllm_omni/model_executor/models/dreamzero/__init__.py b/vllm_omni/model_executor/models/dreamzero/__init__.py index 0897383fc8d..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/dreamzero/__init__.py +++ b/vllm_omni/model_executor/models/dreamzero/__init__.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE - -__all__ = ["DREAMZERO_PIPELINE"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dreamzero/pipeline.py b/vllm_omni/model_executor/models/dreamzero/pipeline.py index bc815dedcf7..58ba5b24d1d 100644 --- a/vllm_omni/model_executor/models/dreamzero/pipeline.py +++ b/vllm_omni/model_executor/models/dreamzero/pipeline.py @@ -10,6 +10,7 @@ DREAMZERO_PIPELINE = PipelineConfig( model_type="dreamzero", + default_deploy_config_name="dreamzero.yaml", model_arch="DreamZeroPipeline", stages=( StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/dynin_omni/__init__.py b/vllm_omni/model_executor/models/dynin_omni/__init__.py index 2a3bae8a9f4..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/dynin_omni/__init__.py +++ b/vllm_omni/model_executor/models/dynin_omni/__init__.py @@ -1,59 +1,8 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from .dynin_omni import DyninOmniForConditionalGeneration -from .dynin_omni_common import ( - get_dynin_magvit_attr, - get_dynin_modeling_attr, - get_dynin_sampling_attr, -) - -if TYPE_CHECKING: - from .dynin_omni_token2audio import DyninOmniToken2Audio - from .dynin_omni_token2image import DyninOmniToken2Image - from .dynin_omni_token2text import DyninOmniToken2Text - - -_STAGE_EXPORTS = { - "DyninOmniToken2Audio": (".dynin_omni_token2audio", "DyninOmniToken2Audio"), - "DyninOmniToken2Image": (".dynin_omni_token2image", "DyninOmniToken2Image"), - "DyninOmniToken2Text": (".dynin_omni_token2text", "DyninOmniToken2Text"), -} - -_MODELING_EXPORTS = {"DyninOmniConfig", "DyninOmniModelLM", "VideoTokenMerger"} -_MAGVIT_EXPORTS = {"VQGANEncoder", "VQGANDecoder", "LFQuantizer", "MAGVITv2"} - - -def __getattr__(name: str) -> Any: - if name in _STAGE_EXPORTS: - module_name, attr_name = _STAGE_EXPORTS[name] - module = __import__(module_name, globals(), locals(), [attr_name], 1) - return getattr(module, attr_name) - - if name in _MODELING_EXPORTS: - return get_dynin_modeling_attr(name) - - if name in _MAGVIT_EXPORTS: - return get_dynin_magvit_attr(name) - - if name == "get_mask_schedule": - return get_dynin_sampling_attr("get_mask_schedule") - - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = [ - "DyninOmniForConditionalGeneration", - "DyninOmniToken2Audio", - "DyninOmniToken2Image", - "DyninOmniToken2Text", - "DyninOmniConfig", - "DyninOmniModelLM", - "VideoTokenMerger", - "VQGANEncoder", - "VQGANDecoder", - "LFQuantizer", - "MAGVITv2", - "get_mask_schedule", -] +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/dynin_omni/pipeline.py b/vllm_omni/model_executor/models/dynin_omni/pipeline.py index db0c51e3569..584b6ab5c97 100644 --- a/vllm_omni/model_executor/models/dynin_omni/pipeline.py +++ b/vllm_omni/model_executor/models/dynin_omni/pipeline.py @@ -23,6 +23,7 @@ DYNIN_OMNI_PIPELINE = PipelineConfig( model_type="dynin_omni", + default_deploy_config_name="dynin_omni.yaml", model_arch="DyninOmniForConditionalGeneration", # Arch-fallback safety net: route by hf_config.architectures when the # auto-detected model_type does not match the registry key exactly. diff --git a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py index cdf0da992fc..69ebd29eeb2 100644 --- a/vllm_omni/model_executor/models/fish_speech/dac_encoder.py +++ b/vllm_omni/model_executor/models/fish_speech/dac_encoder.py @@ -17,6 +17,7 @@ DAC_SAMPLE_RATE, build_dac_codec, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -129,7 +130,7 @@ def encode_reference_audio_codes( (dtype=torch.long). """ if device is None: - device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + device = current_omni_platform.get_torch_device() else: device = torch.device(device) dtype = torch.float32 diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py index c6e9064f133..9627dc7ce13 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_dac_decoder.py @@ -348,18 +348,16 @@ def forward( if req_ids.numel() < 1: continue - flat = req_ids - n = flat.numel() - if n == 0 or n % q != 0: - if n > 0: - logger.warning( - "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.", - n, - q, - ) + n = req_ids.numel() + if n % q != 0: + logger.warning( + "DAC decoder input_ids length %d not divisible by num_codebooks %d; returning empty audio.", + n, + q, + ) continue frames = n // q - codes_qf = flat.reshape(q, frames) + codes_qf = req_ids.reshape(q, frames) parsed_ctx_frames[i] = ctx_frames parsed_total_frames[i] = frames valid_codes_qf.append(codes_qf) diff --git a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py index 46ceee49835..eb056f16a10 100644 --- a/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py +++ b/vllm_omni/model_executor/models/fish_speech/fish_speech_fast_ar.py @@ -191,7 +191,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class _FastARDecoderLayer(nn.Module): - """Transformer decoder layer for Fast AR (SDPA, no KV cache).""" + """Transformer decoder layer for Fast AR. + + The batched ``forward`` runs SDPA over the full sequence and is cache-free, + while ``forward_one`` decodes a single step using a per-call KV cache. + """ def __init__(self, config: FishSpeechFastARConfig, *, prefix: str = "") -> None: super().__init__() @@ -519,7 +523,7 @@ def forward( embed_buf[:bsz, 0, :] = projected # Position 1: embedding of semantic code. - code_embed = self.fast_embeddings(semantic_code.clamp(min=0)) + code_embed = self.fast_embeddings(semantic_code) embed_buf[:bsz, 1, :] = code_embed use_sampling = do_sample and temperature > 0 @@ -542,8 +546,7 @@ def forward( logits = self.fast_output(self.fast_norm(hidden_out)) # Residual codebooks (step >= 1) only have 1024 entries. - if step >= 1: - logits = logits[:, :residual_codebook_size] + logits = logits[:, :residual_codebook_size] if use_sampling: scaled = logits * inv_temperature diff --git a/vllm_omni/model_executor/models/fish_speech/pipeline.py b/vllm_omni/model_executor/models/fish_speech/pipeline.py index baee1b5e72f..1cb113094c2 100644 --- a/vllm_omni/model_executor/models/fish_speech/pipeline.py +++ b/vllm_omni/model_executor/models/fish_speech/pipeline.py @@ -20,6 +20,7 @@ FISH_SPEECH_PIPELINE = PipelineConfig( model_type="fish_qwen3_omni", + default_deploy_config_name="fish_qwen3_omni.yaml", model_arch="FishSpeechSlowARForConditionalGeneration", stages=( StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py index 8b8d8559ead..3918a6f9462 100644 --- a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py +++ b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py @@ -132,13 +132,3 @@ def estimate_fish_voice_clone_prompt_len_from_normalized( [0] * semantic_len, ) return len(prompt_ids) - - -def estimate_fish_voice_clone_prompt_len(tokenizer: Any, text: str, ref_text: str, semantic_len: int) -> int: - normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(text, ref_text) - return estimate_fish_voice_clone_prompt_len_from_normalized( - tokenizer, - normalized_text, - normalized_ref_text, - semantic_len, - ) diff --git a/vllm_omni/model_executor/models/glm_image/__init__.py b/vllm_omni/model_executor/models/glm_image/__init__.py index 6153f07c0a0..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/glm_image/__init__.py +++ b/vllm_omni/model_executor/models/glm_image/__init__.py @@ -1,16 +1,8 @@ -def __getattr__(name: str): - """Lazy import to avoid importing transformers.models.glm_image at module init. - - The AR model depends on ``transformers.models.glm_image`` which is only - available when the model weights are present. Importing it eagerly in - ``__init__.py`` breaks the pipeline registry lookup for environments that - don't have this custom transformers extension installed. - """ - if name == "GlmImageForConditionalGeneration": - from .glm_image_ar import GlmImageForConditionalGeneration - - return GlmImageForConditionalGeneration - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - -__all__ = ["GlmImageForConditionalGeneration"] +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/glm_image/pipeline.py b/vllm_omni/model_executor/models/glm_image/pipeline.py index bd041292c6d..fd74bfad8eb 100644 --- a/vllm_omni/model_executor/models/glm_image/pipeline.py +++ b/vllm_omni/model_executor/models/glm_image/pipeline.py @@ -14,6 +14,7 @@ GLM_IMAGE_PIPELINE = PipelineConfig( model_type="glm_image", + default_deploy_config_name="glm_image.yaml", model_arch="GlmImageForConditionalGeneration", hf_architectures=("GlmImageForConditionalGeneration",), diffusers_class_name="GlmImagePipeline", diff --git a/vllm_omni/model_executor/models/glm_tts/__init__.py b/vllm_omni/model_executor/models/glm_tts/__init__.py index 9618227ca43..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/glm_tts/__init__.py +++ b/vllm_omni/model_executor/models/glm_tts/__init__.py @@ -1,18 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GLM-TTS model support for vLLM-Omni. -GLM-TTS is a two-stage text-to-speech system: - - Stage 0 (LLM): Llama-based AR model generates speech tokens from text - - Stage 1 (DiT): Flow matching model converts speech tokens to mel-spectrogram - -Reference: https://github.com/zai-org/GLM-TTS -""" - -from vllm_omni.model_executor.models.glm_tts.glm_tts import ( - GLMTTSForConditionalGeneration, -) - -__all__ = [ - "GLMTTSForConditionalGeneration", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts.py b/vllm_omni/model_executor/models/glm_tts/glm_tts.py index 6e14840e24b..664a794c4c0 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts.py @@ -53,6 +53,7 @@ from vllm_omni.model_executor.models.common.nucleus_ras_sampling import ras_sample_one as _ras_sample_one from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.glm_tts import GLMTTSConfig from vllm_omni.utils.speaker_cache import get_speaker_cache @@ -411,7 +412,7 @@ def _ensure_cached_runtime_components( # Load voice clone components (speech tokenizer + CampPlus) if load_voice_clone: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() self.processor_device = device campplus_path = getattr(self, "_campplus_path", None) or resolve_glm_tts_campplus_path(model_dir) self._campplus_path = campplus_path diff --git a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py index 5cda6ec45f5..666998e9dda 100644 --- a/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py +++ b/vllm_omni/model_executor/models/glm_tts/glm_tts_dit.py @@ -472,10 +472,6 @@ def __init__( self.num_heads = heads self.spkr_emb_adaln = spkr_emb_adaln - # Frame rates for block-causal mask token→mel conversion - self._input_frame_rate: float = 25.0 - self._mel_framerate: float = 50.0 - # Speaker embedding dimension for AdaLN if spkr_emb_adaln: additional_dim = spkr_dim + 256 if use_wavlm_emb else spkr_dim diff --git a/vllm_omni/model_executor/models/glm_tts/pipeline.py b/vllm_omni/model_executor/models/glm_tts/pipeline.py index f016a993e3d..2a1fedd3ff9 100644 --- a/vllm_omni/model_executor/models/glm_tts/pipeline.py +++ b/vllm_omni/model_executor/models/glm_tts/pipeline.py @@ -12,6 +12,7 @@ GLM_TTS_PIPELINE = PipelineConfig( model_type="glm_tts", + default_deploy_config_name="glm_tts.yaml", model_arch="GLMTTSForConditionalGeneration", stages=( StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/gr00t/__init__.py b/vllm_omni/model_executor/models/gr00t/__init__.py new file mode 100644 index 00000000000..fb6738a2967 --- /dev/null +++ b/vllm_omni/model_executor/models/gr00t/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/gr00t/pipeline.py b/vllm_omni/model_executor/models/gr00t/pipeline.py new file mode 100644 index 00000000000..a2e4c8133e4 --- /dev/null +++ b/vllm_omni/model_executor/models/gr00t/pipeline.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GR00T N1.7 single-stage policy topology.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +GR00T_N1D7_PIPELINE = PipelineConfig( + model_type="Gr00tN1d7", + default_deploy_config_name="Gr00tN1d7.yaml", + model_arch="Gr00tN1d7Pipeline", + hf_architectures=("Gr00tN1d7",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="actions", + model_arch="Gr00tN1d7Pipeline", + ), + ), +) diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py index 1a7c23ebad3..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/__init__.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""vllm-omni integration for boson-ai's higgs-audio v2 (two-stage TTS).""" -from vllm_omni.model_executor.models.higgs_audio_v2.configuration_higgs_audio_v2 import ( - HiggsAudioV2Config, -) - -__all__ = ["HiggsAudioV2Config"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py index dc5eb12b475..b6eef8dc4a5 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/configuration_higgs_audio_v2.py @@ -155,15 +155,3 @@ def __init__( @property def num_real_codes(self) -> int: return self.audio_stream_bos_id # i.e. 1024 for the boson-ai checkpoint - - # The full vocab size of each per-codebook output head. - @property - def codebook_output_size(self) -> int: - return self.codebook_size # i.e. 1026 for the boson-ai checkpoint - - # Canonical MusicGen-style delay pattern used by - # `HiggsAudioV2DelayPatternLogitsProcessor`: codebook k starts emitting - # real codes only after k frames. Kept as a function so subclasses can - # override if upstream ever publishes a different convention. - def default_delay_pattern(self) -> list[int]: - return list(range(self.num_codebooks)) diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py index 6a14707ee5b..91609e1a4ee 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_code2wav.py @@ -41,6 +41,7 @@ load_higgs_audio_codec, ) from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform __all__ = [ "HiggsAudioV2Code2Wav", @@ -173,7 +174,7 @@ def load_codec_from_disk(self, model_dir: str, device: torch.device | None = Non 2. Otherwise fall back to ``<model_dir>/<audio_tokenizer_subdir>``. """ if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() audio_tokenizer_path = self._resolve_audio_tokenizer_path(model_dir) quantizer, fc2, acoustic_decoder, _tokenizer_config = load_higgs_audio_codec(audio_tokenizer_path, device) if len(quantizer.quantizers) != self.num_codebooks: diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py index 5e3883fd2ec..be08163e4d2 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_talker.py @@ -821,12 +821,12 @@ def forward( except Exception: self._last_step_query_start_loc = None - # Voice-clone ref-audio substitution: at the prompt-side audio - # placeholder span we overwrite the text-embed lookups with the - # multi-codebook embedding of the encoded reference clip carried in - # ``model_intermediate_buffer``. No-op for plain-text requests and - # for decode steps (which fall through to the audio-feedback path). if input_ids is not None and inputs_embeds is None: + # Voice-clone ref-audio substitution: at the prompt-side audio + # placeholder span we overwrite the text-embed lookups with the + # multi-codebook embedding of the encoded reference clip carried in + # ``model_intermediate_buffer``. No-op for plain-text requests and + # for decode steps (which fall through to the audio-feedback path). info_dicts = kwargs.get("model_intermediate_buffer") if info_dicts is None: info_dicts = kwargs.get("runtime_additional_information") @@ -836,15 +836,14 @@ def forward( info_dicts, ) - # Audio-feedback substitution: vLLM's AR runner calls forward() - # with ``inputs_embeds=None`` and lets the model embed its own - # input_ids inline. ``embed_input_ids`` is only invoked via the - # supports_mm_inputs/prompt_embeds branches, neither of which fires - # for the talker. So we apply the substitution directly here: at - # decode positions where input_ids == audio_token_id, swap the text - # embedding with embed_audio_codes(last audio frame stashed in - # ``self._audio_state[req_id]`` by :meth:`sample`). - if input_ids is not None and inputs_embeds is None: + # Audio-feedback substitution: vLLM's AR runner calls forward() + # with ``inputs_embeds=None`` and lets the model embed its own + # input_ids inline. ``embed_input_ids`` is only invoked via the + # supports_mm_inputs/prompt_embeds branches, neither of which fires + # for the talker. So we apply the substitution directly here: at + # decode positions where input_ids == audio_token_id, swap the text + # embedding with embed_audio_codes(last audio frame stashed in + # ``self._audio_state[req_id]`` by :meth:`sample`). hidden_states = self._maybe_apply_audio_feedback(hidden_states, input_ids) audio_mask = self.audio_token_mask(input_ids) if input_ids is not None else None @@ -887,7 +886,7 @@ def compute_logits( # single step at a time; cleared in :meth:`sample` after consumption # so a missed call doesn't leak stale tensors. self._last_logits_hidden = hidden_states - return self.logits_processor(self.lm_head, hidden_states, sampling_metadata) + return self.logits_processor(self.lm_head, hidden_states) def audio_codebook_logits( self, @@ -1506,7 +1505,7 @@ def _sample_audio_codes_upstream( sampled[~valid] = torch.argmax(logits[~valid].float(), dim=-1) return sampled - def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> int: + def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None: """Mask non-audio tokens at audio-mode positions, in-place on ``logits``. Heuristic: per-request, find the last token seen so far (last of @@ -1524,7 +1523,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) - the stock sampler. """ if logits is None or logits.ndim != 2: - return 0 + return audio_bos = int(self.config.audio_bos_token_id) audio_id = int(self.config.audio_token_id) # (PyTorch align, corrected): upstream NEVER lets the model @@ -1542,7 +1541,6 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) - prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) output_ids = getattr(sampling_metadata, "output_token_ids", None) num_rows = int(logits.shape[0]) - biased = 0 # Fallback "previous token" source: read each request's LAST input # position from ``self._last_step_input_ids`` using ``query_start_loc`` # from the forward context. The previous version of this code took @@ -1619,7 +1617,6 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) - mask[audio_eos_vocab] = row[audio_eos_vocab] logits[i].copy_(mask) slot_state["should_terminate"] = False - biased += 1 continue # Always force ``audio_token_id`` — upstream mirror. allowed: set[int] = {audio_id, *allowed_extra} @@ -1629,9 +1626,7 @@ def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) - if 0 <= tok < row.shape[-1]: mask[tok] = row[tok] logits[i].copy_(mask) - biased += 1 self._last_first_audio_after_bos = first_after_bos - return biased has_postprocess: bool = True 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 old mode 100644 new mode 100755 index cacf82ee56c..c3f5024b3bc --- 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 @@ -22,6 +22,7 @@ from __future__ import annotations +import json import os import re from typing import Any @@ -29,6 +30,8 @@ import numpy as np import torch +from vllm_omni.platforms import current_omni_platform + __all__ = [ "UnsupportedInputError", "MULTI_SPEAKER_TAG_PATTERN", @@ -168,14 +171,25 @@ def _build_delay_pattern(codes: torch.Tensor) -> torch.Tensor: ) +def _is_higgs_audio_tokenizer_config(config_path: str) -> bool: + try: + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + except (OSError, json.JSONDecodeError): + return False + return config.get("model_type") == "higgs_audio_v2_tokenizer" + + 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")): + config_path = os.path.join(expanded, "config.json") + if os.path.isfile(config_path) and _is_higgs_audio_tokenizer_config(config_path): return expanded nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR) - if os.path.isfile(os.path.join(nested, "config.json")): + nested_config_path = os.path.join(nested, "config.json") + if os.path.isfile(nested_config_path) and _is_higgs_audio_tokenizer_config(nested_config_path): return nested return None @@ -199,7 +213,9 @@ def _resolve_audio_tokenizer_dir() -> str | None: filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json", ) if isinstance(cached_config, str) and os.path.isfile(cached_config): - return os.path.dirname(cached_config) + candidate = _normalize_audio_tokenizer_dir(os.path.dirname(cached_config)) + if candidate is not None: + return candidate from huggingface_hub.constants import HF_HUB_CACHE @@ -238,8 +254,10 @@ def _load_audio_tokenizer(): _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" + audio_tokenizer_dir = _normalize_audio_tokenizer_dir(repo_path) + if audio_tokenizer_dir is None: + raise RuntimeError(f"Downloaded {_K2_OMNIVOICE_REPO} does not contain a valid Higgs audio tokenizer") + device = current_omni_platform.get_torch_device() model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device) return model.eval() diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py b/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py index fb683256992..1d3db52d666 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/pipeline.py @@ -12,6 +12,7 @@ HIGGS_AUDIO_V2_PIPELINE = PipelineConfig( model_type="higgs_audio_v2", + default_deploy_config_name="higgs_audio_v2.yaml", model_arch="HiggsAudioV2ForConditionalGeneration", hf_architectures=("HiggsAudioV2ForConditionalGeneration",), stages=( diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py b/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py index 052d51eef06..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/__init__.py @@ -1,9 +1,8 @@ # 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"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. 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 index 40936f85d68..45b5d68be00 100644 --- 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 @@ -23,10 +23,11 @@ build_higgs_audio_acoustic_decoder, load_higgs_audio_codec, ) -from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform +from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) -from vllm_omni.model_executor.models.output_templates import OmniOutput __all__ = [ "HiggsAudioV3Code2Wav", @@ -126,7 +127,7 @@ def _load_from_tokenizer_repo(self) -> None: else: raise FileNotFoundError(f"No cached snapshot for {tokenizer_id}") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() quantizer, fc2, acoustic_decoder, _cfg = load_higgs_audio_codec(tokenizer_dir, device) self.quantizer = quantizer self.fc2 = fc2 @@ -134,6 +135,81 @@ def _load_from_tokenizer_repo(self) -> None: self._loaded = True logger.info("Loaded HiggsAudioV3Code2Wav from standalone tokenizer repo.") + def _resolve_model_dir(self) -> str | None: + """Resolve ``self._model_path`` to a local checkpoint directory. + + Accepts an already-local path or an HF repo id (resolved via the + local cache). Returns ``None`` when nothing is reachable on disk. + """ + path = self._model_path + if not path: + return None + if os.path.isdir(path): + return path + from huggingface_hub import try_to_load_from_cache + + cached = try_to_load_from_cache(repo_id=path, filename="config.json") + if isinstance(cached, str) and os.path.isfile(cached): + return os.path.dirname(cached) + return None + + def _read_bundled_codec_state(self) -> dict[str, torch.Tensor]: + """Read bundled codec tensors directly from the on-disk checkpoint. + + Returns the ``_CODEC_PREFIX``-stripped tensors, or an empty dict when + the checkpoint cannot be located (so callers can fall back). + """ + model_dir = self._resolve_model_dir() + if model_dir is None: + return {} + import glob + + from safetensors import safe_open + + codec_state: dict[str, torch.Tensor] = {} + for st in sorted(glob.glob(os.path.join(model_dir, "*.safetensors"))): + with safe_open(st, framework="pt") as f: + for key in f.keys(): + if key.startswith(_CODEC_PREFIX): + codec_state[key[len(_CODEC_PREFIX) :]] = f.get_tensor(key) + return codec_state + + def _ensure_codec_loaded(self) -> None: + """Build the codec when vLLM's weight loader skipped ``load_weights``. + + ``load_format=dummy`` (and any loader that bypasses + ``model.load_weights``) never hands us the checkpoint weight iterator, + so the lazily-built codec submodules stay ``None`` and the warmup + ``_dummy_run`` crashes in :meth:`decode_codes`. Mirror fish_speech's + ``_ensure_codec_loaded`` and load the bundled V3 codec directly from + the on-disk checkpoint, falling back to the standalone tokenizer repo. + A no-op once ``load_weights`` has already populated the codec. + """ + if self._loaded: + return + + device: torch.device | None = None + if self.vllm_config is not None: + try: + device = self.vllm_config.device_config.device + except AttributeError: + device = None + + codec_state = self._read_bundled_codec_state() + if codec_state: + self._load_from_bundled_state(codec_state, device) + logger.info( + "Lazily loaded HiggsAudioV3Code2Wav from bundled V3 checkpoint " + "(%d codec keys) because the weight loader skipped load_weights " + "(e.g. load_format=dummy).", + len(codec_state), + ) + return + + # No bundled checkpoint reachable on disk — fall back to the standalone + # tokenizer repo (raises a clear error if that is missing too). + self._load_from_tokenizer_repo() + # ------------------------------------------------------------------ engine hooks def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: if input_ids.numel() == 0: @@ -203,7 +279,7 @@ def _load_from_bundled_state( ) if device is None: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() # The bundled codec may use boson-ai key naming. Try remapping. needs_remap = any(k.startswith("quantizer.vq.layers.") for k in codec_state) @@ -371,7 +447,9 @@ def _load_from_bundled_state( 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.") + # load_weights may have been skipped (e.g. load_format=dummy); + # build the codec lazily from the on-disk checkpoint. + self._ensure_codec_loaded() codes = self._validate_codes(audio_codes) rvq_codes = codes.transpose(0, 1).long() 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 index 9d4e0535601..f51adafa9cc 100644 --- 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 @@ -36,10 +36,10 @@ 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 ( +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.higgs_audio_v3 import ( HiggsAudioV3Config, ) -from vllm_omni.model_executor.models.output_templates import OmniOutput __all__ = ["HiggsAudioV3TalkerForConditionalGeneration"] @@ -531,7 +531,7 @@ def _is_decode_only_graph_batch(self, batch: int) -> bool: 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) + return self.logits_processor(self.lm_head, hidden_states) 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) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py b/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py index 2f35c497e71..9bfc411dfec 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py @@ -24,6 +24,7 @@ HIGGS_AUDIO_V3_PIPELINE = PipelineConfig( model_type="higgs_multimodal_qwen3", + default_deploy_config_name="higgs_multimodal_qwen3.yaml", model_arch="HiggsAudioV3TalkerForConditionalGeneration", hf_architectures=("HiggsMultimodalQwen3ForConditionalGeneration",), stages=( 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 deleted file mode 100644 index 2a21aafcd39..00000000000 --- a/vllm_omni/model_executor/models/higgs_audio_v3/ref_audio_cache.py +++ /dev/null @@ -1,111 +0,0 @@ -# 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/__init__.py b/vllm_omni/model_executor/models/hunyuan_image3/__init__.py index a3618aefc38..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/__init__.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/__init__.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .hunyuan_image3 import HunyuanImage3ForConditionalGeneration - -__all__ = ["HunyuanImage3ForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 8c983cff195..17bbad0014e 100755 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -68,7 +68,6 @@ maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.image import rgba_to_rgb from vllm.multimodal.inputs import ( MultiModalFeatureSpec, MultiModalFieldConfig, @@ -223,14 +222,8 @@ def contains_unexpected_keyword(name, keywords): if self.config.tie_word_embeddings and "lm_head.weight" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - # Loading kv cache scales for compressed-tensors quantization - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight[0] - weight_loader(param, loaded_weight) - continue - + # KV-cache scales are renamed to .attn.{k,v,q}_scale by the + # quant_config mapper in the outer AutoWeightsLoader. is_found = False for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: @@ -889,8 +882,8 @@ def process_image(self, image_input: ImageInput): for image in images: current_info = {} - if self.hf_config.vit["num_channels"] == 3 and image.mode == "RGBA": - image = rgba_to_rgb(image, (255, 255, 255)) + if self.hf_config.vit["num_channels"] == 3 and image.mode != "RGB": + image = image.convert("RGB") # VIT processing vit_pixel_values = self.vision_encoder_processor(image) @@ -2009,10 +2002,9 @@ def embed_input_ids( # the bidirectional MM region and biased AR ratio prediction toward # the first image's bucket. timestep_mask = input_ids == self._timestep_token_id - n_timestep = int(timestep_mask.sum().item()) - if n_timestep > 0: - timestep_input = torch.zeros((n_timestep,), device=inputs_embeds.device, dtype=inputs_embeds.dtype) - inputs_embeds[timestep_mask] = self._timestep_encode(timestep_input) + timestep_input = torch.zeros((), device=inputs_embeds.device, dtype=inputs_embeds.dtype) + timestep_embed = self._timestep_encode(timestep_input).to(inputs_embeds.dtype) + inputs_embeds = torch.where(timestep_mask.unsqueeze(-1), timestep_embed, inputs_embeds) if multimodal_embeddings is None or len(multimodal_embeddings) == 0: return inputs_embeds diff --git a/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py b/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py index 06ebff9c71c..fd53043bdb1 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py @@ -18,6 +18,7 @@ HUNYUAN_IMAGE3_PIPELINE = PipelineConfig( model_type="hunyuan_image_3_moe", + default_deploy_config_name="hunyuan_image_3_moe.yaml", model_arch=_HUNYUAN_IMAGE3_MODEL_ARCH, hf_architectures=_HUNYUAN_IMAGE3_HF_ARCHS, stages=( @@ -50,6 +51,7 @@ HUNYUAN_IMAGE3_AR_PIPELINE = PipelineConfig( model_type="hunyuan_image3_ar", + default_deploy_config_name="hunyuan_image3_ar.yaml", model_arch=_HUNYUAN_IMAGE3_MODEL_ARCH, hf_architectures=(), stages=( @@ -71,6 +73,7 @@ HUNYUAN_IMAGE3_DIT_PIPELINE = PipelineConfig( model_type="hunyuan_image3_dit", + default_deploy_config_name="hunyuan_image3_dit.yaml", model_arch=_HUNYUAN_IMAGE3_MODEL_ARCH, hf_architectures=(), stages=( diff --git a/vllm_omni/model_executor/models/hunyuan_video/__init__.py b/vllm_omni/model_executor/models/hunyuan_video/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/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/vllm_omni/model_executor/models/hunyuan_video/pipeline.py b/vllm_omni/model_executor/models/hunyuan_video/pipeline.py new file mode 100644 index 00000000000..341d14e5dca --- /dev/null +++ b/vllm_omni/model_executor/models/hunyuan_video/pipeline.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""HunyuanVideo-1.5 pipeline topology.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +HUNYUAN_VIDEO_15_PIPELINE = PipelineConfig( + model_type="hunyuan_video_15", + model_arch="HunyuanVideo15Pipeline", + diffusers_class_name="HunyuanVideo15Pipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="video", + model_arch="HunyuanVideo15Pipeline", + ), + ), +) diff --git a/vllm_omni/model_executor/models/indextts2/__init__.py b/vllm_omni/model_executor/models/indextts2/__init__.py new file mode 100644 index 00000000000..fb6738a2967 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/__init__.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/indextts2/bigvgan_cuda_graph.py b/vllm_omni/model_executor/models/indextts2/bigvgan_cuda_graph.py new file mode 100644 index 00000000000..c64a46404a1 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/bigvgan_cuda_graph.py @@ -0,0 +1,284 @@ +"""CUDA graph wrapper for BigVGAN vocoding in IndexTTS2 Stage 1. + +BigVGAN is heavily CPU launch-bound (Snake/AliasFree small-op storm: +~3500 copy_ + ~2000 mul + ~340 cudnn_convolution launches per request, +GPU busy <15%). Capturing the forward for fixed mel-frame buckets and +replaying collapses thousands of kernel launches into one graph launch. + +Dynamic mel inputs are right zero-padded to the nearest captured bucket +and the waveform is trimmed back to ``actual_frames * upsample_factor`` +after replay. Padding leakage into the kept region is bounded by the +conv receptive field at the tail (~20-38 mel frames, ~0.2-0.4s audio) +and is validated via WER/SIM A/B. + +The compile layer (``compile_shapes``) additionally applies +``torch.compile`` to the BigVGAN forward before CUDA-graph capture, +letting Inductor fuse the Snake/alias-free small-op chain. The +three-tier dispatch at runtime is: compiled-graph > plain-graph > eager. + +Modeled on ``CUDAGraphDecoderWrapper`` +(vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py) +and ``CUDAGraphGLMTTSDiTWrapper`` +(vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py). +""" + +from __future__ import annotations + +import math +from collections.abc import Callable + +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +# Mel-frame buckets at 22050 Hz / hop 256 (~86 fps). +# Keep buckets modest so that CUDA-graph private pools do not exhaust +# GPU memory when stage0 and stage1 share the same device (~22 GiB). +# Inputs longer than the largest bucket fall back to eager BigVGAN. +DEFAULT_CAPTURE_SIZES = [128, 256, 384, 512] + + +class CUDAGraphBigVGANWrapper: + """Capture BigVGAN vocoding for fixed mel-frame buckets. + + Supports an optional torch.compile layer on top of plain CUDA graphs. + When ``compile_shapes`` is provided, those buckets are compiled with + ``torch.compile(mode="default", fullgraph=False, dynamic=False)`` and + then captured into CUDA graphs — giving both Inductor operator fusion + and zero-launch-overhead replay. + """ + + def __init__( + self, + model: torch.nn.Module, + *, + enabled: bool = True, + capture_sizes: list[int] | None = None, + compile_shapes: list[int] | None = None, + ) -> None: + self.model = model + self.enabled = bool(enabled) + self.capture_sizes = sorted(set(capture_sizes or DEFAULT_CAPTURE_SIZES)) + self.compile_shapes = sorted(set(compile_shapes or [])) + self.num_mels = int(model.h.num_mels) + self.upsample_factor = int(math.prod(model.h.upsample_rates)) + + # Plain CUDA graph tier + self.graphs: dict[int, torch.cuda.CUDAGraph] = {} + self.static_mel: dict[int, torch.Tensor] = {} + self.static_wav: dict[int, torch.Tensor] = {} + + # Compiled CUDA graph tier + self._compiled_decoder: Callable[[torch.Tensor], torch.Tensor] | None = None + self._compiled_graphs: dict[int, torch.cuda.CUDAGraph] = {} + self._compiled_static_mel: dict[int, torch.Tensor] = {} + self._compiled_static_wav: dict[int, torch.Tensor] = {} + self._compiled_shapes: set[int] = set() + + self._warmed_up = False + self.last_call_info: dict[str, int | str | None] = {} + + def warmup(self, device: torch.device, dtype: torch.dtype) -> None: + if not self.enabled or self._warmed_up or device.type != "cuda" or torch.cuda.is_current_stream_capturing(): + return + self.model.eval() + + # Precompute Snake/SnakeBeta caches so torch.compile sees stable + # buffers instead of lazy init inside the traced graph. + self._precompute_snake_caches() + + logger.info("Starting BigVGAN CUDA graph warmup for mel-frame buckets: %s", self.capture_sizes) + for size in self.capture_sizes: + try: + self._capture(size, device, dtype) + logger.info("Captured BigVGAN CUDA graph for mel frames=%d", size) + except Exception: + logger.warning("Failed to capture BigVGAN CUDA graph for size=%d", size, exc_info=True) + + if self.compile_shapes: + self._warmup_compile_shapes(device, dtype) + + self._warmed_up = True + + def _precompute_snake_caches(self) -> None: + from vllm_omni.model_executor.models.common.snake_activation import Snake, SnakeBeta + + for module in self.model.modules(): + if isinstance(module, (Snake, SnakeBeta)): + module.precompute_exp_cache() + + # ------------------------------------------------------------------ + # Plain CUDA graph capture + # ------------------------------------------------------------------ + + def _capture(self, size: int, device: torch.device, dtype: torch.dtype) -> None: + mel = torch.zeros(1, self.num_mels, size, device=device, dtype=dtype) + self.static_mel[size] = mel + # Eager warm run so cuDNN autotuning happens outside capture. + with torch.no_grad(): + _ = self.model(mel) + torch.accelerator.synchronize(device) + + graph = torch.cuda.CUDAGraph() + with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + self.static_wav[size] = self.model(mel) + self.graphs[size] = graph + + # ------------------------------------------------------------------ + # Compiled CUDA graph capture (torch.compile + capture) + # ------------------------------------------------------------------ + + def _warmup_compile_shapes(self, device: torch.device, dtype: torch.dtype) -> None: + logger.info( + "Starting torch.compile + CUDA Graph warmup for BigVGAN shapes: %s", + self.compile_shapes, + ) + try: + self._compiled_decoder = torch.compile( + self.model.forward, + mode="default", + fullgraph=False, + dynamic=False, + ) + except Exception: + logger.warning("Failed to create torch.compile BigVGAN wrapper", exc_info=True) + self._compiled_decoder = None + return + + for size in self.compile_shapes: + try: + self._capture_compiled(size, device, dtype) + self._compiled_shapes.add(size) + logger.info( + " torch.compile + CUDA Graph ready for mel frames=%d", + size, + ) + except Exception: + logger.warning( + " Failed to capture torch.compile CUDA Graph for size=%d; " + "this shape will use plain CUDA graph/eager fallback", + size, + exc_info=True, + ) + logger.info( + "torch.compile + CUDA Graph warmup complete: %d/%d shapes ready", + len(self._compiled_shapes), + len(self.compile_shapes), + ) + + def _capture_compiled(self, size: int, device: torch.device, dtype: torch.dtype) -> None: + if self._compiled_decoder is None: + raise RuntimeError("Compiled decoder is not initialized") + + mel = torch.zeros(1, self.num_mels, size, device=device, dtype=dtype) + # Warmup iterations let Inductor finish tracing/autotuning before capture. + with torch.inference_mode(): + for _ in range(5): + _ = self._compiled_decoder(mel) + torch.accelerator.synchronize(device) + + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_wav = self._compiled_decoder(mel) + + self._compiled_static_mel[size] = mel + self._compiled_static_wav[size] = static_wav + self._compiled_graphs[size] = graph + + # ------------------------------------------------------------------ + # Dispatch + # ------------------------------------------------------------------ + + def _get_padded_size(self, actual_frames: int) -> int | None: + for size in self.capture_sizes: + if actual_frames <= size: + return size + return None + + def __call__(self, mel: torch.Tensor) -> torch.Tensor: + """Vocode ``mel`` [1, num_mels, T]; falls back to eager when ineligible.""" + actual_frames = int(mel.shape[-1]) if mel.ndim >= 3 else 0 + + # Eligibility guard + if ( + not self.enabled + or not self._warmed_up + or mel.shape[0] != 1 + or mel.device.type != "cuda" + or torch.cuda.is_current_stream_capturing() + ): + self.last_call_info = { + "mode": "eager", + "reason": "ineligible", + "batch": int(mel.shape[0]) if mel.ndim > 0 else 0, + "frames": actual_frames, + "bucket": None, + "padding": None, + } + return self.model(mel) + + # Tier 1: compiled CUDA graph (exact shape match) + if actual_frames in self._compiled_shapes: + return self._replay_compiled(actual_frames, actual_frames, mel, "exact_hit") + + # Tier 1b: compiled CUDA graph (padded shape) + padded = self._get_padded_size(actual_frames) + if padded is not None and padded in self._compiled_shapes: + return self._replay_compiled(actual_frames, padded, mel, "padded_hit") + + # Tier 2: plain CUDA graph + size = padded + if size is None or size not in self.graphs: + self.last_call_info = { + "mode": "eager", + "reason": "no_bucket", + "batch": int(mel.shape[0]), + "frames": actual_frames, + "bucket": size, + "padding": None if size is None else int(size - actual_frames), + } + return self.model(mel) + + buf = self.static_mel[size] + if actual_frames == size: + buf.copy_(mel.to(buf.dtype)) + else: + buf.zero_() + buf[..., :actual_frames].copy_(mel.to(buf.dtype)) + self.graphs[size].replay() + self.last_call_info = { + "mode": "graph", + "reason": "hit", + "batch": int(mel.shape[0]), + "frames": actual_frames, + "bucket": int(size), + "padding": int(size - actual_frames), + } + return self.static_wav[size][..., : actual_frames * self.upsample_factor].clone() + + def _replay_compiled( + self, + actual_frames: int, + bucket: int, + mel: torch.Tensor, + reason: str, + ) -> torch.Tensor: + buf = self._compiled_static_mel[bucket] + if actual_frames == bucket: + buf.copy_(mel.to(buf.dtype)) + else: + buf.zero_() + buf[..., :actual_frames].copy_(mel.to(buf.dtype)) + self._compiled_graphs[bucket].replay() + self.last_call_info = { + "mode": "compile_graph", + "reason": reason, + "batch": int(mel.shape[0]), + "frames": actual_frames, + "bucket": int(bucket), + "padding": int(bucket - actual_frames), + } + return self._compiled_static_wav[bucket][..., : actual_frames * self.upsample_factor].clone() diff --git a/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py b/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py new file mode 100644 index 00000000000..504923b1a1a --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/configuration_indextts2.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import copy + +from transformers import PretrainedConfig + + +class IndexTTS2Config(PretrainedConfig): + model_type = "indextts2" + + # Defaults for the official IndexTTS2 checkpoint. Values from HF + # config.json (passed via ``kwargs``) take precedence over these. + _DEFAULTS = { + "num_attention_heads": 20, + "hidden_size": 1280, + "dataset": { + "bpe_model": "bpe.model", + "sample_rate": 24000, + "squeeze": False, + "mel": { + "sample_rate": 24000, + "n_fft": 1024, + "hop_length": 256, + "win_length": 1024, + "n_mels": 100, + "mel_fmin": 0, + "normalize": False, + }, + }, + "gpt": { + "model_dim": 1280, + "max_mel_tokens": 1815, + "max_text_tokens": 600, + "heads": 20, + "use_mel_codes_as_input": True, + "mel_length_compression": 1024, + "layers": 24, + "number_text_tokens": 12000, + "number_mel_codes": 8194, + "start_mel_token": 8192, + "stop_mel_token": 8193, + "start_text_token": 0, + "stop_text_token": 1, + "train_solo_embeddings": False, + "condition_num_latent": 32, + "condition_type": "conformer_perceiver", + "condition_module": { + "output_size": 512, + "linear_units": 2048, + "attention_heads": 8, + "num_blocks": 6, + "input_layer": "conv2d2", + "perceiver_mult": 2, + }, + "emo_condition_module": { + "output_size": 512, + "linear_units": 1024, + "attention_heads": 4, + "num_blocks": 4, + "input_layer": "conv2d2", + "perceiver_mult": 2, + }, + }, + "semantic_codec": { + "codebook_size": 8192, + "hidden_size": 1024, + "codebook_dim": 8, + "vocos_dim": 384, + "vocos_intermediate_dim": 2048, + "vocos_num_layers": 12, + }, + "s2mel": { + "preprocess_params": { + "sr": 22050, + "spect_params": { + "n_fft": 1024, + "win_length": 1024, + "hop_length": 256, + "n_mels": 80, + "fmin": 0, + "fmax": "None", + }, + }, + "dit_type": "DiT", + "reg_loss_type": "l1", + "style_encoder": {"dim": 192}, + "length_regulator": { + "channels": 512, + "is_discrete": False, + "in_channels": 1024, + "content_codebook_size": 2048, + "sampling_ratios": [1, 1, 1, 1], + "vector_quantize": False, + "n_codebooks": 1, + "quantizer_dropout": 0.0, + "f0_condition": False, + "n_f0_bins": 512, + }, + "DiT": { + "hidden_dim": 512, + "num_heads": 8, + "depth": 13, + "class_dropout_prob": 0.1, + "block_size": 8192, + "in_channels": 80, + "style_condition": True, + "final_layer_type": "wavenet", + "target": "mel", + "content_dim": 512, + "content_codebook_size": 1024, + "content_type": "discrete", + "f0_condition": False, + "n_f0_bins": 512, + "content_codebooks": 1, + "is_causal": False, + "long_skip_connection": True, + "zero_prompt_speech_token": False, + "time_as_token": False, + "style_as_token": False, + "uvit_skip_connection": True, + "add_resblock_in_transformer": False, + }, + "wavenet": { + "hidden_dim": 512, + "num_layers": 8, + "kernel_size": 5, + "dilation_rate": 1, + "p_dropout": 0.2, + "style_condition": True, + }, + }, + "gpt_checkpoint": "gpt.pth", + "w2v_stat": "wav2vec2bert_stats.pt", + "s2mel_checkpoint": "s2mel.pth", + "emo_matrix": "feat2.pt", + "spk_matrix": "feat1.pt", + "emo_num": [3, 17, 2, 8, 4, 5, 10, 24], + "qwen_emo_path": "qwen0.6bemo4-merge/", + "vocoder": {"type": "bigvgan", "name": "nvidia/bigvgan_v2_22khz_80band_256x"}, + "version": 2.0, + # Default-off performance tracing for IndexTTS2 overhead analysis. + } + + def __init__(self, **kwargs): + super().__init__(**kwargs) + defaults = copy.deepcopy(self._DEFAULTS) + for attr, default in defaults.items(): + if not hasattr(self, attr): + setattr(self, attr, default) + # Keep gpt.model_dim in sync with top-level hidden_size. + if isinstance(self.gpt, dict) and self.gpt.get("model_dim") != self.hidden_size: + self.gpt["model_dim"] = self.hidden_size + self.output_sample_rate = int( + self.s2mel["preprocess_params"].get("sr", 22050) if isinstance(self.s2mel, dict) else 22050 + ) diff --git a/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py b/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py new file mode 100644 index 00000000000..4d32b8c1b8b --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/dit_cuda_graph.py @@ -0,0 +1,164 @@ +"""Per-shape CUDA graph capture for the S2Mel DiT transformer core. + +Profiling showed the DiT estimator forward is CPU launch-bound (the GPU +is idle ~85% of CFM wall time). Within one request the CFM runs 12 Euler +steps with *identical* tensor shapes (T_in never changes), so the graph +captured on the first step is replayed by every subsequent step. Across +requests T_in varies, so graphs are cached per shape with a small LRU; +a cache miss costs roughly one extra eager step (warm run + capture). + +Capture safety relies on the IndexTTS2 S2Mel decoder only enabling this +runner for full-mask shapes: + +- ``x_lens == T`` for every sample (the decoder feeds full-length sequences), so + the attention padding mask is all-True and dense flash attention is + mathematically identical. The decoder sets ``_assume_full_mask`` on + the gpt_fast ``Attention`` modules so the backend skips its + ``torch.any(~mask)`` D2H sync, which would abort stream capture. + +Modeled on ``CUDAGraphGLMTTSDiTWrapper`` +(vllm_omni/model_executor/models/glm_tts/glm_tts_dit_wrapper.py). +""" + +from __future__ import annotations + +from collections import OrderedDict +from dataclasses import dataclass + +import torch +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +@dataclass +class _GraphEntry: + graph: torch.cuda.CUDAGraph + static_x: torch.Tensor + static_c: torch.Tensor + static_pos: torch.Tensor + static_mask: torch.Tensor | None + static_out: torch.Tensor + + +class CUDAGraphDiTRunner: + """Replay the gpt_fast Transformer forward via per-shape CUDA graphs.""" + + def __init__(self, transformer: torch.nn.Module, *, max_graphs: int = 4, enabled: bool = True) -> None: + self.transformer = transformer + self.max_graphs = int(max_graphs) + self.enabled = bool(enabled) + self._cache: OrderedDict[tuple, _GraphEntry] = OrderedDict() + self.last_call_info: dict[str, object] = {} + + def __call__( + self, + x_in: torch.Tensor, + c: torch.Tensor, + input_pos: torch.Tensor, + mask: torch.Tensor | None, + ) -> torch.Tensor: + if not self.enabled or x_in.device.type != "cuda" or torch.cuda.is_current_stream_capturing(): + self.last_call_info = { + "mode": "eager", + "reason": "ineligible", + "shape": tuple(x_in.shape), + "cache_size": len(self._cache), + } + return self.transformer(x_in, c, input_pos, mask) + + mask_key = None if mask is None else (tuple(mask.shape), mask.dtype) + device_key = (x_in.device.type, x_in.device.index) + key = ( + device_key, + tuple(x_in.shape), + x_in.dtype, + tuple(c.shape), + c.dtype, + tuple(input_pos.shape), + input_pos.dtype, + mask_key, + ) + entry = self._cache.get(key) + if entry is None: + entry = self._capture(x_in, c, input_pos, mask) + if entry is None: + self.last_call_info = { + "mode": "eager", + "reason": "capture_failed", + "shape": tuple(x_in.shape), + "cache_size": len(self._cache), + } + return self.transformer(x_in, c, input_pos, mask) + self._cache[key] = entry + while len(self._cache) > self.max_graphs: + _, evicted = self._cache.popitem(last=False) + del evicted + # Kernels are only recorded (not executed) during capture, so + # replay once to produce the first real output. + entry.graph.replay() + self.last_call_info = { + "mode": "graph", + "reason": "capture", + "shape": tuple(x_in.shape), + "cache_size": len(self._cache), + } + return entry.static_out.clone() + + self._cache.move_to_end(key) + entry.static_x.copy_(x_in) + entry.static_c.copy_(c) + entry.static_pos.copy_(input_pos) + if entry.static_mask is not None and mask is not None: + entry.static_mask.copy_(mask) + entry.graph.replay() + self.last_call_info = { + "mode": "graph", + "reason": "hit", + "shape": tuple(x_in.shape), + "cache_size": len(self._cache), + } + return entry.static_out.clone() + + def _capture( + self, + x_in: torch.Tensor, + c: torch.Tensor, + input_pos: torch.Tensor, + mask: torch.Tensor | None, + ) -> _GraphEntry | None: + static_x = x_in.clone() + static_c = c.clone() + static_pos = input_pos.clone() + # The mask content is constant (all-True; see module docstring) and + # ignored by Attention via `_assume_full_mask`, but keep a static + # clone so the captured graph never dereferences a freed tensor. + static_mask = mask.contiguous().clone() if mask is not None else None + try: + # Eager warm run: cuDNN/flash-attn autotuning and lazy inits + # must happen outside stream capture. + with torch.no_grad(): + _ = self.transformer(static_x, static_c, static_pos, static_mask) + torch.accelerator.synchronize(x_in.device) + + graph = torch.cuda.CUDAGraph() + with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_out = self.transformer(static_x, static_c, static_pos, static_mask) + except Exception: + logger.warning( + "Disabling S2Mel DiT CUDA graphs after capture failure for shape=%s", + tuple(x_in.shape), + exc_info=True, + ) + self.enabled = False + return None + logger.debug("Captured S2Mel DiT CUDA graph for shape=%s", tuple(x_in.shape)) + return _GraphEntry( + graph=graph, + static_x=static_x, + static_c=static_c, + static_pos=static_pos, + static_mask=static_mask, + static_out=static_out, + ) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/__init__.py b/vllm_omni/model_executor/models/indextts2/gpt/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/__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/indextts2/gpt/conformer/__init__.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/conformer/__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/indextts2/gpt/conformer/attention.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/attention.py new file mode 100644 index 00000000000..7b34ab38175 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/conformer/attention.py @@ -0,0 +1,288 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2019 Shigeki Karita +# 2020 Mobvoi Inc (Binbin Zhang) +# 2022 Xingchen Song (sxc19@mails.tsinghua.edu.cn) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Multi-Head Attention layer definition.""" + +import math + +import torch +from torch import nn + + +class MultiHeadedAttention(nn.Module): + """Multi-Head Attention layer. + + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, n_head: int, n_feat: int, dropout_rate: float): + """Construct an MultiHeadedAttention object.""" + super().__init__() + assert n_feat % n_head == 0 + # We assume d_v always equals d_k + self.d_k = n_feat // n_head + self.h = n_head + self.linear_q = nn.Linear(n_feat, n_feat) + self.linear_k = nn.Linear(n_feat, n_feat) + self.linear_v = nn.Linear(n_feat, n_feat) + self.linear_out = nn.Linear(n_feat, n_feat) + self.dropout = nn.Dropout(p=dropout_rate) + + def forward_qkv( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Transform query, key and value. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + + Returns: + torch.Tensor: Transformed query tensor, size + (#batch, n_head, time1, d_k). + torch.Tensor: Transformed key tensor, size + (#batch, n_head, time2, d_k). + torch.Tensor: Transformed value tensor, size + (#batch, n_head, time2, d_k). + + """ + n_batch = query.size(0) + q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) + k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) + v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) + q = q.transpose(1, 2) # (batch, head, time1, d_k) + k = k.transpose(1, 2) # (batch, head, time2, d_k) + v = v.transpose(1, 2) # (batch, head, time2, d_k) + + return q, k, v + + def forward_attention( + self, value: torch.Tensor, scores: torch.Tensor, mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool) + ) -> torch.Tensor: + """Compute attention context vector. + + Args: + value (torch.Tensor): Transformed value, size + (#batch, n_head, time2, d_k). + scores (torch.Tensor): Attention score, size + (#batch, n_head, time1, time2). + mask (torch.Tensor): Mask, size (#batch, 1, time2) or + (#batch, time1, time2), (0, 0, 0) means fake mask. + + Returns: + torch.Tensor: Transformed value (#batch, time1, d_model) + weighted by the attention score (#batch, time1, time2). + + """ + n_batch = value.size(0) + # NOTE(xcsong): When will `if mask.size(2) > 0` be True? + # 1. onnx(16/4) [WHY? Because we feed real cache & real mask for the + # 1st chunk to ease the onnx export.] + # 2. pytorch training + if mask.size(2) > 0: # time2 > 0 + mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) + # For last chunk, time2 might be larger than scores.size(-1) + mask = mask[:, :, :, : scores.size(-1)] # (batch, 1, *, time2) + scores = scores.masked_fill(mask, -float("inf")) + attn = torch.softmax(scores, dim=-1).masked_fill(mask, 0.0) # (batch, head, time1, time2) + # NOTE(xcsong): When will `if mask.size(2) > 0` be False? + # 1. onnx(16/-1, -1/-1, 16/0) + # 2. jit (16/-1, -1/-1, 16/0, 16/4) + else: + attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) + + p_attn = self.dropout(attn) + x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) + x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) # (batch, time1, d_model) + + return self.linear_out(x) # (batch, time1, d_model) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute scaled dot product attention. + + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2). + 1.When applying cross attention between decoder and encoder, + the batch padding mask for input is in (#batch, 1, T) shape. + 2.When applying self attention of encoder, + the mask is in (#batch, T, T) shape. + 3.When applying self attention of decoder, + the mask is in (#batch, L, L) shape. + 4.If the different position in decoder see different block + of the encoder, such as Mocha, the passed in mask could be + in (#batch, L, T) shape. But there is no such case in current + Wenet. + cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + + + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + + """ + q, k, v = self.forward_qkv(query, key, value) + + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + new_cache = torch.cat((k, v), dim=-1) + + # Build boolean attention mask for SDPA. + # Input mask convention: True = attend, False = ignore. + # SDPA boolean mask convention: True = attend, False = masked to -inf. + if mask.size(2) > 0: + attn_mask = mask.unsqueeze(1) # (batch, 1, *, time2) + attn_mask = attn_mask[:, :, :, : k.size(2)] + else: + attn_mask = None + + # q/k/v: (batch, head, time, d_k) + x = torch.nn.functional.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=self.dropout.p if self.training else 0.0, + scale=1.0 / math.sqrt(self.d_k), + ) + n_batch = query.size(0) + x = x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) + + return self.linear_out(x), new_cache + + +class RelPositionMultiHeadedAttention(MultiHeadedAttention): + """Multi-Head Attention layer with relative position encoding. + Paper: https://arxiv.org/abs/1901.02860 + Args: + n_head (int): The number of heads. + n_feat (int): The number of features. + dropout_rate (float): Dropout rate. + """ + + def __init__(self, n_head, n_feat, dropout_rate): + """Construct an RelPositionMultiHeadedAttention object.""" + super().__init__(n_head, n_feat, dropout_rate) + # linear transformation for positional encoding + self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) + # these two learnable bias are used in matrix c and matrix d + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + self.pos_bias_u = nn.Parameter(torch.Tensor(self.h, self.d_k)) + self.pos_bias_v = nn.Parameter(torch.Tensor(self.h, self.d_k)) + torch.nn.init.xavier_uniform_(self.pos_bias_u) + torch.nn.init.xavier_uniform_(self.pos_bias_v) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + pos_emb: torch.Tensor = torch.empty(0), + cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute 'Scaled Dot Product Attention' with rel. positional encoding. + Args: + query (torch.Tensor): Query tensor (#batch, time1, size). + key (torch.Tensor): Key tensor (#batch, time2, size). + value (torch.Tensor): Value tensor (#batch, time2, size). + mask (torch.Tensor): Mask tensor (#batch, 1, time2) or + (#batch, time1, time2), (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): Positional embedding tensor + (#batch, time2, size). + cache (torch.Tensor): Cache tensor (1, head, cache_t, d_k * 2), + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + Returns: + torch.Tensor: Output tensor (#batch, time1, d_model). + torch.Tensor: Cache tensor (1, head, cache_t + time1, d_k * 2) + where `cache_t == chunk_size * num_decoding_left_chunks` + and `head * d_k == size` + """ + q, k, v = self.forward_qkv(query, key, value) + q = q.transpose(1, 2) # (batch, time1, head, d_k) + + # NOTE(xcsong): + # when export onnx model, for 1st chunk, we feed + # cache(1, head, 0, d_k * 2) (16/-1, -1/-1, 16/0 mode) + # or cache(1, head, real_cache_t, d_k * 2) (16/4 mode). + # In all modes, `if cache.size(0) > 0` will always be `True` + # and we will always do splitting and + # concatnation(this will simplify onnx export). Note that + # it's OK to concat & split zero-shaped tensors(see code below). + # when export jit model, for 1st chunk, we always feed + # cache(0, 0, 0, 0) since jit supports dynamic if-branch. + # >>> a = torch.ones((1, 2, 0, 4)) + # >>> b = torch.ones((1, 2, 3, 4)) + # >>> c = torch.cat((a, b), dim=2) + # >>> torch.equal(b, c) # True + # >>> d = torch.split(a, 2, dim=-1) + # >>> torch.equal(d[0], d[1]) # True + if cache.size(0) > 0: + key_cache, value_cache = torch.split(cache, cache.size(-1) // 2, dim=-1) + k = torch.cat([key_cache, k], dim=2) + v = torch.cat([value_cache, v], dim=2) + # NOTE(xcsong): We do cache slicing in encoder.forward_chunk, since it's + # non-trivial to calculate `next_cache_start` here. + new_cache = torch.cat((k, v), dim=-1) + + n_batch_pos = pos_emb.size(0) + p = self.linear_pos(pos_emb).view(n_batch_pos, -1, self.h, self.d_k) + p = p.transpose(1, 2) # (batch, head, time1, d_k) + + # (batch, head, time1, d_k) + q_with_bias_u = (q + self.pos_bias_u).transpose(1, 2) + # (batch, head, time1, d_k) + q_with_bias_v = (q + self.pos_bias_v).transpose(1, 2) + + # compute attention score + # first compute matrix a and matrix c + # as described in https://arxiv.org/abs/1901.02860 Section 3.3 + # (batch, head, time1, time2) + matrix_ac = torch.matmul(q_with_bias_u, k.transpose(-2, -1)) + + # compute matrix b and matrix d + # (batch, head, time1, time2) + matrix_bd = torch.matmul(q_with_bias_v, p.transpose(-2, -1)) + + scores = (matrix_ac + matrix_bd) / math.sqrt(self.d_k) # (batch, head, time1, time2) + + return self.forward_attention(v, scores, mask), new_cache diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer/embedding.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/embedding.py new file mode 100644 index 00000000000..9c52df58668 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/conformer/embedding.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2020 Mobvoi Inc. (authors: Binbin Zhang, Di Wu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) + +"""Positional Encoding Module.""" + +import math + +import torch +import torch.nn.functional as F + + +class PositionalEncoding(torch.nn.Module): + """Positional encoding. + + :param int d_model: embedding dim + :param float dropout_rate: dropout rate + :param int max_len: maximum input length + + PE(pos, 2i) = sin(pos/(10000^(2i/dmodel))) + PE(pos, 2i+1) = cos(pos/(10000^(2i/dmodel))) + """ + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000, reverse: bool = False): + """Construct an PositionalEncoding object.""" + super().__init__() + self.d_model = d_model + self.xscale = math.sqrt(self.d_model) + self.dropout = torch.nn.Dropout(p=dropout_rate) + self.max_len = max_len + + pe = torch.zeros(self.max_len, self.d_model) + position = torch.arange(0, self.max_len).unsqueeze(1) + div_term = torch.exp(torch.arange(0, self.d_model, 2) * -(math.log(10000.0) / self.d_model)) + pe[:, 0::2] = torch.sin(position * div_term) + pe[:, 1::2] = torch.cos(position * div_term) + pe = pe.unsqueeze(0) + self.register_buffer("pe", pe) + + def forward(self, x: torch.Tensor, offset: int | torch.Tensor = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Add positional encoding. + + Args: + x (torch.Tensor): Input. Its shape is (batch, time, ...) + offset (int, torch.tensor): position offset + + Returns: + torch.Tensor: Encoded tensor. Its shape is (batch, time, ...) + torch.Tensor: for compatibility to RelPositionalEncoding + """ + + self.pe = self.pe.to(x.device) + pos_emb = self.position_encoding(offset, x.size(1), False) + x = x * self.xscale + pos_emb + return self.dropout(x), self.dropout(pos_emb) + + def position_encoding(self, offset: int | torch.Tensor, size: int, apply_dropout: bool = True) -> torch.Tensor: + """For getting encoding in a streaming fashion + + Attention!!!!! + we apply dropout only once at the whole utterance level in a none + streaming way, but will call this function several times with + increasing input size in a streaming scenario, so the dropout will + be applied several times. + + Args: + offset (int or torch.tensor): start offset + size (int): required size of position encoding + + Returns: + torch.Tensor: Corresponding encoding + """ + # How to subscript a Union type: + # https://github.com/pytorch/pytorch/issues/69434 + if isinstance(offset, int): + assert offset + size < self.max_len + pos_emb = self.pe[:, offset : offset + size] + elif isinstance(offset, torch.Tensor) and offset.dim() == 0: # scalar + assert offset + size < self.max_len + pos_emb = self.pe[:, offset : offset + size] + else: # for batched streaming decoding on GPU + assert torch.max(offset) + size < self.max_len + index = offset.unsqueeze(1) + torch.arange(0, size).to(offset.device) # B X T + flag = index > 0 + # remove negative offset + index = index * flag + pos_emb = F.embedding(index, self.pe[0]) # B X T X d_model + + if apply_dropout: + pos_emb = self.dropout(pos_emb) + return pos_emb + + +class RelPositionalEncoding(PositionalEncoding): + """Relative positional encoding module. + See : Appendix B in https://arxiv.org/abs/1901.02860 + Args: + d_model (int): Embedding dimension. + dropout_rate (float): Dropout rate. + max_len (int): Maximum input length. + """ + + def __init__(self, d_model: int, dropout_rate: float, max_len: int = 5000): + """Initialize class.""" + super().__init__(d_model, dropout_rate, max_len, reverse=True) + + def forward(self, x: torch.Tensor, offset: int | torch.Tensor = 0) -> tuple[torch.Tensor, torch.Tensor]: + """Compute positional encoding. + Args: + x (torch.Tensor): Input tensor (batch, time, `*`). + Returns: + torch.Tensor: Encoded tensor (batch, time, `*`). + torch.Tensor: Positional embedding tensor (1, time, `*`). + """ + self.pe = self.pe.to(x.device) + x = x * self.xscale + pos_emb = self.position_encoding(offset, x.size(1), False) + return self.dropout(x), self.dropout(pos_emb) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer/subsampling.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer/subsampling.py new file mode 100644 index 00000000000..f8e0b7aa04c --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/conformer/subsampling.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2021 Mobvoi Inc (Binbin Zhang, Di Wu) +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Modified from ESPnet(https://github.com/espnet/espnet) + + +"""Subsampling layer definition.""" + +import torch + + +class BaseSubsampling(torch.nn.Module): + def __init__(self): + super().__init__() + self.right_context = 0 + self.subsampling_rate = 1 + + def position_encoding(self, offset: int | torch.Tensor, size: int) -> torch.Tensor: + return self.pos_enc.position_encoding(offset, size) + + +class Conv2dSubsampling2(BaseSubsampling): + """Convolutional 2D subsampling (to 1/2 length). + + Args: + idim (int): Input dimension. + odim (int): Output dimension. + dropout_rate (float): Dropout rate. + + """ + + def __init__(self, idim: int, odim: int, dropout_rate: float, pos_enc_class: torch.nn.Module): + super().__init__() + self.conv = torch.nn.Sequential( + torch.nn.Conv2d(1, odim, 3, 2), + torch.nn.ReLU(), + ) + self.out = torch.nn.Sequential(torch.nn.Linear(odim * ((idim - 1) // 2), odim)) + self.pos_enc = pos_enc_class + self.subsampling_rate = 2 + self.right_context = 2 + + def forward( + self, x: torch.Tensor, x_mask: torch.Tensor, offset: int | torch.Tensor = 0 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = x.unsqueeze(1) # (b, c=1, t, f) + x = self.conv(x) + b, c, t, f = x.size() + x = self.out(x.transpose(1, 2).contiguous().view(b, t, c * f)) + x, pos_emb = self.pos_enc(x, offset) + return x, pos_emb, x_mask[:, :, 2::2] diff --git a/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py b/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py new file mode 100644 index 00000000000..67eebbab3fd --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/conformer_encoder.py @@ -0,0 +1,422 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.indextts2.gpt.conformer.attention import ( + MultiHeadedAttention, + RelPositionMultiHeadedAttention, +) +from vllm_omni.model_executor.models.indextts2.gpt.conformer.embedding import ( + PositionalEncoding, + RelPositionalEncoding, +) +from vllm_omni.model_executor.models.indextts2.gpt.conformer.subsampling import ( + Conv2dSubsampling2, +) +from vllm_omni.model_executor.models.indextts2.utils.common import make_pad_mask + + +class PositionwiseFeedForward(torch.nn.Module): + """Positionwise feed forward layer. + + FeedForward are applied on each position of the sequence. + The output dim is same with the input dim. + + Args: + idim (int): Input dimension. + hidden_units (int): The number of hidden units. + dropout_rate (float): Dropout rate. + activation (torch.nn.Module): Activation function + """ + + def __init__( + self, idim: int, hidden_units: int, dropout_rate: float, activation: torch.nn.Module = torch.nn.ReLU() + ): + """Construct a PositionwiseFeedForward object.""" + super().__init__() + self.w_1 = torch.nn.Linear(idim, hidden_units) + self.activation = activation + self.dropout = torch.nn.Dropout(dropout_rate) + self.w_2 = torch.nn.Linear(hidden_units, idim) + + def forward(self, xs: torch.Tensor) -> torch.Tensor: + """Forward function. + + Args: + xs: input tensor (B, L, D) + Returns: + output tensor, (B, L, D) + """ + return self.w_2(self.dropout(self.activation(self.w_1(xs)))) + + +class ConvolutionModule(nn.Module): + """ConvolutionModule in Conformer model.""" + + def __init__(self, channels: int, kernel_size: int = 15, activation: nn.Module = nn.ReLU(), bias: bool = True): + """Construct an ConvolutionModule object. + Args: + channels (int): The number of channels of conv layers. + kernel_size (int): Kernel size of conv layers. + causal (int): Whether use causal convolution or not + """ + super().__init__() + + self.pointwise_conv1 = nn.Conv1d( + channels, + 2 * channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + # self.lorder is used to distinguish if it's a causal convolution, + # if self.lorder > 0: it's a causal convolution, the input will be + # padded with self.lorder frames on the left in forward. + # else: it's a symmetrical convolution + # kernel_size should be an odd number for none causal convolution + assert (kernel_size - 1) % 2 == 0 + padding = (kernel_size - 1) // 2 + self.lorder = 0 + + self.depthwise_conv = nn.Conv1d( + channels, + channels, + kernel_size, + stride=1, + padding=padding, + groups=channels, + bias=bias, + ) + + self.use_layer_norm = True + self.norm = nn.LayerNorm(channels) + + self.pointwise_conv2 = nn.Conv1d( + channels, + channels, + kernel_size=1, + stride=1, + padding=0, + bias=bias, + ) + self.activation = activation + + def forward( + self, + x: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + cache: torch.Tensor = torch.zeros((0, 0, 0)), + ) -> tuple[torch.Tensor, torch.Tensor]: + """Compute convolution module. + Args: + x (torch.Tensor): Input tensor (#batch, time, channels). + mask_pad (torch.Tensor): used for batch padding (#batch, 1, time), + (0, 0, 0) means fake mask. + cache (torch.Tensor): left context cache, it is only + used in causal convolution (#batch, channels, cache_t), + (0, 0, 0) meas fake cache. + Returns: + torch.Tensor: Output tensor (#batch, time, channels). + """ + # exchange the temporal dimension and the feature dimension + x = x.transpose(1, 2) # (#batch, channels, time) + + # mask batch padding + if mask_pad.size(2) > 0: # time > 0 + x.masked_fill_(~mask_pad, 0.0) + + if self.lorder > 0: + if cache.size(2) == 0: # cache_t == 0 + x = nn.functional.pad(x, (self.lorder, 0), "constant", 0.0) + else: + assert cache.size(0) == x.size(0) # equal batch + assert cache.size(1) == x.size(1) # equal channel + x = torch.cat((cache, x), dim=2) + assert x.size(2) > self.lorder + new_cache = x[:, :, -self.lorder :] + else: + # It's better we just return None if no cache is required, + # However, for JIT export, here we just fake one tensor instead of + # None. + new_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + + # GLU mechanism + x = self.pointwise_conv1(x) # (batch, 2*channel, dim) + x = nn.functional.glu(x, dim=1) # (batch, channel, dim) + + # 1D Depthwise Conv + x = self.depthwise_conv(x) + if self.use_layer_norm: + x = x.transpose(1, 2) + x = self.activation(self.norm(x)) + if self.use_layer_norm: + x = x.transpose(1, 2) + x = self.pointwise_conv2(x) + # mask batch padding + if mask_pad.size(2) > 0: # time > 0 + x.masked_fill_(~mask_pad, 0.0) + + return x.transpose(1, 2), new_cache + + +class ConformerEncoderLayer(nn.Module): + """Encoder layer module. + Args: + size (int): Input dimension. + self_attn (torch.nn.Module): Self-attention module instance. + `MultiHeadedAttention` or `RelPositionMultiHeadedAttention` + instance can be used as the argument. + feed_forward (torch.nn.Module): Feed-forward module instance. + `PositionwiseFeedForward` instance can be used as the argument. + feed_forward_macaron (torch.nn.Module): Additional feed-forward module + instance. + `PositionwiseFeedForward` instance can be used as the argument. + conv_module (torch.nn.Module): Convolution module instance. + `ConvlutionModule` instance can be used as the argument. + dropout_rate (float): Dropout rate. + normalize_before (bool): + True: use layer_norm before each sub-block. + False: use layer_norm after each sub-block. + """ + + def __init__( + self, + size: int, + self_attn: torch.nn.Module, + feed_forward: nn.Module | None = None, + feed_forward_macaron: nn.Module | None = None, + conv_module: nn.Module | None = None, + dropout_rate: float = 0.1, + normalize_before: bool = True, + ): + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.feed_forward_macaron = feed_forward_macaron + self.conv_module = conv_module + self.norm_ff = nn.LayerNorm(size, eps=1e-5) + self.norm_mha = nn.LayerNorm(size, eps=1e-5) + if feed_forward_macaron is not None: + self.norm_ff_macaron = nn.LayerNorm(size, eps=1e-5) + self.ff_scale = 0.5 + else: + self.ff_scale = 1.0 + if self.conv_module is not None: + self.norm_conv = nn.LayerNorm(size, eps=1e-5) + self.norm_final = nn.LayerNorm(size, eps=1e-5) + self.dropout = nn.Dropout(dropout_rate) + self.size = size + self.normalize_before = normalize_before + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor, + pos_emb: torch.Tensor, + mask_pad: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool), + att_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + cnn_cache: torch.Tensor = torch.zeros((0, 0, 0, 0)), + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute encoded features. + + Args: + x (torch.Tensor): (#batch, time, size) + mask (torch.Tensor): Mask tensor for the input (#batch, time，time), + (0, 0, 0) means fake mask. + pos_emb (torch.Tensor): positional encoding, must not be None + for ConformerEncoderLayer. + mask_pad (torch.Tensor): batch padding mask used for conv module. + (#batch, 1，time), (0, 0, 0) means fake mask. + att_cache (torch.Tensor): Cache tensor of the KEY & VALUE + (#batch=1, head, cache_t1, d_k * 2), head * d_k == size. + cnn_cache (torch.Tensor): Convolution cache in conformer layer + (#batch=1, size, cache_t2) + Returns: + torch.Tensor: Output tensor (#batch, time, size). + torch.Tensor: Mask tensor (#batch, time, time). + torch.Tensor: att_cache tensor, + (#batch=1, head, cache_t1 + time, d_k * 2). + torch.Tensor: cnn_cache tensor (#batch, size, cache_t2). + """ + + if self.feed_forward_macaron is not None: + residual = x + if self.normalize_before: + x = self.norm_ff_macaron(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward_macaron(x)) + + residual = x + if self.normalize_before: + x = self.norm_mha(x) + x_att, new_att_cache = self.self_attn(x, x, x, mask, pos_emb, att_cache) + x = residual + self.dropout(x_att) + + new_cnn_cache = torch.zeros((0, 0, 0), dtype=x.dtype, device=x.device) + if self.conv_module is not None: + residual = x + if self.normalize_before: + x = self.norm_conv(x) + x, new_cnn_cache = self.conv_module(x, mask_pad, cnn_cache) + x = residual + self.dropout(x) + + residual = x + if self.normalize_before: + x = self.norm_ff(x) + x = residual + self.ff_scale * self.dropout(self.feed_forward(x)) + + if self.conv_module is not None: + x = self.norm_final(x) + + return x, mask, new_att_cache, new_cnn_cache + + +class BaseEncoder(torch.nn.Module): + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.0, + input_layer: str = "conv2d", + pos_enc_layer_type: str = "abs_pos", + normalize_before: bool = True, + concat_after: bool = False, + ): + super().__init__() + self._output_size = output_size + + pos_enc_map = {"abs_pos": PositionalEncoding, "rel_pos": RelPositionalEncoding} + if pos_enc_layer_type not in pos_enc_map: + raise ValueError("unknown pos_enc_layer: " + pos_enc_layer_type) + pos_enc_class = pos_enc_map[pos_enc_layer_type] + + subsampling_map = {"conv2d2": Conv2dSubsampling2} + if input_layer not in subsampling_map: + raise ValueError("unknown input_layer: " + input_layer) + subsampling_class = subsampling_map[input_layer] + + self.embed = subsampling_class( + input_size, + output_size, + dropout_rate, + pos_enc_class(output_size, dropout_rate), + ) + + self.normalize_before = normalize_before + self.after_norm = torch.nn.LayerNorm(output_size, eps=1e-5) + + def output_size(self) -> int: + return self._output_size + + def forward( + self, + xs: torch.Tensor, + xs_lens: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Embed positions in tensor. + + Args: + xs: padded input tensor (B, T, D) + xs_lens: input length (B) + + Returns: + encoder output tensor xs, and subsampled masks + xs: padded output tensor (B, T' ~= T/subsample_rate, D) + masks: torch.Tensor batch padding mask after subsample + (B, 1, T' ~= T/subsample_rate) + """ + T = xs.size(1) + masks = ~make_pad_mask(xs_lens, T).unsqueeze(1) # (B, 1, T) + xs, pos_emb, masks = self.embed(xs, masks) + chunk_masks = masks + mask_pad = masks # (B, 1, T/subsample_rate) + for layer in self.encoders: + xs, chunk_masks, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad) + if self.normalize_before: + xs = self.after_norm(xs) + # Here we assume the mask is not changed in encoder layers, so just + # return the masks before encoder layers, and the masks will be used + # for cross attention with decoder later + return xs, masks + + +class ConformerEncoder(BaseEncoder): + """Conformer encoder module.""" + + def __init__( + self, + input_size: int, + output_size: int = 256, + attention_heads: int = 4, + linear_units: int = 2048, + num_blocks: int = 6, + dropout_rate: float = 0.0, + input_layer: str = "conv2d", + pos_enc_layer_type: str = "rel_pos", + normalize_before: bool = True, + concat_after: bool = False, + macaron_style: bool = False, + use_cnn_module: bool = True, + cnn_module_kernel: int = 15, + ): + super().__init__( + input_size, + output_size, + attention_heads, + linear_units, + num_blocks, + dropout_rate, + input_layer, + pos_enc_layer_type, + normalize_before, + concat_after, + ) + + activation = torch.nn.SiLU() + + # self-attention module definition + if pos_enc_layer_type != "rel_pos": + encoder_selfattn_layer = MultiHeadedAttention + else: + encoder_selfattn_layer = RelPositionMultiHeadedAttention + encoder_selfattn_layer_args = ( + attention_heads, + output_size, + dropout_rate, + ) + + # feed-forward module definition + positionwise_layer = PositionwiseFeedForward + positionwise_layer_args = ( + output_size, + linear_units, + dropout_rate, + activation, + ) + # convolution module definition + convolution_layer = ConvolutionModule + convolution_layer_args = ( + output_size, + cnn_module_kernel, + activation, + ) + + self.encoders = torch.nn.ModuleList( + [ + ConformerEncoderLayer( + output_size, + encoder_selfattn_layer(*encoder_selfattn_layer_args), + positionwise_layer(*positionwise_layer_args), + positionwise_layer(*positionwise_layer_args) if macaron_style else None, + convolution_layer(*convolution_layer_args) if use_cnn_module else None, + dropout_rate, + normalize_before, + ) + for _ in range(num_blocks) + ] + ) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/embeddings.py b/vllm_omni/model_executor/models/indextts2/gpt/embeddings.py new file mode 100644 index 00000000000..bc9e8e466e8 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/embeddings.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Learned position embeddings for IndexTTS2 GPT model.""" + +import torch +import torch.nn as nn + + +class LearnedPositionEmbeddings(nn.Module): + def __init__(self, seq_len: int, model_dim: int, init: float = 0.02): + super().__init__() + self.emb = nn.Embedding(seq_len, model_dim) + self.emb.weight.data.normal_(mean=0.0, std=init) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + sl = x.shape[1] + return self.emb(torch.arange(0, sl, device=x.device)) + + def get_fixed_embedding(self, ind: int, dev: torch.device) -> torch.Tensor: + return self.emb(torch.tensor([ind], device=dev)).unsqueeze(0) diff --git a/vllm_omni/model_executor/models/indextts2/gpt/perceiver.py b/vllm_omni/model_executor/models/indextts2/gpt/perceiver.py new file mode 100644 index 00000000000..98c7becfab8 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/gpt/perceiver.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Adapted from lucidrains/naturalspeech2-pytorch (PerceiverResampler) +# https://github.com/lucidrains/naturalspeech2-pytorch/blob/659bec7f/ +# naturalspeech2_pytorch/naturalspeech2_pytorch.py#L532 + +import torch +import torch.nn.functional as F +from einops import rearrange, repeat +from torch import nn + + +def exists(val): + return val is not None + + +def default(val, d): + if exists(val): + return val + return d() if callable(d) else d + + +class Attend(nn.Module): + def __init__(self, dropout=0.0): + super().__init__() + self.dropout = dropout + + def forward(self, q, k, v, mask=None): + attn_mask = None + if exists(mask): + attn_mask = rearrange(mask, "b j -> b 1 1 j") + + return F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=self.dropout if self.training else 0.0, + ) + + +class RMSNorm(nn.Module): + def __init__(self, dim, scale=True): + super().__init__() + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(dim)) if scale else None + + def forward(self, x): + gamma = default(self.gamma, 1) + return F.normalize(x, dim=-1) * self.scale * gamma + + +class GEGLU(nn.Module): + def forward(self, x): + x, gate = x.chunk(2, dim=-1) + return F.gelu(gate) * x + + +def FeedForward(dim, mult=4): + dim_inner = int(dim * mult * 2 / 3) + return nn.Sequential(nn.Linear(dim, dim_inner * 2), GEGLU(), nn.Linear(dim_inner, dim)) + + +class PerceiverResampler(nn.Module): + def __init__( + self, + dim, + depth=2, + dim_context=None, + num_latents=32, + dim_head=64, + heads=8, + ff_mult=4, + ): + super().__init__() + dim_context = default(dim_context, dim) + + self.proj_context = nn.Linear(dim_context, dim) if dim_context != dim else nn.Identity() + + self.latents = nn.Parameter(torch.randn(num_latents, dim)) + nn.init.normal_(self.latents, std=0.02) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + Attention(dim=dim, dim_head=dim_head, heads=heads, cross_attn_include_queries=True), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + self.norm = RMSNorm(dim) + + def forward(self, x, mask=None): + batch = x.shape[0] + + x = self.proj_context(x) + + latents = repeat(self.latents, "n d -> b n d", b=batch) + + for attn, ff in self.layers: + latents = attn(latents, x, mask=mask) + latents + latents = ff(latents) + latents + + return self.norm(latents) + + +class Attention(nn.Module): + def __init__( + self, + dim, + *, + dim_context=None, + dim_head=64, + heads=8, + dropout=0.0, + cross_attn_include_queries=False, + ): + super().__init__() + self.scale = dim_head**-0.5 + self.heads = heads + self.cross_attn_include_queries = cross_attn_include_queries + + dim_inner = dim_head * heads + dim_context = default(dim_context, dim) + + self.attend = Attend(dropout=dropout) + self.to_q = nn.Linear(dim, dim_inner, bias=False) + self.to_kv = nn.Linear(dim_context, dim_inner * 2, bias=False) + self.to_out = nn.Linear(dim_inner, dim, bias=False) + + def forward(self, x, context=None, mask=None): + h, has_context = self.heads, exists(context) + + context = default(context, x) + + if has_context and self.cross_attn_include_queries: + context = torch.cat((x, context), dim=-2) + + q, k, v = (self.to_q(x), *self.to_kv(context).chunk(2, dim=-1)) + q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) + + out = self.attend(q, k, v, mask=mask) + + out = rearrange(out, "b h n d -> b n (h d)") + return self.to_out(out) diff --git a/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py b/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py new file mode 100644 index 00000000000..a374f1b4570 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/indextts2_s2mel_decoder.py @@ -0,0 +1,894 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""IndexTTS2 Stage 1: S2Mel decoder + BigVGAN vocoder. + +Receives mel_codes + latent from Stage 0 (GPT AR talker), runs flow matching +to synthesize mel spectrogram, then BigVGAN to produce waveform audio. +""" + +from __future__ import annotations + +import logging +import math +import os +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.sequence import IntermediateTensors + +from vllm_omni.diffusion.model_loader.hub_prefetch import _repo_prefetch_lock +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .configuration_indextts2 import IndexTTS2Config +from .preprocess_utils import load_semantic_codec, resolve_model_file +from .s2mel.modules.commons import AttrDict, MyModel + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Lazy loaders for external models +# --------------------------------------------------------------------------- + +_bigvgan_models: dict[tuple[str, str], nn.Module] = {} + + +def _patch_bigvgan_compat(cls): + """Monkey-patch BigVGAN._from_pretrained for huggingface_hub>=1.0 compat.""" + if getattr(cls._from_pretrained, "_hf1_patched", False): + return + orig = cls._from_pretrained.__func__ + + @classmethod + def _compat(klass, *, proxies=None, resume_download=False, **kw): + return orig(klass, proxies=proxies, resume_download=resume_download, **kw) + + _compat._hf1_patched = True + cls._from_pretrained = _compat + + +def _strip_weight_norm(root: nn.Module) -> list[str]: + """Bake legacy weight_norm (weight_g/weight_v) into plain weights. + + The DiT final-layer WaveNet (and any other S2Mel submodule) keeps + torch.nn.utils.weight_norm forward pre-hooks after loading, so every + CFM estimator step recomputes the normalization (~217 calls/request + in profiling). Inference never updates these weights, so we fold the + parametrization once at load time. Hooks live on the innermost conv + (e.g. SConv1d.conv.conv), hence the generic module walk instead of + the per-class remove_weight_norm() helpers. + + Returns the qualified names (relative to ``root``) of the plain + weight parameters created by the fold, so callers can mark them as + checkpoint-initialized for vLLM's loader validation. + """ + from torch.nn.utils import remove_weight_norm + + folded: list[str] = [] + for mod_name, module in root.named_modules(): + for hook in list(module._forward_pre_hooks.values()): + if hook.__class__.__name__ == "WeightNorm": + remove_weight_norm(module, name=hook.name) + folded.append(f"{mod_name}.{hook.name}" if mod_name else hook.name) + return folded + + +def _load_bigvgan(vocoder_name: str, device: torch.device, dtype: torch.dtype = torch.float32): + cache_key = (vocoder_name, str(device), str(dtype)) + if cache_key in _bigvgan_models: + return _bigvgan_models[cache_key] + + lock = _repo_prefetch_lock(vocoder_name) if not os.path.isdir(vocoder_name) else None + if lock is not None: + lock.__enter__() + try: + try: + from .s2mel.modules import bigvgan as bigvgan_mod + + _patch_bigvgan_compat(bigvgan_mod.BigVGAN) + bigvgan_model = bigvgan_mod.BigVGAN.from_pretrained(vocoder_name) + except (ImportError, ModuleNotFoundError): + import bigvgan + + _patch_bigvgan_compat(bigvgan.BigVGAN) + bigvgan_model = bigvgan.BigVGAN.from_pretrained(vocoder_name) + finally: + if lock is not None: + lock.__exit__(None, None, None) + # One-time weight cast (vs per-call autocast, which re-casts every conv's + # fp32 weights on each forward and is slower than fp32 on BigVGAN). + bigvgan_model = bigvgan_model.to(device=device, dtype=dtype).eval() + bigvgan_model.remove_weight_norm() + for p in bigvgan_model.parameters(): + p.requires_grad_(False) + _bigvgan_models[cache_key] = bigvgan_model + return bigvgan_model + + +# --------------------------------------------------------------------------- +# Stage 1 Decoder +# --------------------------------------------------------------------------- + + +class IndexTTS2S2MelDecoder(nn.Module): + """S2Mel + BigVGAN decoder for IndexTTS2 Stage 1. + + Receives from Stage 0: + - mel_codes: [T] mel code sequence + - latent: [T, 1280] hidden states from GPT AR + - S_ref: speaker semantic embeddings (from Wav2Vec2 → RepCodec) + - ref_mel: [80, T_ref] reference mel spectrogram + - style: [192] CAMPPlus style vector + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + self.config: IndexTTS2Config = vllm_config.model_config.hf_config # type: ignore[assignment] + + # --- Flags for vLLM-Omni framework --- + 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 + + # --- Build S2Mel model (CFM + LengthRegulator + gpt_layer) --- + s2mel_cfg = self.config.s2mel + s2mel_args = AttrDict(s2mel_cfg) + # Ensure nested dicts are also AttrDicts + for key in ["DiT", "wavenet", "length_regulator", "style_encoder", "preprocess_params"]: + if key in s2mel_args and isinstance(s2mel_args[key], dict): + s2mel_args[key] = AttrDict(s2mel_args[key]) + + self.s2mel = MyModel(s2mel_args, use_emovec=False, use_gpt_latent=True) + + # Diffusion config — configurable via deploy YAML hf_overrides + self.diffusion_steps: int = getattr(self.config, "diffusion_steps", 25) + self.inference_cfg_rate: float = getattr(self.config, "inference_cfg_rate", 0.7) + self.mel_code_to_frame_ratio: float = 1.72 + + # Run the DiT estimator under bf16 autocast (flash attention + faster + # GEMMs); the CFM Euler solver stays float32. Disable via deploy YAML + # hf_overrides `s2mel_dit_bf16: false` if quality regresses. + self.s2mel_dit_bf16: bool = getattr(self.config, "s2mel_dit_bf16", True) + # Run BigVGAN vocoding under bf16 autocast (conv-heavy, ~2x faster). + # Disable via deploy YAML hf_overrides `s2mel_vocoder_bf16: false`. + self.s2mel_vocoder_bf16: bool = getattr(self.config, "s2mel_vocoder_bf16", True) + # Capture BigVGAN into bucketed CUDA graphs (collapses the Snake/conv + # small-op launch storm). Enable via `s2mel_vocoder_cuda_graph: true`. + self.s2mel_vocoder_cuda_graph: bool = getattr(self.config, "s2mel_vocoder_cuda_graph", False) + self._vocoder_graph: Any = None + # Capture the DiT transformer core into per-shape CUDA graphs. + # Enable via `s2mel_dit_cuda_graph: true`. + self.s2mel_dit_cuda_graph: bool = getattr(self.config, "s2mel_dit_cuda_graph", False) + self.s2mel_dit_cuda_graph_max_graphs: int = int(getattr(self.config, "s2mel_dit_cuda_graph_max_graphs", 8)) + self.s2mel_vocoder_capture_sizes: list[int] | None = getattr(self.config, "s2mel_vocoder_capture_sizes", None) + self.s2mel_vocoder_compile_shapes: list[int] | None = getattr(self.config, "s2mel_vocoder_compile_shapes", None) + logger.info( + "[S2Mel] dit_bf16=%s, vocoder_bf16=%s, " + "vocoder_cuda_graph=%s, dit_cuda_graph=%s, dit_graph_lru=%d, vocoder_buckets=%s, " + "vocoder_compile_shapes=%s", + self.s2mel_dit_bf16, + self.s2mel_vocoder_bf16, + self.s2mel_vocoder_cuda_graph, + self.s2mel_dit_cuda_graph, + self.s2mel_dit_cuda_graph_max_graphs, + self.s2mel_vocoder_capture_sizes, + self.s2mel_vocoder_compile_shapes, + ) + + # ------------------------------------------------------------------ + # vLLM 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) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> torch.Tensor | OmniOutput: + """Run S2Mel flow matching + BigVGAN vocoding.""" + model_intermediate_buffer = kwargs.get("model_intermediate_buffer") + if model_intermediate_buffer is None: + model_intermediate_buffer = runtime_additional_information + if runtime_additional_information is not None and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + + request_infos: list[dict[str, Any]] = [] + if model_intermediate_buffer: + request_infos = [info for info in model_intermediate_buffer if isinstance(info, dict)] + additional_information: dict[str, Any] = request_infos[0] if request_infos else {} + + device = input_ids.device + model_dtype = self.s2mel.models["gpt_layer"][0].weight.dtype + + logger.debug( + "[S2Mel forward] model_intermediate_buffer len=%d, keys=%s", + len(model_intermediate_buffer) if model_intermediate_buffer else 0, + list(additional_information.keys()) if additional_information else [], + ) + for k, v in additional_information.items(): + if isinstance(v, torch.Tensor): + logger.debug("[S2Mel forward] %s: shape=%s, dtype=%s", k, v.shape, v.dtype) + + # Extract Stage 0 outputs and pad multiple ready Stage-1 requests into + # one S2Mel batch. The runner can split list-valued multimodal outputs + # back to per-request payloads, so never silently drop batched items. + if len(request_infos) > 1: + mel_codes, latent, code_lens, code_lens_list, s_ref, ref_mel, style = self._build_batched_inputs( + request_infos, + device=device, + model_dtype=model_dtype, + ) + code_lens_raw = code_lens + else: + # Extract Stage 0 outputs and cast to model dtype + latent = self._get_tensor(additional_information, "latent", device, model_dtype) + mel_codes = self._get_tensor(additional_information, "mel_codes", device) + code_lens_raw = additional_information.get("code_lens") + s_ref = self._get_tensor(additional_information, "S_ref", device) + ref_mel = self._get_tensor(additional_information, "ref_mel", device, model_dtype) + style = self._get_tensor(additional_information, "style", device, model_dtype) + + if mel_codes is None or latent is None: + logger.warning("S2Mel decoder received empty mel_codes or latent") + return OmniOutput( + text_hidden_states=torch.zeros(1, device=device), + multimodal_outputs={"audio": torch.zeros(1, device=device), "sr": 22050}, + ) + + logger.debug( + "[S2Mel] extracted tensors — mel_codes=%s latent=%s code_lens=%s s_ref=%s ref_mel=%s style=%s", + mel_codes.shape if mel_codes is not None else None, + latent.shape if latent is not None else None, + code_lens_raw, + s_ref.shape if isinstance(s_ref, torch.Tensor) else None, + ref_mel.shape if isinstance(ref_mel, torch.Tensor) else None, + style.shape if isinstance(style, torch.Tensor) else None, + ) + + # Ensure batch dimension + if mel_codes.ndim == 1: + mel_codes = mel_codes.unsqueeze(0) + if latent.ndim == 2: + latent = latent.unsqueeze(0) + + # Compute code_lens + if code_lens_raw is not None: + if isinstance(code_lens_raw, torch.Tensor): + if "code_lens_list" not in locals(): + code_lens_list = [int(x) for x in code_lens_raw.detach().cpu().reshape(-1).tolist()] + code_lens = code_lens_raw.to(device=device, dtype=torch.long) + else: + code_lens_list = [int(x) for x in code_lens_raw] + code_lens = torch.tensor(code_lens_list, device=device, dtype=torch.long) + else: + # Strip stop token (8193) and compute actual length + stop_token = self.config.gpt.get("stop_mel_token", 8193) + code_lens_list = [] + for i in range(mel_codes.shape[0]): + stop_mask = (mel_codes[i] == stop_token).nonzero(as_tuple=False) + if stop_mask.numel() > 0: + code_lens_list.append(stop_mask[0].item()) + else: + code_lens_list.append(mel_codes.shape[1]) + code_lens = torch.tensor(code_lens_list, device=device, dtype=torch.long) + + # Trim to actual length + max_len = max(code_lens_list) + mel_codes = mel_codes[:, :max_len] + latent = latent[:, :max_len, :] + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "[S2Mel] after trim — mel_codes=%s latent=%s code_lens=%s max_len=%d", + mel_codes.shape, + latent.shape, + code_lens.tolist(), + max_len, + ) + + # --- S2Mel pipeline --- + # 1. Project GPT latent: 1280 → 1024 + latent = latent.to(device=device, dtype=model_dtype) + latent = self.s2mel.forward_gpt(latent).to(dtype=model_dtype) # [B, T, 1024] + logger.debug("[S2Mel] step1 forward_gpt → latent=%s dtype=%s", latent.shape, latent.dtype) + + # 2. Embed mel codes via semantic codec vq2emb + semantic_codec = load_semantic_codec(self.model_path, self.config.semantic_codec, device) + codebook_size = self.config.semantic_codec.get("codebook_size", 8192) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "mel_codes debug: shape=%s, min=%d, max=%d", + mel_codes.shape, + mel_codes.min().item(), + mel_codes.max().item(), + ) + if getattr(self.config, "s2mel_validate_code_range", False) and torch.any( + (mel_codes < 0) | (mel_codes >= codebook_size) + ): + raise ValueError(f"IndexTTS2 generated mel code outside semantic codebook range [0, {codebook_size - 1}]") + with torch.no_grad(): + S_infer = semantic_codec.quantizer.vq2emb(mel_codes.unsqueeze(1)) # [B, T, 1024] + S_infer = S_infer.transpose(1, 2).to(dtype=model_dtype) # [B, T, 1024] + S_infer = S_infer + latent + logger.debug("[S2Mel] step2 vq2emb → S_infer=%s dtype=%s", S_infer.shape, S_infer.dtype) + + # 3. Length regulate: codes → mel frames + target_lengths_list = [int(length * self.mel_code_to_frame_ratio) for length in code_lens_list] + target_lengths = torch.tensor(target_lengths_list, device=device, dtype=torch.long) + if logger.isEnabledFor(logging.DEBUG): + logger.debug("[S2Mel] step3 length_regulator target_lengths=%s", target_lengths_list) + cond = self.s2mel.models["length_regulator"](S_infer, ylens=target_lengths, n_quantizers=3, f0=None)[ + 0 + ] # [B, T_mel, 512] + logger.debug("[S2Mel] step3 length_regulator → cond=%s", cond.shape) + + # 4. Reference prompt conditioning + if s_ref is not None and ref_mel is not None: + if ref_mel.ndim == 2: + ref_mel = ref_mel.unsqueeze(0) + + if s_ref.ndim == 3: + # Official infer_v2.py uses the quantized embedding returned by + # semantic_codec.quantize() directly as S_ref. + s_ref_emb = s_ref.to(device=device, dtype=model_dtype) + else: + # Backward-compatible path for older payloads that carried + # codebook indices instead of quantized embeddings. + s_ref_codes = s_ref.long() + if s_ref_codes.ndim == 1: + s_ref_codes = s_ref_codes.unsqueeze(0) # [1, T] + codebook_size = self.config.semantic_codec.get("codebook_size", 8192) + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "S_ref debug: shape=%s, min=%d, max=%d, codebook_size=%d", + s_ref_codes.shape, + s_ref_codes.min().item(), + s_ref_codes.max().item(), + codebook_size, + ) + s_ref_codes = s_ref_codes.clamp(0, codebook_size - 1) + with torch.no_grad(): + s_ref_emb = semantic_codec.quantizer.vq2emb(s_ref_codes.unsqueeze(1)) # [B, 1024, T] + s_ref_emb = s_ref_emb.transpose(1, 2).to(dtype=model_dtype) # [B, T, 1024] + + ref_target_lengths = self._infer_ref_lengths(ref_mel) + prompt_condition = self.s2mel.models["length_regulator"]( + s_ref_emb, ylens=ref_target_lengths, n_quantizers=3, f0=None + )[0] + cat_condition = torch.cat([prompt_condition, cond], dim=1) + logger.debug( + "[S2Mel] step4 ref conditioning — s_ref_emb=%s prompt_condition=%s cat_condition=%s", + s_ref_emb.shape, + prompt_condition.shape, + cat_condition.shape, + ) + else: + prompt_condition = None + cat_condition = cond + ref_mel = torch.zeros(cond.shape[0], 80, 0, device=device, dtype=model_dtype) + + if style is None: + style = torch.zeros(cat_condition.shape[0], 192, device=device, dtype=model_dtype) + if style.ndim == 1: + style = style.unsqueeze(0) + + logger.debug( + "[S2Mel] step5 CFM input — cat_condition=%s ref_mel=%s style=%s steps=%d cfg_rate=%.2f", + cat_condition.shape, + ref_mel.shape, + style.shape, + self.diffusion_steps, + self.inference_cfg_rate, + ) + + # 5. Flow matching inference (Euler ODE steps, CFG rate configurable) + # CFM ODE solver runs in float32 — fp16 Euler steps diverge to noise. + # The DiT estimator itself may run under bf16 autocast (see below). + cfm = self.s2mel.models["cfm"] + cfm.estimator_autocast_dtype = torch.bfloat16 if self.s2mel_dit_bf16 else None + cond = cond.float() + if prompt_condition is not None: + prompt_condition = prompt_condition.float() + cat_condition = cat_condition.float() + ref_mel = ref_mel.float() + style = style.float() + + if prompt_condition is None: + ref_lengths = torch.zeros(cat_condition.shape[0], device=device, dtype=torch.long) + else: + ref_lengths = self._infer_ref_lengths(ref_mel) + ref_lengths_list = [int(length) for length in ref_lengths.tolist()] + if len(ref_lengths_list) == 1 and cat_condition.shape[0] > 1: + ref_lengths_list = ref_lengths_list * int(cat_condition.shape[0]) + ref_lengths = ref_lengths.expand(int(cat_condition.shape[0])) + max_ref_len = int(ref_mel.size(-1)) + ref_len = ref_lengths_list[0] if ref_lengths_list else max_ref_len + x_lens = target_lengths + ref_lengths + + estimator = cfm.estimator + required_cache_batch = max(2, int(cat_condition.shape[0]) * (2 if self.inference_cfg_rate > 0 else 1)) + if ( + estimator.transformer.freqs_cis is None + or getattr(estimator.transformer, "max_batch_size", -1) < required_cache_batch + ): + estimator.setup_caches(max_batch_size=required_cache_batch, max_seq_length=16384) + logger.info("[S2Mel] DiT caches initialized (freqs_cis + causal_mask)") + + cat_len = int(cat_condition.size(1)) + x_lens_list = [target_len + ref_len_i for target_len, ref_len_i in zip(target_lengths_list, ref_lengths_list)] + same_ref_len = len(set(ref_lengths_list)) <= 1 + same_target_len = len(set(target_lengths_list)) <= 1 + dit_graph_full_mask = same_ref_len and same_target_len and all(length == cat_len for length in x_lens_list) + if dit_graph_full_mask: + self._enable_dit_cuda_graph(estimator) + with torch.no_grad(): + mel = cfm.inference( + cat_condition, + x_lens, + ref_mel, + style, + None, # f0 + self.diffusion_steps, + inference_cfg_rate=self.inference_cfg_rate, + ) + # Strip reference portion + mel = mel[:, :, ref_len:] + else: + # Batched Stage-1 requests are padded to a shared length. The DiT + # CUDA graph fast path assumes a full mask, so rebuild each request + # with its exact reference/target lengths before enabling the graph. + mels_list: list[torch.Tensor] = [] + for i in range(cat_condition.shape[0]): + self._enable_dit_cuda_graph(estimator) + ref_len_i = ref_lengths_list[i] + target_len_i = target_lengths_list[i] + if prompt_condition is not None: + prompt_i = prompt_condition[i : i + 1, :ref_len_i] + target_i = cond[i : i + 1, :target_len_i] + cond_i = torch.cat([prompt_i, target_i], dim=1) + else: + cond_i = cat_condition[i : i + 1, :target_len_i] + x_lens_i = torch.tensor([ref_len_i + target_len_i], device=device, dtype=torch.long) + ref_mel_i = ref_mel[i : i + 1, :, :ref_len_i] if ref_mel.ndim == 3 else ref_mel[:, :ref_len_i] + style_i = style[i : i + 1] if style.ndim == 2 else style + with torch.no_grad(): + mel_i = cfm.inference( + cond_i, + x_lens_i, + ref_mel_i, + style_i, + None, + self.diffusion_steps, + inference_cfg_rate=self.inference_cfg_rate, + ) + mels_list.append(mel_i[0, :, ref_len_i:]) + mel = mels_list + dit_graph_info = getattr(getattr(estimator, "_cuda_graph_runner", None), "last_call_info", None) + mel_shape = [tuple(m.shape) for m in mel] if isinstance(mel, list) else tuple(mel.shape) + logger.debug("[S2Mel] step5 CFM done → mel=%s (stripped ref %d frames)", mel_shape, ref_len) + + # 6. BigVGAN vocoding + vocoder_cfg = self.config.vocoder + vocoder_name = vocoder_cfg.get("name", "nvidia/bigvgan_v2_22khz_80band_256x") + voc_dtype = torch.bfloat16 if (self.s2mel_vocoder_bf16 and device.type == "cuda") else torch.float32 + bigvgan = _load_bigvgan(vocoder_name, device, voc_dtype) + vocode = self._get_vocoder_graph(bigvgan, device, voc_dtype) + upsample_factor = int(math.prod(bigvgan.h.upsample_rates)) + with torch.no_grad(): + wavs: list[torch.Tensor] = [] + vocoder_infos: list[dict[str, Any]] = [] + if isinstance(mel, list): + # Per-request CFM output: variable-length mels. Batch them + # for BigVGAN by padding to the max length. + batch_size = len(mel) + if batch_size == 1: + mel_i = mel[0:1].to(voc_dtype) if mel[0].ndim == 3 else mel[0].unsqueeze(0).to(voc_dtype) + wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) + wavs.append(torch.clamp(wav_i, -1.0, 1.0)) + last_call_info = getattr(vocode, "last_call_info", None) + if isinstance(last_call_info, dict): + vocoder_infos.append(dict(last_call_info)) + else: + for i in range(batch_size): + mel_i = mel[i].unsqueeze(0).to(voc_dtype) if mel[i].ndim == 2 else mel[i : i + 1].to(voc_dtype) + wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) + wavs.append(torch.clamp(wav_i, -1.0, 1.0)) + last_call_info = getattr(vocode, "last_call_info", None) + if isinstance(last_call_info, dict): + vocoder_infos.append(dict(last_call_info)) + else: + batch_size = mel.shape[0] + if batch_size == 1: + mel_i = mel[:, :, : target_lengths_list[0]].to(voc_dtype) + wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) + wavs.append(torch.clamp(wav_i, -1.0, 1.0)) + last_call_info = getattr(vocode, "last_call_info", None) + if isinstance(last_call_info, dict): + vocoder_infos.append(dict(last_call_info)) + else: + for i in range(batch_size): + mel_i = mel[i : i + 1, :, : target_lengths_list[i]].to(voc_dtype) + wav_i = vocode(mel_i).float().squeeze(0).squeeze(0) + wavs.append(torch.clamp(wav_i, -1.0, 1.0)) + last_call_info = getattr(vocode, "last_call_info", None) + if isinstance(last_call_info, dict): + vocoder_infos.append(dict(last_call_info)) + wav: torch.Tensor | list[torch.Tensor] = wavs[0] if len(wavs) == 1 else wavs + + # Keep the public vLLM-Omni audio contract as normalized float. This is + # numerically equivalent to official infer_v2.py before its int16 save + # step, while avoiding double scaling in OpenAI/soundfile encoders. + if isinstance(wav, torch.Tensor) and wav.ndim == 0: + wav = wav.unsqueeze(0) + + if logger.isEnabledFor(logging.DEBUG): + logger.debug( + "[S2Mel] step6 BigVGAN done → wav=%s min=%.1f max=%.1f sr=22050", + wav.shape if isinstance(wav, torch.Tensor) else [w.shape for w in wav], + wav.min().item() if isinstance(wav, torch.Tensor) else min(w.min().item() for w in wav), + wav.max().item() if isinstance(wav, torch.Tensor) else max(w.max().item() for w in wav), + ) + + audio_cpu = [w.cpu() for w in wav] if isinstance(wav, list) else wav.cpu() + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "audio": audio_cpu, + "sr": torch.tensor(22050, dtype=torch.int32), + }, + ) + + def compute_logits(self, hidden_states: Any, sampling_metadata: Any = None) -> None: + return None + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _build_batched_inputs( + self, + infos: list[dict[str, Any]], + *, + device: torch.device, + model_dtype: torch.dtype, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + list[int], + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + stop_token = self.config.gpt.get("stop_mel_token", 8193) + mel_items: list[torch.Tensor] = [] + latent_items: list[torch.Tensor] = [] + code_lens: list[int] = [] + s_ref_items: list[torch.Tensor | None] = [] + ref_mel_items: list[torch.Tensor | None] = [] + style_items: list[torch.Tensor | None] = [] + + for idx, info in enumerate(infos): + mel = self._get_tensor(info, "mel_codes", device) + latent = self._get_tensor(info, "latent", device, model_dtype) + if mel is None or latent is None: + raise ValueError(f"S2Mel batched input {idx} is missing mel_codes or latent") + + mel = mel.to(device=device, dtype=torch.long).reshape(-1) + if latent.ndim == 3 and latent.shape[0] == 1: + latent = latent[0] + elif latent.ndim != 2: + latent = latent.reshape(-1, latent.shape[-1]) + + raw_len = info.get("code_lens") + if isinstance(raw_len, torch.Tensor) and raw_len.numel() > 0: + code_len = int(raw_len.reshape(-1)[0].item()) + elif raw_len is not None: + code_len = int(raw_len[0] if isinstance(raw_len, list) else raw_len) + else: + stop_mask = (mel == stop_token).nonzero(as_tuple=False) + code_len = int(stop_mask[0].item()) if stop_mask.numel() > 0 else int(mel.shape[0]) + code_len = max(0, min(code_len, int(mel.shape[0]), int(latent.shape[0]))) + mel_items.append(mel[:code_len]) + latent_items.append(latent[:code_len]) + code_lens.append(code_len) + + s_ref = self._get_tensor(info, "S_ref", device) + if isinstance(s_ref, torch.Tensor) and s_ref.ndim == 3 and s_ref.shape[0] == 1: + s_ref = s_ref[0] + s_ref_items.append(s_ref) + + ref_mel = self._get_tensor(info, "ref_mel", device, model_dtype) + if isinstance(ref_mel, torch.Tensor) and ref_mel.ndim == 3 and ref_mel.shape[0] == 1: + ref_mel = ref_mel[0] + ref_mel_items.append(ref_mel) + + style = self._get_tensor(info, "style", device, model_dtype) + if isinstance(style, torch.Tensor) and style.ndim == 2 and style.shape[0] == 1: + style = style[0] + style_items.append(style) + + max_code_len = max(code_lens) if code_lens else 0 + latent_dim = latent_items[0].shape[-1] + # Padding reaches vq2emb before length masking, so it must be a valid + # semantic-code index. Actual sequence lengths are still carried by + # code_lens/target_lengths; stop_token (8193) is outside the codebook. + mel_batch = torch.zeros((len(mel_items), max_code_len), dtype=torch.long, device=device) + latent_batch = torch.zeros(len(latent_items), max_code_len, latent_dim, device=device, dtype=model_dtype) + for i, (mel, latent) in enumerate(zip(mel_items, latent_items)): + mel_batch[i, : mel.shape[0]] = mel + latent_batch[i, : latent.shape[0]] = latent.to(dtype=model_dtype) + + s_ref_batch = self._stack_optional_sequence(s_ref_items, device=device, dtype=model_dtype) + ref_mel_batch = self._stack_optional_ref_mel(ref_mel_items, device=device, dtype=model_dtype) + style_batch = self._stack_optional_style(style_items, device=device, dtype=model_dtype) + return ( + mel_batch, + latent_batch, + torch.tensor(code_lens, device=device, dtype=torch.long), + code_lens, + s_ref_batch, + ref_mel_batch, + style_batch, + ) + + @staticmethod + def _stack_optional_sequence( + items: list[torch.Tensor | None], + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if not items or any(item is None for item in items): + return None + tensors = [item.to(device=device) for item in items if item is not None] + max_len = max(int(t.shape[0]) for t in tensors) + if tensors[0].ndim == 1: + out = torch.zeros(len(tensors), max_len, device=device, dtype=tensors[0].dtype) + for i, t in enumerate(tensors): + out[i, : t.shape[0]] = t + return out + out = torch.zeros(len(tensors), max_len, tensors[0].shape[-1], device=device, dtype=dtype) + for i, t in enumerate(tensors): + out[i, : t.shape[0], :] = t.to(dtype=dtype) + return out + + @staticmethod + def _stack_optional_ref_mel( + items: list[torch.Tensor | None], + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if not items or any(item is None for item in items): + return None + tensors = [item.to(device=device, dtype=dtype) for item in items if item is not None] + max_len = max(int(t.shape[-1]) for t in tensors) + out = torch.zeros(len(tensors), tensors[0].shape[-2], max_len, device=device, dtype=dtype) + for i, t in enumerate(tensors): + out[i, :, : t.shape[-1]] = t + return out + + @staticmethod + def _stack_optional_style( + items: list[torch.Tensor | None], + *, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor | None: + if not items or any(item is None for item in items): + return None + return torch.stack( + [item.to(device=device, dtype=dtype).reshape(-1) for item in items if item is not None], + dim=0, + ) + + @staticmethod + def _infer_ref_lengths(ref_mel: torch.Tensor) -> torch.Tensor: + if ref_mel.ndim == 2: + return torch.tensor([ref_mel.size(-1)], device=ref_mel.device, dtype=torch.long) + nonzero = ref_mel.abs().sum(dim=1) > 0 + lengths = nonzero.long().sum(dim=1) + return torch.clamp(lengths, min=1).to(device=ref_mel.device, dtype=torch.long) + + @staticmethod + def _get_tensor( + info: dict, + key: str, + device: torch.device, + dtype: torch.dtype | None = None, + ) -> torch.Tensor | None: + """Extract tensor from additional_information, handling nested dicts.""" + val = info.get(key) + if val is None: + for sub_key in ["codes", "hidden_states", "meta"]: + sub = info.get(sub_key) + if isinstance(sub, dict) and key in sub: + val = sub[key] + break + if val is None: + return None + if isinstance(val, torch.Tensor): + return val.to(device=device, dtype=dtype) if dtype else val.to(device=device) + if isinstance(val, list): + if val and isinstance(val[0], torch.Tensor): + return val[0].to(device=device, dtype=dtype) if dtype else val[0].to(device=device) + return None + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights from s2mel.pth checkpoint. + + Bypasses vllm's default .pt iterator (same reason as talker: raw + tensor .pt files in model dir). Loads s2mel.pth directly and + flattens the nested ``state["net"]`` structure: + + - ``net["cfm"][key]`` → ``s2mel.models.cfm.{key}`` + - ``net["length_regulator"][key]`` → ``s2mel.models.length_regulator.{key}`` + - ``net["gpt_layer"][key]`` → ``s2mel.models.gpt_layer.{key}`` + """ + _ = weights + + ckpt_path = resolve_model_file(self.model_path, self.config.s2mel_checkpoint) + if ckpt_path is None: + raise FileNotFoundError(f"IndexTTS2 S2Mel checkpoint {self.config.s2mel_checkpoint!r} not found") + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + net = state["net"] if "net" in state else state + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + prefix_map = { + "cfm.": "s2mel.models.cfm.", + "length_regulator.": "s2mel.models.length_regulator.", + "gpt_layer.": "s2mel.models.gpt_layer.", + } + + # Flatten nested ModuleDict state: {module: {param: tensor}} → flat iter + flat_items: Iterable[tuple[str, torch.Tensor]] + if net and isinstance(next(iter(net.values())), dict): + flat_items = ( + (f"{mod}.{param}", tensor) for mod, mod_state in net.items() for param, tensor in mod_state.items() + ) + else: + flat_items = net.items() + + for name, loaded_weight in flat_items: + mapped_name = name + for old_prefix, new_prefix in prefix_map.items(): + if name.startswith(old_prefix): + mapped_name = new_prefix + name[len(old_prefix) :] + break + + if mapped_name not in params_dict: + logger.debug("Skipping unrecognized S2Mel weight: %s → %s", name, mapped_name) + continue + + param = params_dict[mapped_name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(mapped_name) + + logger.info("Loaded %d weights for IndexTTS2S2MelDecoder", len(loaded_params)) + required_prefixes = [ + "s2mel.models.cfm.", + "s2mel.models.length_regulator.", + "s2mel.models.gpt_layer.", + ] + missing_prefixes = [ + prefix for prefix in required_prefixes if not any(name.startswith(prefix) for name in loaded_params) + ] + if missing_prefixes: + raise RuntimeError( + "IndexTTS2 S2Mel checkpoint did not load required parameter groups: " + ", ".join(missing_prefixes) + ) + + stripped = _strip_weight_norm(self.s2mel) + if stripped: + # The fold replaces weight_g/weight_v with a plain `weight` + # parameter; register it as loaded so vLLM's initialization + # check does not flag it as missing from the checkpoint. + for name in stripped: + loaded_params.add(f"s2mel.{name}") + logger.info("Folded %d weight_norm parametrizations in S2Mel", len(stripped)) + + # Keep the CFM ODE solver in float32 once after loading instead of + # recursively re-casting the module on every request. + cfm = self.s2mel.models["cfm"] + cfm.float() + cfm.estimator_autocast_dtype = torch.bfloat16 if self.s2mel_dit_bf16 else None + + self._warmup_external_models() + return loaded_params + + def _warmup_external_models(self) -> None: + """Preload vocoder + semantic codec to eliminate first-request latency.""" + device = next((p.device for p in self.s2mel.parameters()), torch.device("cpu")) + try: + vocoder_name = self.config.vocoder.get("name", "nvidia/bigvgan_v2_22khz_80band_256x") + voc_dtype = torch.bfloat16 if (self.s2mel_vocoder_bf16 and device.type == "cuda") else torch.float32 + bigvgan = _load_bigvgan(vocoder_name, device, voc_dtype) + load_semantic_codec(self.model_path, self.config.semantic_codec, device) + # Warm up the CUDA graph + torch.compile wrapper at load time so the + # ~95s compile cost is paid during startup rather than on the first + # request. Mirrors the Qwen3-TTS pattern (enable_cudagraph in + # load_weights). _get_vocoder_graph constructs the wrapper and calls + # warmup() on first access. + if self.s2mel_vocoder_cuda_graph and device.type == "cuda": + self._get_vocoder_graph(bigvgan, device, voc_dtype) + logger.info("BigVGAN + semantic codec preloaded") + except Exception as e: + logger.warning("Failed to preload Stage 1 models: %s", e) + + def _get_vocoder_graph(self, bigvgan: nn.Module, device: torch.device, dtype: torch.dtype): + """Return the CUDA-graph vocoder wrapper, or the eager model when disabled. + + Warmup is lazy (first request) following the GLM-TTS DiT wrapper + pattern; on warmup failure the wrapper is disabled and all + subsequent calls fall back to the eager model. + """ + if not self.s2mel_vocoder_cuda_graph or device.type != "cuda": + return bigvgan + if self._vocoder_graph is None: + from vllm_omni.model_executor.models.indextts2.bigvgan_cuda_graph import ( + CUDAGraphBigVGANWrapper, + ) + + self._vocoder_graph = CUDAGraphBigVGANWrapper( + bigvgan, + capture_sizes=self.s2mel_vocoder_capture_sizes, + compile_shapes=self.s2mel_vocoder_compile_shapes, + ) + try: + self._vocoder_graph.warmup(device, dtype) + except Exception: + logger.warning("Disabling BigVGAN CUDA graphs after warmup failure", exc_info=True) + self._vocoder_graph.enabled = False + return self._vocoder_graph + + def _enable_dit_cuda_graph(self, estimator: nn.Module) -> None: + """Attach a per-shape CUDA graph runner to the DiT transformer core. + + Sets `_assume_full_mask` on the attention modules: this decoder + always passes x_lens == T, so the padding mask is all-True and + dense attention is identical — the flag removes the backend's + per-step ``torch.any(~mask)`` D2H sync, which is also a + stream-capture blocker. + """ + if not self.s2mel_dit_cuda_graph: + return + if getattr(estimator, "_cuda_graph_runner", None) is not None: + return + from vllm_omni.model_executor.models.indextts2.dit_cuda_graph import CUDAGraphDiTRunner + + for layer in estimator.transformer.layers: + layer.attention._assume_full_mask = True + estimator._cuda_graph_runner = CUDAGraphDiTRunner( + estimator.transformer, + max_graphs=self.s2mel_dit_cuda_graph_max_graphs, + ) + logger.info( + "S2Mel DiT CUDA graph runner enabled (per-shape capture, LRU=%d)", + self.s2mel_dit_cuda_graph_max_graphs, + ) diff --git a/vllm_omni/model_executor/models/indextts2/indextts2_talker.py b/vllm_omni/model_executor/models/indextts2/indextts2_talker.py new file mode 100644 index 00000000000..2b1244cccb3 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/indextts2_talker.py @@ -0,0 +1,1062 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""IndexTTS2 Stage 0: GPT-2 AR Talker with vLLM-native PagedAttention. + +Predicts mel codes autoregressively and collects hidden_states as latent +for Stage 1 (S2Mel + BigVGAN). +""" + +from __future__ import annotations + +import json +import math +import re +from collections.abc import Iterable +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +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.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.gpt2 import GPT2Block +from vllm.model_executor.models.utils import ( + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) +from vllm.sequence import IntermediateTensors + +from vllm_omni.data_entry_keys import OmniPayload +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.utils.speaker_cache import get_speaker_cache + +from .configuration_indextts2 import IndexTTS2Config +from .gpt.conformer_encoder import ConformerEncoder +from .gpt.embeddings import LearnedPositionEmbeddings +from .gpt.perceiver import PerceiverResampler +from .preprocess_utils import ( + compute_fbank, + load_campplus, + load_qwen_emotion, + load_reference_audio, + load_semantic_codec, + load_wav2vec2, + resolve_model_file, + wav2vec_extract, +) + +logger = init_logger(__name__) + + +def _find_most_similar_cosine(query: torch.Tensor, matrix: torch.Tensor) -> int: + sims = F.cosine_similarity(query.float(), matrix.float(), dim=1) + return int(torch.argmax(sims).item()) + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": 0, + "intermediate_tensors": 0, + "inputs_embeds": 0, + } +) +class IndexTTS2TalkerForConditionalGeneration(nn.Module): + """vLLM-native GPT-2 AR talker for IndexTTS2. + + Stage 0 of the two-stage pipeline. Predicts mel codes (8194 vocab) + and accumulates hidden_states as latent for Stage 1 S2Mel decoder. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + self.config: IndexTTS2Config = vllm_config.model_config.hf_config # type: ignore[assignment] + gpt_cfg = self.config.gpt + + self.model_dim = gpt_cfg["model_dim"] + self.num_layers = gpt_cfg["layers"] + self.num_heads = gpt_cfg["heads"] + self.max_mel_tokens = gpt_cfg["max_mel_tokens"] + self.max_text_tokens = gpt_cfg["max_text_tokens"] + self.number_mel_codes = gpt_cfg["number_mel_codes"] + self.start_mel_token = gpt_cfg["start_mel_token"] + self.stop_mel_token = gpt_cfg["stop_mel_token"] + self.number_text_tokens = gpt_cfg.get("number_text_tokens", 12000) + self.condition_num_latent = gpt_cfg.get("condition_num_latent", 32) + + # --- Flags for vLLM-Omni framework --- + self.have_multimodal_outputs = True + self.has_preprocess = True + self.has_postprocess = True + self.requires_raw_input_tokens = True + self.enable_update_additional_information = True + # S2Mel consumes the complete code sequence only at request end, but + # the connector accumulator needs to observe each appendable AR delta. + # Keep runner request-end gating disabled and emit per-step rows from + # make_omni_output; the full-payload builder reconstructs the sequence. + self.omni_payload_at_request_end = False + self.omni_request_end_token_ids = () + self._speaker_cache = get_speaker_cache() + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("codes", "mel"), + ("hidden_states", "latent"), + ("meta", "mel_start_offset"), + ("meta", "latent_acc"), + ("meta", "mel_code_count"), + } + + # --- GPT-2 transformer (vLLM-native with PagedAttention) --- + # Build a GPT2Config-like object for GPT2Block + from transformers import GPT2Config as HFGpt2Config + + max_seq_len = self.max_mel_tokens + self.max_text_tokens + self.condition_num_latent + 4 + gpt2_config = HFGpt2Config( + vocab_size=self.number_mel_codes, + n_positions=max_seq_len, + n_ctx=max_seq_len, + n_embd=self.model_dim, + n_layer=self.num_layers, + n_head=self.num_heads, + activation_function="gelu_new", + layer_norm_epsilon=1e-5, + ) + # Store for GPT2Block construction + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + self.start_layer, self.end_layer, self.h = make_layers( + gpt2_config.n_layer, + lambda prefix: GPT2Block(gpt2_config, cache_config, quant_config, prefix=prefix), + prefix=f"{prefix}.h", + ) + self.ln_f = nn.LayerNorm(self.model_dim, eps=gpt2_config.layer_norm_epsilon) + self.final_norm = nn.LayerNorm(self.model_dim, eps=gpt2_config.layer_norm_epsilon) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states"], self.model_dim + ) + + # --- Embeddings --- + self.text_embedding = nn.Embedding(self.number_text_tokens + 1, self.model_dim) + self.mel_embedding = nn.Embedding(self.number_mel_codes, self.model_dim) + self.text_pos_embedding = LearnedPositionEmbeddings(self.max_text_tokens + 2, self.model_dim) + self.mel_pos_embedding = LearnedPositionEmbeddings(self.max_mel_tokens + 2 + 1, self.model_dim) + self.speed_emb = nn.Embedding(2, self.model_dim) + self.emo_layer = nn.Linear(self.model_dim, self.model_dim) + self.emovec_layer = nn.Linear(1024, self.model_dim) + + # --- Mel head (logits) --- + if get_pp_group().is_last_rank: + self.mel_head = ParallelLMHead( + self.number_mel_codes, + self.model_dim, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "mel_head"), + ) + else: + self.mel_head = None # type: ignore[assignment] + self.logits_processor = LogitsProcessor(self.number_mel_codes) + + # --- Conditioning modules (speaker / emotion) --- + cond_cfg = gpt_cfg.get("condition_module", {}) + emo_cond_cfg = gpt_cfg.get("emo_condition_module", {}) + + self.conditioning_encoder = ConformerEncoder( + input_size=1024, + output_size=cond_cfg.get("output_size", 512), + linear_units=cond_cfg.get("linear_units", 2048), + attention_heads=cond_cfg.get("attention_heads", 8), + num_blocks=cond_cfg.get("num_blocks", 6), + input_layer=cond_cfg.get("input_layer", "conv2d2"), + ) + self.perceiver_encoder = PerceiverResampler( + self.model_dim, + dim_context=cond_cfg.get("output_size", 512), + ff_mult=cond_cfg.get("perceiver_mult", 2), + heads=cond_cfg.get("attention_heads", 8), + num_latents=self.condition_num_latent, + ) + + self.emo_conditioning_encoder = ConformerEncoder( + input_size=1024, + output_size=emo_cond_cfg.get("output_size", 512), + linear_units=emo_cond_cfg.get("linear_units", 1024), + attention_heads=emo_cond_cfg.get("attention_heads", 4), + num_blocks=emo_cond_cfg.get("num_blocks", 4), + input_layer=emo_cond_cfg.get("input_layer", "conv2d2"), + ) + self.emo_perceiver_encoder = PerceiverResampler( + 1024, + dim_context=emo_cond_cfg.get("output_size", 512), + ff_mult=emo_cond_cfg.get("perceiver_mult", 2), + heads=emo_cond_cfg.get("attention_heads", 4), + num_latents=1, + ) + + # Padding masks for perceiver cross-attention + self.cond_mask_pad = nn.ConstantPad1d((self.condition_num_latent, 0), True) + self.emo_cond_mask_pad = nn.ConstantPad1d((1, 0), True) + + # --- Lazy-loaded external models --- + self._w2v_stat: torch.Tensor | None = None + self._emo_matrix: torch.Tensor | None = None + self._spk_matrix: torch.Tensor | None = None + self._text_tokenizer: Any = None + + # Initialize embeddings per GPT-2 convention + for emb in [self.text_embedding, self.mel_embedding]: + emb.weight.data.normal_(mean=0.0, std=0.02) + self.speed_emb.weight.data.normal_(mean=0.0, std=0.0) + + # ------------------------------------------------------------------ + # vLLM required hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.mel_embedding(input_ids) + + def _extract_mel_offsets(self, positions: torch.Tensor, kwargs: dict[str, Any]) -> torch.Tensor: + """Build per-token mel_start_offset from model_intermediate_buffer.""" + info_dicts = kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + if not info_dicts: + logger.warning_once("[OFFSET] no model_intermediate_buffer available; defaulting mel offsets to 0") + return torch.zeros_like(positions) + offsets = [] + for info in info_dicts: + if isinstance(info, dict): + meta = info.get("meta", {}) + val = meta.get("mel_start_offset", 0) + offsets.append(int(val) if not isinstance(val, int) else val) + else: + offsets.append(0) + return torch.tensor(offsets, dtype=positions.dtype, device=positions.device) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor | IntermediateTensors: + """AR transformer forward. + + During prefill: inputs_embeds is set by preprocess(), input_ids ignored. + During decode: inputs_embeds is mel_embedding(input_ids) + mel_pos. + """ + if get_pp_group().is_first_rank: + if inputs_embeds is None: + logger.warning("[FORWARD] inputs_embeds=None — fallback decode path") + mel_start_offset = self._extract_mel_offsets(positions, kwargs) + mel_pos = positions - mel_start_offset + mel_pos = torch.clamp(mel_pos, min=0) + inputs_embeds = self.mel_embedding(input_ids) + self.mel_pos_embedding.emb(mel_pos) + hidden_states = inputs_embeds + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + + for layer in self.h[self.start_layer : self.end_layer]: + hidden_states = layer(hidden_states) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + return self.final_norm(self.ln_f(hidden_states)) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: Any = None, + ) -> torch.Tensor | None: + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None or self.mel_head is None: + return None + return self.logits_processor(self.mel_head, hidden_states) + + # ------------------------------------------------------------------ + # Omni multimodal output plumbing + # ------------------------------------------------------------------ + + def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: + """Collect mel_codes and hidden_states from intermediate buffer.""" + 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") or [] + if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs: + logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer") + + device = hidden.device + latent_dim = hidden.shape[-1] + _zero_codes = torch.zeros(0, 1, dtype=torch.long, device=device) + _zero_latent = torch.zeros(0, latent_dim, dtype=hidden.dtype, device=device) + + # Per-request lists: to_payload_element dispatches lists via + # element[idx], giving correct per-request mapping even during + # mixed prefill+decode steps. Zero-length placeholders for prefill + # requests are harmless — torch.cat ignores them in accumulation. + # + # Export semantics are appendable deltas for the existing full-payload + # accumulator: codes.mel is a 2-D row [1, 1], hidden_states.latent is + # the aligned one-row [1, D] slice from meta.latent_acc. + mel_codes_list: list[torch.Tensor] = [] + latent_list: list[torch.Tensor] = [] + meta_list: list[dict[str, Any]] = [] + + for info in info_dicts: + if not isinstance(info, dict): + mel_codes_list.append(_zero_codes) + latent_list.append(_zero_latent) + meta_list.append({}) + continue + codes = info.get("codes", {}) + mc = codes.get("mel") + info_meta = info.get("meta", {}) + # mel_len now comes from the int counter in meta, not from + # code_seq.shape[0] (codes.mel is a single-token delta). + mel_len = int(info_meta.get("mel_code_count", 0)) + if isinstance(mc, torch.Tensor) and mc.numel() > 0 and mel_len > 0: + code_seq = mc.reshape(-1) + mel_codes_list.append(code_seq[-1:].reshape(1, 1).contiguous()) + else: + mel_codes_list.append(_zero_codes) + + lat_acc = info_meta.get("latent_acc") + if isinstance(lat_acc, torch.Tensor) and lat_acc.numel() > 0 and mel_len > 0: + lat_seq = lat_acc.reshape(-1, lat_acc.shape[-1]) if lat_acc.ndim != 2 else lat_acc + idx = min(max(mel_len - 1, 0), int(lat_seq.shape[0]) - 1) + latent_list.append(lat_seq[idx : idx + 1].contiguous()) + else: + hs = info.get("hidden_states", {}) + lat = hs.get("latent") + if isinstance(lat, torch.Tensor) and lat.numel() > 0: + lat_seq = lat.reshape(-1, lat.shape[-1]) if lat.ndim != 2 else lat + latent_list.append(lat_seq[-1:].contiguous()) + else: + latent_list.append(_zero_latent) + + req_meta: dict[str, Any] = {} + if mel_len == 1: + s_ref = info_meta.get("S_ref") + ref_mel = info_meta.get("ref_mel") + style = info_meta.get("style") + if isinstance(s_ref, torch.Tensor): + req_meta["S_ref"] = s_ref + if isinstance(ref_mel, torch.Tensor): + req_meta["ref_mel"] = ref_mel + if isinstance(style, torch.Tensor): + req_meta["style"] = style + meta_list.append(req_meta) + + if not any(t.numel() > 0 for t in mel_codes_list): + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + mm: OmniPayload = { + "codes": {"mel": mel_codes_list}, + "hidden_states": {"latent": latent_list}, + } + if any(meta_list): + mm["meta"] = meta_list + + return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm) + + # ------------------------------------------------------------------ + # preprocess / postprocess + # ------------------------------------------------------------------ + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build prompt embeddings for prefill; compute mel embeddings for decode. + + Prefill layout:: + + [conds(32) + emo_vec(1) + duration(2)] [text_emb + text_pos] [start_mel + mel_pos(0)] + + Decode: ``mel_embedding(token) + mel_pos_embedding(step)``. + """ + span_len = int(input_ids.shape[0]) + meta = info_dict.get("meta", {}) + + # --- Decode path: span_len==1 and mel_start_offset already set --- + if span_len == 1 and isinstance(meta, dict) and meta.get("mel_start_offset") is not None: + return self._preprocess_decode(input_ids, info_dict, meta) + + # --- Prefill path --- + return self._preprocess_prefill(input_ids, input_embeds, info_dict, span_len) + + def _preprocess_decode( + self, + input_ids: torch.Tensor, + info_dict: dict[str, Any], + meta: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Lightweight decode step: embed token + track code count.""" + # Track code count via int counter instead of accumulating a full tensor + # (the connector reconstructs the full sequence from per-step deltas). + prev_count = meta.get("mel_code_count", 0) + if not isinstance(prev_count, int): + prev_count = int(prev_count) + decode_step = prev_count + 1 + + # Only keep the current token; no O(T) history accumulation. + new_codes = input_ids.detach() + + mel_pos = input_ids.new_full((1,), decode_step) + embeds = self.mel_embedding(input_ids) + self.mel_pos_embedding.emb(mel_pos) + + return ( + input_ids, + embeds, + { + "codes": {"mel": new_codes}, + "meta": {"mel_code_count": decode_step}, + }, + ) + + def _preprocess_prefill( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + info_dict: dict[str, Any], + span_len: int, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Full conditioning pipeline — runs once during prefill.""" + additional_information = info_dict.get("additional_information") + if isinstance(additional_information, dict): + merged: dict[str, Any] = {k: v for k, v in info_dict.items() if k != "additional_information"} + for k, v in additional_information.items(): + merged.setdefault(k, v) + info_dict = merged + + device = input_ids.device + + def _first(key, default=None): + v = info_dict.get(key) + return v[0] if isinstance(v, list) and v else default + + text = _first("text") + if not text: + raise ValueError("Missing additional_information.text for IndexTTS2 talker.") + + voice_path = _first("voice") + if voice_path is None: + raise ValueError("IndexTTS2 requires voice (reference audio path) in additional_information.") + + emo_audio_path = _first("emo_audio") + emo_text = _first("emo_text") + use_emo_text = bool(_first("use_emo_text", False)) + emo_vector = _first("emo_vector") + emo_alpha = float(_first("emo_alpha", 1.0)) + use_random = bool(_first("use_random", False)) + _raw_emo_voice = _first("emo_voice_name") + emo_voice_name = str(_raw_emo_voice).strip().lower() if _raw_emo_voice else None + + # --- Speaker cache lookup --- + _voice_name = _first("voice_name") + _voice_created_at = int(_first("voice_created_at", 0)) + _speaker_cache_key = None + _cached = None + if _voice_name: + _speaker_cache_key = self._speaker_cache.make_cache_key(_voice_name, "indextts2", _voice_created_at) + _cached = self._speaker_cache.get(_speaker_cache_key) + + # --- Load audio and extract features --- + # Keep the speaker-cache hit path truly lazy: cached speaker artifacts are + # sufficient for conditioning and same-speaker emotion reuse, so avoid + # loading/resampling reference audio and touching the Wav2Vec2 singleton + # unless a miss (or an uncached separate emotion audio) actually needs it. + wav_16k: torch.Tensor | None = None + w2v_model = None + w2v_proc = None + + if _cached is not None: + s_ref = _cached["S_ref"].to(device) + style = _cached["style"].to(device) + ref_mel = _cached["ref_mel"].to(device) + spk_cond_emb = _cached["spk_cond_emb"].to(device) + else: + wav_16k, wav_22k = load_reference_audio(voice_path, device, mode="speaker") + w2v_model, w2v_proc = load_wav2vec2(self.model_path, device) + self._ensure_w2v_stat_loaded(device) + spk_cond_emb = wav2vec_extract(wav_16k, w2v_model, w2v_proc, device, self._w2v_stat) + + semantic_codec = load_semantic_codec(self.model_path, self.config.semantic_codec, device) + with torch.no_grad(): + _, s_ref = semantic_codec.quantize(spk_cond_emb) # [B, T, 1024] quantized embeddings + + campplus = load_campplus(self.model_path, device) + fbank = compute_fbank(wav_16k, device) + with torch.no_grad(): + style = campplus(fbank) # [1, 192] + + ref_mel = self._compute_mel_22k(wav_22k, device) # [1, 80, T_ref] + + if _speaker_cache_key is not None: + self._speaker_cache.put( + _speaker_cache_key, + { + "S_ref": s_ref.detach().cpu().contiguous(), + "style": style.detach().cpu().contiguous(), + "ref_mel": ref_mel.detach().cpu().contiguous(), + "spk_cond_emb": spk_cond_emb.detach().cpu().contiguous(), + }, + ) + + model_dtype = next(self.conditioning_encoder.parameters()).dtype + spk_cond_emb = spk_cond_emb.to(device=device, dtype=model_dtype) + spk_lens = torch.tensor([spk_cond_emb.shape[1]], device=device, dtype=torch.long) + speech_cond, mask = self.conditioning_encoder(spk_cond_emb, spk_lens) + conds_mask = self.cond_mask_pad(mask.squeeze(1)) + conds = self.perceiver_encoder(speech_cond, conds_mask) # [1, 32, D] + + emo_vec = self._compute_emotion_vector( + wav_16k=wav_16k, + speaker_audio_path=voice_path, + emo_audio_path=emo_audio_path, + emo_voice_name=emo_voice_name, + main_text=text, + use_emo_text=use_emo_text, + emo_text=emo_text, + emo_vector=emo_vector, + emo_alpha=emo_alpha, + use_random=use_random, + style=style, + spk_cond_emb=spk_cond_emb, + w2v_model=w2v_model, + w2v_proc=w2v_proc, + device=device, + ) + + speed_zero = torch.zeros(1, device=device, dtype=torch.long) + speed_one = torch.ones(1, device=device, dtype=torch.long) + duration_emb = self.speed_emb(speed_zero) + duration_emb_half = self.speed_emb(speed_one) + + conds_with_emo = conds + emo_vec.unsqueeze(1) # [1, 32, D] + conds_prefix = torch.cat( + [conds_with_emo, duration_emb_half.unsqueeze(1), duration_emb.unsqueeze(1)], + dim=1, + ) # [1, 34, D] + + text_tokens = self._tokenize_text(text, device) # [1, L] + text_emb = self.text_embedding(text_tokens) + self.text_pos_embedding(text_tokens) + + start_mel = self.mel_embedding(torch.tensor([[self.start_mel_token]], device=device, dtype=torch.long)) + start_mel_pos = self.mel_pos_embedding.get_fixed_embedding(0, device) + start_mel = start_mel + start_mel_pos # [1, 1, D] + + inputs_embeds = torch.cat([conds_prefix, text_emb, start_mel], dim=1).squeeze(0) + mel_start_offset = int(conds_prefix.shape[1]) + int(text_emb.shape[1]) + prefill_mel_index = int(inputs_embeds.shape[0]) - 1 + + if int(inputs_embeds.shape[0]) < span_len: + pad_n = span_len - int(inputs_embeds.shape[0]) + pad_emb = torch.zeros(pad_n, self.model_dim, device=device, dtype=inputs_embeds.dtype) + inputs_embeds = torch.cat([inputs_embeds, pad_emb], dim=0) + elif int(inputs_embeds.shape[0]) > span_len: + inputs_embeds = inputs_embeds[:span_len] + + input_ids_out = torch.full_like(input_ids, self.start_mel_token) + + info_update: dict[str, Any] = { + "meta": { + "mel_start_offset": mel_start_offset, + "prefill_mel_index": prefill_mel_index, + "mel_code_count": 0, + "S_ref": s_ref.cpu().contiguous(), + "ref_mel": ref_mel.cpu().contiguous(), + "style": style.cpu().contiguous(), + "latent_acc": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device), + }, + "codes": {"mel": torch.zeros(0, dtype=torch.long, device=device)}, + "hidden_states": {"latent": torch.zeros(0, self.model_dim, dtype=inputs_embeds.dtype, device=device)}, + } + + return input_ids_out, inputs_embeds, info_update + + def postprocess( + self, + hidden_states: torch.Tensor, + multimodal_outputs: Any = None, + **kwargs: Any, + ) -> dict[str, Any]: + """Store current decode-step hidden state for Stage 1. + + Only the current row is kept; the connector's full-payload accumulator + reconstructs the complete latent sequence from per-step deltas emitted + by make_omni_output. + """ + if hidden_states.shape[0] != 1: + meta = kwargs.get("meta", {}) + prefill_mel_index = int(meta.get("prefill_mel_index", hidden_states.shape[0] - 1)) + prefill_mel_index = min(max(prefill_mel_index, 0), hidden_states.shape[0] - 1) + return {"meta": {"latent_acc": hidden_states[prefill_mel_index : prefill_mel_index + 1].detach()}} + + return {"meta": {"latent_acc": hidden_states.detach()}} + + def _compute_mel_22k(self, wav_22k: torch.Tensor, device: torch.device) -> torch.Tensor: + """Compute 80-band mel spectrogram at 22.05kHz for Stage 1.""" + s2mel_cfg = self.config.s2mel.get("preprocess_params", {}) + sr = s2mel_cfg.get("sr", 22050) + spect = s2mel_cfg.get("spect_params", {}) + n_fft = spect.get("n_fft", 1024) + hop_length = spect.get("hop_length", 256) + win_length = spect.get("win_length", 1024) + n_mels = spect.get("n_mels", 80) + + # mel_spectrogram is pure torch with per-device caches — run on the + # wav's device (GPU after load_reference_audio) instead of CPU. + wav = wav_22k.float() + if wav.ndim == 1: + wav = wav.unsqueeze(0) + + from .s2mel.modules.audio import mel_spectrogram as mel_fn + + mel = mel_fn(wav, n_fft, n_mels, sr, hop_length, win_length, 0, None) + return mel.to(device=device) # [1, 80, T] + + # ------------------------------------------------------------------ + # Emotion conditioning + # ------------------------------------------------------------------ + + def _compute_emotion_vector( + self, + *, + wav_16k: torch.Tensor | None, + speaker_audio_path: Any, + emo_audio_path: str | None, + emo_voice_name: str | None, + main_text: str, + use_emo_text: bool, + emo_text: str | None, + emo_vector: Any, + emo_alpha: float, + use_random: bool, + style: torch.Tensor, + spk_cond_emb: torch.Tensor, + w2v_model: Any, + w2v_proc: Any, + device: torch.device, + ) -> torch.Tensor: + """Full emotion control logic aligned with official infer_v2.py. Returns [1, D].""" + # --- Mutual exclusion: text/vector guidance clears emotion reference audio --- + # Mirrors official infer_v2.py: use_emo_text is evaluated before explicit + # vectors, so text guidance wins if both are supplied by an internal caller. + if use_emo_text: + if emo_text is None: + emo_text = main_text + emo_vector = self._predict_emotion_from_text(emo_text, device) + if use_emo_text or emo_vector is not None: + emo_audio_path = None + + # --- Validate emo_vector (user-provided or text-predicted) --- + if emo_vector is not None: + if isinstance(emo_vector, torch.Tensor): + emo_vector = emo_vector.tolist() + elif isinstance(emo_vector, np.ndarray): + emo_vector = emo_vector.tolist() + if not isinstance(emo_vector, list) or len(emo_vector) != len(self._DESIRED_ORDER): + raise ValueError( + "IndexTTS2 emo_vector must contain 8 values in order: " + "happy, angry, sad, afraid, disgusted, melancholic, surprised, calm" + ) + if any(isinstance(v, bool) for v in emo_vector): + raise ValueError("IndexTTS2 emo_vector values must be numeric") + try: + emo_vector = [float(v) for v in emo_vector] + except (TypeError, ValueError) as exc: + raise ValueError("IndexTTS2 emo_vector values must be numeric") from exc + if not all(math.isfinite(v) for v in emo_vector): + raise ValueError("IndexTTS2 emo_vector values must be finite") + if not all(0.0 <= v <= 1.2 for v in emo_vector): + raise ValueError("IndexTTS2 emo_vector values must be between 0.0 and 1.2") + + # Alpha pre-scaling (official infer_v2.py:414-418): + # when emo_alpha != 1.0, scale emo_vector down before overlay. + emo_vector_scale = max(0.0, min(1.0, emo_alpha)) + if emo_vector_scale != 1.0: + emo_vector = [int(x * emo_vector_scale * 10000) / 10000 for x in emo_vector] + + # --- Path B: emo_vector (8-dim) → spk_matrix cosine lookup --- + emovec_mat = None + weight_sum = 0.0 + if emo_vector is not None: + weight_sum = sum(emo_vector) + vec = torch.tensor(emo_vector, device=device, dtype=torch.float32) + emovec_mat = self._compute_emo_from_distribution(vec, style, device, use_random=use_random) + + # --- Audio-based emotion with alpha blending --- + if emo_audio_path is None: + # No separate emotion audio: reuse speaker w2v-bert features instead of + # reloading + re-extracting the same audio (saves ~0.5s per prefill). + # Official code resamples orig→16k directly for emotion vs orig→22k→16k + # for speaker; the resulting w2v features are near-identical. + effective_alpha = 1.0 + emo_cond_emb = spk_cond_emb + else: + effective_alpha = emo_alpha + # emotion audio cache (enabled when emo_voice_name is provided) + _emo_cache_key = None + _emo_cached = None + if emo_voice_name: + _emo_cache_key = self._speaker_cache.make_cache_key(emo_voice_name, "indextts2_emo", 0) + _emo_cached = self._speaker_cache.get(_emo_cache_key) + + if _emo_cached is not None: + emo_cond_emb = _emo_cached["spk_cond_emb"].to(device) + logger.debug("[EMO] cache HIT for %s", emo_voice_name) + else: + emo_wav_16k, _ = load_reference_audio(emo_audio_path, device, mode="emotion") + if w2v_model is None or w2v_proc is None: + w2v_model, w2v_proc = load_wav2vec2(self.model_path, device) + self._ensure_w2v_stat_loaded(device) + emo_cond_emb = wav2vec_extract(emo_wav_16k, w2v_model, w2v_proc, device, self._w2v_stat) + if _emo_cache_key is not None: + self._speaker_cache.put( + _emo_cache_key, + { + "spk_cond_emb": emo_cond_emb.detach().cpu().contiguous(), + }, + ) + logger.debug("[EMO] cache MISS, cached for %s", emo_voice_name) + + emovec = self._merge_emovec(spk_cond_emb, emo_cond_emb, effective_alpha) + + # --- emo_vector overlay (official infer_v2.py:560-561) --- + # emovec_mat is raw feat2.pt weighted sum [1, model_dim], NOT projected. + # Mix: emovec_mat contributes the emotion prototype; (1-sum)*emovec + # preserves the audio-based emotion as a "base color". + if emovec_mat is not None: + emovec = emovec_mat.to(dtype=emovec.dtype) + (1.0 - weight_sum) * emovec + + return emovec + + _MELANCHOLIC_WORDS = frozenset( + { + "低落", + "melancholy", + "melancholic", + "depression", + "depressed", + "gloomy", + } + ) + + _DESIRED_ORDER = ["高兴", "愤怒", "悲伤", "恐惧", "反感", "低落", "惊讶", "自然"] + + def _predict_emotion_from_text(self, text: str, device: torch.device) -> list[float] | None: + """Use QwenEmotion CausalLM to predict 8-dim emotion vector from text (aligned with official).""" + model, tokenizer = load_qwen_emotion( + self.model_path, + device, + trust_remote_code=self.vllm_config.model_config.trust_remote_code, + ) + if model is None or tokenizer is None: + return None + + messages = [ + {"role": "system", "content": "文本情感分类"}, + {"role": "user", "content": text}, + ] + prompt_text = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + model_inputs = tokenizer([prompt_text], return_tensors="pt").to(model.device) + + with torch.no_grad(): + generated_ids = model.generate( + **model_inputs, + max_new_tokens=512, + pad_token_id=tokenizer.eos_token_id, + ) + output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist() + + try: + index = len(output_ids) - output_ids[::-1].index(151668) + except ValueError: + index = 0 + content_str = tokenizer.decode(output_ids[index:], skip_special_tokens=True) + + try: + content = json.loads(content_str) + except json.JSONDecodeError: + content = { + m.group(1): float(m.group(2)) for m in re.finditer(r'([^\s":.,]+?)"?\s*:\s*([\d.]+)', content_str) + } + + text_lower = text.lower() + if any(w in text_lower for w in self._MELANCHOLIC_WORDS): + content["悲伤"], content["低落"] = content.get("低落", 0.0), content.get("悲伤", 0.0) + + max_score, min_score = 1.2, 0.0 + emo_vector = [] + for cn_key in self._DESIRED_ORDER: + val = content.get(cn_key, 0.0) + if isinstance(val, str): + try: + val = float(val) + except ValueError: + val = 0.0 + emo_vector.append(max(min_score, min(max_score, val))) + + if all(v <= 0.0 for v in emo_vector): + emo_vector[7] = 1.0 # default calm + + return emo_vector + + def _compute_emo_from_distribution( + self, + emo_dist: torch.Tensor, + style: torch.Tensor, + device: torch.device, + use_random: bool = False, + ) -> torch.Tensor: + """From 8-dim emotion distribution + CAMPPlus style, build emovec_mat via spk_matrix cosine lookup. + + Returns [1, model_dim] (feat2.pt rows are already model_dim). + """ + self._ensure_matrix_loaded("_emo_matrix", "emo_matrix", "Emotion matrix") + self._ensure_matrix_loaded("_spk_matrix", "spk_matrix", "Speaker matrix") + if self._emo_matrix is None or self._spk_matrix is None: + return torch.zeros(1, self.model_dim, device=device) + + weight_vector = emo_dist.to(device=device, dtype=torch.float32) + if weight_vector.ndim == 2: + weight_vector = weight_vector.squeeze(0) # [8] + + if use_random: + indices = [int(torch.randint(0, sub.shape[0], (1,), device=device).item()) for sub in self._spk_matrix] + else: + indices = [_find_most_similar_cosine(style, sub.to(device)) for sub in self._spk_matrix] + emo_rows = torch.cat( + [sub.to(device)[idx].unsqueeze(0) for idx, sub in zip(indices, self._emo_matrix)], dim=0 + ) # [8, model_dim] + + return (weight_vector.unsqueeze(1) * emo_rows).sum(0).unsqueeze(0) # [1, model_dim] + + def _compute_audio_emo_vec(self, cond_emb: torch.Tensor) -> torch.Tensor: + """Conformer → Perceiver → emovec_layer → emo_layer. Input: [1, T, 1024], output: [1, D].""" + emo_dtype = next(self.emo_conditioning_encoder.parameters()).dtype + cond_emb = cond_emb.to(device=cond_emb.device, dtype=emo_dtype) + cond_lens = torch.tensor([cond_emb.shape[1]], device=cond_emb.device, dtype=torch.long) + emo_cond, emo_mask = self.emo_conditioning_encoder(cond_emb, cond_lens) + emo_conds_mask = self.emo_cond_mask_pad(emo_mask.squeeze(1)) + emo_percept = self.emo_perceiver_encoder(emo_cond, emo_conds_mask) # [1, 1, 1024] + emo_vec_syn = self.emovec_layer(emo_percept.squeeze(1)) # [1, D] + return self.emo_layer(emo_vec_syn) # [1, D] + + def _merge_emovec( + self, + spk_cond_emb: torch.Tensor, + emo_cond_emb: torch.Tensor, + alpha: float, + ) -> torch.Tensor: + """Alpha-blend speaker and emotion audio vectors: base + alpha * (emo - base).""" + if emo_cond_emb is spk_cond_emb or alpha == 1.0: + return self._compute_audio_emo_vec(emo_cond_emb) + if alpha == 0.0: + return self._compute_audio_emo_vec(spk_cond_emb) + base_vec = self._compute_audio_emo_vec(spk_cond_emb) + emo_vec = self._compute_audio_emo_vec(emo_cond_emb) + return base_vec + alpha * (emo_vec - base_vec) # [1, D] + + # ------------------------------------------------------------------ + # Text tokenizer + # ------------------------------------------------------------------ + + def _tokenize_text(self, text: str, device: torch.device) -> torch.Tensor: + """Tokenize text using the BPE tokenizer, add start/stop tokens.""" + if self._text_tokenizer is None: + from .tokenizer import IndexTTS2Tokenizer + + bpe_path = resolve_model_file(self.model_path, "bpe.model") + if bpe_path is None: + raise FileNotFoundError(f"BPE model not found in {self.model_path}") + self._text_tokenizer = IndexTTS2Tokenizer(bpe_path, model_dir=self.model_path) + + token_ids = self._text_tokenizer.encode(text, add_special_tokens=False) + + # Wrap with start/stop text tokens + start_text = 0 + stop_text = 1 + token_ids = [start_text] + token_ids + [stop_text] + return torch.tensor([token_ids], device=device, dtype=torch.long) + + # ------------------------------------------------------------------ + # Lazy loading helpers + # ------------------------------------------------------------------ + + def _ensure_w2v_stat_loaded(self, device: torch.device) -> None: + if self._w2v_stat is not None: + return + stat_path = resolve_model_file(self.model_path, self.config.w2v_stat) + if stat_path is not None: + raw = torch.load(stat_path, map_location="cpu", weights_only=True) + if isinstance(raw, dict): + self._w2v_stat = torch.stack([raw["mean"], raw["var"]]) # [2, 1024] + else: + self._w2v_stat = raw + else: + logger.warning("Wav2Vec2-BERT stat file not found in %s", self.model_path) + + def _ensure_matrix_loaded(self, attr: str, cfg_key: str, label: str) -> None: + if getattr(self, attr) is not None: + return + path = resolve_model_file(self.model_path, getattr(self.config, cfg_key)) + if path is not None: + raw = torch.load(path, map_location="cpu", weights_only=True) + setattr(self, attr, torch.split(raw, list(self.config.emo_num))) + else: + logger.warning("%s not found in %s", label, self.model_path) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights from IndexTTS2 checkpoint (gpt.pth). + + Bypasses vllm's default .pt iterator because the model directory + contains raw-tensor files (feat1.pt, feat2.pt) that are not state + dicts and would crash ``pt_weights_iterator``. + + Weight mapping from gpt.pth checkpoint → vLLM model params: + gpt.h.{i}.* → h.{i}.* (strip ``gpt.`` prefix) + gpt.ln_f.* → ln_f.* + final_norm.* → final_norm.* + text_head.* → (skipped) + <everything else> → identity + """ + _ = weights # don't iterate — raw .pt files crash the default iterator + + ckpt_path = resolve_model_file(self.model_path, self.config.gpt_checkpoint) + if ckpt_path is None: + raise FileNotFoundError(f"IndexTTS2 GPT checkpoint {self.config.gpt_checkpoint!r} not found") + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + + prefix_map = { + "gpt.h.": "h.", + "gpt.ln_f.": "ln_f.", + "final_norm.": "final_norm.", + "text_head.": "_skip_text_head.", + } + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + for name, loaded_weight in state.items(): + # Skip attention masks + if ".attn.bias" in name or ".attn.masked_bias" in name: + continue + # Skip wpe (null position embeddings) + if ".wpe." in name or ".wte." in name: + continue + + # Remap checkpoint name to model name + mapped_name = name + for old_prefix, new_prefix in prefix_map.items(): + if name.startswith(old_prefix): + mapped_name = new_prefix + name[len(old_prefix) :] + break + + # Skip unmapped / intentionally skipped weights + if mapped_name.startswith("_skip_"): + continue + + if is_pp_missing_parameter(mapped_name, self): + continue + + if mapped_name not in params_dict: + logger.debug("Skipping unrecognized weight: %s → %s", name, mapped_name) + continue + + param = params_dict[mapped_name] + + # GPT-2 Conv1D → Linear transpose + for conv1d_name in ["c_attn", "c_proj", "c_fc"]: + if conv1d_name not in mapped_name: + continue + if not mapped_name.endswith(".weight"): + continue + loaded_weight = loaded_weight.t() + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(mapped_name) + + logger.info( + "Loaded %d weights for IndexTTS2TalkerForConditionalGeneration", + len(loaded_params), + ) + required_prefixes = [ + "text_embedding.", + "mel_embedding.", + "text_pos_embedding.", + "mel_pos_embedding.", + "speed_emb.", + "emo_layer.", + "emovec_layer.", + "conditioning_encoder.", + "perceiver_encoder.", + "emo_conditioning_encoder.", + "emo_perceiver_encoder.", + "h.", + "ln_f.", + "final_norm.", + ] + if self.mel_head is not None: + required_prefixes.append("mel_head.") + missing_prefixes = [ + prefix for prefix in required_prefixes if not any(name.startswith(prefix) for name in loaded_params) + ] + if missing_prefixes: + raise RuntimeError( + "IndexTTS2 GPT checkpoint did not load required parameter groups: " + ", ".join(missing_prefixes) + ) + + # Ensure all sub-modules on correct device (some nn.Parameter created + # with torch.Tensor() end up on CPU even when model is on CUDA). + target_device = next((p.device for p in self.h.parameters()), torch.device("cpu")) + for mod in [ + self.conditioning_encoder, + self.emo_conditioning_encoder, + self.perceiver_encoder, + self.emo_perceiver_encoder, + ]: + mod.to(target_device) + + self._warmup_external_models(target_device) + return loaded_params + + def _warmup_external_models(self, device: torch.device) -> None: + """Preload external models to eliminate first-request latency.""" + try: + load_wav2vec2(self.model_path, device) + load_semantic_codec(self.model_path, self.config.semantic_codec, device) + load_campplus(self.model_path, device) + self._ensure_w2v_stat_loaded(device) + logger.info("External models preloaded: Wav2Vec2, RepCodec, CAMPPlus") + except Exception as e: + logger.warning("Failed to preload external models: %s", e) diff --git a/vllm_omni/model_executor/models/indextts2/pipeline.py b/vllm_omni/model_executor/models/indextts2/pipeline.py new file mode 100644 index 00000000000..50f866ffd2d --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/pipeline.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""IndexTTS2 pipeline: GPT AR talker (text → mel codes) → S2Mel + BigVGAN (mel → audio). + +Two-stage non-streaming pipeline. S2Mel flow matching (25 Euler steps) requires +the full mel code sequence, so async_chunk is not applicable. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.indextts2" + +INDEXTTS2_PIPELINE = PipelineConfig( + model_type="indextts2", + default_deploy_config_name="indextts2.yaml", + model_arch="IndexTTS2TalkerForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="indextts2_talker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + extras={"skip_tokenizer_init": True, "tokenizer": "gpt2"}, + custom_process_next_stage_input_func=f"{_PROC}.talker2s2mel_full_payload", + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [8193], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="indextts2_s2mel_decoder", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="IndexTTS2S2MelDecoder", + sync_process_input_func=f"{_PROC}.talker2s2mel_token_only", + extras={"skip_tokenizer_init": True, "tokenizer": "gpt2"}, + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/indextts2/preprocess_utils.py b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py new file mode 100644 index 00000000000..452bce9bb2b --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/preprocess_utils.py @@ -0,0 +1,324 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""External model loading, audio I/O, and emotion conditioning for IndexTTS2.""" + +from __future__ import annotations + +import os +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from transformers.utils.hub import cached_file +from vllm.logger import init_logger + +from vllm_omni.diffusion.model_loader.hub_prefetch import _repo_prefetch_lock + +logger = init_logger(__name__) + +# --------------------------------------------------------------------------- +# Lazy external model caches — keyed by (model_path, device) for multi-device +# --------------------------------------------------------------------------- + +_wav2vec2_cache: dict[tuple[str, str], tuple] = {} +_semantic_codec_cache: dict[tuple[str, str], Any] = {} +_campplus_cache: dict[tuple[str, str], Any] = {} +_qwen_emotion_cache: dict[tuple[str, str], tuple] = {} + + +def _freeze(model): + model.eval() + for p in model.parameters(): + p.requires_grad_(False) + return model + + +def resolve_model_file(model_path: str, filename: str) -> str | None: + """Resolve an IndexTTS2 asset from a local model dir or HF repo id.""" + local_path = os.path.join(model_path, filename) + if os.path.isfile(local_path): + return local_path + try: + return cached_file(model_path, filename) + except Exception as e: + logger.warning("Could not resolve IndexTTS2 asset %s from %s: %s", filename, model_path, e) + return None + + +def load_wav2vec2(model_path: str, device: torch.device): + cache_key = (model_path, str(device)) + if cache_key in _wav2vec2_cache: + return _wav2vec2_cache[cache_key] + from transformers import AutoFeatureExtractor, Wav2Vec2BertModel + + model, processor = None, None + # Try local subfolder first (no lock needed), then HF hub with lock. + try: + src_args = dict(pretrained_model_name_or_path=model_path, subfolder="wav2vec2bert") + processor = AutoFeatureExtractor.from_pretrained(**src_args) + model = Wav2Vec2BertModel.from_pretrained(**src_args) + except Exception: + try: + with _repo_prefetch_lock("facebook/w2v-bert-2.0"): + src_args = dict(pretrained_model_name_or_path="facebook/w2v-bert-2.0") + processor = AutoFeatureExtractor.from_pretrained(**src_args) + model = Wav2Vec2BertModel.from_pretrained(**src_args) + except Exception: + pass + if model is None or processor is None: + raise RuntimeError(f"Could not load Wav2Vec2-BERT from {model_path}/wav2vec2bert or facebook/w2v-bert-2.0") + _freeze(model.to(device=device, dtype=torch.float32)) + _wav2vec2_cache[cache_key] = (model, processor) + return model, processor + + +def load_semantic_codec(model_path: str, config: dict, device: torch.device): + cache_key = (model_path, str(device)) + if cache_key in _semantic_codec_cache: + return _semantic_codec_cache[cache_key] + from .utils.maskgct.repcodec_model import RepCodec + + codec = RepCodec( + codebook_size=config.get("codebook_size", 8192), + hidden_size=config.get("hidden_size", 1024), + codebook_dim=config.get("codebook_dim", 8), + ) + ckpt_path = resolve_model_file(model_path, "semantic_codec.pth") + if ckpt_path is not None: + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + codec.load_state_dict(state, strict=False) + else: + import safetensors.torch + from huggingface_hub import hf_hub_download + + ckpt_path = hf_hub_download("amphion/MaskGCT", filename="semantic_codec/model.safetensors") + safetensors.torch.load_model(codec, ckpt_path) + _freeze(codec.to(device=device, dtype=torch.float32)) + _semantic_codec_cache[cache_key] = codec + return codec + + +def load_campplus(model_path: str, device: torch.device): + cache_key = (model_path, str(device)) + if cache_key in _campplus_cache: + return _campplus_cache[cache_key] + from .utils.campplus.dtdnn import CAMPPlus + + campplus = CAMPPlus(feat_dim=80, embedding_size=192) + ckpt_path = resolve_model_file(model_path, "campplus.pth") + if ckpt_path is None: + from huggingface_hub import hf_hub_download + + ckpt_path = hf_hub_download("funasr/campplus", filename="campplus_cn_common.bin") + state = torch.load(ckpt_path, map_location="cpu", weights_only=True) + campplus.load_state_dict(state, strict=False) + _freeze(campplus.to(device=device, dtype=torch.float32)) + _campplus_cache[cache_key] = campplus + return campplus + + +def load_qwen_emotion(model_path: str, device: torch.device, *, trust_remote_code: bool = True): + cache_key = (model_path, str(device)) + if cache_key in _qwen_emotion_cache: + return _qwen_emotion_cache[cache_key] + from transformers import AutoModelForCausalLM, AutoTokenizer + + subfolder = "qwen0.6bemo4-merge" + try: + tokenizer = AutoTokenizer.from_pretrained( + model_path, + subfolder=subfolder, + trust_remote_code=trust_remote_code, + ) + model = AutoModelForCausalLM.from_pretrained( + model_path, + subfolder=subfolder, + torch_dtype="float16", + device_map={"": str(device)}, + trust_remote_code=trust_remote_code, + ) + except Exception as e: + logger.warning( + "Could not load QwenEmotion from %s/%s: %s. use_emo_text will fall back to default emotion.", + model_path, + subfolder, + e, + ) + return None, None + _freeze(model) + _qwen_emotion_cache[cache_key] = (model, tokenizer) + return model, tokenizer + + +# --------------------------------------------------------------------------- +# Audio utilities +# --------------------------------------------------------------------------- + + +def compute_fbank(wav_16k: torch.Tensor, device: torch.device) -> torch.Tensor: + """Compute 80-dim fbank features for CAMPPlus. Input: [T] at 16kHz.""" + try: + import torchaudio.compliance.kaldi as kaldi + + # kaldi.fbank is pure torch — keep on the wav's device (GPU when + # available) instead of forcing a CPU round-trip (~100ms for 15s). + wav = wav_16k.unsqueeze(0).float() + fbank = kaldi.fbank( + wav, + num_mel_bins=80, + dither=0, + sample_frequency=16000, + ) + fbank = fbank - fbank.mean(dim=0, keepdim=True) + return fbank.unsqueeze(0).to(device=device) # [1, T, 80] + except (ImportError, RuntimeError, OSError): + logger.warning("torchaudio not available, returning zero fbank") + return torch.zeros(1, 100, 80, device=device) + + +def wav2vec_extract( + wav_16k: torch.Tensor, + model: Any, + processor: Any, + device: torch.device, + w2v_stat: torch.Tensor | None = None, +) -> torch.Tensor: + """Extract Wav2Vec2-BERT features. Returns [1, T, 1024].""" + wav_np = wav_16k.cpu().numpy() + inputs = processor(wav_np, sampling_rate=16000, return_tensors="pt") + input_features = inputs["input_features"].to(device=device, dtype=torch.float32) + attention_mask = inputs.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(device=device) + with torch.no_grad(): + # bf16 autocast roughly halves the w2v-bert forward (~600M params, + # fp32 weights). Features are cast back to fp32 before stat + # normalization; downstream consumers are fp32. + if input_features.is_cuda: + with torch.autocast("cuda", dtype=torch.bfloat16): + outputs = model( + input_features=input_features, + attention_mask=attention_mask, + output_hidden_states=True, + ) + else: + outputs = model( + input_features=input_features, + attention_mask=attention_mask, + output_hidden_states=True, + ) + feat = outputs.hidden_states[17].float() + if w2v_stat is not None: + mean = w2v_stat[0:1, :].to(device=device) + std = torch.sqrt(w2v_stat[1:2, :].to(device=device)) + feat = (feat - mean) / std + return feat + + +def load_reference_audio( + audio_path: str | tuple | list, + device: torch.device, + max_audio_length_seconds: float | None = 15, + mode: str = "speaker", +) -> tuple[torch.Tensor, torch.Tensor]: + """Load reference audio and resample to 16kHz and 22.05kHz. + + Accepts either a file path (str) or a pre-loaded (wav_list, sr) tuple + from the serving layer. + + ``mode`` mirrors official IndexTTS2 v2: + - speaker: librosa default path first normalizes to 22.05kHz, truncates, + then derives the 16kHz wav2vec/CAMPPlus input from that 22.05kHz signal. + - emotion: librosa loads directly at 16kHz, then truncates. + """ + if mode not in {"speaker", "emotion"}: + raise ValueError("IndexTTS2 audio mode must be 'speaker' or 'emotion'") + + wav, sr = _load_audio_1d(audio_path) + # Resample on GPU when available: CPU sinc resampling for awkward ratios + # (e.g. 16000->22050, gcd=50) costs hundreds of ms; on GPU it is <5ms. + # Downstream consumers move tensors to CPU explicitly where needed. + if device.type == "cuda": + wav = wav.to(device) + if mode == "speaker": + wav_22k = _resample(wav, sr, 22050) + wav_22k = _truncate_audio(wav_22k, 22050, max_audio_length_seconds) + wav_16k = _resample(wav_22k, 22050, 16000) + return wav_16k, wav_22k + + wav_16k = _resample(wav, sr, 16000) + wav_16k = _truncate_audio(wav_16k, 16000, max_audio_length_seconds) + wav_22k = _resample(wav_16k, 16000, 22050) + return wav_16k, wav_22k + + +def _load_audio_1d(audio_path: str | tuple | list) -> tuple[torch.Tensor, int]: + if isinstance(audio_path, (list, tuple)) and len(audio_path) == 2: + wav_data, sr = audio_path + if isinstance(wav_data, np.ndarray): + wav = torch.from_numpy(wav_data).float() + elif isinstance(wav_data, list): + # np.asarray + from_numpy is ~10x faster than torch.tensor(list) + # for the long float lists the serving layer passes through msgspec. + wav = torch.from_numpy(np.asarray(wav_data, dtype=np.float32)) + elif isinstance(wav_data, torch.Tensor): + wav = wav_data.float() + else: + raise TypeError(f"Unsupported audio data type: {type(wav_data)}") + return _mono_1d(wav), int(sr) + + try: + import torchaudio + + wav, sr = torchaudio.load(audio_path) + return _mono_1d(wav, channels_first=True), int(sr) + except (ImportError, RuntimeError, OSError): + pass + + import soundfile as sf + + audio_np, sr = sf.read(audio_path, dtype="float32", always_2d=False) + return _mono_1d(torch.from_numpy(audio_np).float(), channels_first=False), int(sr) + + +def _mono_1d(wav: torch.Tensor, channels_first: bool | None = None) -> torch.Tensor: + if wav.ndim == 1: + return wav.contiguous() + if channels_first is True: + return wav.mean(dim=0).contiguous() + if channels_first is False: + return wav.mean(dim=-1).contiguous() + if wav.shape[0] <= 8 and wav.shape[-1] > wav.shape[0]: + return wav.mean(dim=0).contiguous() + return wav.mean(dim=-1).contiguous() + + +def _truncate_audio( + wav: torch.Tensor, + sample_rate: int, + max_audio_length_seconds: float | None, +) -> torch.Tensor: + if max_audio_length_seconds is None: + return wav + max_audio_samples = int(max_audio_length_seconds * sample_rate) + if max_audio_samples > 0 and wav.shape[-1] > max_audio_samples: + return wav[..., :max_audio_samples] + return wav + + +def _resample(wav: torch.Tensor, orig_sr: int, target_sr: int) -> torch.Tensor: + if orig_sr == target_sr: + return wav + try: + import torchaudio + + return torchaudio.functional.resample(wav, orig_sr, target_sr) + except (ImportError, RuntimeError, OSError): + return ( + F.interpolate( + wav.unsqueeze(0).unsqueeze(0), scale_factor=target_sr / orig_sr, mode="linear", align_corners=False + ) + .squeeze(0) + .squeeze(0) + ) diff --git a/vllm_omni/model_executor/models/indextts2/prompt_utils.py b/vllm_omni/model_executor/models/indextts2/prompt_utils.py new file mode 100644 index 00000000000..c33202785b8 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/prompt_utils.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Prompt helpers for IndexTTS2 talker prefill.""" + +from __future__ import annotations + +from functools import lru_cache + +from vllm_omni.model_executor.models.indextts2.preprocess_utils import resolve_model_file +from vllm_omni.model_executor.models.indextts2.tokenizer import IndexTTS2Tokenizer + +_CONDITIONING_PREFIX_TOKENS = 34 +_TEXT_WRAPPER_TOKENS = 2 +_START_MEL_TOKENS = 1 + + +def _resolve_bpe_model_path(model_id_or_path: str) -> str: + path = resolve_model_file(model_id_or_path, "bpe.model") + if path is None: + raise FileNotFoundError(f"Could not resolve bpe.model for {model_id_or_path!r}") + return path + + +@lru_cache(maxsize=16) +def _get_text_tokenizer(model_id_or_path: str) -> IndexTTS2Tokenizer: + return IndexTTS2Tokenizer(_resolve_bpe_model_path(model_id_or_path), model_dir=model_id_or_path) + + +def estimate_indextts2_prefill_prompt_len(model_id_or_path: str, text: str) -> int: + """Return the placeholder prompt length expected by the IndexTTS2 talker. + + The model-side prefill replaces placeholder embeddings with: + 34 conditioning tokens, text tokens wrapped with start/stop markers, and + one start-mel token. + """ + tokenizer = _get_text_tokenizer(model_id_or_path) + text_token_count = len(tokenizer.encode(text, add_special_tokens=False)) + _TEXT_WRAPPER_TOKENS + return _CONDITIONING_PREFIX_TOKENS + text_token_count + _START_MEL_TOKENS + + +def build_indextts2_prefill_prompt_ids( + model_id_or_path: str, + text: str, + *, + placeholder_token_id: int = 1, +) -> list[int]: + return [placeholder_token_id] * estimate_indextts2_prefill_prompt_len(model_id_or_path, text) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/__init__.py b/vllm_omni/model_executor/models/indextts2/s2mel/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/__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/indextts2/s2mel/modules/__init__.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/__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/indextts2/s2mel/modules/audio.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/audio.py new file mode 100644 index 00000000000..21bb45d735e --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/audio.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import torch + + +def dynamic_range_compression_torch(x, compression=1, clip_val=1e-5): + return torch.log(torch.clamp(x, min=clip_val) * compression) + + +def spectral_normalize_torch(magnitudes): + return dynamic_range_compression_torch(magnitudes) + + +def _hz_to_mel(freqs: np.ndarray | float) -> np.ndarray | float: + """librosa's default Slaney mel scale (htk=False).""" + scalar = not hasattr(freqs, "__len__") + freqs = np.atleast_1d(np.asarray(freqs, dtype=np.float64)) + f_min = 0.0 + f_sp = 200.0 / 3 + mels = (freqs - f_min) / f_sp + + min_log_hz = 1000.0 + min_log_mel = (min_log_hz - f_min) / f_sp + logstep = np.log(6.4) / 27.0 + log_t = freqs >= min_log_hz + mels[log_t] = min_log_mel + np.log(freqs[log_t] / min_log_hz) / logstep + return float(mels[0]) if scalar else mels + + +def _mel_to_hz(mels: np.ndarray | float) -> np.ndarray | float: + """Inverse of librosa's default Slaney mel scale (htk=False).""" + scalar = not hasattr(mels, "__len__") + mels = np.atleast_1d(np.asarray(mels, dtype=np.float64)) + f_min = 0.0 + f_sp = 200.0 / 3 + freqs = f_min + f_sp * mels + + min_log_hz = 1000.0 + min_log_mel = (min_log_hz - f_min) / f_sp + logstep = np.log(6.4) / 27.0 + log_t = mels >= min_log_mel + freqs[log_t] = min_log_hz * np.exp(logstep * (mels[log_t] - min_log_mel)) + return float(freqs[0]) if scalar else freqs + + +def _mel_filterbank(sr: int, n_fft: int, n_mels: int, fmin: float, fmax: float | None) -> np.ndarray: + """Create a librosa.filters.mel-compatible Slaney filterbank.""" + if fmax is None: + fmax = float(sr) / 2.0 + mels = np.linspace(_hz_to_mel(fmin), _hz_to_mel(fmax), n_mels + 2) + freqs = _mel_to_hz(mels) + fft_freqs = np.fft.rfftfreq(n_fft, d=1.0 / sr) + + fdiff = np.diff(freqs) + ramps = np.subtract.outer(freqs, fft_freqs) + weights = np.zeros((n_mels, len(fft_freqs)), dtype=np.float32) + for i in range(n_mels): + lower = -ramps[i] / fdiff[i] + upper = ramps[i + 2] / fdiff[i + 1] + weights[i] = np.maximum(0, np.minimum(lower, upper)) + + enorm = 2.0 / (freqs[2 : n_mels + 2] - freqs[:n_mels]) + weights *= enorm[:, np.newaxis] + return weights + + +mel_basis = {} +hann_window = {} + + +def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False): + global mel_basis, hann_window # pylint: disable=global-statement + if f"{sampling_rate}_{fmax}_{y.device}" not in mel_basis: + mel = _mel_filterbank(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax) + mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)] = ( + torch.from_numpy(mel).float().to(y.device) + ) + hann_window[str(sampling_rate) + "_" + str(y.device)] = torch.hann_window(win_size).to(y.device) + + y = torch.nn.functional.pad( + y.unsqueeze(1), (int((n_fft - hop_size) / 2), int((n_fft - hop_size) / 2)), mode="reflect" + ) + y = y.squeeze(1) + + spec = torch.view_as_real( + torch.stft( + y, + n_fft, + hop_length=hop_size, + win_length=win_size, + window=hann_window[str(sampling_rate) + "_" + str(y.device)], + center=center, + pad_mode="reflect", + normalized=False, + onesided=True, + return_complex=True, + ) + ) + + spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9)) + + spec = torch.matmul(mel_basis[str(sampling_rate) + "_" + str(fmax) + "_" + str(y.device)], spec) + spec = spectral_normalize_torch(spec) + + return spec diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py new file mode 100644 index 00000000000..5876bb14742 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/bigvgan.py @@ -0,0 +1,359 @@ +# Copyright (c) 2024 NVIDIA CORPORATION. +# Licensed under the MIT license. + +# Adapted from https://github.com/jik876/hifi-gan under the MIT license. +# LICENSE is in incl_licenses directory. + +import json +import os + +import torch +import torch.nn as nn +from huggingface_hub import PyTorchModelHubMixin, hf_hub_download +from torch.nn import Conv1d, ConvTranspose1d +from torch.nn.utils import remove_weight_norm, weight_norm + +from vllm_omni.model_executor.models.common.alias_free_activation import AliasFreeActivation1d +from vllm_omni.model_executor.models.common.snake_activation import Snake, SnakeBeta +from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import AttrDict + +# --------------------------------------------------------------------------- +# Helpers (inlined from env.py / utils.py) +# --------------------------------------------------------------------------- + + +def load_hparams_from_json(path) -> AttrDict: + with open(path) as f: + data = f.read() + return AttrDict(json.loads(data)) + + +def init_weights(m, mean=0.0, std=0.01): + classname = m.__class__.__name__ + if classname.find("Conv") != -1: + m.weight.data.normal_(mean, std) + + +def get_padding(kernel_size, dilation=1): + return int((kernel_size * dilation - dilation) / 2) + + +# --------------------------------------------------------------------------- +# AMP Residual Blocks +# --------------------------------------------------------------------------- + + +class AMPBlock1(torch.nn.Module): + """AMPBlock with Snake/SnakeBeta activations and extra convs2 (dilation=1).""" + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs1 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + self.convs1.apply(init_weights) + + self.convs2 = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=1, + padding=get_padding(kernel_size, 1), + ) + ) + for _ in range(len(dilation)) + ] + ) + self.convs2.apply(init_weights) + + self.num_layers = len(self.convs1) + len(self.convs2) + + if activation == "snake": + self.activations = nn.ModuleList( + [ + AliasFreeActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": + self.activations = nn.ModuleList( + [ + AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + acts1, acts2 = self.activations[::2], self.activations[1::2] + for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2): + xt = a1(x) + xt = c1(xt) + xt = a2(xt) + xt = c2(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs1: + remove_weight_norm(layer) + for layer in self.convs2: + remove_weight_norm(layer) + + +class AMPBlock2(torch.nn.Module): + """AMPBlock with Snake/SnakeBeta activations, no extra convs2.""" + + def __init__( + self, + h: AttrDict, + channels: int, + kernel_size: int = 3, + dilation: tuple = (1, 3, 5), + activation: str = None, + ): + super().__init__() + + self.h = h + + self.convs = nn.ModuleList( + [ + weight_norm( + Conv1d( + channels, + channels, + kernel_size, + stride=1, + dilation=d, + padding=get_padding(kernel_size, d), + ) + ) + for d in dilation + ] + ) + self.convs.apply(init_weights) + + self.num_layers = len(self.convs) + + if activation == "snake": + self.activations = nn.ModuleList( + [ + AliasFreeActivation1d(activation=Snake(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + elif activation == "snakebeta": + self.activations = nn.ModuleList( + [ + AliasFreeActivation1d(activation=SnakeBeta(channels, alpha_logscale=h.snake_logscale)) + for _ in range(self.num_layers) + ] + ) + else: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + def forward(self, x): + for c, a in zip(self.convs, self.activations): + xt = a(x) + xt = c(xt) + x = xt + x + + return x + + def remove_weight_norm(self): + for layer in self.convs: + remove_weight_norm(layer) + + +# --------------------------------------------------------------------------- +# BigVGAN Generator +# --------------------------------------------------------------------------- + + +class BigVGAN( + torch.nn.Module, + PyTorchModelHubMixin, + library_name="bigvgan", + repo_url="https://github.com/NVIDIA/BigVGAN", + docs_url="https://github.com/NVIDIA/BigVGAN/blob/main/README.md", + pipeline_tag="audio-to-audio", + license="mit", + tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"], +): + """BigVGAN neural vocoder with anti-aliased periodic activation.""" + + def __init__(self, h: AttrDict, use_cuda_kernel: bool = False): + super().__init__() + self.h = h + + self.num_kernels = len(h.resblock_kernel_sizes) + self.num_upsamples = len(h.upsample_rates) + + # Pre-conv + self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)) + + # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default + if h.resblock == "1": + resblock_class = AMPBlock1 + elif h.resblock == "2": + resblock_class = AMPBlock2 + else: + raise ValueError(f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}") + + # Transposed conv-based upsamplers. does not apply anti-aliasing + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)): + self.ups.append( + nn.ModuleList( + [ + weight_norm( + ConvTranspose1d( + h.upsample_initial_channel // (2**i), + h.upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + ] + ) + ) + + # Residual blocks using anti-aliased multi-periodicity composition modules (AMP) + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = h.upsample_initial_channel // (2 ** (i + 1)) + for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)): + self.resblocks.append(resblock_class(h, ch, k, d, activation=h.activation)) + + # Post-conv + activation_post = ( + Snake(ch, alpha_logscale=h.snake_logscale) + if h.activation == "snake" + else (SnakeBeta(ch, alpha_logscale=h.snake_logscale) if h.activation == "snakebeta" else None) + ) + if activation_post is None: + raise NotImplementedError( + "activation incorrectly specified. check the config file and look for 'activation'." + ) + + self.activation_post = AliasFreeActivation1d(activation=activation_post) + + # Whether to use bias for the final conv_post + self.use_bias_at_final = h.get("use_bias_at_final", True) + self.conv_post = weight_norm(Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)) + + # Weight initialization + for i in range(len(self.ups)): + self.ups[i].apply(init_weights) + self.conv_post.apply(init_weights) + + # Final tanh activation + self.use_tanh_at_final = h.get("use_tanh_at_final", True) + + def forward(self, x): + # Pre-conv + x = self.conv_pre(x) + + for i in range(self.num_upsamples): + # Upsampling + for i_up in range(len(self.ups[i])): + x = self.ups[i][i_up](x) + # AMP blocks + xs = None + for j in range(self.num_kernels): + if xs is None: + xs = self.resblocks[i * self.num_kernels + j](x) + else: + xs += self.resblocks[i * self.num_kernels + j](x) + x = xs / self.num_kernels + + # Post-conv + x = self.activation_post(x) + x = self.conv_post(x) + # Final tanh activation + if self.use_tanh_at_final: + x = torch.tanh(x) + else: + x = torch.clamp(x, min=-1.0, max=1.0) + + return x + + def remove_weight_norm(self): + try: + for up_layer in self.ups: + for sub in up_layer: + remove_weight_norm(sub) + for block in self.resblocks: + block.remove_weight_norm() + remove_weight_norm(self.conv_pre) + remove_weight_norm(self.conv_post) + except ValueError: + pass + + @classmethod + def _from_pretrained(cls, *, model_id: str, map_location: str = "cpu", **kwargs): + def _resolve(filename): + if os.path.isdir(model_id): + return os.path.join(model_id, filename) + return hf_hub_download( + repo_id=model_id, + filename=filename, + **{ + k: kwargs[k] + for k in ( + "revision", + "cache_dir", + "force_download", + "proxies", + "resume_download", + "token", + "local_files_only", + ) + if k in kwargs + }, + ) + + h = load_hparams_from_json(_resolve("config.json")) + model = cls(h) + + checkpoint_dict = torch.load(_resolve("bigvgan_generator.pt"), map_location=map_location) + try: + model.load_state_dict(checkpoint_dict["generator"], strict=False) + except RuntimeError: + model.remove_weight_norm() + model.load_state_dict(checkpoint_dict["generator"], strict=False) + + return model diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py new file mode 100644 index 00000000000..dd86f58475b --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/commons.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + + +class AttrDict(dict): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.__dict__ = self + + +def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): + n_channels_int = n_channels if isinstance(n_channels, int) else n_channels[0] + in_act = input_a + input_b + t_act_part, s_act_part = torch.split(in_act, n_channels_int, dim=1) + t_act = torch.tanh(t_act_part) + s_act = torch.sigmoid(s_act_part) + acts = t_act * s_act + return acts + + +def sequence_mask(length, max_length=None): + if max_length is None: + max_length = length.max() + x = torch.arange(max_length, dtype=length.dtype, device=length.device) + return x.unsqueeze(0) < length.unsqueeze(1) + + +class MyModel(nn.Module): + def __init__(self, args, use_emovec=False, use_gpt_latent=False): + super().__init__() + from vllm_omni.model_executor.models.indextts2.s2mel.modules.flow_matching import CFM + from vllm_omni.model_executor.models.indextts2.s2mel.modules.length_regulator import ( + InterpolateRegulator, + ) + + length_regulator = InterpolateRegulator( + channels=args.length_regulator.channels, + sampling_ratios=args.length_regulator.sampling_ratios, + is_discrete=args.length_regulator.is_discrete, + in_channels=getattr(args.length_regulator, "in_channels", None), + vector_quantize=getattr(args.length_regulator, "vector_quantize", False), + codebook_size=args.length_regulator.content_codebook_size, + n_codebooks=getattr(args.length_regulator, "n_codebooks", 1), + quantizer_dropout=getattr(args.length_regulator, "quantizer_dropout", 0.0), + f0_condition=getattr(args.length_regulator, "f0_condition", False), + n_f0_bins=getattr(args.length_regulator, "n_f0_bins", 512), + ) + + if use_gpt_latent: + self.models = nn.ModuleDict( + { + "cfm": CFM(args), + "length_regulator": length_regulator, + "gpt_layer": torch.nn.Sequential( + torch.nn.Linear(1280, 256), + torch.nn.Linear(256, 128), + torch.nn.Linear(128, 1024), + ), + } + ) + else: + self.models = nn.ModuleDict({"cfm": CFM(args), "length_regulator": length_regulator}) + + def forward_gpt(self, x): + x = self.models["gpt_layer"](x) + return x + + def enable_torch_compile(self): + if "cfm" in self.models: + self.models["cfm"].enable_torch_compile() diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py new file mode 100644 index 00000000000..f73fdc800e0 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/diffusion_transformer.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import torch +from torch import nn +from torch.nn.utils import weight_norm + +from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import sequence_mask +from vllm_omni.model_executor.models.indextts2.s2mel.modules.gpt_fast.model import ModelArgs, Transformer +from vllm_omni.model_executor.models.indextts2.s2mel.modules.wavenet import WN + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +################################################################################# +# Embedding Layers for Timesteps and Class Labels # +################################################################################# + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__(self, hidden_size, frequency_embedding_size=256): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + self.frequency_embedding_size = frequency_embedding_size + self.max_period = 10000 + self.scale = 1000 + + half = frequency_embedding_size // 2 + freqs = torch.exp(-math.log(self.max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half) + self.register_buffer("freqs", freqs) + + def timestep_embedding(self, t): + """ + Create sinusoidal timestep embeddings. + :param t: a 1-D Tensor of N indices, one per batch element. + These may be fractional. + :param dim: the dimension of the output. + :param max_period: controls the minimum frequency of the embeddings. + :return: an (N, D) Tensor of positional embeddings. + """ + # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py + + args = self.scale * t[:, None].float() * self.freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if self.frequency_embedding_size % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + t_freq = self.timestep_embedding(t) + t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype)) + return t_emb + + +class FinalLayer(nn.Module): + """ + The final layer of DiT. + """ + + def __init__(self, hidden_size, patch_size, out_channels): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + self.linear = weight_norm(nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)) + self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) + + def forward(self, x, c): + shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class DiT(torch.nn.Module): + def __init__(self, args): + super().__init__() + self.time_as_token = getattr(args.DiT, "time_as_token", False) + self.style_as_token = getattr(args.DiT, "style_as_token", False) + self.uvit_skip_connection = getattr(args.DiT, "uvit_skip_connection", False) + model_args = ModelArgs( + block_size=16384, # args.DiT.block_size, + n_layer=args.DiT.depth, + n_head=args.DiT.num_heads, + dim=args.DiT.hidden_dim, + head_dim=args.DiT.hidden_dim // args.DiT.num_heads, + vocab_size=1024, + uvit_skip_connection=self.uvit_skip_connection, + time_as_token=self.time_as_token, + ) + self.transformer = Transformer(model_args) + self.in_channels = args.DiT.in_channels + self.out_channels = args.DiT.in_channels + self.num_heads = args.DiT.num_heads + + self.x_embedder = weight_norm(nn.Linear(args.DiT.in_channels, args.DiT.hidden_dim, bias=True)) + + self.content_type = args.DiT.content_type # 'discrete' or 'continuous' + self.content_codebook_size = args.DiT.content_codebook_size # for discrete content + self.content_dim = args.DiT.content_dim # for continuous content + self.cond_projection = nn.Linear(args.DiT.content_dim, args.DiT.hidden_dim, bias=True) # continuous content + + self.is_causal = args.DiT.is_causal + + self.t_embedder = TimestepEmbedder(args.DiT.hidden_dim) + + input_pos = torch.arange(16384) + self.register_buffer("input_pos", input_pos) + + self.final_layer_type = args.DiT.final_layer_type # mlp or wavenet + if self.final_layer_type == "wavenet": + self.t_embedder2 = TimestepEmbedder(args.wavenet.hidden_dim) + self.conv1 = nn.Linear(args.DiT.hidden_dim, args.wavenet.hidden_dim) + self.conv2 = nn.Conv1d(args.wavenet.hidden_dim, args.DiT.in_channels, 1) + self.wavenet = WN( + hidden_channels=args.wavenet.hidden_dim, + kernel_size=args.wavenet.kernel_size, + dilation_rate=args.wavenet.dilation_rate, + n_layers=args.wavenet.num_layers, + gin_channels=args.wavenet.hidden_dim, + p_dropout=args.wavenet.p_dropout, + causal=False, + ) + self.final_layer = FinalLayer(args.wavenet.hidden_dim, 1, args.wavenet.hidden_dim) + self.res_projection = nn.Linear( + args.DiT.hidden_dim, args.wavenet.hidden_dim + ) # residual connection from transformer output to final output + self.wavenet_style_condition = args.wavenet.style_condition + assert args.DiT.style_condition == args.wavenet.style_condition + else: + self.final_mlp = nn.Sequential( + nn.Linear(args.DiT.hidden_dim, args.DiT.hidden_dim), + nn.SiLU(), + nn.Linear(args.DiT.hidden_dim, args.DiT.in_channels), + ) + self.transformer_style_condition = args.DiT.style_condition + + self.class_dropout_prob = args.DiT.class_dropout_prob + self.content_mask_embedder = nn.Embedding(1, args.DiT.hidden_dim) + + self.long_skip_connection = args.DiT.long_skip_connection + self.skip_linear = nn.Linear(args.DiT.hidden_dim + args.DiT.in_channels, args.DiT.hidden_dim) + + self.cond_x_merge_linear = nn.Linear( + args.DiT.hidden_dim + + args.DiT.in_channels * 2 + + args.style_encoder.dim * self.transformer_style_condition * (not self.style_as_token), + args.DiT.hidden_dim, + ) + if self.style_as_token: + self.style_in = nn.Linear(args.style_encoder.dim, args.DiT.hidden_dim) + + def setup_caches(self, max_batch_size, max_seq_length): + self.transformer.setup_caches(max_batch_size, max_seq_length, use_kv_cache=False) + + def forward(self, x, prompt_x, x_lens, t, style, cond, mask_content=False, pre_mask=None): + """ + x (torch.Tensor): random noise + prompt_x (torch.Tensor): reference mel + zero mel + shape: (batch_size, 80, 795+1068) + x_lens (torch.Tensor): mel frames output + shape: (batch_size, mel_timesteps) + t (torch.Tensor): radshape: + shape: (batch_size) + style (torch.Tensor): reference global style + shape: (batch_size, 192) + cond (torch.Tensor): semantic info of reference audio and altered audio + shape: (batch_size, mel_timesteps(795+1069), 512) + pre_mask (torch.Tensor | None): precomputed (x_mask, x_mask_expanded) to + avoid recomputation across ODE steps. Tuple of (x_mask, x_mask_expanded). + """ + class_dropout = False + if self.training and torch.rand(1) < self.class_dropout_prob: + class_dropout = True + if not self.training and mask_content: + class_dropout = True + cond_in_module = self.cond_projection + + B, _, T = x.size() + + t1 = self.t_embedder(t) # (N, D) # t1 [2, 512] + cond = cond_in_module(cond) # cond [2,1863,512]->[2,1863,512] + + x = x.transpose(1, 2) # [2,1863,80] + prompt_x = prompt_x.transpose(1, 2) # [2,1863,80] + + x_in = torch.cat([x, prompt_x, cond], dim=-1) # 80+80+512=672 [2, 1863, 672] + + if self.transformer_style_condition and not self.style_as_token: # True and True + x_in = torch.cat([x_in, style[:, None, :].expand(-1, T, -1)], dim=-1) # [2, 1863, 864] + + if class_dropout: # False + x_in[..., self.in_channels :] = x_in[..., self.in_channels :] * 0 # 80维后全置为0 + + x_in = self.cond_x_merge_linear(x_in) # (N, T, D) [2, 1863, 512] + + if self.style_as_token: # False + style = self.style_in(style) + style = torch.zeros_like(style) if class_dropout else style + x_in = torch.cat([style.unsqueeze(1), x_in], dim=1) + + if self.time_as_token: # False + x_in = torch.cat([t1.unsqueeze(1), x_in], dim=1) + + if pre_mask is not None: + x_mask, x_mask_expanded = pre_mask + else: + x_mask = ( + sequence_mask(x_lens + self.style_as_token + self.time_as_token, max_length=x_in.size(1)) + .to(x.device) + .unsqueeze(1) + ) + input_pos = self.input_pos[: x_in.size(1)] + x_mask_expanded = x_mask[:, None, :].expand(-1, -1, x_in.size(1), -1) if not self.is_causal else None + + input_pos = self.input_pos[: x_in.size(1)] # (T,) range（0，1863） + # Route through per-shape CUDA graph runner if enabled. + graph_runner = getattr(self, "_cuda_graph_runner", None) + if graph_runner is not None: + x_res = graph_runner(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded) + else: + x_res = self.transformer(x_in, t1.unsqueeze(1), input_pos, x_mask_expanded) # [2, 1863, 512] + x_res = x_res[:, 1:] if self.time_as_token else x_res + x_res = x_res[:, 1:] if self.style_as_token else x_res + + if self.long_skip_connection: # True + x_res = self.skip_linear(torch.cat([x_res, x], dim=-1)) + if self.final_layer_type == "wavenet": + x = self.conv1(x_res) + x = x.transpose(1, 2) + t2 = self.t_embedder2(t) + x = self.wavenet(x, x_mask, g=t2.unsqueeze(2)).transpose(1, 2) + self.res_projection( + x_res + ) # long residual connection + x = self.final_layer(x, t1).transpose(1, 2) + x = self.conv2(x) + else: + x = self.final_mlp(x_res) + x = x.transpose(1, 2) + # x [2,80,1863] + return x diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/encodec.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/encodec.py new file mode 100644 index 00000000000..a41141dcbc6 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/encodec.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + +"""Convolutional layers wrappers and utilities.""" + +import math +import typing as tp +import warnings + +import torch +from torch import nn +from torch.nn import functional as F +from torch.nn.utils import weight_norm + + +def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module: + if norm == "weight_norm": + return weight_norm(module) + return module + + +def get_norm_module(module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs) -> nn.Module: + return nn.Identity() + + +def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0) -> int: + length = x.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + +def pad1d(x: torch.Tensor, paddings: tuple[int, int], mode: str = "zero", value: float = 0.0): + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + if mode == "reflect": + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + else: + return F.pad(x, paddings, mode, value) + + +class NormConv1d(nn.Module): + def __init__(self, *args, causal: bool = False, norm: str = "none", norm_kwargs: dict[str, tp.Any] = {}, **kwargs): + super().__init__() + self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) + self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x): + x = self.conv(x) + x = self.norm(x) + return x + + +class SConv1d(nn.Module): + """Conv1d with builtin handling of asymmetric or causal padding and normalization.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + norm: str = "none", + norm_kwargs: dict[str, tp.Any] = {}, + pad_mode: str = "reflect", + **kwargs, + ): + super().__init__() + if stride > 1 and dilation > 1: + warnings.warn( + "SConv1d has been initialized with stride > 1 and dilation > 1" + f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." + ) + self.conv = NormConv1d( + in_channels, + out_channels, + kernel_size, + stride, + dilation=dilation, + groups=groups, + bias=bias, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + self.causal = causal + self.pad_mode = pad_mode + + def forward(self, x): + B, C, T = x.shape + kernel_size = self.conv.conv.kernel_size[0] + stride = self.conv.conv.stride[0] + dilation = self.conv.conv.dilation[0] + kernel_size = (kernel_size - 1) * dilation + 1 + padding_total = kernel_size - stride + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + if self.causal: + x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) + else: + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode) + return self.conv(x) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py new file mode 100644 index 00000000000..5192867af37 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/flow_matching.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from abc import ABC + +import torch + +from vllm_omni.model_executor.models.indextts2.s2mel.modules.diffusion_transformer import DiT + + +class BASECFM(torch.nn.Module, ABC): + def __init__( + self, + args, + ): + super().__init__() + self.sigma_min = 1e-6 + + self.estimator = None + + self.in_channels = args.DiT.in_channels + + if hasattr(args.DiT, "zero_prompt_speech_token"): + self.zero_prompt_speech_token = args.DiT.zero_prompt_speech_token + else: + self.zero_prompt_speech_token = False + + @torch.inference_mode() + def inference(self, mu, x_lens, prompt, style, f0, n_timesteps, temperature=1.0, inference_cfg_rate=0.5): + """Forward diffusion + + Args: + mu (torch.Tensor): semantic info of reference audio and altered audio + shape: (batch_size, mel_timesteps(795+1069), 512) + x_lens (torch.Tensor): mel frames output + shape: (batch_size, mel_timesteps) + prompt (torch.Tensor): reference mel + shape: (batch_size, 80, 795) + style (torch.Tensor): reference global style + shape: (batch_size, 192) + f0 (None): unused, reserved for f0 conditioning + n_timesteps (int): number of diffusion steps + temperature (float, optional): temperature for scaling noise. Defaults to 1.0. + + Returns: + sample (torch.Tensor): generated mel-spectrogram + shape: (batch_size, 80, mel_timesteps) + """ + B, T = mu.size(0), mu.size(1) + z = torch.randn([B, self.in_channels, T], device=mu.device, dtype=mu.dtype) * temperature + t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + return self.solve_euler(z, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate) + + def solve_euler(self, x, x_lens, prompt, mu, style, f0, t_span, inference_cfg_rate=0.5): + """ + Fixed euler solver for ODEs. + Args: + x (torch.Tensor): random noise + t_span (torch.Tensor): n_timesteps interpolated + shape: (n_timesteps + 1,) + mu (torch.Tensor): semantic info of reference audio and altered audio + shape: (batch_size, mel_timesteps(795+1069), 512) + x_lens (torch.Tensor): mel frames output + shape: (batch_size, mel_timesteps) + prompt (torch.Tensor): reference mel + shape: (batch_size, 80, 795) + style (torch.Tensor): reference global style + shape: (batch_size, 192) + """ + t = t_span[0] + + # Optional bf16 autocast for the DiT estimator (set by the decoder via + # `estimator_autocast_dtype`). The Euler solver state stays float32 — + # only the estimator forward runs in reduced precision, which restores + # flash-attention eligibility and halves GEMM cost. + autocast_dtype = getattr(self, "estimator_autocast_dtype", None) + + def _run_estimator(*args, **kw): + if autocast_dtype is None: + return self.estimator(*args, **kw) + with torch.autocast(x.device.type, dtype=autocast_dtype): + out = self.estimator(*args, **kw) + return out.float() + + prompt_len = prompt.size(-1) + prompt_x = torch.zeros_like(x) + prompt_x[..., :prompt_len] = prompt[..., :prompt_len] + x[..., :prompt_len] = 0 + if self.zero_prompt_speech_token: + mu[..., :prompt_len] = 0 + + # Precompute mask once — T and x_lens are constant across all ODE steps. + # x_mask [B, 1, T] broadcasts to [2B, 1, T] when CFG doubles batch. + pre_mask = self._precompute_mask(x, x_lens) + + if inference_cfg_rate > 0: + batch = x.shape[0] + cfg_prompt_x = prompt_x.new_zeros((2 * batch, *prompt_x.shape[1:])) + cfg_prompt_x[:batch].copy_(prompt_x) + cfg_style = style.new_zeros((2 * batch, *style.shape[1:])) + cfg_style[:batch].copy_(style) + cfg_mu = mu.new_zeros((2 * batch, *mu.shape[1:])) + cfg_mu[:batch].copy_(mu) + cfg_x = x.new_empty((2 * batch, *x.shape[1:])) + cfg_t = x.new_empty((2 * batch,)) + cfg_pre_mask = self._repeat_pre_mask_for_cfg(pre_mask) + + for step in range(1, len(t_span)): + dt = t_span[step] - t_span[step - 1] + if inference_cfg_rate > 0: + cfg_x[:batch].copy_(x) + cfg_x[batch:].copy_(x) + cfg_t.fill_(t) + + stacked_dphi_dt = _run_estimator( + cfg_x, + cfg_prompt_x, + x_lens, + cfg_t, + cfg_style, + cfg_mu, + pre_mask=cfg_pre_mask, + ) + + dphi_dt, cfg_dphi_dt = stacked_dphi_dt.chunk(2, dim=0) + dphi_dt = (1.0 + inference_cfg_rate) * dphi_dt - inference_cfg_rate * cfg_dphi_dt + else: + dphi_dt = _run_estimator(x, prompt_x, x_lens, t.unsqueeze(0), style, mu, pre_mask=pre_mask) + + x = x + dt * dphi_dt + t = t + dt + x[:, :, :prompt_len] = 0 + + return x + + @staticmethod + def _repeat_pre_mask_for_cfg(pre_mask): + """Duplicate precomputed masks when CFG doubles the batch. + + The original IndexTTS2 inference path used B=1, where a two-element + timestep tensor happened to match CFG. Stage-1 batching makes the + estimator input batch 2B, so all batch-indexed masks must be repeated + consistently with ``torch.cat([cond, uncond], dim=0)``. + """ + if not isinstance(pre_mask, tuple): + return pre_mask + repeated = [] + for mask in pre_mask: + if isinstance(mask, torch.Tensor): + repeated.append(torch.cat([mask, mask], dim=0)) + else: + repeated.append(mask) + return tuple(repeated) + + def _precompute_mask(self, x: torch.Tensor, x_lens: torch.Tensor): + """Precompute padding mask for DiT — reused across all ODE steps.""" + from .commons import sequence_mask + + estimator = self._eager_estimator if hasattr(self, "_eager_estimator") else self.estimator + T = x.size(2) + style_offset = getattr(estimator, "style_as_token", False) + time_offset = getattr(estimator, "time_as_token", False) + T_in = T + int(style_offset) + int(time_offset) + is_causal = getattr(estimator, "is_causal", False) + + x_mask = sequence_mask(x_lens + int(style_offset) + int(time_offset), max_length=T_in).to(x.device).unsqueeze(1) + if is_causal: + return (x_mask, None) + x_mask_expanded = x_mask[:, None, :].expand(-1, -1, T_in, -1) + return (x_mask, x_mask_expanded) + + +class CFM(BASECFM): + def __init__(self, args): + super().__init__(args) + if args.dit_type == "DiT": + self.estimator = DiT(args) + object.__setattr__(self, "_eager_estimator", self.estimator) + else: + raise NotImplementedError(f"Unknown diffusion type {args.dit_type}") + self._compiled = False + + def enable_torch_compile(self): + """Enable torch.compile optimization for the estimator model. + + This method applies torch.compile to the estimator (DiT model) for significant + performance improvements during inference. It also configures distributed + training optimizations if applicable. + """ + if torch.distributed.is_initialized(): + torch._inductor.config.reorder_for_compute_comm_overlap = True + self.estimator = torch.compile( + self._eager_estimator, + fullgraph=True, + dynamic=True, + ) + self._compiled = True + + def disable_torch_compile(self): + self.estimator = self._eager_estimator + self._compiled = False diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/__init__.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/__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/indextts2/s2mel/modules/gpt_fast/model.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py new file mode 100644 index 00000000000..c410b63878a --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/gpt_fast/model.py @@ -0,0 +1,303 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +import math +from dataclasses import dataclass + +import torch +import torch.nn as nn +from torch import Tensor +from torch.nn import functional as F + +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata +from vllm_omni.diffusion.attention.layer import Attention as DiffusionAttention +from vllm_omni.diffusion.layers.norm import RMSNorm + + +def find_multiple(n: int, k: int) -> int: + if n % k == 0: + return n + return n + k - (n % k) + + +class AdaptiveLayerNorm(nn.Module): + r"""Adaptive Layer Normalization""" + + def __init__(self, d_model, norm) -> None: + super().__init__() + self.project_layer = nn.Linear(d_model, 2 * d_model) + self.norm = norm + self.d_model = d_model + self.eps = getattr(self.norm, "eps", getattr(self.norm, "variance_epsilon", 1e-5)) + + def forward(self, input: Tensor, embedding: Tensor = None) -> Tensor: + if embedding is None: + return self.norm(input) + weight, bias = torch.split( + self.project_layer(embedding), + split_size_or_sections=self.d_model, + dim=-1, + ) + return weight * self.norm(input) + bias + + +@dataclass +class ModelArgs: + block_size: int = 2048 + vocab_size: int = 32000 + n_layer: int = 32 + n_head: int = 32 + dim: int = 4096 + intermediate_size: int = None + n_local_heads: int = -1 + head_dim: int = 64 + rope_base: float = 10000 + norm_eps: float = 1e-5 + has_cross_attention: bool = False + context_dim: int = 0 + uvit_skip_connection: bool = False + time_as_token: bool = False + + def __post_init__(self): + if self.n_local_heads == -1: + self.n_local_heads = self.n_head + if self.intermediate_size is None: + hidden_dim = 4 * self.dim + n_hidden = int(2 * hidden_dim / 3) + self.intermediate_size = find_multiple(n_hidden, 256) + + +class Transformer(nn.Module): + def __init__(self, config: ModelArgs) -> None: + super().__init__() + self.config = config + + self.layers = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layer)) + self.norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + + self.freqs_cis: Tensor | None = None + self.mask_cache: Tensor | None = None + self.max_batch_size = -1 + self.max_seq_length = -1 + + def setup_caches(self, max_batch_size, max_seq_length, use_kv_cache=False): + if self.max_seq_length >= max_seq_length and self.max_batch_size >= max_batch_size: + return + max_seq_length = find_multiple(max_seq_length, 8) + self.max_seq_length = max_seq_length + self.max_batch_size = max_batch_size + dtype = self.norm.project_layer.weight.dtype + device = self.norm.project_layer.weight.device + + self.freqs_cis = precompute_freqs_cis( + self.config.block_size, self.config.head_dim, self.config.rope_base, dtype + ).to(device) + self.causal_mask = torch.tril(torch.ones(self.max_seq_length, self.max_seq_length, dtype=torch.bool)).to(device) + self.use_kv_cache = use_kv_cache + self.uvit_skip_connection = self.config.uvit_skip_connection + if self.uvit_skip_connection: + self.layers_emit_skip = [i for i in range(self.config.n_layer) if i < self.config.n_layer // 2] + self.layers_receive_skip = [i for i in range(self.config.n_layer) if i > self.config.n_layer // 2] + else: + self.layers_emit_skip = [] + self.layers_receive_skip = [] + + def forward( + self, + x: Tensor, + c: Tensor, + input_pos: Tensor | None = None, + mask: Tensor | None = None, + context: Tensor | None = None, + context_input_pos: Tensor | None = None, + cross_attention_mask: Tensor | None = None, + ) -> Tensor: + assert self.freqs_cis is not None, "Caches must be initialized first" + if mask is None: # in case of non-causal model + if not self.training and self.use_kv_cache: + mask = self.causal_mask[None, None, input_pos] + else: + mask = self.causal_mask[None, None, input_pos] + mask = mask[..., input_pos] + freqs_cis = self.freqs_cis[input_pos] + if context is not None: + context_freqs_cis = self.freqs_cis[context_input_pos] + else: + context_freqs_cis = None + skip_in_x_list = [] + for i, layer in enumerate(self.layers): + if self.uvit_skip_connection and i in self.layers_receive_skip: + skip_in_x = skip_in_x_list.pop(-1) + else: + skip_in_x = None + x = layer(x, c, input_pos, freqs_cis, mask, context, context_freqs_cis, cross_attention_mask, skip_in_x) + if self.uvit_skip_connection and i in self.layers_emit_skip: + skip_in_x_list.append(x) + x = self.norm(x, c) + return x + + +class TransformerBlock(nn.Module): + def __init__(self, config: ModelArgs) -> None: + super().__init__() + self.attention = Attention(config) + self.feed_forward = FeedForward(config) + self.ffn_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + self.attention_norm = AdaptiveLayerNorm(config.dim, RMSNorm(config.dim, eps=config.norm_eps)) + + if config.uvit_skip_connection: + self.skip_in_linear = nn.Linear(config.dim * 2, config.dim) + self.uvit_skip_connection = True + else: + self.uvit_skip_connection = False + + self.time_as_token = config.time_as_token + + def forward( + self, + x: Tensor, + c: Tensor, + input_pos: Tensor, + freqs_cis: Tensor, + mask: Tensor, + context: Tensor | None = None, + context_freqs_cis: Tensor | None = None, + cross_attention_mask: Tensor | None = None, + skip_in_x: Tensor | None = None, + ) -> Tensor: + c = None if self.time_as_token else c + if self.uvit_skip_connection and skip_in_x is not None: + x = self.skip_in_linear(torch.cat([x, skip_in_x], dim=-1)) + h = x + self.attention(self.attention_norm(x, c), freqs_cis, mask, input_pos) + out = h + self.feed_forward(self.ffn_norm(h, c)) + return out + + +class Attention(nn.Module): + def __init__(self, config: ModelArgs, is_cross_attention: bool = False): + super().__init__() + assert config.dim % config.n_head == 0 + + total_head_dim = (config.n_head + 2 * config.n_local_heads) * config.head_dim + # key, query, value projections for all heads, but in a batch + if is_cross_attention: + self.wq = nn.Linear(config.dim, config.n_head * config.head_dim, bias=False) + self.wkv = nn.Linear(config.context_dim, 2 * config.n_local_heads * config.head_dim, bias=False) + else: + self.wqkv = nn.Linear(config.dim, total_head_dim, bias=False) + self.wo = nn.Linear(config.head_dim * config.n_head, config.dim, bias=False) + self.kv_cache = None + + self.n_head = config.n_head + self.head_dim = config.head_dim + self.n_local_heads = config.n_local_heads + self.dim = config.dim + + # vLLM-Omni diffusion attention backend. + # Auto-selects FlashAttention / SageAttention / SDPA based on + # DIFFUSION_ATTENTION_BACKEND env var or platform default. + self.diff_attn = DiffusionAttention( + num_heads=config.n_head, + head_size=config.head_dim, + softmax_scale=1.0 / math.sqrt(config.head_dim), + causal=config.time_as_token, + num_kv_heads=config.n_local_heads, + role="cross" if is_cross_attention else "self", + ) + + def forward( + self, + x: Tensor, + freqs_cis: Tensor, + mask: Tensor, + input_pos: Tensor | None = None, + context: Tensor | None = None, + context_freqs_cis: Tensor | None = None, + ) -> Tensor: + bsz, seqlen, _ = x.shape + + kv_size = self.n_local_heads * self.head_dim + if context is None: + q, k, v = self.wqkv(x).split([kv_size, kv_size, kv_size], dim=-1) + context_seqlen = seqlen + else: + q = self.wq(x) + k, v = self.wkv(context).split([kv_size, kv_size], dim=-1) + context_seqlen = context.shape[1] + + q = q.view(bsz, seqlen, self.n_head, self.head_dim) + k = k.view(bsz, context_seqlen, self.n_local_heads, self.head_dim) + v = v.view(bsz, context_seqlen, self.n_local_heads, self.head_dim) + + q = apply_rotary_emb(q, freqs_cis) + k = apply_rotary_emb(k, context_freqs_cis if context_freqs_cis is not None else freqs_cis) + + # vLLM-Omni diffusion attention backend. + # Expects (B, S, H, D) layout — skip the transpose. + # GQA handled internally via num_kv_heads init param. + # KV cache is not used by the S2Mel DiT, so skip it here. + # + # Mask format: diffusion_transformer.py passes a 4D padding mask + # (B, 1, T, T). The FlashAttention backend requires 2D (B, T). + # Squeeze to 2D for universal compatibility. + # + # Fast path: when the caller guarantees no padding (x_lens == T, + # the S2Mel CFM always runs full-length sequences), the mask is + # all-True and dense attention is mathematically identical. Skip + # building AttentionMetadata entirely — this avoids the backend's + # `torch.any(~mask)` D2H sync per ODE step and keeps the forward + # CUDA-graph capturable. + if getattr(self, "_assume_full_mask", False): + attn_metadata = None + else: + attn_mask = mask + if attn_mask is not None and attn_mask.dim() == 4: + # (B, 1, T, T) → (B, T) — take the first query row; + # all query positions share the same padding pattern. + attn_mask = attn_mask[:, 0, 0, :] + attn_metadata = AttentionMetadata(attn_mask=attn_mask) if attn_mask is not None else None + y = self.diff_attn(q, k, v, attn_metadata=attn_metadata) + y = y.flatten(2, 3) + + y = self.wo(y) + return y + + +class FeedForward(nn.Module): + def __init__(self, config: ModelArgs) -> None: + super().__init__() + self.w1 = nn.Linear(config.dim, config.intermediate_size, bias=False) + self.w3 = nn.Linear(config.dim, config.intermediate_size, bias=False) + self.w2 = nn.Linear(config.intermediate_size, config.dim, bias=False) + + def forward(self, x: Tensor) -> Tensor: + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +def precompute_freqs_cis(seq_len: int, n_elem: int, base: int = 10000, dtype: torch.dtype = torch.bfloat16) -> Tensor: + freqs = 1.0 / (base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)) + t = torch.arange(seq_len, device=freqs.device) + freqs = torch.outer(t, freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) + cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1) + return cache.to(dtype=dtype) + + +def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor: + xshaped = x.float().reshape(*x.shape[:-1], -1, 2) + freqs_cis = freqs_cis.view(1, xshaped.size(1), 1, xshaped.size(3), 2) + x_out2 = torch.stack( + [ + xshaped[..., 0] * freqs_cis[..., 0] - xshaped[..., 1] * freqs_cis[..., 1], + xshaped[..., 1] * freqs_cis[..., 0] + xshaped[..., 0] * freqs_cis[..., 1], + ], + -1, + ) + + x_out2 = x_out2.flatten(3) + return x_out2.type_as(x) diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/length_regulator.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/length_regulator.py new file mode 100644 index 00000000000..1f2757781f9 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/length_regulator.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import torch +import torch.nn as nn +from torch.nn import functional as F + +from vllm_omni.model_executor.models.indextts2.s2mel.modules.commons import sequence_mask + +# f0_bin = 256 +f0_max = 1100.0 +f0_min = 50.0 +f0_mel_min = 1127 * np.log(1 + f0_min / 700) +f0_mel_max = 1127 * np.log(1 + f0_max / 700) + + +def f0_to_coarse(f0, f0_bin): + f0_mel = 1127 * (1 + f0 / 700).log() + a = (f0_bin - 2) / (f0_mel_max - f0_mel_min) + b = f0_mel_min * a - 1.0 + f0_mel = torch.where(f0_mel > 0, f0_mel * a - b, f0_mel) + f0_coarse = torch.round(f0_mel).long() + f0_coarse = f0_coarse * (f0_coarse > 0) + f0_coarse = f0_coarse + ((f0_coarse < 1) * 1) + f0_coarse = f0_coarse * (f0_coarse < f0_bin) + f0_coarse = f0_coarse + ((f0_coarse >= f0_bin) * (f0_bin - 1)) + return f0_coarse + + +class InterpolateRegulator(nn.Module): + def __init__( + self, + channels: int, + sampling_ratios: tuple, + is_discrete: bool = False, + in_channels: int = None, # only applies to continuous input + vector_quantize: bool = False, # whether to use vector quantization, only applies to continuous input + codebook_size: int = 1024, # for discrete only + out_channels: int = None, + groups: int = 1, + n_codebooks: int = 1, # number of codebooks + quantizer_dropout: float = 0.0, # dropout for quantizer + f0_condition: bool = False, + n_f0_bins: int = 512, + ): + super().__init__() + self.sampling_ratios = sampling_ratios + out_channels = out_channels or channels + model = nn.ModuleList([]) + if len(sampling_ratios) > 0: + self.interpolate = True + for _ in sampling_ratios: + module = nn.Conv1d(channels, channels, 3, 1, 1) + norm = nn.GroupNorm(groups, channels) + act = nn.Mish() + model.extend([module, norm, act]) + else: + self.interpolate = False + model.append(nn.Conv1d(channels, out_channels, 1, 1)) + self.model = nn.Sequential(*model) + self.embedding = nn.Embedding(codebook_size, channels) + self.is_discrete = is_discrete + + self.mask_token = nn.Parameter(torch.zeros(1, channels)) + + self.n_codebooks = n_codebooks + if n_codebooks > 1: + self.extra_codebooks = nn.ModuleList( + [nn.Embedding(codebook_size, channels) for _ in range(n_codebooks - 1)] + ) + self.extra_codebook_mask_tokens = nn.ParameterList( + [nn.Parameter(torch.zeros(1, channels)) for _ in range(n_codebooks - 1)] + ) + self.quantizer_dropout = quantizer_dropout + + if f0_condition: + self.f0_embedding = nn.Embedding(n_f0_bins, channels) + self.f0_condition = f0_condition + self.n_f0_bins = n_f0_bins + self.f0_bins = torch.arange(2, 1024, 1024 // n_f0_bins) + self.f0_mask = nn.Parameter(torch.zeros(1, channels)) + else: + self.f0_condition = False + + if not is_discrete: + self.content_in_proj = nn.Linear(in_channels, channels) + if vector_quantize: + from dac.nn.quantize import VectorQuantize + + self.vq = VectorQuantize(channels, codebook_size, 8) + + def forward(self, x, ylens=None, n_quantizers=None, f0=None): + # apply token drop + if self.training: + n_quantizers = torch.ones((x.shape[0],)) * self.n_codebooks + dropout = torch.randint(1, self.n_codebooks + 1, (x.shape[0],)) + n_dropout = int(x.shape[0] * self.quantizer_dropout) + n_quantizers[:n_dropout] = dropout[:n_dropout] + n_quantizers = n_quantizers.to(x.device) + # decide whether to drop for each sample in batch + else: + n_quantizers = torch.ones((x.shape[0],), device=x.device) * ( + self.n_codebooks if n_quantizers is None else n_quantizers + ) + if self.is_discrete: + if self.n_codebooks > 1: + assert len(x.size()) == 3 + x_emb = self.embedding(x[:, 0]) + for i, emb in enumerate(self.extra_codebooks): + x_emb = x_emb + (n_quantizers > i + 1)[..., None, None] * emb(x[:, i + 1]) + # add mask token if not using this codebook + # x_emb = x_emb + (n_quantizers <= i+1)[..., None, None] * self.extra_codebook_mask_tokens[i] + x = x_emb + elif self.n_codebooks == 1: + if len(x.size()) == 2: + x = self.embedding(x) + else: + x = self.embedding(x[:, 0]) + else: + x = self.content_in_proj(x) + # x in (B, T, D) + mask = sequence_mask(ylens).unsqueeze(-1) + if self.interpolate: + x = F.interpolate(x.transpose(1, 2).contiguous(), size=ylens.max(), mode="nearest") + else: + x = x.transpose(1, 2).contiguous() + mask = mask[:, : x.size(2), :] + ylens = ylens.clamp(max=x.size(2)).long() + if self.f0_condition: + if f0 is None: + x = x + self.f0_mask.unsqueeze(-1) + else: + # quantized_f0 = torch.bucketize(f0, self.f0_bins.to(f0.device)) # (N, T) + quantized_f0 = f0_to_coarse(f0, self.n_f0_bins) + quantized_f0 = quantized_f0.clamp(0, self.n_f0_bins - 1).long() + f0_emb = self.f0_embedding(quantized_f0) + f0_emb = F.interpolate(f0_emb.transpose(1, 2).contiguous(), size=ylens.max(), mode="nearest") + x = x + f0_emb + out = self.model(x).transpose(1, 2).contiguous() + if hasattr(self, "vq"): + ( + out_q, + commitment_loss, + codebook_loss, + codes, + out, + ) = self.vq(out.transpose(1, 2)) + out_q = out_q.transpose(1, 2) + return out_q * mask, ylens, codes, commitment_loss, codebook_loss + olens = ylens + return out * mask, olens, None, None, None diff --git a/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py b/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py new file mode 100644 index 00000000000..4c5381766ef --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/s2mel/modules/wavenet.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm_omni.model_executor.models.indextts2.s2mel.modules.encodec import SConv1d + +from . import commons + + +class WN(torch.nn.Module): + def __init__( + self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0, causal=False + ): + super().__init__() + conv1d_type = SConv1d + assert kernel_size % 2 == 1 + self.hidden_channels = hidden_channels + self.kernel_size = (kernel_size,) + self.dilation_rate = dilation_rate + self.n_layers = n_layers + self.gin_channels = gin_channels + self.p_dropout = p_dropout + + self.in_layers = torch.nn.ModuleList() + self.res_skip_layers = torch.nn.ModuleList() + self.drop = nn.Dropout(p_dropout) + + if gin_channels != 0: + self.cond_layer = conv1d_type(gin_channels, 2 * hidden_channels * n_layers, 1, norm="weight_norm") + + for i in range(n_layers): + dilation = dilation_rate**i + padding = int((kernel_size * dilation - dilation) / 2) + in_layer = conv1d_type( + hidden_channels, + 2 * hidden_channels, + kernel_size, + dilation=dilation, + padding=padding, + norm="weight_norm", + causal=causal, + ) + self.in_layers.append(in_layer) + + # last one is not necessary + if i < n_layers - 1: + res_skip_channels = 2 * hidden_channels + else: + res_skip_channels = hidden_channels + + res_skip_layer = conv1d_type(hidden_channels, res_skip_channels, 1, norm="weight_norm", causal=causal) + self.res_skip_layers.append(res_skip_layer) + + def forward(self, x, x_mask, g=None, **kwargs): + output = torch.zeros_like(x) + + if g is not None: + g = self.cond_layer(g) + + for i in range(self.n_layers): + x_in = self.in_layers[i](x) + if g is not None: + cond_offset = i * 2 * self.hidden_channels + g_l = g[:, cond_offset : cond_offset + 2 * self.hidden_channels, :] + else: + g_l = torch.zeros_like(x_in) + + acts = commons.fused_add_tanh_sigmoid_multiply(x_in, g_l, self.hidden_channels) + acts = self.drop(acts) + + res_skip_acts = self.res_skip_layers[i](acts) + if i < self.n_layers - 1: + res_acts = res_skip_acts[:, : self.hidden_channels, :] + x = (x + res_acts) * x_mask + output = output + res_skip_acts[:, self.hidden_channels :, :] + else: + output = output + res_skip_acts + return output * x_mask + + def remove_weight_norm(self): + if self.gin_channels != 0: + torch.nn.utils.remove_weight_norm(self.cond_layer) + for conv_layer in self.in_layers: + torch.nn.utils.remove_weight_norm(conv_layer) + for conv_layer in self.res_skip_layers: + torch.nn.utils.remove_weight_norm(conv_layer) diff --git a/vllm_omni/model_executor/models/indextts2/tokenizer.py b/vllm_omni/model_executor/models/indextts2/tokenizer.py new file mode 100644 index 00000000000..be20d02be93 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/tokenizer.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os + +from transformers import PreTrainedTokenizer +from vllm.logger import init_logger + +from vllm_omni.model_executor.models.indextts2.utils.front import TextTokenizer + +logger = init_logger(__name__) + + +class IndexTTS2Tokenizer(PreTrainedTokenizer): + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): + vocab_file = kwargs.pop("vocab_file", None) + if vocab_file is None: + vocab_file = os.path.join(pretrained_model_name_or_path, "bpe.model") + logger.info(f"IndexTTS2Tokenizer.from_pretrained: {pretrained_model_name_or_path}") + return cls(vocab_file, **kwargs) + + def __init__(self, vocab_file: str, **kwargs): + logger.info(f"IndexTTS2Tokenizer.__init__: vocab_file={vocab_file}") + self.vocab_file = vocab_file + self._tok = TextTokenizer(vocab_file) + logger.info("IndexTTS2Tokenizer initialized successfully") + super().__init__(**kwargs) + + @property + def vocab_size(self): + return self._tok.vocab_size + + def get_vocab(self): + return self._tok.get_vocab() + + def _tokenize(self, text): + return self._tok.tokenize(text) + + def _convert_token_to_id(self, token): + return self._tok.convert_tokens_to_ids(token)[0] + + def convert_tokens_to_string(self, tokens): + return self._tok.decode(self._tok.convert_tokens_to_ids(tokens)) + + def split_segments( + self, + tokenized: list[str], + max_text_tokens_per_segment: int = 120, + quick_streaming_tokens: int = 0, + ) -> list[list[str]]: + return self._tok.split_segments( + tokenized, + max_text_tokens_per_segment=max_text_tokens_per_segment, + quick_streaming_tokens=quick_streaming_tokens, + ) diff --git a/vllm_omni/model_executor/models/indextts2/utils/__init__.py b/vllm_omni/model_executor/models/indextts2/utils/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/__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/indextts2/utils/campplus/__init__.py b/vllm_omni/model_executor/models/indextts2/utils/campplus/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/campplus/__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/indextts2/utils/campplus/dtdnn.py b/vllm_omni/model_executor/models/indextts2/utils/campplus/dtdnn.py new file mode 100644 index 00000000000..71766d04afb --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/campplus/dtdnn.py @@ -0,0 +1,117 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +from collections import OrderedDict + +import torch.nn.functional as F +from torch import nn + +from .layers import ( + BasicResBlock, + CAMDenseTDNNBlock, + DenseLayer, + StatsPool, + TDNNLayer, + TransitLayer, + get_nonlinear, +) + + +class FCM(nn.Module): + def __init__(self, block=BasicResBlock, num_blocks=[2, 2], m_channels=32, feat_dim=80): + super().__init__() + self.in_planes = m_channels + self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(m_channels) + + self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=2) + self.layer2 = self._make_layer(block, m_channels, num_blocks[1], stride=2) + + self.conv2 = nn.Conv2d(m_channels, m_channels, kernel_size=3, stride=(2, 1), padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(m_channels) + self.out_channels = m_channels * (feat_dim // 8) + + def _make_layer(self, block, planes, num_blocks, stride): + strides = [stride] + [1] * (num_blocks - 1) + layers = [] + for stride in strides: + layers.append(block(self.in_planes, planes, stride)) + self.in_planes = planes * block.expansion + return nn.Sequential(*layers) + + def forward(self, x): + x = x.unsqueeze(1) + out = F.relu(self.bn1(self.conv1(x))) + out = self.layer1(out) + out = self.layer2(out) + out = F.relu(self.bn2(self.conv2(out))) + + shape = out.shape + out = out.reshape(shape[0], shape[1] * shape[2], shape[3]) + return out + + +class CAMPPlus(nn.Module): + def __init__( + self, + feat_dim=80, + embedding_size=512, + growth_rate=32, + bn_size=4, + init_channels=128, + config_str="batchnorm-relu", + memory_efficient=True, + ): + super().__init__() + + self.head = FCM(feat_dim=feat_dim) + channels = self.head.out_channels + + self.xvector = nn.Sequential( + OrderedDict( + [ + ( + "tdnn", + TDNNLayer(channels, init_channels, 5, stride=2, dilation=1, padding=-1, config_str=config_str), + ), + ] + ) + ) + channels = init_channels + for i, (num_layers, kernel_size, dilation) in enumerate(zip((12, 24, 16), (3, 3, 3), (1, 2, 2))): + block = CAMDenseTDNNBlock( + num_layers=num_layers, + in_channels=channels, + out_channels=growth_rate, + bn_channels=bn_size * growth_rate, + kernel_size=kernel_size, + dilation=dilation, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.xvector.add_module(f"block{i + 1}", block) + channels = channels + num_layers * growth_rate + self.xvector.add_module( + f"transit{i + 1}", TransitLayer(channels, channels // 2, bias=False, config_str=config_str) + ) + channels //= 2 + + self.xvector.add_module("out_nonlinear", get_nonlinear(config_str, channels)) + + self.xvector.add_module("stats", StatsPool()) + self.xvector.add_module("dense", DenseLayer(channels * 2, embedding_size, config_str="batchnorm_")) + + for m in self.modules(): + if isinstance(m, (nn.Conv1d, nn.Linear)): + nn.init.kaiming_normal_(m.weight.data) + if m.bias is not None: + nn.init.zeros_(m.bias) + + def forward(self, x): + x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T) + x = self.head(x) + x = self.xvector(x) + return x diff --git a/vllm_omni/model_executor/models/indextts2/utils/campplus/layers.py b/vllm_omni/model_executor/models/indextts2/utils/campplus/layers.py new file mode 100644 index 00000000000..4919cf16e68 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/campplus/layers.py @@ -0,0 +1,221 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0) + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint as cp +from torch import nn + + +def get_nonlinear(config_str, channels): + nonlinear = nn.Sequential() + for name in config_str.split("-"): + if name == "relu": + nonlinear.add_module("relu", nn.ReLU(inplace=True)) + elif name == "prelu": + nonlinear.add_module("prelu", nn.PReLU(channels)) + elif name == "batchnorm": + nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels)) + elif name == "batchnorm_": + nonlinear.add_module("batchnorm", nn.BatchNorm1d(channels, affine=False)) + else: + raise ValueError(f"Unexpected module ({name}).") + return nonlinear + + +def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2): + mean = x.mean(dim=dim) + std = x.std(dim=dim, unbiased=unbiased) + stats = torch.cat([mean, std], dim=-1) + if keepdim: + stats = stats.unsqueeze(dim=dim) + return stats + + +class StatsPool(nn.Module): + def forward(self, x): + return statistics_pooling(x) + + +class TDNNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + kernel_size, + stride=1, + padding=0, + dilation=1, + bias=False, + config_str="batchnorm-relu", + ): + super().__init__() + if padding < 0: + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.linear = nn.Conv1d( + in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + x = self.linear(x) + x = self.nonlinear(x) + return x + + +class CAMLayer(nn.Module): + def __init__(self, bn_channels, out_channels, kernel_size, stride, padding, dilation, bias, reduction=2): + super().__init__() + self.linear_local = nn.Conv1d( + bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1) + self.relu = nn.ReLU(inplace=True) + self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + y = self.linear_local(x) + context = x.mean(-1, keepdim=True) + self.seg_pooling(x) + context = self.relu(self.linear1(context)) + m = self.sigmoid(self.linear2(context)) + return y * m + + def seg_pooling(self, x, seg_len=100, stype="avg"): + if stype == "avg": + seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + elif stype == "max": + seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True) + else: + raise ValueError("Wrong segment pooling type.") + shape = seg.shape + seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1) + seg = seg[..., : x.shape[-1]] + return seg + + +class CAMDenseTDNNLayer(nn.Module): + def __init__( + self, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + assert kernel_size % 2 == 1, f"Expect equal paddings, but got even kernel size ({kernel_size})" + padding = (kernel_size - 1) // 2 * dilation + self.memory_efficient = memory_efficient + self.nonlinear1 = get_nonlinear(config_str, in_channels) + self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False) + self.nonlinear2 = get_nonlinear(config_str, bn_channels) + self.cam_layer = CAMLayer( + bn_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias + ) + + def bn_function(self, x): + return self.linear1(self.nonlinear1(x)) + + def forward(self, x): + if self.training and self.memory_efficient: + x = cp.checkpoint(self.bn_function, x) + else: + x = self.bn_function(x) + x = self.cam_layer(self.nonlinear2(x)) + return x + + +class CAMDenseTDNNBlock(nn.ModuleList): + def __init__( + self, + num_layers, + in_channels, + out_channels, + bn_channels, + kernel_size, + stride=1, + dilation=1, + bias=False, + config_str="batchnorm-relu", + memory_efficient=False, + ): + super().__init__() + for i in range(num_layers): + layer = CAMDenseTDNNLayer( + in_channels=in_channels + i * out_channels, + out_channels=out_channels, + bn_channels=bn_channels, + kernel_size=kernel_size, + stride=stride, + dilation=dilation, + bias=bias, + config_str=config_str, + memory_efficient=memory_efficient, + ) + self.add_module(f"tdnnd{i + 1}", layer) + + def forward(self, x): + for layer in self: + x = torch.cat([x, layer(x)], dim=1) + return x + + +class TransitLayer(nn.Module): + def __init__(self, in_channels, out_channels, bias=True, config_str="batchnorm-relu"): + super().__init__() + self.nonlinear = get_nonlinear(config_str, in_channels) + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + + def forward(self, x): + x = self.nonlinear(x) + x = self.linear(x) + return x + + +class DenseLayer(nn.Module): + def __init__(self, in_channels, out_channels, bias=False, config_str="batchnorm-relu"): + super().__init__() + self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias) + self.nonlinear = get_nonlinear(config_str, out_channels) + + def forward(self, x): + if len(x.shape) == 2: + x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1) + else: + x = self.linear(x) + x = self.nonlinear(x) + return x + + +class BasicResBlock(nn.Module): + expansion = 1 + + def __init__(self, in_planes, planes, stride=1): + super().__init__() + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, stride=(stride, 1), padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + + self.shortcut = nn.Sequential() + if stride != 1 or in_planes != self.expansion * planes: + self.shortcut = nn.Sequential( + nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=(stride, 1), bias=False), + nn.BatchNorm2d(self.expansion * planes), + ) + + def forward(self, x): + out = F.relu(self.bn1(self.conv1(x))) + out = self.bn2(self.conv2(out)) + out += self.shortcut(x) + out = F.relu(out) + return out diff --git a/vllm_omni/model_executor/models/indextts2/utils/common.py b/vllm_omni/model_executor/models/indextts2/utils/common.py new file mode 100644 index 00000000000..c969dc15903 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/common.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import re + +import torch + + +def tokenize_by_CJK_char(line: str, do_upper_case=True) -> str: + """ + Tokenize a line of text with CJK char. + + Note: All return characters will be upper case. + + Example: + input = "你好世界是 hello world 的中文" + output = "你 好 世 界 是 HELLO WORLD 的 中 文" + + Args: + line: + The input text. + + Return: + A new string tokenize by CJK char. + """ + # CJK ranges from nltk/tokenize/util.py (alvations/nltk@79eed6d) + CJK_RANGE_PATTERN = ( + r"([\u1100-\u11ff\u2e80-\ua4cf\ua840-\uD7AF\uF900-\uFAFF\uFE30-\uFE4F\uFF65-\uFFDC\U00020000-\U0002FFFF])" + ) + chars = re.split(CJK_RANGE_PATTERN, line.strip()) + return " ".join([w.strip().upper() if do_upper_case else w.strip() for w in chars if w.strip()]) + + +def de_tokenized_by_CJK_char(line: str, do_lower_case=False) -> str: + """ + Example: + input = "你 好 世 界 是 HELLO WORLD 的 中 文" + output = "你好世界是 hello world 的中文" + + do_lower_case: + input = "SEE YOU!" + output = "see you!" + """ + # replace english words in the line with placeholders + english_word_pattern = re.compile(r"([A-Z]+(?:[\s'-][A-Z-]+)*)", re.IGNORECASE) + english_sents = english_word_pattern.findall(line) + for i, sent in enumerate(english_sents): + line = line.replace(sent, f"<sent_{i}>") + + words = line.split() + # restore english sentences + sent_placeholder_pattern = re.compile(r"(<sent_(\d+)>)") + for i in range(len(words)): + all_matches = sent_placeholder_pattern.findall(words[i]) + if all_matches: + # restore the english word + for h, j in all_matches: + placeholder_index = int(j) + words[i] = words[i].replace(h, english_sents[placeholder_index]) + if do_lower_case: + words[i] = words[i].lower() + return "".join(words) + + +def make_pad_mask(lengths: torch.Tensor, max_len: int = 0) -> torch.Tensor: + """Make mask tensor containing indices of padded part. + + See description of make_non_pad_mask. + + Args: + lengths (torch.Tensor): Batch of lengths (B,). + Returns: + torch.Tensor: Mask tensor containing indices of padded part. + + Examples: + >>> lengths = [5, 3, 2] + >>> make_pad_mask(lengths) + masks = [[0, 0, 0, 0 ,0], + [0, 0, 0, 1, 1], + [0, 0, 1, 1, 1]] + """ + batch_size = lengths.size(0) + max_len = max_len if max_len > 0 else lengths.max().item() + seq_range = torch.arange(0, max_len, dtype=torch.int64, device=lengths.device) + seq_range_expand = seq_range.unsqueeze(0).expand(batch_size, max_len) + seq_length_expand = lengths.unsqueeze(-1) + mask = seq_range_expand >= seq_length_expand + return mask diff --git a/vllm_omni/model_executor/models/indextts2/utils/front.py b/vllm_omni/model_executor/models/indextts2/utils/front.py new file mode 100644 index 00000000000..3c3bd4cb7c5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/front.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import warnings +from typing import overload + +from sentencepiece import SentencePieceProcessor + +from vllm_omni.model_executor.models.indextts2.utils.common import ( + de_tokenized_by_CJK_char, + tokenize_by_CJK_char, +) + + +class TextTokenizer: + def __init__(self, vocab_file: str): + self.vocab_file = vocab_file + + if self.vocab_file is None: + raise ValueError("vocab_file is None") + if not os.path.exists(self.vocab_file): + raise ValueError(f"vocab_file {self.vocab_file} does not exist") + self.sp_model = SentencePieceProcessor(model_file=self.vocab_file) + self.pre_tokenizers = [tokenize_by_CJK_char] + + @property + def vocab_size(self): + return self.sp_model.GetPieceSize() + + @property + def unk_token(self): + return "<unk>" + + @property + def pad_token(self): + return None + + @property + def bos_token(self): + return "<s>" + + @property + def eos_token(self): + return "</s>" + + @property + def pad_token_id(self): + return -1 + + @property + def bos_token_id(self): + return 0 + + @property + def eos_token_id(self): + return 1 + + @property + def unk_token_id(self): + return self.sp_model.unk_id() + + @property + def special_tokens_map(self): + return { + "unk_token": self.unk_token, + "pad_token": self.pad_token, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + } + + def get_vocab(self): + vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)} + return vocab + + @overload + def convert_ids_to_tokens(self, ids: int) -> str: ... + + @overload + def convert_ids_to_tokens(self, ids: list[int]) -> list[str]: ... + + def convert_ids_to_tokens(self, ids: list[int] | int): + return self.sp_model.IdToPiece(ids) + + def convert_tokens_to_ids(self, tokens: list[str] | str) -> list[int]: + if isinstance(tokens, str): + tokens = [tokens] + return [self.sp_model.PieceToId(token) for token in tokens] + + def tokenize(self, text: str) -> list[str]: + return self.encode(text, out_type=str) + + def encode(self, text: str, **kwargs): + if len(text) == 0: + return [] + if len(text.strip()) == 1: + return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs) + for pre_tokenizer in self.pre_tokenizers: + text = pre_tokenizer(text) + return self.sp_model.Encode(text, out_type=kwargs.pop("out_type", int), **kwargs) + + def decode(self, ids: list[int] | int, do_lower_case=False, **kwargs): + if isinstance(ids, int): + ids = [ids] + decoded = self.sp_model.Decode(ids, out_type=kwargs.pop("out_type", str), **kwargs) + return de_tokenized_by_CJK_char(decoded, do_lower_case=do_lower_case) + + @staticmethod + def split_segments_by_token( + tokenized_str: list[str], + split_tokens: list[str], + max_text_tokens_per_segment: int, + quick_streaming_tokens: int = 0, + ) -> list[list[str]]: + if len(tokenized_str) == 0: + return [] + segments: list[list[str]] = [] + current_segment = [] + current_segment_tokens_len = 0 + i = 0 + while i < len(tokenized_str): + token = tokenized_str[i] + current_segment.append(token) + current_segment_tokens_len += 1 + if not ("," in split_tokens or "▁," in split_tokens) and ( + "," in current_segment or "▁," in current_segment + ): + sub_segments = TextTokenizer.split_segments_by_token( + current_segment, + [",", "▁,"], + max_text_tokens_per_segment=max_text_tokens_per_segment, + quick_streaming_tokens=quick_streaming_tokens, + ) + elif "-" not in split_tokens and "-" in current_segment: + sub_segments = TextTokenizer.split_segments_by_token( + current_segment, + ["-"], + max_text_tokens_per_segment=max_text_tokens_per_segment, + quick_streaming_tokens=quick_streaming_tokens, + ) + elif current_segment_tokens_len <= max_text_tokens_per_segment: + if token in split_tokens and current_segment_tokens_len > 2: + if i < len(tokenized_str) - 1: + if tokenized_str[i + 1] in ["'", "▁'"]: + current_segment.append(tokenized_str[i + 1]) + i += 1 + segments.append(current_segment) + current_segment = [] + current_segment_tokens_len = 0 + i += 1 + continue + else: + sub_segments = [] + for j in range(0, len(current_segment), max_text_tokens_per_segment): + if j + max_text_tokens_per_segment < len(current_segment): + sub_segments.append(current_segment[j : j + max_text_tokens_per_segment]) + else: + sub_segments.append(current_segment[j:]) + warnings.warn( + f"The tokens length of segment exceeds limit: {max_text_tokens_per_segment}, " + f"Tokens in segment: {current_segment}." + "Maybe unexpected behavior", + RuntimeWarning, + ) + segments.extend(sub_segments) + current_segment = [] + current_segment_tokens_len = 0 + i += 1 + if current_segment_tokens_len > 0: + assert current_segment_tokens_len <= max_text_tokens_per_segment + segments.append(current_segment) + merged_segments = [] + total_token = 0 + for segment in segments: + total_token += len(segment) + if len(segment) == 0: + continue + if len(merged_segments) == 0: + merged_segments.append(segment) + elif ( + len(merged_segments[-1]) + len(segment) <= max_text_tokens_per_segment + and total_token > quick_streaming_tokens + ): + merged_segments[-1] = merged_segments[-1] + segment + elif len(merged_segments[-1]) + len(segment) <= max_text_tokens_per_segment / 2: + merged_segments[-1] = merged_segments[-1] + segment + else: + merged_segments.append(segment) + return merged_segments + + punctuation_marks_tokens = [ + ".", + "!", + "?", + "▁.", + "▁?", + "▁...", + ] + + def split_segments( + self, tokenized: list[str], max_text_tokens_per_segment=120, quick_streaming_tokens=0 + ) -> list[list[str]]: + return TextTokenizer.split_segments_by_token( + tokenized, + self.punctuation_marks_tokens, + max_text_tokens_per_segment=max_text_tokens_per_segment, + quick_streaming_tokens=quick_streaming_tokens, + ) diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/__init__.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/maskgct/__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/indextts2/utils/maskgct/quantize.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/quantize.py new file mode 100644 index 00000000000..e4fcc6c79f8 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/maskgct/quantize.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +# +# Merged from factorized_vector_quantize.py and residual_vq.py. +# Only FactorizedVectorQuantize ("fvq") quantizer type is retained +# as it is the only one used by RepCodec in IndexTTS2. + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch.nn.utils import weight_norm + + +def WNConv1d(*args, **kwargs): + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +class FactorizedVectorQuantize(nn.Module): + def __init__( + self, + input_dim, + codebook_size, + codebook_dim, + commitment=0.005, + codebook_loss_weight=1.0, + use_l2_normalize=True, + ): + super().__init__() + self.input_dim = input_dim + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + self.commitment = commitment + self.codebook_loss_weight = codebook_loss_weight + self.use_l2_normalize = use_l2_normalize + + if self.input_dim != self.codebook_dim: + self.in_project = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1) + self.out_project = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1) + else: + self.in_project = nn.Identity() + self.out_project = nn.Identity() + + self.codebook = nn.Embedding(self.codebook_size, self.codebook_dim) + + def forward(self, z): + z_e = self.in_project(z) + z_q, indices = self.decode_latents(z_e) + + if self.training: + commit_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2]) * self.commitment + codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2]) * self.codebook_loss_weight + else: + commit_loss = torch.zeros(z.shape[0], device=z.device) + codebook_loss = torch.zeros(z.shape[0], device=z.device) + + z_q = z_e + (z_q - z_e).detach() + z_q = self.out_project(z_q) + + return z_q, commit_loss, codebook_loss, indices, z_e + + def embed_code(self, embed_id): + return F.embedding(embed_id, self.codebook.weight) + + def decode_code(self, embed_id): + return self.embed_code(embed_id).transpose(1, 2) + + def decode_latents(self, latents): + encodings = rearrange(latents, "b d t -> (b t) d") + codebook = self.codebook.weight + + if self.use_l2_normalize: + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0)) + z_q = self.decode_code(indices) + return z_q, indices + + def vq2emb(self, vq, out_proj=True): + emb = self.decode_code(vq) + if out_proj: + emb = self.out_project(emb) + return emb + + +class ResidualVQ(nn.Module): + """Residual Vector Quantization (SoundStream). + + Only supports FactorizedVectorQuantize ("fvq") quantizer type. + """ + + def __init__( + self, + input_dim: int = 256, + num_quantizers: int = 8, + codebook_size: int = 1024, + codebook_dim: int = 256, + quantizer_type: str = "fvq", + quantizer_dropout: float = 0.5, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + self.quantizer_dropout = quantizer_dropout + + self.quantizers = nn.ModuleList( + [ + FactorizedVectorQuantize( + input_dim=input_dim, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + **kwargs, + ) + for _ in range(num_quantizers) + ] + ) + + def forward(self, z, n_quantizers: int = None): + quantized_out = 0.0 + residual = z + + all_commit_losses = [] + all_codebook_losses = [] + all_indices = [] + all_quantized = [] + + if n_quantizers is None: + n_quantizers = self.num_quantizers + + if self.training: + n_q = torch.ones((z.shape[0],)) * self.num_quantizers + 1 + dropout = torch.randint(1, self.num_quantizers + 1, (z.shape[0],)) + n_dropout = int(z.shape[0] * self.quantizer_dropout) + n_q[:n_dropout] = dropout[:n_dropout] + n_quantizers = n_q.to(z.device) + + for i, quantizer in enumerate(self.quantizers): + if not self.training and i >= n_quantizers: + break + + z_q_i, commit_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(residual) + + mask = torch.full((z.shape[0],), fill_value=i, device=z.device) < n_quantizers + quantized_out = quantized_out + z_q_i * mask[:, None, None] + residual = residual - z_q_i + + commit_loss_i = (commit_loss_i * mask).mean() + codebook_loss_i = (codebook_loss_i * mask).mean() + + all_commit_losses.append(commit_loss_i) + all_codebook_losses.append(codebook_loss_i) + all_indices.append(indices_i) + all_quantized.append(z_q_i) + + ( + all_commit_losses, + all_codebook_losses, + all_indices, + all_quantized, + ) = map( + torch.stack, + ( + all_commit_losses, + all_codebook_losses, + all_indices, + all_quantized, + ), + ) + + return ( + quantized_out, + all_indices, + all_commit_losses, + all_codebook_losses, + all_quantized, + ) + + def vq2emb(self, vq, n_quantizers=None): + quantized_out = 0.0 + if n_quantizers is None: + n_quantizers = self.num_quantizers + for idx, quantizer in enumerate(self.quantizers): + if idx >= n_quantizers: + break + quantized_out += quantizer.vq2emb(vq[idx]) + return quantized_out diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/repcodec_model.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/repcodec_model.py new file mode 100644 index 00000000000..fbccdb36324 --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/maskgct/repcodec_model.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import torch.nn as nn +from torch.nn import functional as F + +from vllm_omni.model_executor.models.indextts2.utils.maskgct.quantize import ResidualVQ +from vllm_omni.model_executor.models.indextts2.utils.maskgct.vocos import VocosBackbone + + +def init_weights(m): + if isinstance(m, nn.Conv1d): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + +class RepCodec(nn.Module): + def __init__( + self, + codebook_size=8192, + hidden_size=1024, + codebook_dim=8, + vocos_dim=384, + vocos_intermediate_dim=2048, + vocos_num_layers=12, + num_quantizers=1, + downsample_scale=1, + cfg=None, + ): + super().__init__() + if cfg is not None: + _g = lambda attr, default: getattr(cfg, attr, default) # noqa: E731 + codebook_size = _g("codebook_size", codebook_size) + codebook_dim = _g("codebook_dim", codebook_dim) + hidden_size = _g("hidden_size", hidden_size) + vocos_dim = _g("vocos_dim", vocos_dim) + vocos_intermediate_dim = _g("vocos_intermediate_dim", vocos_intermediate_dim) + vocos_num_layers = _g("vocos_num_layers", vocos_num_layers) + num_quantizers = _g("num_quantizers", num_quantizers) + downsample_scale = _g("downsample_scale", downsample_scale) + + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + self.hidden_size = hidden_size + self.vocos_dim = vocos_dim + self.vocos_intermediate_dim = vocos_intermediate_dim + self.vocos_num_layers = vocos_num_layers + self.num_quantizers = num_quantizers + self.downsample_scale = downsample_scale + + if self.downsample_scale is not None and self.downsample_scale > 1: + self.down = nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1) + self.up = nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=1, padding=1) + + self.encoder = nn.Sequential( + VocosBackbone( + input_channels=self.hidden_size, + dim=self.vocos_dim, + intermediate_dim=self.vocos_intermediate_dim, + num_layers=self.vocos_num_layers, + adanorm_num_embeddings=None, + ), + nn.Linear(self.vocos_dim, self.hidden_size), + ) + self.decoder = nn.Sequential( + VocosBackbone( + input_channels=self.hidden_size, + dim=self.vocos_dim, + intermediate_dim=self.vocos_intermediate_dim, + num_layers=self.vocos_num_layers, + adanorm_num_embeddings=None, + ), + nn.Linear(self.vocos_dim, self.hidden_size), + ) + + self.quantizer = ResidualVQ( + input_dim=hidden_size, + num_quantizers=num_quantizers, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + quantizer_type="fvq", + quantizer_dropout=0.0, + commitment=0.15, + codebook_loss_weight=1.0, + use_l2_normalize=True, + ) + + self.reset_parameters() + + def quantize(self, x): + if self.downsample_scale is not None and self.downsample_scale > 1: + x = x.transpose(1, 2) + x = self.down(x) + x = F.gelu(x) + x = x.transpose(1, 2) + + x = self.encoder(x.transpose(1, 2)).transpose(1, 2) + + ( + quantized_out, + all_indices, + all_commit_losses, + all_codebook_losses, + _, + ) = self.quantizer(x) + + if all_indices.shape[0] == 1: + return all_indices.squeeze(0), quantized_out.transpose(1, 2) + return all_indices, quantized_out.transpose(1, 2) + + def reset_parameters(self): + self.apply(init_weights) diff --git a/vllm_omni/model_executor/models/indextts2/utils/maskgct/vocos.py b/vllm_omni/model_executor/models/indextts2/utils/maskgct/vocos.py new file mode 100644 index 00000000000..ce735a9adcc --- /dev/null +++ b/vllm_omni/model_executor/models/indextts2/utils/maskgct/vocos.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Copyright (c) 2024 Amphion. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +import torch +from torch import nn + + +class ConvNeXtBlock(nn.Module): + """ConvNeXt Block adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal.""" + + def __init__( + self, + dim: int, + intermediate_dim: int, + layer_scale_init_value: float, + adanorm_num_embeddings: int | None = None, + ): + super().__init__() + self.dwconv = nn.Conv1d(dim, dim, kernel_size=7, padding=3, groups=dim) + self.adanorm = adanorm_num_embeddings is not None + if adanorm_num_embeddings: + raise ValueError("AdaLayerNorm is not supported in inference-only mode") + else: + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, intermediate_dim) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(intermediate_dim, dim) + self.gamma = ( + nn.Parameter(layer_scale_init_value * torch.ones(dim), requires_grad=True) + if layer_scale_init_value > 0 + else None + ) + + def forward(self, x: torch.Tensor, cond_embedding_id: torch.Tensor | None = None) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x = x.transpose(1, 2) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.transpose(1, 2) + x = residual + x + return x + + +class VocosBackbone(nn.Module): + """Vocos backbone module built with ConvNeXt blocks.""" + + def __init__( + self, + input_channels: int, + dim: int, + intermediate_dim: int, + num_layers: int, + layer_scale_init_value: float | None = None, + adanorm_num_embeddings: int | None = None, + ): + super().__init__() + self.input_channels = input_channels + self.embed = nn.Conv1d(input_channels, dim, kernel_size=7, padding=3) + self.adanorm = adanorm_num_embeddings is not None + if adanorm_num_embeddings: + raise ValueError("AdaLayerNorm is not supported in inference-only mode") + else: + self.norm = nn.LayerNorm(dim, eps=1e-6) + layer_scale_init_value = layer_scale_init_value or 1 / num_layers + self.convnext = nn.ModuleList( + [ + ConvNeXtBlock( + dim=dim, + intermediate_dim=intermediate_dim, + layer_scale_init_value=layer_scale_init_value, + adanorm_num_embeddings=adanorm_num_embeddings, + ) + for _ in range(num_layers) + ] + ) + self.final_layer_norm = nn.LayerNorm(dim, eps=1e-6) + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, (nn.Conv1d, nn.Linear)): + nn.init.trunc_normal_(m.weight, std=0.02) + nn.init.constant_(m.bias, 0) + + def forward(self, x: torch.Tensor, **kwargs) -> torch.Tensor: + x = self.embed(x) + x = self.norm(x.transpose(1, 2)) + x = x.transpose(1, 2) + for conv_block in self.convnext: + x = conv_block(x) + x = self.final_layer_norm(x.transpose(1, 2)) + return x diff --git a/vllm_omni/model_executor/models/lance/__init__.py b/vllm_omni/model_executor/models/lance/__init__.py index 2ecd5df1f05..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/lance/__init__.py +++ b/vllm_omni/model_executor/models/lance/__init__.py @@ -1,9 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Lance (ByteDance) — vLLM-Omni model_executor package. -Single-stage topology: a self-contained diffusion stage (``LancePipeline`` -defined in ``vllm_omni.diffusion.models.lance``) that owns the Qwen2-MoT LLM, -the Qwen2.5-VL ViT, the Wan2.2 VAE and the tokenizer — directly analogous to -``bagel_single_stage``. The ``PipelineConfig`` lives in :mod:`.pipeline`. -""" +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/lance/pipeline.py b/vllm_omni/model_executor/models/lance/pipeline.py index 777f8942dd9..126a23acb58 100644 --- a/vllm_omni/model_executor/models/lance/pipeline.py +++ b/vllm_omni/model_executor/models/lance/pipeline.py @@ -15,6 +15,7 @@ LANCE_PIPELINE = PipelineConfig( model_type="lance", + default_deploy_config_name="lance.yaml", model_arch="LancePipeline", hf_architectures=(), stages=( diff --git a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py index 80e7c226343..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py @@ -1,17 +1,8 @@ -from transformers import AutoTokenizer - -from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer -from vllm_omni.transformers_utils.configs.mammoth_moda2 import ( # noqa: F401 registers AutoConfig - Mammothmoda2Config, - Mammothmoda2Qwen2_5_VLConfig, -) - -from .mammoth_moda2 import MammothModa2ARForConditionalGeneration - -# AutoConfig.register is done inside transformers_utils/configs/mammoth_moda2.py -AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) -AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer) - -__all__ = [ - "MammothModa2ARForConditionalGeneration", -] +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py index 6311216dc0c..2fe616cfd8c 100644 --- a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py +++ b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py @@ -483,19 +483,15 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ] params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() + if self.quant_config is not None: + cache_scale_mapper = self.quant_config.get_cache_scale_mapper() + if cache_scale_mapper is not None: + weights = cache_scale_mapper.apply(weights) for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue - for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py new file mode 100644 index 00000000000..830c63aff32 --- /dev/null +++ b/vllm_omni/model_executor/models/mammoth_moda2/pipeline.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MammothModa2 pipeline topology (frozen). + +Stage 0: AR — multimodal understanding + latent generation +Stage 1: DiT — latent → image + +For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2" + +MAMMOTH_MODA2_PIPELINE = PipelineConfig( + model_type="mammoth_moda2", + default_deploy_config_name="mammoth_moda2.yaml", + model_arch="MammothModa2ForConditionalGeneration", + hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=False, + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="latent", + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="image", + owns_tokenizer=False, + requires_multimodal_data=False, + engine_output_type="image", + custom_process_input_func=f"{_PROC}.ar2dit", + ), + ), +) + +# Single-stage AR variant for understanding / summarization tasks +MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig( + model_type="mammoth_moda2_ar", + default_deploy_config_name="mammoth_moda2_ar.yaml", + model_arch="MammothModa2ForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + ), + ), +) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py index bd178afe667..26482c6bce0 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_code2wav.py @@ -23,6 +23,7 @@ from vllm_omni.model_executor.models.mimo_audio.cuda_graph_decoder_wrapper import CUDAGraphMiMoDecoderWrapper from vllm_omni.model_executor.models.mimo_audio.modeling_audio_tokenizer import MiMoAudioTokenizer from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) @@ -442,8 +443,7 @@ def __init__(self, vllm_config: VllmConfig, prefix: str = ""): self.logits_processor = LogitsProcessor(config.vocab_size) - device_str = "cuda" if torch.cuda.is_available() else "cpu" - self.device = torch.device(device_str) + self.device = current_omni_platform.get_torch_device() self.sample_rate = getattr(config, "audio_sample_rate", 24000) self.tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path, trust_remote_code=True) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index a640ade3aaf..48715fb98e6 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -1153,17 +1153,14 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters(remove_duplicate=False)) loaded_params: set[str] = set() + if self.quant_config is not None: + cache_scale_mapper = self.quant_config.get_cache_scale_mapper() + if cache_scale_mapper is not None: + weights = cache_scale_mapper.apply(weights) + for name, loaded_weight in weights: if name.startswith("model."): name = "model." + name - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - # Loading kv cache quantization scales - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if name.startswith("input_local_transformer.") or name.startswith("local_transformer."): diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py index 126c901763c..9e54979deb1 100644 --- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py +++ b/vllm_omni/model_executor/models/mimo_audio/pipeline.py @@ -23,6 +23,7 @@ MIMO_AUDIO_PIPELINE = PipelineConfig( model_type="mimo_audio", + default_deploy_config_name="mimo_audio.yaml", # HF ``architectures: ["MiMoAudioModel"]`` is also the registry key in # ``model_executor/models/registry.py``; it resolves to the internal # class ``MiMoAudioForConditionalGeneration`` in ``mimo_audio.py``. @@ -56,7 +57,6 @@ final_output=True, final_output_type="audio", engine_output_type="audio", - custom_process_input_func=f"{_PROC}.llm2code2wav", sync_process_input_func=f"{_PROC}.llm2code2wav_token_only", sampling_constraints={"detokenize": False}, ), diff --git a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py index 4cd086e6426..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py @@ -1,20 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 -# Copyright 2025 The vLLM-Omni team. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .ming_flash_omni import MingFlashOmniForConditionalGeneration -from .ming_flash_omni_talker import MingFlashOmniTalkerForConditionalGeneration -from .ming_flash_omni_thinker import ( - MingFlashOmniThinkerDummyInputsBuilder, - MingFlashOmniThinkerForConditionalGeneration, - MingFlashOmniThinkerMultiModalProcessor, - MingFlashOmniThinkerProcessingInfo, -) - -__all__ = [ - "MingFlashOmniForConditionalGeneration", - "MingFlashOmniTalkerForConditionalGeneration", - "MingFlashOmniThinkerForConditionalGeneration", - "MingFlashOmniThinkerProcessingInfo", - "MingFlashOmniThinkerMultiModalProcessor", - "MingFlashOmniThinkerDummyInputsBuilder", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. 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 5a1af2fc23c..b70fdee8a02 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 @@ -32,7 +32,10 @@ from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, @@ -341,7 +344,7 @@ def __init__( prefix=f"{prefix}.experts", ) - self.experts.expert_mapping = FusedMoE.make_expert_params_mapping( + self.experts.expert_mapping = fused_moe_make_expert_params_mapping( self.experts, ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", diff --git a/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py b/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py index 7ebbfd9d869..da8c24843b8 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/pipeline.py @@ -6,8 +6,8 @@ Stage 1: Talker — text -> audio waveform via CFM + AudioVAE. The thinker -> talker bridge passes the detokenized text rather than -hidden states through `ming_flash_omni.thinker2talker`; the talker has a -self-contained Qwen2 LLM that retokenizes the string itself. +hidden states; the talker has a self-contained Qwen2 LLM that retokenizes +the string itself. """ from vllm_omni.config.stage_config import ( @@ -21,6 +21,7 @@ MING_FLASH_OMNI_PIPELINE = PipelineConfig( model_type="ming_flash_omni", + default_deploy_config_name="ming_flash_omni.yaml", model_arch="MingFlashOmniForConditionalGeneration", # Upstream HF config applies model_type="bailingmm_moe_v2_lite" # (the thinker sub-config name) rather than "ming_flash_omni". @@ -57,7 +58,6 @@ hf_config_name="talker_config", engine_output_type="audio", tokenizer_subdir="talker/llm", - custom_process_input_func=f"{_PROC}.thinker2talker", sync_process_input_func=f"{_PROC}.thinker2talker_token_only", ), ), @@ -67,6 +67,7 @@ # Standalone TTS variant: talker only. MING_FLASH_OMNI_TTS_PIPELINE = PipelineConfig( model_type="ming_flash_omni_tts", + default_deploy_config_name="ming_flash_omni_tts.yaml", model_arch="MingFlashOmniTalkerForConditionalGeneration", stages=( StagePipelineConfig( @@ -88,6 +89,7 @@ # Thinker-only variant: multimodal understanding with text output MING_FLASH_OMNI_THINKER_ONLY_PIPELINE = PipelineConfig( model_type="ming_flash_omni_thinker_only", + default_deploy_config_name="ming_flash_omni_thinker_only.yaml", model_arch="MingFlashOmniForConditionalGeneration", stages=( StagePipelineConfig( @@ -110,6 +112,7 @@ # Thinker + image-generation (diffusion) variant: text-to-image / img2img. MING_FLASH_OMNI_IMAGE_PIPELINE = PipelineConfig( model_type="ming_flash_omni_image", + default_deploy_config_name="ming_flash_omni_image.yaml", model_arch="MingFlashOmniForConditionalGeneration", stages=( StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/ming_flash_omni/projectors.py b/vllm_omni/model_executor/models/ming_flash_omni/projectors.py index 42e53d1c635..03527ba6327 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/projectors.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/projectors.py @@ -108,20 +108,6 @@ def __init__( layers.append(nn.Linear(llm_dim, llm_dim)) self.proj = nn.Sequential(*layers) - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Project audio features with temporal downsampling. - - Args: - x: [B, T, audio_dim] audio encoder output (channel-last). - - Returns: - [B, T', llm_dim] projected features (channel-last), - where T' = (T - ds_kernel_size + 2*(ds_kernel_size//2)) // ds_stride + 1. - """ - # Conv1d expects [B, C, T], so transpose input - x = x.transpose(-1, -2) # [B, audio_dim, T] - return self.proj(x) - def forward_packed( self, packed: torch.Tensor, @@ -155,19 +141,6 @@ def forward_packed( packed_proj = mlp(packed_proj) return packed_proj, proj_lens - def compute_output_length(self, input_length: torch.Tensor) -> torch.Tensor: - """Compute output sequence length after Conv1d downsampling. - - Args: - input_length: Original mel spectrogram lengths. - - Returns: - Output lengths after both convolutions. - """ - length = (input_length - 3 + 2 * 1) // 2 + 1 - length = (length - self.ds_kernel_size + 2 * (self.ds_kernel_size // 2)) // self.ds_stride + 1 - return length - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() diff --git a/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py b/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py index 436b92f1428..e0e8e73e669 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/text_processing.py @@ -15,19 +15,6 @@ logger = init_logger(__name__) -# Ported from -# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/toolkit.py -_TOKENIZE_PATTERN = re.compile(r"(?:[a-zA-Z]\.)+|[a-zA-Z]+(?:['\-][a-zA-Z]+)*|\d+(?:\.\d+)?|[\u4e00-\u9fff]|\s+|\S") - - -def tokenize_mixed_text(text: str) -> list[str]: - return re.findall(_TOKENIZE_PATTERN, text) - - -def tokenize_mixed_text_iterator(text: str): - for match in _TOKENIZE_PATTERN.finditer(text): - yield match.group(0) - # Ported from # https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/text_segment_cut.py @@ -218,84 +205,6 @@ def cut_text_by_semantic_length( return [f.replace(_DOT_PLACEHOLDER, ".") for f in result_fragments] -# Streaming sentence boundary detection - -_RE_CJK = re.compile(r"[\u4e00-\u9fff]") -_RE_DIGIT_LAST = re.compile(r"[0-9]") - - -# Left for reference for now -# Alternative to `cut_text_by_semantic_length` -def detect_sentence_boundaries( - text: str, - max_length: int = 50, -) -> list[str]: - """Accumulate tokens and flush at sentence boundaries. - - Ported from the streaming sentence detection loop from the Ming repo - TTS branch, but operates on the full text since we have it available upfront. - """ - sentences: list[str] = [] - streaming_text: list[str] = [] - count = 0 - - for ele in tokenize_mixed_text_iterator(text): - if len(ele) == 0: - continue - - should_process = False - min_tokens = 12 if count == 0 else 8 - - if ele[-1] in "！？。，!?": - if len(streaming_text) >= min_tokens: - should_process = True - streaming_text.append(ele) - - elif ele[-1] == ".": - if ( - len(streaming_text) >= min_tokens - and streaming_text - and not _RE_DIGIT_LAST.search(streaming_text[-1][-1]) - ): - should_process = True - streaming_text.append(ele) - - elif ele[-1] == "\n": - if streaming_text: - joined = "".join(streaming_text) - if _RE_CJK.search(joined): - if _RE_CJK.search(streaming_text[-1][-1]): - ele = "，" - streaming_text.append(ele) - else: - if len(ele) > 1 and re.search(r"[a-zA-Z]", ele[-2]): - ele = ele[:-1] + "." - else: - ele = ele[:-1] - streaming_text.append(ele) - - if len(streaming_text) >= min_tokens: - should_process = True - else: - streaming_text.append(ele) - continue - - if should_process: - joined = "".join(streaming_text) - fragments = cut_text_by_semantic_length(joined, max_length) - sentences.extend(fragments) - streaming_text = [] - count += 1 - - # Flush remaining - if streaming_text and re.search(r"[a-zA-Z\u4e00-\u9fff1-9]", "".join(streaming_text)): - joined = "".join(streaming_text) - fragments = cut_text_by_semantic_length(joined, max_length) - sentences.extend(fragments) - - return sentences - - # number normalization for English. Ported from # https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/front/number_en.py diff --git a/vllm_omni/model_executor/models/ming_tts/__init__.py b/vllm_omni/model_executor/models/ming_tts/__init__.py index 35afa2ce51a..da407694069 100644 --- a/vllm_omni/model_executor/models/ming_tts/__init__.py +++ b/vllm_omni/model_executor/models/ming_tts/__init__.py @@ -1,13 +1,7 @@ # 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", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/ming_tts/audio_prep.py b/vllm_omni/model_executor/models/ming_tts/audio_prep.py index b63e70cc8a0..378cd222de7 100644 --- a/vllm_omni/model_executor/models/ming_tts/audio_prep.py +++ b/vllm_omni/model_executor/models/ming_tts/audio_prep.py @@ -14,7 +14,6 @@ PATCH_SIZE, SAMPLE_RATE, VAE_PATCH_SIZE, - VISION_START_TOKEN_ID, MingTTSConfig, ) @@ -232,14 +231,17 @@ def _find_audio_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfi 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: +def _find_speaker_placeholder_positions(input_ids: torch.Tensor, cfg: MingTTSConfig) -> list[int]: + # The speaker embedding is written at the slot AFTER the placeholder marker: + # dense ``<|vision_start|>``+1 (the ``<|vision_pad|>`` slot), moe ``<spk>``+1 + # (the ``<audioPatch>`` slot). See upstream prepare_input_embed. + placeholder_token_id = int(cfg.speaker_placeholder_token_id) + placeholder_pos = (input_ids == placeholder_token_id).nonzero(as_tuple=True)[0] + if placeholder_pos.numel() == 0: return [] slots = [] - for pos in vision_start_pos: + for pos in placeholder_pos: slot = int(pos.item()) + 1 if slot < int(input_ids.shape[0]): slots.append(slot) 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 index b42c4fdca62..fdac2a30e2d 100644 --- a/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py +++ b/vllm_omni/model_executor/models/ming_tts/config_ming_tts.py @@ -22,6 +22,7 @@ DEFAULT_SIGMA, DEFAULT_TEMPERATURE, HISTORY_PATCH_SIZE, + INITIAL_LATENT_CHUNK_SIZE, KEY_CFG, KEY_CHUNK_ID, KEY_DECODE_STEP, @@ -30,7 +31,6 @@ KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS, KEY_NEXT_EMBEDS, - KEY_PROMPT_LATENT_TAIL, KEY_PROMPT_LATENTS, KEY_REQUEST_ID, KEY_SIGMA, @@ -43,6 +43,12 @@ LLM_HIDDEN_SIZE, LLM_VOCAB_SIZE, MAX_DECODE_STEPS, + MOE_AUDIO_DUMMY_TOKEN_ID, + MOE_AUDIO_END_TOKEN_ID, + MOE_AUDIO_EOS_TOKEN_ID, + MOE_AUDIO_START_TOKEN_ID, + MOE_SPK_TOKEN_ID, + MOE_TEXT_EOS_TOKEN_ID, PATCH_SIZE, SAMPLE_RATE, STOP_HEAD_MIN_STEPS, @@ -90,10 +96,126 @@ def get_text_config(self, decoder: bool = False, **kwargs: Any) -> Qwen2Config: return self.llm_config +class BailingMoeConfig(PretrainedConfig): + """Stage-0 LLM config for the Ming MoE variant. + + Ported from upstream inclusionAI ``configuration_bailing_moe.py``; the + community vLLM ``BailingMoeModel`` is duck-typed over these attributes. + Only the fields the backbone reads are kept; unknown keys flow through + ``**kwargs`` to ``PretrainedConfig``. + """ + + model_type = "bailing_moe" + + def __init__( + self, + vocab_size: int = 126464, + hidden_size: int = 2048, + intermediate_size: int | None = None, + num_hidden_layers: int = 28, + num_attention_heads: int = 16, + num_key_value_heads: int = 0, + hidden_act: str = "silu", + use_qkv_bias: bool = False, + use_bias: bool = True, + rms_norm_eps: float = 1e-05, + norm_head: bool = False, + tie_word_embeddings: bool = False, + max_position_embeddings: int = 32768, + rope_theta: float = 10000.0, + use_cache: bool = True, + rope_scaling: Any = None, + pad_token_id: int = 126081, + num_experts: int = 64, + num_shared_experts: int = 2, + num_experts_per_tok: int = 6, + norm_topk_prob: bool = True, + moe_intermediate_size: int | None = None, + first_k_dense_replace: int = 0, + head_dim: int | None = None, + output_router_logits: bool = False, + multi_gate: bool = False, + image_patch_token: int = 126346, + use_grouped_gemm: bool = False, + **kwargs: Any, + ) -> None: + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.use_qkv_bias = use_qkv_bias + self.use_bias = use_bias + self.norm_head = norm_head + self.rms_norm_eps = rms_norm_eps + self.max_position_embeddings = max_position_embeddings + self.rope_theta = rope_theta + self.use_cache = use_cache + self.head_dim = head_dim or self.hidden_size // self.num_attention_heads + self.rope_scaling = rope_scaling + self.num_experts = num_experts + self.num_shared_experts = num_shared_experts + self.num_experts_per_tok = num_experts_per_tok + self.norm_topk_prob = norm_topk_prob + self.moe_intermediate_size = moe_intermediate_size + self.first_k_dense_replace = first_k_dense_replace + self.output_router_logits = output_router_logits + self.multi_gate = multi_gate + self.image_patch_token = image_patch_token + self.use_grouped_gemm = use_grouped_gemm + super().__init__(pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs) + + +def _coerce_bailing_moe_config(value: Any) -> BailingMoeConfig: + if isinstance(value, BailingMoeConfig): + cfg = value + elif isinstance(value, PretrainedConfig): + cfg = BailingMoeConfig.from_dict(value.to_dict()) + elif isinstance(value, dict): + cfg = BailingMoeConfig.from_dict(dict(value)) + else: + raise TypeError(f"Unsupported llm_config type for Ming MoE config: {type(value)!r}") + # The multimodal bailing_moe declares 3D mrope (rope_scaling type "3D"). + # Ming-TTS Stage-0 carries no vision/video tokens, so mrope degenerates to + # standard 1D RoPE (all three position dims are equal). Strip rope_scaling + # so vLLM's get_rope uses plain RoPE — it has no "3D" scaling type. + cfg.rope_scaling = None + cfg.rope_parameters = {"rope_type": "default", "rope_theta": cfg.rope_theta} + return cfg + + +class MingMoeConfig(PretrainedConfig): + # Upstream inclusionAI Ming-omni-tts-16.8B-A3B HF config reports model_type="bailingmm". + model_type = "bailingmm" + + def __init__( + self, + llm_config: BailingMoeConfig | 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_bailing_moe_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) -> BailingMoeConfig: + 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().""" + model_variant: str = "dense" # "dense" (qwen2 backbone) | "moe" (bailing_moe backbone) + llm_hidden_size: int = LLM_HIDDEN_SIZE llm_vocab_size: int = LLM_VOCAB_SIZE llm_config: dict[str, Any] = field(default_factory=dict) @@ -118,6 +240,7 @@ class MingTTSConfig: max_decode_steps: int = MAX_DECODE_STEPS latent_chunk_size: int = LATENT_CHUNK_SIZE + initial_latent_chunk_size: int = INITIAL_LATENT_CHUNK_SIZE latent_left_context: int = LATENT_LEFT_CONTEXT text_eos_token_id: int = TEXT_EOS_TOKEN_ID @@ -125,6 +248,7 @@ class MingTTSConfig: 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 + speaker_placeholder_token_id: int = VISION_START_TOKEN_ID # dense <|vision_start|>; moe <spk> @classmethod def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig: @@ -145,7 +269,12 @@ def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig: enc_hop_size = _nested_get(atc, "enc_kwargs", "hop_size", default=AUDIO_FRAME_HOP) + # Variant is decided by the Stage-0 LLM backbone family, not the wrapper + # model_type: dense -> qwen2, MoE -> bailing_moe. + model_variant = "moe" if str(llm_dict.get("model_type", "")) == "bailing_moe" else "dense" + cfg = cls( + model_variant=model_variant, 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, @@ -159,6 +288,15 @@ def from_hf_config(cls, hf_config: PretrainedConfig) -> MingTTSConfig: sample_rate=atc_sample_rate, audio_frame_hop=enc_hop_size if enc_hop_size is not None else AUDIO_FRAME_HOP, ) + if model_variant == "moe": + # The bailing tokenizer uses a different vocab; override the dense + # (Qwen2) special-token defaults with the bailing token IDs. + cfg.text_eos_token_id = MOE_TEXT_EOS_TOKEN_ID + cfg.audio_dummy_token_id = MOE_AUDIO_DUMMY_TOKEN_ID + cfg.audio_start_token_id = MOE_AUDIO_START_TOKEN_ID + cfg.audio_end_token_id = MOE_AUDIO_END_TOKEN_ID + cfg.audio_eos_token_id = MOE_AUDIO_EOS_TOKEN_ID + cfg.speaker_placeholder_token_id = MOE_SPK_TOKEN_ID return cfg def validate(self) -> None: @@ -184,7 +322,6 @@ def validate(self) -> None: "KEY_MAX_DECODE_STEPS", "KEY_MIN_DECODE_STEPS", "KEY_NEXT_EMBEDS", - "KEY_PROMPT_LATENT_TAIL", "KEY_PROMPT_LATENTS", "KEY_REQUEST_ID", "KEY_SIGMA", @@ -197,7 +334,9 @@ def validate(self) -> None: "LLM_HIDDEN_SIZE", "LLM_VOCAB_SIZE", "MAX_DECODE_STEPS", + "BailingMoeConfig", "MingDenseConfig", + "MingMoeConfig", "MingTTSConfig", "PATCH_SIZE", "SAMPLE_RATE", diff --git a/vllm_omni/model_executor/models/ming_tts/constants.py b/vllm_omni/model_executor/models/ming_tts/constants.py index 2f59e209401..e0508bda767 100644 --- a/vllm_omni/model_executor/models/ming_tts/constants.py +++ b/vllm_omni/model_executor/models/ming_tts/constants.py @@ -15,6 +15,22 @@ TEXT_EOS_TOKEN_ID = 151669 # <text_eos> +# --------------------------------------------------------------------------- +# MoE (bailing_moe / Ming-omni-tts-16.8B-A3B) token IDs +# (confirmed from the bailing tokenizer_config.json — different vocab than the +# dense Qwen2 tokenizer above). The bailing tokenizer has NO <text_eos>; the +# audio AR loop terminates purely on the stop-head (see upstream +# modeling_bailingmm.sample), so <end_of_audio> doubles as the AR stop token. +# --------------------------------------------------------------------------- + +MOE_AUDIO_DUMMY_TOKEN_ID = 126357 # <audioPatch> +MOE_AUDIO_START_TOKEN_ID = 126358 # <audio> +MOE_AUDIO_END_TOKEN_ID = 126359 # </audio> +MOE_AUDIO_EOS_TOKEN_ID = 126356 # <end_of_audio> +MOE_TEXT_EOS_TOKEN_ID = 126356 # no <text_eos> in bailing; reuse <end_of_audio> as AR stop +MOE_SPK_TOKEN_ID = 126368 # <spk> — speaker-embedding placeholder (dense uses <|vision_start|>) + + # --------------------------------------------------------------------------- # Architectural constants (confirmed from original config.json) # --------------------------------------------------------------------------- @@ -42,6 +58,7 @@ # Connector / Stage-2 streaming defaults (runtime tuning) LATENT_CHUNK_SIZE = 25 +INITIAL_LATENT_CHUNK_SIZE = 4 LATENT_LEFT_CONTEXT = 0 MAX_DECODE_STEPS = 200 @@ -55,7 +72,6 @@ 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" diff --git a/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py b/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py new file mode 100644 index 00000000000..7a7280034a9 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_tts/fm/cfm_cudagraph.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CUDAGraph-accelerated CFM (flow-matching) diffusion head for Ming-TTS. + +The per-AR-step flow head runs ``steps`` ODE iterations, each a tiny DiT forward +(history+patch tokens, CFG cond+null batched). In eager mode these are dozens of +micro kernels per step and the wall time is dominated by kernel-launch / Python +overhead, not GPU math. Capturing the whole sampling loop + Aggregator + stop +head into one CUDA graph removes that overhead. + +Adapted from adaspeech/ming_flash ``CFMGraphExecutor``. Captured for batch=1 +(Ming Stage-0 runs ``max_num_seqs: 1``); other batch sizes fall back to eager. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.common.ming.fm import build_timesteps + + +class CFMSampler(nn.Module): + """Unrolled CFM ODE solver, graph-capturable (no ``torch.randn``/``.item()``). + + All randomness (``y0``, ``sde_rnd``) and the timestep schedule ``t`` are + passed in as float32 tensors so the loop is deterministic during replay and + the integration matches the eager float32 path. Equivalent to eager + ``Solver.integrate`` for ``temperature==0`` (the TTS default). + """ + + def __init__(self, dit_model: nn.Module, steps: int = 10): + super().__init__() + self.model = dit_model + self.steps = steps + + @torch.no_grad() + def forward(self, llm_cond, lat_cond, y0, t, sde_args, sde_rnd): + # y0 / t / sde_rnd / sde_args are float32; integrate in float32 and cast + # only the DiT inputs to the model dtype, mirroring eager FlowLoss.sample. + model_dtype = llm_cond.dtype + for step in range(self.steps): + dt = t[step + 1] - t[step] + pred_cfg = self.model.forward_with_cfg( + x=y0.to(model_dtype), + t=t[step].to(model_dtype), + c=llm_cond, + latent_history=lat_cond, + cfg_scale=sde_args[0], + patch_size=y0.shape[1], + ) + pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) + flow = (pred + (pred - null_pred) * sde_args[0]).float() + y0 = y0 + flow * dt + if step < self.steps - 1: + y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] + return y0 + + +class CFMGraphExecutor: + """Capture CFM sampling + Aggregator + stop head into one CUDA graph. + + First ``execute`` captures the graph; later calls copy inputs into the + static placeholders and ``replay()``. Runtime controls (cfg/sigma/temp) are + fed via the ``sde_args`` placeholder, so they are tunable without recapture. + """ + + def __init__( + self, + cfm_sampler: CFMSampler, + aggregator: nn.Module, + stop_head: nn.Module, + *, + patch_size: int, + latent_dim: int, + steps: int = 10, + ): + self.cfm_sampler = cfm_sampler + self.aggregator = aggregator + self.stop_head = stop_head + self.patch_size = int(patch_size) + self.latent_dim = int(latent_dim) + self.steps = int(steps) + self.initialized = False + self.graph: torch.cuda.CUDAGraph | None = None + + def _timesteps(self, device): + # Raw EPSS schedule + sway in float32, identical to eager build_timesteps. + return build_timesteps(self.steps, device=device, dtype=torch.float32, use_epss=True, sway_sampling_coef=-1.0) + + @torch.no_grad() + def execute(self, llm_cond, his_lat, cfg_strength, sigma, temperature): + bsz = llm_cond.shape[0] + device = llm_cond.device + # ODE state / noise / timesteps are float32 to match the eager path. + y0 = torch.randn((bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32) + t = self._timesteps(device) + sde_rnd = torch.randn((self.steps, bsz, self.patch_size, self.latent_dim), device=device, dtype=torch.float32) + + if not self.initialized: + self._initialize_graph(llm_cond, his_lat, y0, t, sde_rnd) + + self._llm_cond_ph.copy_(llm_cond) + self._his_lat_ph.copy_(his_lat) + self._y0_ph.copy_(y0) + self._t_ph.copy_(t) + self._sde_args_ph[0] = float(cfg_strength) + self._sde_args_ph[1] = float(sigma) + self._sde_args_ph[2] = float(temperature) + self._sde_rnd_ph.copy_(sde_rnd) + + self.graph.replay() + + return self._gen_lat_ph.clone(), self._embeds_ph.clone(), self._stop_ph.clone() + + def _initialize_graph(self, llm_cond, his_lat, y0, t, sde_rnd): + model_dtype = llm_cond.dtype + self._llm_cond_ph = torch.empty_like(llm_cond) + self._his_lat_ph = torch.empty_like(his_lat) + self._y0_ph = torch.empty_like(y0) # float32 + self._t_ph = torch.empty_like(t) # float32 + self._sde_args_ph = torch.empty(3, device=llm_cond.device, dtype=torch.float32) + self._sde_rnd_ph = torch.empty_like(sde_rnd) # float32 + + self._llm_cond_ph.copy_(llm_cond) + self._his_lat_ph.copy_(his_lat) + self._y0_ph.copy_(y0) + self._t_ph.copy_(t) + self._sde_args_ph[0], self._sde_args_ph[1], self._sde_args_ph[2] = 2.0, 0.25, 0.0 + self._sde_rnd_ph.copy_(sde_rnd) + + def _run(): + # Solver runs in float32; cast the final latent to the model dtype + # for the Aggregator / stop head (matches eager ``.to(dtype)``). + gen_lat = self.cfm_sampler( + self._llm_cond_ph, self._his_lat_ph, self._y0_ph, self._t_ph, self._sde_args_ph, self._sde_rnd_ph + ).to(model_dtype) + embeds = self.aggregator(gen_lat) + stop = self.stop_head(self._llm_cond_ph[:, -1, :]).softmax(dim=-1) + return gen_lat, embeds, stop + + # Warm up on a side stream to stabilize allocations before capture. + warmup = torch.cuda.Stream() + warmup.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup): + for _ in range(3): + _run() + torch.cuda.current_stream().wait_stream(warmup) + + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self._gen_lat_ph, self._embeds_ph, self._stop_ph = _run() + + self.initialized = True diff --git a/vllm_omni/model_executor/models/ming_tts/fm/dit.py b/vllm_omni/model_executor/models/ming_tts/fm/dit.py index 9aa05d4a25e..3eb13e98797 100644 --- a/vllm_omni/model_executor/models/ming_tts/fm/dit.py +++ b/vllm_omni/model_executor/models/ming_tts/fm/dit.py @@ -87,12 +87,17 @@ def forward(self, x, t, c, latent_history, mask=None): 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) + # Always run the classifier-free-guidance path (cond + uncond). TTS + # inference always uses cfg>1, and branching on ``cfg_scale`` is avoided + # because under CUDA-graph capture it is a CUDA tensor — ``bool(cfg_scale)`` + # would force a device sync and break capture. ``expand`` (not ``repeat``) + # avoids a fresh allocation each step. + 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]) + t = t.reshape(1) + t = t.expand(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 index 667aded27d7..70a4141861b 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts.py +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts.py @@ -16,7 +16,6 @@ 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, @@ -32,7 +31,6 @@ KEY_MAX_DECODE_STEPS, KEY_MIN_DECODE_STEPS, KEY_NEXT_EMBEDS, - KEY_PROMPT_LATENT_TAIL, KEY_REQUEST_ID, KEY_SIGMA, KEY_SPEAKER_EMBEDDING, @@ -175,7 +173,6 @@ def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tenso 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 @@ -185,7 +182,7 @@ def _prefill_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tenso 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) + speaker_slots = _find_speaker_placeholder_positions(input_ids, self.ming_config) if len(speaker_slots) < len(speaker_embeddings): raise RuntimeError( "Could not locate enough speaker placeholder slots: " @@ -262,7 +259,6 @@ def _copy_runtime_controls(update: dict[str, Any], info_dict: dict[str, Any]) -> __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_llm.py b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py index 92daad967e0..39ddd7b9f4a 100644 --- a/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py +++ b/vllm_omni/model_executor/models/ming_tts/ming_tts_llm.py @@ -3,7 +3,7 @@ # Adopted from https://github.com/inclusionAI/Ming-omni-tts/blob/main/modeling_bailingmm.py from __future__ import annotations -import warnings +import os from collections.abc import Iterable from typing import Any @@ -52,6 +52,9 @@ logger = init_logger(__name__) +# CFM ODE integration steps; must match CFM.sample(steps=...) default (10). +_CFM_STEPS = 10 + class MingLLMModel(nn.Module): hf_to_vllm_mapper = WeightsMapper( @@ -67,7 +70,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): 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")) + if self.ming_config.model_variant == "moe": + from vllm.model_executor.models.bailing_moe import BailingMoeModel + + # BailingMoeModel reads ``vllm_config.model_config.hf_config`` directly + # (no get_text_config()), so re-wrap with the nested bailing_moe config. + llm_config = vllm_config.model_config.hf_config.get_text_config() + llm_vllm_config = vllm_config.with_hf_config(llm_config) + self.model = BailingMoeModel(vllm_config=llm_vllm_config, prefix=maybe_prefix(prefix, "model")) + else: + 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, @@ -87,6 +99,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._pending_postprocess_updates: dict[str, dict[str, Any]] = {} self._last_ming_next_token_ids = None self._last_text_mode = False + # CUDAGraph for the flow-matching (CFM) diffusion head — the TTFP/latency + # hot spot. Built lazily on the first batch=1 CUDA decode; set + # MING_CFM_CUDAGRAPH=0 to force the eager flow head. + self._cfm_graph = None + self._cfm_graph_enabled = os.environ.get("MING_CFM_CUDAGRAPH", "1") == "1" def embed_input_ids( self, input_ids: torch.Tensor, inputs_embeds: torch.Tensor | None = None, **_: Any @@ -119,11 +136,14 @@ def forward( min_stop_step=int(self.ming_config.stop_head_min_steps), default_max_decode_steps=self.ming_config.max_decode_steps, ) + # Positional call: Qwen2Model names the 2nd arg ``positions`` while + # BailingMoeModel names it ``position_ids`` — both share the signature + # (input_ids, positions, intermediate_tensors, inputs_embeds). backbone_out = self.model( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, + input_ids, + positions, + intermediate_tensors, + inputs_embeds, ) if isinstance(backbone_out, IntermediateTensors): self._last_ming_next_token_ids = None @@ -295,6 +315,55 @@ def sample(self, logits, sampling_metadata): logprobs_tensors=None, ) + def _maybe_build_cfm_graph(self, z_diff_cond: torch.Tensor): + """Lazily build the CFM CUDAGraph executor (batch=1, CUDA only).""" + if not self._cfm_graph_enabled: + return None + if z_diff_cond.shape[0] != 1 or not z_diff_cond.is_cuda: + return None + if self._cfm_graph is not None: + return self._cfm_graph + try: + from .fm.cfm_cudagraph import CFMGraphExecutor, CFMSampler + + # TODO(perf): + # (1) compile per-block (dit_model.blocks) instead of the + # whole forward for tighter fusion / fewer graph breaks; + # (2) move the compile + a synthetic warmup forward to a post-weight-load hook + dit_model = self.flowloss.cfm.model + if not getattr(dit_model, "_ming_compiled", False): + try: + # Scoped options (avoid mutating global inductor config). + dit_model.forward = torch.compile( + dit_model.forward, + fullgraph=False, + dynamic=False, + options={ + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + ) + dit_model._ming_compiled = True + logger.info("Ming CFM DiT torch.compile enabled (inductor cudagraphs off).") + except Exception as exc: + logger.warning("Ming CFM torch.compile failed (%s); using uncompiled DiT.", exc) + + sampler = CFMSampler(self.flowloss.cfm.model, steps=_CFM_STEPS) + self._cfm_graph = CFMGraphExecutor( + sampler, + self.linear_proj_audio, + self.stop_head, + patch_size=self.ming_config.patch_size, + latent_dim=self.ming_config.latent_dim, + steps=_CFM_STEPS, + ) + logger.info("Ming CFM CUDAGraph enabled (steps=%d, patch=%d).", _CFM_STEPS, self.ming_config.patch_size) + except Exception as exc: + logger.warning("Ming CFM CUDAGraph unavailable (%s); using eager flow head.", exc) + self._cfm_graph_enabled = False + return None + return self._cfm_graph + def _decode_one_step( self, *, @@ -318,23 +387,45 @@ def _decode_one_step( 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, - ) + + # Fast path: CUDAGraph-captured flow head (CFM sampling + Aggregator + + # stop head). Falls back permanently to the eager path on any failure. + sampled_token_latent = next_embeds = stop_probs = None + # cfg < 1e-5 disables CFG; the eager flow head handles this with a + # dedicated unconditional branch (c=zeros) that the captured graph does + # not replicate, so fall back to eager there to match FlowLoss.sample. + graph_exec = self._maybe_build_cfm_graph(z_diff_cond) if cfg_scale >= 1e-5 else None + if graph_exec is not None: + try: + sampled_token_latent, next_embeds, stop_full = graph_exec.execute( + z_diff_cond, latent_history, cfg_scale, sigma, temperature + ) + stop_probs = stop_full[:, 1] + except Exception as exc: + logger.warning("Ming CFM CUDAGraph failed (%s); falling back to eager flow head.", exc) + self._cfm_graph = None + self._cfm_graph_enabled = False + sampled_token_latent = next_embeds = stop_probs = None + + if sampled_token_latent is None: + 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, + ) + # 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] + 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(): @@ -346,10 +437,10 @@ def _decode_one_step( 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) + _warn_missing_prefix("flowloss", params_dict, loaded_params, prefix="flowloss.") + _warn_missing_prefix("linear_proj_audio", params_dict, loaded_params, prefix="linear_proj_audio.") + _warn_missing_prefix("stop_head", params_dict, loaded_params, prefix="stop_head.") + _warn_missing_prefix("spk_head", params_dict, loaded_params, prefix="spk_head.") return loaded_params @@ -383,7 +474,6 @@ def _warn_missing_prefix( 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: @@ -392,9 +482,7 @@ def _warn_missing_prefix( 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) + raise RuntimeError(msg) __all__ = [ diff --git a/vllm_omni/model_executor/models/ming_tts/patch_emission.py b/vllm_omni/model_executor/models/ming_tts/patch_emission.py index c0a25a1c1f9..89acdf23129 100644 --- a/vllm_omni/model_executor/models/ming_tts/patch_emission.py +++ b/vllm_omni/model_executor/models/ming_tts/patch_emission.py @@ -127,7 +127,7 @@ def _validate_ming_decode_window( 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) + min_required = max(min_stop_step + 2, min_steps) if max_steps < min_required: req_id = info.get(KEY_REQUEST_ID, f"idx={i}") raise ValueError( @@ -148,7 +148,7 @@ def _resolve_ming_stop_decision( 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) + min_required_decode_steps = max(min_stop_step + 2, min_decode_steps) if max_decode_steps < min_required_decode_steps: raise RuntimeError( "Invalid Ming decode window: " diff --git a/vllm_omni/model_executor/models/ming_tts/pipeline.py b/vllm_omni/model_executor/models/ming_tts/pipeline.py index 31ec98e5a05..86eecd41ee9 100644 --- a/vllm_omni/model_executor/models/ming_tts/pipeline.py +++ b/vllm_omni/model_executor/models/ming_tts/pipeline.py @@ -12,8 +12,12 @@ MING_TTS_PIPELINE = PipelineConfig( model_type="ming_tts", + default_deploy_config_name="ming_tts.yaml", model_arch="MingTTSForConditionalGeneration", - hf_architectures=("MingTTSForConditionalGeneration",), + hf_architectures=("MingTTSForConditionalGeneration", "BailingMMNativeForConditionalGeneration"), + # The dense (0.5B) and MoE (16.8B) share architectures=["BailingMMNativeForConditionalGeneration"] + # Here we disambiguate by the model_type the upstream HF config reports + hf_config_predicate=lambda c: str(getattr(c, "model_type", "")) == "dense", stages=( StagePipelineConfig( stage_id=0, @@ -52,3 +56,18 @@ ), ), ) + + +# MoE variant (inclusionAI/Ming-omni-tts-16.8B-A3B) +# Same two-stage topology and model class as the above, +# but with a distinct model_type reported by the upstream HF config. +# Keep this for auto-detection of deploy config yaml. +MING_TTS_MOE_PIPELINE = PipelineConfig( + model_type="ming_tts_moe", + default_deploy_config_name="ming_tts_moe.yaml", + model_arch="MingTTSForConditionalGeneration", + hf_architectures=("MingTTSForConditionalGeneration", "BailingMMNativeForConditionalGeneration"), + # Disambiguate from the dense-0.5B ming_tts pipeline + hf_config_predicate=lambda c: str(getattr(c, "model_type", "")) == "bailingmm", + stages=MING_TTS_PIPELINE.stages, +) diff --git a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py index c732e07a045..36817022a7f 100644 --- a/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py +++ b/vllm_omni/model_executor/models/ming_tts/prompt_assembly.py @@ -176,6 +176,84 @@ def build_dense_prompt_token_ids( ) +def build_moe_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]: + """Stage-0 prompt for the Ming MoE (bailing_moe) variant. + + Mirrors the ``model_type != 'dense'`` branch of upstream + ``BailingMMNativeForConditionalGeneration.prepare_input_embed``: no system + message, ``<role>HUMAN/ASSISTANT</role>`` turns, and a ``<spk><audioPatch></spk>`` + speaker placeholder (vs. the dense ``<|vision_*|>`` placeholder). + """ + speaker_prompt = [] + for idx in range(int(speaker_count)): + speaker_prompt.extend( + tokenizer.encode(f" speaker_{idx + 1}:") + + tokenizer.encode("<spk>") + + tokenizer.encode("<audioPatch>") + + tokenizer.encode("</spk>\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("<audioPatch>") + if audio_patch_token_id == tokenizer.unk_token_id: + raise ValueError("Ming tokenizer is missing required <audioPatch> 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("<role>HUMAN</role>") + + tokenizer.encode(prompt) + + speaker_prompt + + text_input_prefix + + prompt_text_tokens + + tokenizer.encode(text) + + tokenizer.encode("<role>ASSISTANT</role>") + + instruction_prompt + + tokenizer.encode("<audio>") + + prompt_latent_tokens + ) + + +def build_prompt_token_ids(model_variant: str, tokenizer: Any, **kwargs: Any) -> list[int]: + """Dispatch to the dense or MoE Stage-0 prompt builder by model variant.""" + if model_variant == "moe": + return build_moe_prompt_token_ids(tokenizer, **kwargs) + return build_dense_prompt_token_ids(tokenizer, **kwargs) + + +def detect_model_variant(tokenizer: Any) -> str: + """Infer the Ming variant from the tokenizer's special tokens. + + The bailing (MoE) tokenizer carries ``<role>`` chat markers; the dense + Qwen2 tokenizer does not. Used to pick the prompt format when the caller + does not pass an explicit variant. + """ + try: + role_id = tokenizer.convert_tokens_to_ids("<role>") + except Exception: + return "dense" + unk = getattr(tokenizer, "unk_token_id", None) + if role_id is None or (unk is not None and role_id == unk): + return "dense" + return "moe" + + def build_ming_dense_prompt( tokenizer: Any, *, @@ -189,7 +267,10 @@ def build_ming_dense_prompt( speaker_embedding: Any = None, use_zero_spk_emb: bool = False, request_id: str | None = None, + model_variant: str | None = None, ) -> dict[str, Any]: + if model_variant is None: + model_variant = detect_model_variant(tokenizer) instruction_text = create_instruction(instruction) speaker_embeddings = coerce_speaker_embeddings(speaker_embedding, use_zero_spk_emb=use_zero_spk_emb) effective_runtime_controls = resolve_effective_runtime_controls(text=text, runtime_controls=runtime_controls) @@ -217,7 +298,8 @@ def build_ming_dense_prompt( prompt_latent_value, patch_size=PATCH_SIZE, latent_dim=LATENT_DIM ) - prompt_token_ids = build_dense_prompt_token_ids( + prompt_token_ids = build_prompt_token_ids( + model_variant, tokenizer, prompt=prompt, text=text, @@ -266,7 +348,10 @@ def build_ming_dense_prompt( "DEFAULT_PROMPT", "build_dense_prompt_token_ids", "build_ming_dense_prompt", + "build_moe_prompt_token_ids", + "build_prompt_token_ids", "create_instruction", + "detect_model_variant", "estimate_decode_step_window_for_duration", "estimate_decode_steps_for_duration", "parse_duration_seconds", diff --git a/vllm_omni/model_executor/models/ming_tts/validation.py b/vllm_omni/model_executor/models/ming_tts/validation.py index 59e1fdb8f4a..3f3da6cf0b2 100644 --- a/vllm_omni/model_executor/models/ming_tts/validation.py +++ b/vllm_omni/model_executor/models/ming_tts/validation.py @@ -14,6 +14,9 @@ LATENT_DIM, LLM_HIDDEN_SIZE, LLM_VOCAB_SIZE, + MOE_AUDIO_DUMMY_TOKEN_ID, + MOE_AUDIO_EOS_TOKEN_ID, + MOE_TEXT_EOS_TOKEN_ID, PATCH_SIZE, SAMPLE_RATE, ) @@ -70,19 +73,23 @@ def _nested_get(obj: Any, *keys: str, default: Any = None) -> Any: def validate_ming_tts_config(cfg: Any) -> None: """Run before GPU allocation/weight loading. Raises ValueError on mismatches.""" - if cfg.audio_dummy_token_id != 151705: + is_moe = getattr(cfg, "model_variant", "dense") == "moe" + exp_audio_dummy = MOE_AUDIO_DUMMY_TOKEN_ID if is_moe else 151705 + exp_audio_eos = MOE_AUDIO_EOS_TOKEN_ID if is_moe else 151704 + exp_text_eos = MOE_TEXT_EOS_TOKEN_ID if is_moe else 151669 + if cfg.audio_dummy_token_id != exp_audio_dummy: raise ValueError( - f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected 151705 (<audioPatch>). " + f"audio_dummy_token_id={cfg.audio_dummy_token_id}, expected {exp_audio_dummy} (<audioPatch>). " "Wrong tokenizer/checkpoint?" ) - if cfg.audio_eos_token_id != 151704: + if cfg.audio_eos_token_id != exp_audio_eos: raise ValueError( - f"audio_eos_token_id={cfg.audio_eos_token_id}, expected 151704 (<end_of_audio>). " + f"audio_eos_token_id={cfg.audio_eos_token_id}, expected {exp_audio_eos} (<end_of_audio>). " "Wrong tokenizer/checkpoint?" ) - if cfg.text_eos_token_id != 151669: + if cfg.text_eos_token_id != exp_text_eos: raise ValueError( - f"text_eos_token_id={cfg.text_eos_token_id}, expected 151669 (<text_eos>). Wrong tokenizer/checkpoint?" + f"text_eos_token_id={cfg.text_eos_token_id}, expected {exp_text_eos}. Wrong tokenizer/checkpoint?" ) if cfg.audio_tokenizer_config is None: @@ -102,13 +109,16 @@ def validate_ming_tts_config(cfg: Any) -> None: f"history_patch_size mismatch: got {cfg.history_patch_size}, expected {HISTORY_PATCH_SIZE}. " "Check ditar_config.history_patch_size." ) - if cfg.llm_hidden_size != LLM_HIDDEN_SIZE: - raise ValueError( - f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. " - "Check llm_config.hidden_size." - ) - if cfg.llm_vocab_size != LLM_VOCAB_SIZE: - raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.") + # Absolute hidden/vocab pins are dense-only; the MoE backbone carries its + # own (larger) dims, cross-checked against llm_config.hidden_size below. + if not is_moe: + if cfg.llm_hidden_size != LLM_HIDDEN_SIZE: + raise ValueError( + f"llm_hidden_size mismatch: got {cfg.llm_hidden_size}, expected {LLM_HIDDEN_SIZE}. " + "Check llm_config.hidden_size." + ) + if cfg.llm_vocab_size != LLM_VOCAB_SIZE: + raise ValueError(f"llm_vocab_size mismatch: got {cfg.llm_vocab_size}, expected {LLM_VOCAB_SIZE}.") if cfg.sample_rate != SAMPLE_RATE: raise ValueError(f"sample_rate mismatch: got {cfg.sample_rate}, expected {SAMPLE_RATE}.") @@ -162,6 +172,9 @@ def validate_ming_tts_config(cfg: Any) -> None: if cfg.latent_chunk_size <= 0: raise ValueError(f"latent_chunk_size must be > 0, got {cfg.latent_chunk_size}.") + initial_latent_chunk_size = getattr(cfg, "initial_latent_chunk_size", 0) + if initial_latent_chunk_size < 0: + raise ValueError(f"initial_latent_chunk_size must be >= 0, got {initial_latent_chunk_size}.") if cfg.latent_left_context < 0: raise ValueError(f"latent_left_context must be >= 0, got {cfg.latent_left_context}.") if cfg.max_decode_steps <= 0: diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py index 7284c5b7ff1..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/__init__.py @@ -1,9 +1,8 @@ -from .minicpmo_4_5_omni import MiniCPMO45OmniForConditionalGeneration -from .minicpmo_4_5_omni_llm import MiniCPMO45OmniLLMForConditionalGeneration -from .minicpmo_4_5_omni_tts import MiniCPMO45OmniTTSForConditionalGeneration +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = [ - "MiniCPMO45OmniForConditionalGeneration", - "MiniCPMO45OmniLLMForConditionalGeneration", - "MiniCPMO45OmniTTSForConditionalGeneration", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py index 7746b423e9e..d9b48807ae4 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni.py @@ -291,7 +291,7 @@ def forward( device = inputs_embeds.device else: num_tokens = 1 - device = torch.device("cuda") + device = current_omni_platform.get_torch_device() hidden_dim = self.config.hidden_size if hasattr(self.config, "hidden_size") else 2560 # Profile/dummy run: both input_ids and inputs_embeds are None. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index b930578d69b..3019a5b3deb 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -760,7 +760,12 @@ def preprocess( ) -AutoImageProcessor.register("MiniCPMVImageProcessor", MiniCPMVImageProcessor) +# transformers >= 5.x requires the config *class* (not a string) as the +# first arg to AutoImageProcessor.register — it does `key.__module__` +# internally. Passing "MiniCPMVImageProcessor" raises: +# AttributeError: 'str' object has no attribute '__module__' +# See also processing_gr00t_n1d7.py for the same transformers API change. +AutoImageProcessor.register(MiniCPMOConfig, MiniCPMVImageProcessor, exist_ok=True) # ============== SigLIP Vision Transformer Classes ============== diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index c7208752ba8..d14fe14cffd 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -15,6 +15,7 @@ import io import logging import os +import sys from collections.abc import Iterable import numpy as np @@ -25,13 +26,35 @@ from vllm.config import VllmConfig from vllm.model_executor.models.interfaces import SupportsPP -try: - from stepaudio2 import Token2wav as _Token2wav +from vllm_omni.platforms import current_omni_platform - _stepaudio2_available = True -except ImportError: - _Token2wav = None - _stepaudio2_available = False +# Preserve the established external vocoder on CUDA. Ascend uses the in-tree +# adapter because ``stepaudio2-minicpmo`` hard-codes CUDA device placement. +if current_omni_platform.is_npu(): + try: + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_token2wav import ( + MiniCPMO45Token2wav as _Token2wav, + ) + + _token2wav_backend = "step_audio2_core" + except ImportError: + try: + from stepaudio2 import Token2wav as _Token2wav + + _token2wav_backend = "stepaudio2_pkg" + except ImportError: + _Token2wav = None + _token2wav_backend = None +else: + try: + from stepaudio2 import Token2wav as _Token2wav + + _token2wav_backend = "stepaudio2_pkg" + except ImportError: + _Token2wav = None + _token2wav_backend = None + +_stepaudio2_available = _Token2wav is not None logger = logging.getLogger(__name__) @@ -99,14 +122,25 @@ def _lazy_init_tts(self): return try: model_path = self.vllm_config.model_config.model - import os - import sys if model_path not in sys.path: sys.path.insert(0, model_path) + from transformers import AutoImageProcessor from transformers.dynamic_module_utils import get_class_from_dynamic_module - MiniCPMTTS = get_class_from_dynamic_module("modeling_minicpmo.MiniCPMTTS", model_path) + # openbmb/MiniCPM-o-4_5/processing_minicpmo.py registers via a + # string: AutoImageProcessor.register("MiniCPMVImageProcessor", ...), + # which crashes on transformers>=5 (register reads key.__module__). + # Loading MiniCPMTTS imports that module, so no-op the string form + # (unused by the standalone talker) while it runs, then restore. + original_register = AutoImageProcessor.register + AutoImageProcessor.register = ( # type: ignore[method-assign] + lambda key, *a, **k: None if isinstance(key, str) else original_register(key, *a, **k) + ) + try: + MiniCPMTTS = get_class_from_dynamic_module("modeling_minicpmo.MiniCPMTTS", model_path) + finally: + AutoImageProcessor.register = original_register # type: ignore[method-assign] # MiniCPMTTS.__init__ reads `config.top_p / top_k / repetition_penalty` # directly (modeling_minicpmo.py L4112-4114), but the model repo's @@ -119,6 +153,11 @@ def _lazy_init_tts(self): if not hasattr(self._tts_config, _attr): setattr(self._tts_config, _attr, _default) + # The copied Hugging Face flash_attention_2 setting is not valid + # for this standalone MiniCPMTTS path. Use PyTorch SDPA on every + # backend until a dedicated flash-attention implementation exists. + self._tts_config.attn_implementation = "sdpa" + prev_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.float32) try: @@ -157,7 +196,11 @@ def _lazy_init_tts(self): finally: torch.set_default_dtype(prev_dtype2) self.tts_obj.audio_tokenizer = self.audio_tokenizer - logger.info("Loaded Token2wav from %s", token2wav_dir) + logger.info( + "Loaded Token2wav from %s (backend=%s)", + token2wav_dir, + _token2wav_backend, + ) # Only mark init as complete after every step succeeds, so a # partial failure leaves the next call free to retry the full # init instead of short-circuiting back to a silent empty path. @@ -187,19 +230,25 @@ def generate_speech( tts = self.tts_obj device = tts.emb_text.weight.device - dtype = tts.emb_text.weight.dtype + # MiniCPMTTS AR backbone uses FlashAttention (fp16/bf16 only). The + # submodule is constructed under float32 default dtype during lazy init, + # so pin the condition embeddings to bfloat16 explicitly rather than + # inheriting the (float32) parameter dtype — a float32 condition breaks + # the CUDA FA2 path and wastes memory on the NPU sdpa path. + ar_dtype = torch.bfloat16 llm_embeds = tts.emb_text(tts_token_ids.to(device)) - hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=dtype)) + hidden_embeds = tts.projector_semantic(tts_hidden_states.to(device=device, dtype=ar_dtype)) if getattr(tts.config, "normalize_projected_hidden", False): hidden_embeds = F.normalize(hidden_embeds, p=2, dim=-1) - tts_embeds = llm_embeds + hidden_embeds + tts_embeds = (llm_embeds + hidden_embeds).to(dtype=ar_dtype) text_eos = tts.emb_text(torch.tensor([tts.config.text_eos_token_id], device=device, dtype=torch.long)) audio_bos = tts.emb_text(torch.tensor([tts.audio_bos_token_id], device=device, dtype=torch.long)) - spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=tts_embeds.dtype) + spk_embeds = torch.zeros(0, tts.config.hidden_size, device=device, dtype=ar_dtype) inputs_embeds = torch.cat([spk_embeds, tts_embeds, text_eos, audio_bos], dim=0).unsqueeze(0) + inputs_embeds = inputs_embeds.to(dtype=ar_dtype) logger.info("generate_speech: inputs_embeds shape=%s", list(inputs_embeds.shape)) # Scale max_new_token with input text length to avoid mid-stream truncation on long @@ -247,7 +296,15 @@ def _patched_save(uri, src, sample_rate, **kw): prev_dtype = torch.get_default_dtype() torch.set_default_dtype(torch.float32) try: - with torch.amp.autocast("cuda", enabled=False): + # Vocoder path is float32; use the platform abstraction because + # torch.amp.autocast validates unsupported device types even when + # autocast is disabled. + autocast_device = device.type if isinstance(device, torch.device) else str(device) + with current_omni_platform.create_autocast_context( + device_type=autocast_device, + dtype=torch.float32, + enabled=False, + ): token_list = generated_tokens.squeeze(0).tolist() num_tokens = len(token_list) @@ -379,7 +436,7 @@ def _dummy_hidden_states( break else: num_tokens = 1 - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() hidden_size = int(getattr(self, "_hidden_size", 768) or 768) return torch.zeros((num_tokens, hidden_size), device=device, dtype=torch.bfloat16) @@ -416,7 +473,17 @@ def forward( return None, None def compute_logits(self, hidden_states, *args, **kwargs): - return torch.zeros(1, 2, device=hidden_states.device if isinstance(hidden_states, torch.Tensor) else "cuda") + # Placeholder logits: one row per sampled request (the scheduler + # indexes sampled_token_ids by req_index). Hardcoding a single row + # breaks batched/concurrent decoding with IndexError. The values are + # discarded — real output is the waveform via multimodal_outputs. + if isinstance(hidden_states, torch.Tensor): + device = hidden_states.device + num_reqs = hidden_states.shape[0] if hidden_states.ndim >= 1 else 1 + else: + device = current_omni_platform.get_torch_device() + num_reqs = 1 + return torch.zeros(num_reqs, 2, device=device) def sample(self, logits, sampling_metadata): return None @@ -440,10 +507,31 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): logger.warning("TTS missing keys (%d): %s", len(missing), missing[:5]) if unexpected: logger.warning("TTS unexpected keys (%d): %s", len(unexpected), unexpected[:5]) - self.tts_obj = self.tts_obj.to("cuda") + # Move the AR backbone to the active device (cuda / npu / …) and + # cast to bfloat16: MiniCPMTTS AR uses FlashAttention (fp16/bf16 + # only) and is built under a float32 default dtype during lazy + # init, so an uncast float32 backbone breaks CUDA FA2 and wastes + # memory on the NPU sdpa path. Detach the Token2wav vocoder first + # so the cast does not drag it onto the accelerator or downcast + # its float32 flow/HiFT weights: it manages its own device + # placement (e.g. CPU HiFT on NPU) and may not be an nn.Module. + device = current_omni_platform.get_torch_device() + audio_tok = getattr(self.tts_obj, "audio_tokenizer", None) + if audio_tok is not None: + self.tts_obj.audio_tokenizer = None + try: + self.tts_obj = self.tts_obj.to(device=device, dtype=torch.bfloat16) + finally: + if audio_tok is not None: + self.tts_obj.audio_tokenizer = audio_tok + self.audio_tokenizer = audio_tok self.emb_text = self.tts_obj.emb_text self.projector_semantic = self.tts_obj.projector_semantic - logger.info("Loaded %d TTS weights, moved to cuda", len(tts_weights)) + logger.info( + "Loaded %d TTS weights, moved to %s (bfloat16)", + len(tts_weights), + device, + ) return loaded diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py new file mode 100644 index 00000000000..c069b8d09c2 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiniCPM-o 4.5 Token2wav adapter over in-tree ``StepAudio2Token2WavCore``. + +``minicpmo_4_5_omni_tts`` historically depended on the external +``from stepaudio2 import Token2wav`` entry point (``stepaudio2-minicpmo``). +That package hard-codes ``.cuda()`` and duplicates the flow/HiFT stack that +vLLM-Omni already vendors for Step-Audio2 (which also carries the Ascend/NPU +fixes: CPU HiFT, DiT mask expand, MATH SDPA, compile disable, PE buffer +extension). + +This module exposes the same call surface MiniCPM expects +(``__call__`` / ``set_stream_cache`` / ``stream``) while delegating the +actual vocoder work to +``vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav.StepAudio2Token2WavCore``. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch + +from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import ( + StepAudio2Token2WavCore, + _StreamState, +) + + +def _resolve_device(device: str | torch.device | None) -> str: + if device is not None: + return str(device) + try: + from vllm_omni.platforms import current_omni_platform + + plat_dev = getattr(current_omni_platform, "device_type", None) + if plat_dev: + return str(plat_dev) + except Exception: + pass + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch, "npu") and torch.npu.is_available(): + return "npu" + return "cpu" + + +class MiniCPMO45Token2wav: + """Token2wav-compatible facade around ``StepAudio2Token2WavCore``. + + Matches the MiniCPM-o / ``stepaudio2-minicpmo`` vocoder API used by + ``MiniCPMO45OmniTTSForConditionalGeneration.generate_speech``: + + * ``Token2wav(model_path, float16=False, n_timesteps=10)`` + * ``__(tokens, prompt_wav) -> wav_bytes`` + * ``set_stream_cache(prompt_wav) -> (stream_cache, hift_cache_dict)`` + * ``stream(tokens, prompt_wav, last_chunk=..., return_waveform=...)`` + """ + + def __init__( + self, + model_path: str, + float16: bool = False, + n_timesteps: int = 10, + device: str | torch.device | None = None, + ): + self.float16 = float16 + self.n_timesteps = n_timesteps + self.device = _resolve_device(device) + self._core = StepAudio2Token2WavCore( + model_path=model_path, + float16=float16, + device=self.device, + n_timesteps=n_timesteps, + ) + # Eager-load so construction failures surface at init time (same as + # the external Token2wav package), not on the first request. + self._core._ensure_models_loaded() + + # Mutable streaming fields expected by MiniCPM's long-form path. + self.stream_cache: Any | None = None + self.hift_cache_dict: dict[str, torch.Tensor] = {} + + def __call__(self, generated_speech_tokens, prompt_wav) -> bytes: + """One-shot tokens → 24 kHz WAV bytes.""" + return self._core.forward( + generated_speech_tokens, + prompt_wav, + return_bytes=True, + ) + + def set_stream_cache(self, prompt_wav: str): + """Initialise flow + HiFT caches for chunked vocoding. + + Returns ``(stream_cache, hift_cache_dict)`` so callers can assign them + back onto this object the way ``stepaudio2.Token2wav`` does. + """ + state = _StreamState() + self._core.setup_stream_for(prompt_wav, state) + self.stream_cache = state.stream_cache + self.hift_cache_dict = state.hift_cache_dict + return self.stream_cache, self.hift_cache_dict + + def stream( + self, + generated_speech_tokens, + prompt_wav: str, + last_chunk: bool = False, + return_waveform: bool = False, + ): + """Process one streaming chunk; updates ``stream_cache`` / ``hift_cache_dict``.""" + if self.stream_cache is None: + raise ValueError("stream_cache is not set") + + state = _StreamState() + state.setup_done = True + state.stream_cache = self.stream_cache + state.hift_cache_dict = self.hift_cache_dict + + speech = self._core.stream_chunk_for( + list(generated_speech_tokens), + prompt_wav, + last_chunk, + state, + ) + self.stream_cache = state.stream_cache + self.hift_cache_dict = state.hift_cache_dict + + wav_np = speech.detach().float().cpu().numpy() + if return_waveform: + return wav_np + + wav_np = np.clip(wav_np, -1.0, 1.0) + wav_int16 = (wav_np * 32767.0).astype("<i2") + return wav_int16.tobytes() diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 701e014e69f..f966555d02c 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -22,6 +22,7 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( model_type="minicpmo_4_5", + default_deploy_config_name="minicpmo_4_5.yaml", model_arch="MiniCPMO45OmniForConditionalGeneration", # MiniCPM-o 4.5's HF config.json reports `model_type="minicpmo"` and # `architectures=["MiniCPMO"]` — both shared verbatim with older MiniCPM-o diff --git a/vllm_omni/model_executor/models/moss_tts/__init__.py b/vllm_omni/model_executor/models/moss_tts/__init__.py index 23a73b3fe88..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/moss_tts/__init__.py +++ b/vllm_omni/model_executor/models/moss_tts/__init__.py @@ -1,13 +1,8 @@ -from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_codec import ( - MossTTSCodecDecoder, -) -from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import ( - MossTTSDelayTalkerForGeneration, - MossTTSRealtimeTalkerForGeneration, -) +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = [ - "MossTTSDelayTalkerForGeneration", - "MossTTSRealtimeTalkerForGeneration", - "MossTTSCodecDecoder", -] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py index a2cf75b07b1..b34b5b94431 100644 --- a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer.py @@ -45,6 +45,8 @@ def __init__( causal_transformer_context_duration: float = 10.0, encoder_kwargs: list[dict[str, Any]] | None = None, decoder_kwargs: list[dict[str, Any]] | None = None, + number_channels: int = 1, + enable_channel_interleave: bool = True, quantizer_type: str = "rlfq", quantizer_kwargs: dict[str, Any] | None = None, **kwargs: Any, @@ -54,6 +56,8 @@ def __init__( self.sampling_rate = sampling_rate self.downsample_rate = downsample_rate self.causal_transformer_context_duration = causal_transformer_context_duration + self.number_channels = number_channels + self.enable_channel_interleave = enable_channel_interleave self.encoder_kwargs = encoder_kwargs or _default_encoder_kwargs() self.decoder_kwargs = decoder_kwargs or _default_decoder_kwargs() if quantizer_kwargs is None: @@ -78,7 +82,7 @@ def num_quantizers(self) -> int: @property def codebook_size(self) -> int: - return int(self.quantizer_kwargs.get("codebook_size", 4096)) + return int(self.quantizer_kwargs.get("codebook_size", 1024)) @property def frame_rate(self) -> float: @@ -201,13 +205,21 @@ class _Attention(nn.Module): """Causal multi-head self-attention with RoPE, no streaming KV cache.""" def __init__( - self, embed_dim: int, num_heads: int, causal: bool, max_period: float, device=None, dtype=None + self, + embed_dim: int, + num_heads: int, + causal: bool, + max_period: float, + context: int | None = None, + device=None, + dtype=None, ) -> None: super().__init__() self.embed_dim = embed_dim self.num_heads = num_heads self.causal = causal self.max_period = max_period + self.context = context kw = {"device": device, "dtype": dtype} self.in_proj = nn.Linear(embed_dim, 3 * embed_dim, bias=False, **kw) self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False, **kw) @@ -234,7 +246,38 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: qkv = self.in_proj(x).reshape(B, T, 3, H, D).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] # each (B, H, T, D) q, k = _apply_rope(q, k, self.max_period) - out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal) + if self.context is not None and self.context < T: + # Local-windowed causal attention: query i may only see keys in + # [i - context + 1, i]. Matches upstream's per-stage receptive + # field — wider context windows are not numerically equivalent. + # + # A dense (T, T) mask is O(T^2) memory (e.g. 4 GiB at T=65536) — + # for long utterances decoded in a single non-streaming call this + # OOMs. Process queries in blocks against only their reachable + # key range so peak mask memory is O(block * (block + context)) + # regardless of T. + block = 4096 + if T <= block: + positions = torch.arange(T, device=x.device) + delta = positions.view(-1, 1) - positions.view(1, -1) + mask = (delta >= 0) & (delta < self.context) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + else: + outs = [] + for start in range(0, T, block): + end = min(start + block, T) + k_start = max(0, start - self.context + 1) + q_pos = torch.arange(start, end, device=x.device) + k_pos = torch.arange(k_start, end, device=x.device) + delta = q_pos.view(-1, 1) - k_pos.view(1, -1) + blk_mask = (delta >= 0) & (delta < self.context) + out_blk = F.scaled_dot_product_attention( + q[:, :, start:end], k[:, :, k_start:end], v[:, :, k_start:end], attn_mask=blk_mask + ) + outs.append(out_blk) + out = torch.cat(outs, dim=2) + else: + out = F.scaled_dot_product_attention(q, k, v, is_causal=self.causal) out = out.transpose(1, 2).reshape(B, T, self.embed_dim) return self.out_proj(out) @@ -248,6 +291,7 @@ def __init__( causal: bool, max_period: float, layer_scale: float | None, + context: int | None = None, device=None, dtype=None, ) -> None: @@ -255,7 +299,7 @@ def __init__( kw = {"device": device, "dtype": dtype} self.norm1 = nn.LayerNorm(d_model, eps=1e-5, **kw) self.norm2 = nn.LayerNorm(d_model, eps=1e-5, **kw) - self.attn = _Attention(d_model, num_heads, causal, max_period, device, dtype) + self.attn = _Attention(d_model, num_heads, causal, max_period, context, device, dtype) self.ff1 = nn.Linear(d_model, dim_feedforward, bias=False, **kw) self.ff2 = nn.Linear(dim_feedforward, d_model, bias=False, **kw) ls_kw = cast(dict[str, object], kw) @@ -282,13 +326,16 @@ def __init__( causal: bool, max_period: float, layer_scale: float | None, + context: int | None = None, device=None, dtype=None, ) -> None: super().__init__() self.layers = nn.ModuleList( [ - _TransformerLayer(d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, device, dtype) + _TransformerLayer( + d_model, num_heads, dim_feedforward, causal, max_period, layer_scale, context, device, dtype + ) for _ in range(num_layers) ] ) @@ -307,11 +354,12 @@ def __init__( ) -> None: super().__init__() self.downsample_ratio: int = 1 - self.in_proj = nn.Linear(input_dimension, d_model, bias=False) if d_model != input_dimension else nn.Identity() + # Upstream always materializes a learned Linear here regardless of + # whether input_dimension == d_model — substituting Identity when + # the dims happen to match silently drops a trained projection. + self.in_proj = nn.Linear(input_dimension, d_model, bias=False) self.transformer = _Transformer(d_model=d_model, **kwargs) - self.out_proj = ( - nn.Linear(d_model, output_dimension, bias=False) if d_model != output_dimension else nn.Identity() - ) + self.out_proj = nn.Linear(d_model, output_dimension, bias=False) def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: x = self.in_proj(x.transpose(1, 2)) # (B, D, T) → (B, T, d_model) @@ -456,10 +504,21 @@ def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: def _build_modules( - specs: list[dict[str, Any]], is_downsample: bool, context_duration: float, sampling_rate: float -) -> nn.ModuleList: + specs: list[dict[str, Any]], is_downsample: bool, default_context_duration: float, frame_rate: float +) -> tuple[nn.ModuleList, float]: + """Build encoder/decoder stage modules. + + ``frame_rate`` tracks the running frames-per-second at the current stage + (updated after each module by its downsample/upsample ratio, mirroring + upstream's ``current_frame_rate``). Each ``Transformer`` module's local + attention window (``context``, in frames) is derived from its own + ``context_duration`` (seconds) at the rate *at that point in the stack* — + not a global constant — so it must be computed here, not by the caller. + Returns the final rate too, since the decoder's starting rate is the + encoder's bottleneck rate, not the raw sampling rate. + """ modules: list[nn.Module] = [] - rate = float(sampling_rate) + rate = float(frame_rate) for spec in specs: spec = dict(spec) if spec["module_type"] == "PatchedPretransform": @@ -471,16 +530,18 @@ def _build_modules( spec.pop("positional_embedding", None) spec.pop("norm", None) spec.pop("causal", None) + context_duration = float(spec.pop("context_duration", default_context_duration)) m = _ProjectedTransformer( module_type="Transformer", causal=True, max_period=spec.pop("max_period", 10000), layer_scale=spec.pop("layer_scale", None), + context=int(round(rate * context_duration)), **spec, ) modules.append(m) rate = rate / m.downsample_ratio if is_downsample else rate * m.downsample_ratio - return nn.ModuleList(modules) + return nn.ModuleList(modules), rate class MossAudioTokenizerModel(PreTrainedModel): @@ -494,12 +555,27 @@ def __init__(self, config: MossAudioTokenizerConfig) -> None: super().__init__(config) self.sampling_rate = config.sampling_rate self.downsample_rate = config.downsample_rate + # Real v1 checkpoints store these as an explicit JSON `null` rather + # than omitting the key, so `getattr(..., default)` doesn't apply — + # normalize None to the same defaults here. + self.number_channels = getattr(config, "number_channels", 1) or 1 + self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True) + if self.enable_channel_interleave is None: + self.enable_channel_interleave = True ctx = config.causal_transformer_context_duration - sr = float(config.sampling_rate) - self.encoder = _build_modules(config.encoder_kwargs, is_downsample=True, context_duration=ctx, sampling_rate=sr) - self.decoder = _build_modules( - copy.deepcopy(config.decoder_kwargs), is_downsample=False, context_duration=ctx, sampling_rate=sr + channel_interleave_factor = ( + self.number_channels if (self.enable_channel_interleave and self.number_channels > 1) else 1 + ) + frame_rate = float(config.sampling_rate) * channel_interleave_factor + self.encoder, bottleneck_rate = _build_modules( + config.encoder_kwargs, is_downsample=True, default_context_duration=ctx, frame_rate=frame_rate + ) + self.decoder, _ = _build_modules( + copy.deepcopy(config.decoder_kwargs), + is_downsample=False, + default_context_duration=ctx, + frame_rate=bottleneck_rate, ) kw = dict(config.quantizer_kwargs) @@ -512,15 +588,29 @@ def batch_encode( wav_list: list[torch.Tensor], num_quantizers: int | None = None, ) -> MossAudioTokenizerEncoderOutput: - """Encode a list of 1-D waveform tensors → RVQ codes (NQ, B, T).""" + """Encode a list of waveform tensors → RVQ codes (NQ, B, T). + + Each element of ``wav_list`` is ``(T,)``/``(1, T)`` for a mono codec + (``number_channels == 1``) or ``(number_channels, T)`` otherwise. + """ device = wav_list[0].device B = len(wav_list) - max_len = max(w.shape[-1] for w in wav_list) - x = torch.zeros(B, 1, max_len, device=device) + C = self.number_channels + normalized: list[torch.Tensor] = [] lengths = torch.zeros(B, device=device, dtype=torch.long) for i, w in enumerate(wav_list): - x[i, 0, : w.shape[-1]] = w - lengths[i] = w.shape[-1] + if C == 1: + w_i = w.unsqueeze(0) if w.dim() == 1 else w + else: + if w.dim() != 2 or w.shape[0] != C: + raise ValueError(f"Expected wav_list[{i}] to have shape ({C}, T), got {tuple(w.shape)}.") + w_i = w + normalized.append(w_i) + lengths[i] = w_i.shape[-1] + max_len = int(lengths.max().item()) + x = torch.zeros(B, C, max_len, device=device, dtype=wav_list[0].dtype) + for i, w_i in enumerate(normalized): + x[i, :, : w_i.shape[-1]] = w_i return self._encode(x, lengths, num_quantizers) @torch.no_grad() @@ -541,6 +631,45 @@ def batch_decode( lengths[i] = c.shape[-1] return self._decode(codes, lengths) + def _flatten_channels_for_codec( + self, + input_values: torch.Tensor, + input_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Pad to a ``downsample_rate`` multiple, then interleave channels into one stream. + + Mirrors upstream's ``_flatten_channels_for_codec``: for a stereo + codec, ``(B, C, T)`` becomes ``(B, 1, T*C)`` with samples interleaved + frame-by-frame across channels, so the mono encoder/decoder stack can + process a multi-channel waveform unchanged. + """ + if input_values.shape[-1] % self.downsample_rate != 0: + pad_length = self.downsample_rate - (input_values.shape[-1] % self.downsample_rate) + input_values = F.pad(input_values, (0, pad_length)) + if self.number_channels > 1 and self.enable_channel_interleave: + input_values = input_values.transpose(1, 2).contiguous().view(input_values.shape[0], 1, -1) + input_lengths = input_lengths * self.number_channels + return input_values, input_lengths + + def _restore_channels_from_codec( + self, + output_values: torch.Tensor, + output_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Inverse of ``_flatten_channels_for_codec``: de-interleave back to ``(B, C, T)``.""" + if self.number_channels == 1 or not self.enable_channel_interleave: + return output_values.float(), output_lengths + output_values = ( + output_values.squeeze(1) + .contiguous() + .view(output_values.shape[0], -1, self.number_channels) + .transpose(1, 2) + .contiguous() + .float() + ) + output_lengths = torch.div(output_lengths, self.number_channels, rounding_mode="floor") + return output_values, output_lengths + def _encode( self, x: torch.Tensor, @@ -549,9 +678,7 @@ def _encode( ) -> MossAudioTokenizerEncoderOutput: if x.dim() == 2: x = x.unsqueeze(1) - T = x.shape[-1] - if T % self.downsample_rate: - x = F.pad(x, (0, self.downsample_rate - T % self.downsample_rate)) + x, lengths = self._flatten_channels_for_codec(x, lengths) e, e_len = x, lengths for m in self.encoder: e, e_len = m(e, e_len) @@ -567,6 +694,7 @@ def _decode( d, d_len = z, lengths for m in self.decoder: d, d_len = m(d, d_len) + d, d_len = self._restore_channels_from_codec(d, d_len) return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_len) diff --git a/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py new file mode 100644 index 00000000000..f3ba18c3677 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/audio_tokenizer_v2.py @@ -0,0 +1,1959 @@ +# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch MossAudioTokenizer model.""" + +from __future__ import annotations + +import copy +import math +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from typing import cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.modeling_utils import PreTrainedAudioTokenizerBase +from transformers.utils import ModelOutput, auto_docstring, logging + +from .configuration_moss_audio_tokenizer_v2 import MossAudioTokenizerConfig + +logger = logging.get_logger(__name__) + + +# ============================================================================= +# Output Classes +# ============================================================================= + + +@dataclass +@auto_docstring +class MossAudioTokenizerEncoderOutput(ModelOutput): + r""" + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*): + Discrete audio codes computed using the encoder and quantizer. + audio_codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio codes. + encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, hidden_size, sequence_length)`, *optional*): + Hidden states from the encoder before quantization. + """ + + audio_codes: torch.Tensor | None = None + audio_codes_lengths: torch.Tensor | None = None + encoder_hidden_states: torch.Tensor | None = None + + +@dataclass +@auto_docstring +class MossAudioTokenizerDecoderOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Decoded audio waveform. + audio_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio. + """ + + audio: torch.Tensor | None = None + audio_lengths: torch.Tensor | None = None + + +@dataclass +@auto_docstring +class MossAudioTokenizerOutput(ModelOutput): + r""" + audio (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Decoded audio waveform. + audio_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio. + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*): + Discrete audio codes computed using the encoder and quantizer. + audio_codes_lengths (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Valid lengths for each sample's audio codes. + """ + + audio: torch.Tensor | None = None + audio_lengths: torch.Tensor | None = None + audio_codes: torch.Tensor | None = None + audio_codes_lengths: torch.Tensor | None = None + + +# ============================================================================= +# Streaming Module Base Classes +# ============================================================================= + + +@dataclass +class StreamingState: + """Base state for streaming modules.""" + + batch_size: int + device: torch.device + + def __post_init__(self): + self.exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=self.device) + + def set_exec_mask(self, exec_mask: torch.Tensor): + self.exec_mask[:] = exec_mask + + def reset(self, reset_mask: torch.Tensor) -> None: + self.exec_mask[:] = torch.where(reset_mask, torch.ones_like(self.exec_mask), self.exec_mask) + + def __enter__(self): + # ExitStack expects a context manager; returning self is conventional and useful for debugging. + return self + + def __exit__(self, exc_type, exc_value, traceback) -> None: + pass + + +class StreamingModule(nn.Module): + """Base class for streaming components.""" + + def __init__(self) -> None: + super().__init__() + self._streaming_state: StreamingState | None = None + self._streaming_detached: bool = False + self._cached_children: list[tuple[str, StreamingModule]] | None = None + + @property + def is_streaming(self): + return self._streaming_state is not None + + def _apply_named_streaming(self, fn): + def _handle_module(prefix: str, module: nn.Module): + if isinstance(module, StreamingModule): + if module._streaming_detached and prefix != "": + return + if self._cached_children is None: + raise RuntimeError("Internal error: _cached_children should be initialized before traversal.") + self._cached_children.append((prefix, module)) + for name, child in module.named_children(): + new_prefix = f"{prefix}.{name}" if prefix else name + _handle_module(new_prefix, child) + + if self._cached_children is None: + self._cached_children = [] + _handle_module("", self) + for name, child in self._cached_children: + fn(name, child) + + def _start_streaming(self, batch_size: int, exit_stack: ExitStack): + def _start_streaming_fn(name: str, module: StreamingModule): + if module._streaming_state is not None: + raise RuntimeError(f"{name} is already streaming!") + state = module._init_streaming_state(batch_size) + exit_stack.enter_context(state) + module._streaming_state = state + + self._apply_named_streaming(_start_streaming_fn) + + def _stop_streaming(self) -> None: + def _stop_streaming_fn(name: str, module: StreamingModule): + module._streaming_state = None + + self._apply_named_streaming(_stop_streaming_fn) + + def _init_streaming_state(self, batch_size: int) -> StreamingState: + device = next(iter(self.parameters())).device + return StreamingState(batch_size, device) + + def streaming(self, batch_size: int) -> ExitStack: + """Context manager to enter streaming mode.""" + exit_stack = ExitStack() + self._start_streaming(batch_size, exit_stack) + exit_stack.callback(self._stop_streaming) + return exit_stack + + def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: + """Set the active slot mask on every live streaming module.""" + + def _set_exec_mask(_: str, module: StreamingModule) -> None: + state = module._streaming_state + if state is not None: + state.set_exec_mask(exec_mask.to(state.device)) + + self._apply_named_streaming(_set_exec_mask) + + +class StreamingContainer(StreamingModule): + """Container for streaming modules.""" + + pass + + +# ============================================================================= +# Normalization Layers +# ============================================================================= + + +class MossAudioTokenizerRMSNorm(nn.Module): + """Root Mean Square Layer Normalization.""" + + def __init__( + self, + dim: int, + eps: float = 1e-5, + dtype: torch.dtype | None = None, + device=None, + ): + super().__init__() + self.eps = eps + self.dtype = dtype + self.alpha = nn.Parameter(torch.full((1, 1, dim), 1.0, requires_grad=True, device=device, dtype=dtype)) + + def forward(self, x: torch.Tensor): + x_dtype = x.dtype + if self.dtype is not None: + x = x.to(self.dtype) + var = self.eps + torch.mean(x**2, dim=2, keepdim=True) + y = (x * (self.alpha.to(var) * torch.rsqrt(var))).to(x_dtype) + return y + + +class MossAudioTokenizerLayerScale(nn.Module): + """Layer scale from Touvron et al. 2021.""" + + def __init__( + self, + channels: int, + init: float = 1e-4, + channel_last: bool = True, + device=None, + dtype=None, + ): + super().__init__() + self.channel_last = channel_last + self.scale = nn.Parameter(torch.full((channels,), init, requires_grad=True, device=device, dtype=dtype)) + + def forward(self, x: torch.Tensor): + if self.channel_last: + return self.scale * x + else: + return self.scale[:, None] * x + + +def create_norm_fn(norm_type: str, dim: int, **kwargs) -> nn.Module: + """Create normalization module.""" + if norm_type == "layer_norm": + return nn.LayerNorm(dim, eps=1e-5, **kwargs) + elif norm_type in {"rms_norm"}: + return MossAudioTokenizerRMSNorm(dim, eps=1e-5, **kwargs) + elif norm_type in {"rms_norm_f32"}: + kwargs.pop("dtype", None) + return MossAudioTokenizerRMSNorm(dim, eps=1e-8, dtype=torch.float, **kwargs) + else: + raise ValueError(f"Unknown norm type: {norm_type}") + + +# ============================================================================= +# Rotary Position Embedding +# ============================================================================= + + +def apply_rope( + q: torch.Tensor, + k: torch.Tensor, + offset: torch.Tensor, + max_period: float = 10_000, + time_before_heads: bool = False, +): + """Apply rotary position embedding.""" + if time_before_heads: + B, T, H, D = q.shape + else: + B, H, T, D = q.shape + if k.shape != q.shape: + raise ValueError(f"Expected k.shape == q.shape, got k={tuple(k.shape)} q={tuple(q.shape)}") + if D <= 0 or (D % 2) != 0: + raise ValueError(f"RoPE requires an even last dimension, got D={D}") + + ds = torch.arange(D // 2, device=q.device, dtype=torch.float32) + freqs = torch.exp(ds * (-math.log(max_period) * 2 / D)) + ts = offset.float().view(-1, 1) + torch.arange(T, device=q.device, dtype=torch.float32) + + if time_before_heads: + ts = ts.view(B, -1, 1, 1) + else: + ts = ts.view(B, 1, -1, 1) + + dims = q.shape[:-1] + q = q.view(*dims, D // 2, 2) + k = k.view(*dims, D // 2, 2) + + qr, qi = q[..., 0].float(), q[..., 1].float() + kr, ki = k[..., 0].float(), k[..., 1].float() + + rotr = torch.cos(freqs * ts) + roti = torch.sin(freqs * ts) + + qor = qr * rotr - qi * roti + qoi = qr * roti + qi * rotr + kor = kr * rotr - ki * roti + koi = kr * roti + ki * rotr + + dtype = q.dtype + qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1) + ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1) + + return qo.view(*dims, D), ko.view(*dims, D) + + +class MossAudioTokenizerRotaryEmbedding(nn.Module): + """Rotary positional embedding (RoPE).""" + + def __init__(self, max_period: float = 10000.0): + super().__init__() + self.max_period = max_period + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + offset: torch.Tensor, + time_before_heads: bool = False, + ): + return apply_rope(q, k, offset, self.max_period, time_before_heads) + + +# ============================================================================= +# Gating Modules +# ============================================================================= + + +class MossAudioTokenizerActivationGating(nn.Module): + """Gating FFN layer with activation.""" + + def __init__(self, dim: int, dim_feedforward: int, activation, **factory_kwargs): + super().__init__() + if dim_feedforward == 4 * dim: + hidden = (21 * dim) // 8 + else: + hidden = (2 * dim_feedforward) // 3 + + self.linear_in = nn.Linear(dim, 2 * hidden, bias=False, **factory_kwargs) + self.linear_out = nn.Linear(hidden, dim, bias=False, **factory_kwargs) + self.activation = activation + + def forward(self, x: torch.Tensor): + x = self.linear_in(x) + B, T, _ = x.shape + x = x.view(B, T, 2, -1) + x = self.activation(x[..., 0, :]) * x[..., 1, :] + x = self.linear_out(x) + return x + + +def _get_activation(name: str): + if name in ["sigmoid", "tanh", "relu"]: + return getattr(torch, name) + elif name in ["leaky_relu", "elu", "gelu", "silu", "mish", "softsign"]: + return getattr(F, name) + elif name == "identity": + return nn.Identity() + else: + raise ValueError(f"Unknown activation {name}") + + +def make_gating(name: str, dim: int, dim_feedforward: int, **factory_kwargs) -> nn.Module: + return MossAudioTokenizerActivationGating(dim, dim_feedforward, _get_activation(name), **factory_kwargs) + + +# ============================================================================= +# Positional Embeddings +# ============================================================================= + + +def create_sin_embedding( + positions: torch.Tensor, + dim: int, + max_period: float = 10000, + dtype: torch.dtype = torch.float32, +) -> torch.Tensor: + """Create sinusoidal positional embedding with shape [B, T, C].""" + if dim % 2 != 0: + raise ValueError(f"Sinusoidal embedding requires even dim, got dim={dim}") + half_dim = dim // 2 + if half_dim <= 1: + raise ValueError(f"Sinusoidal embedding requires dim >= 4, got dim={dim}") + positions = positions.to(dtype) + adim = torch.arange(half_dim, device=positions.device, dtype=dtype).view(1, 1, -1) + max_period_tensor = torch.full([], max_period, device=positions.device, dtype=dtype) + phase = positions / (max_period_tensor ** (adim / (half_dim - 1))) + return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1) + + +# ============================================================================= +# KV Cache for Attention +# ============================================================================= + + +class KVCacheResult: + """Container for KV cache results that supports tuple unpacking.""" + + __slots__ = ("keys", "values", "positions") + + def __init__(self, keys: torch.Tensor, values: torch.Tensor, positions: torch.Tensor): + self.keys = keys + self.values = values + self.positions = positions + + def __iter__(self): + """Allow unpacking as (keys, values, positions).""" + return iter((self.keys, self.values, self.positions)) + + @staticmethod + def from_kv(keys: torch.Tensor, values: torch.Tensor) -> KVCacheResult: + B, H, T, D = keys.shape + positions = torch.arange(T, device=keys.device, dtype=torch.long) + return KVCacheResult(keys, values, positions.expand(B, -1)) + + +class RingKVCache: + """Efficient streaming KVCache compatible with CUDA Graph.""" + + def __init__( + self, + batch_size: int, + num_heads: int, + dim_per_head: int, + capacity: int, + respect_exec_mask: bool = True, + device: torch.device = torch.device("cuda"), + dtype: torch.dtype = torch.bfloat16, + ): + self.capacity = capacity + self.cache = torch.zeros( + (2, batch_size, num_heads, capacity, dim_per_head), + device=device, + dtype=dtype, + ) + self.respect_exec_mask = respect_exec_mask + if self.respect_exec_mask: + self.end_offset = torch.zeros(batch_size, device=device, dtype=torch.long) + else: + self.end_offset = torch.zeros(1, device=device, dtype=torch.long) + + def reset(self, reset_mask: torch.Tensor) -> None: + self.end_offset[:] = torch.where(reset_mask, torch.zeros_like(self.end_offset), self.end_offset) + + def complete(self, k: torch.Tensor, v: torch.Tensor, exec_mask: torch.Tensor) -> KVCacheResult: + B, H, T, D = k.shape + if T <= 0: + raise ValueError(f"Expected T > 0, got T={T}") + + indexes = torch.arange(T, device=self.end_offset.device, dtype=self.end_offset.dtype) + indexes = indexes + self.end_offset.view(-1, 1) + indexes = indexes % self.capacity + + if self.respect_exec_mask: + this_indexes = indexes.view(B, 1, T, 1).expand(-1, H, T, D) + self.cache[0].scatter_(2, this_indexes, k) + self.cache[1].scatter_(2, this_indexes, v) + else: + self.cache[0].index_copy_(2, indexes[0], k) + self.cache[1].index_copy_(2, indexes[0], v) + + keys = self.cache[0] + values = self.cache[1] + + indexes = torch.arange(self.capacity, device=self.end_offset.device, dtype=torch.long) + last_offset = self.end_offset.view(-1, 1) + T - 1 + end_index = last_offset % self.capacity + delta = indexes - end_index + + positions = torch.where( + delta <= 0, + last_offset + delta, + last_offset + delta - self.capacity, + ) + + if self.respect_exec_mask: + self.end_offset[:] = torch.where(exec_mask, self.end_offset + T, self.end_offset) + else: + self.end_offset.add_(T) + + invalid = indexes >= self.end_offset.view(-1, 1) + positions = torch.where(invalid, torch.full_like(positions, -1), positions) + + return KVCacheResult(keys, values, positions) + + +# ============================================================================= +# Multi-Head Attention +# ============================================================================= + + +@dataclass +class MHAState(StreamingState): + kv_cache: RingKVCache | None + offset: torch.Tensor + offset_cpu: int + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offset[:] = torch.where(reset_mask, torch.zeros_like(self.offset), self.offset) + if self.kv_cache is not None: + self.kv_cache.reset(reset_mask) + self.offset_cpu = 0 + + +def apply_weights_per_step( + modules: nn.ModuleList, + schedule: list[int] | None, + x: torch.Tensor, + offset: int | None, +) -> torch.Tensor: + """Apply different weights for each time step.""" + if len(modules) == 1: + return modules[0](x) + + if offset is None: + raise ValueError("offset must be provided when using per-step weights (len(modules) > 1).") + ys = [] + B, T, C = x.shape + for t in range(T): + module_index = t + offset + if schedule is not None: + if module_index >= len(schedule) or module_index < 0: + raise ValueError( + f"weights_per_step_schedule is too short for module_index={module_index} (len={len(schedule)})." + ) + module_index = schedule[module_index] + if module_index >= len(modules) or module_index < 0: + raise ValueError(f"module_index={module_index} out of range for len(modules)={len(modules)}.") + y = modules[module_index](x[:, t : t + 1]) + ys.append(y) + return torch.cat(ys, 1) + + +class MossAudioTokenizerMultiheadAttention(StreamingModule): + """Multi-head attention with streaming support.""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + causal: bool = False, + context: int | None = None, + rope: MossAudioTokenizerRotaryEmbedding | None = None, + weights_per_step: int = 0, + weights_per_step_schedule: list[int] | None = None, + device=None, + dtype=None, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": dtype} + + self.embed_dim = embed_dim + self.causal = causal + self.context = context + self.rope = rope + self.num_heads = num_heads + self.weights_per_step = weights_per_step + self.weights_per_step_schedule = weights_per_step_schedule + + out_dim = 3 * embed_dim + mult = 1 + if weights_per_step: + mult = max(weights_per_step_schedule) + 1 if weights_per_step_schedule else weights_per_step + self.mult = mult + + self.out_projs = nn.ModuleList( + [nn.Linear(embed_dim, embed_dim, bias=False, **factory_kwargs) for _ in range(mult)] + ) + self.in_projs = nn.ModuleList( + [nn.Linear(embed_dim, out_dim, bias=False, **factory_kwargs) for _ in range(mult)] + ) + + self._register_load_state_dict_pre_hook(self._load_hook, with_module=True) + + @staticmethod + def _load_hook(module, state_dict, prefix, *_): + mappings = { + "in_proj_weight": "in_projs.{i}.weight", + "in_proj.weight": "in_projs.{i}.weight", + "out_proj.weight": "out_projs.{i}.weight", + } + mult = module.mult + for suffix in ["", "_scb"]: + for source, target in mappings.items(): + this_source = prefix + source + suffix + if this_source in state_dict: + weight = state_dict[this_source] + _, *OD = weight.shape + weight = weight.view(mult, -1, *OD) + for i in range(mult): + state_dict[prefix + target.format(i=i) + suffix] = weight[i] + state_dict.pop(this_source) + + def _init_streaming_state(self, batch_size: int) -> MHAState: + in_proj = cast(nn.Linear, self.in_projs[0]) + device = cast(torch.device, in_proj.weight.device) + dtype = cast(torch.dtype, in_proj.weight.dtype) + + dim_per_head = self.embed_dim // self.num_heads + if self.context is None: + capacity = self.weights_per_step if self.weights_per_step else 1024 + else: + capacity = self.context + + kv_cache = RingKVCache( + batch_size, + self.num_heads, + dim_per_head, + capacity, + respect_exec_mask=not self.weights_per_step, + device=cast(torch.device, device), + dtype=cast(torch.dtype, dtype), + ) + return MHAState( + batch_size, + cast(torch.device, device), + kv_cache, + offset=torch.zeros(batch_size, device=cast(torch.device, device), dtype=torch.long), + offset_cpu=0, + ) + + def _complete_kv(self, k, v) -> KVCacheResult: + state = cast(MHAState | None, self._streaming_state) + if state is None: + return KVCacheResult.from_kv(k, v) + if state.kv_cache is None: + return KVCacheResult.from_kv(k, v) + return state.kv_cache.complete(k, v, state.exec_mask) + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor): + state = cast(MHAState | None, self._streaming_state) + B, T = query.shape[:2] + + if state is None: + offset = torch.zeros(B, device=query.device, dtype=torch.long) + offset_cpu = 0 + else: + offset = state.offset + offset_cpu = state.offset_cpu + + projected = apply_weights_per_step(self.in_projs, self.weights_per_step_schedule, query, offset_cpu) + dim_per_head = self.embed_dim // self.num_heads + projected = projected.reshape(B, T, 3, self.num_heads, dim_per_head).permute(2, 0, 3, 1, 4) + q, k, v = projected[0], projected[1], projected[2] + + if self.rope: + q, k = self.rope(q, k, offset, time_before_heads=False) + + k, v, pos_k = self._complete_kv(k, v) + pos_k = pos_k[:, None] + + if self.causal: + pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(-1, 1) + delta = pos_q - pos_k + attn_bias = (pos_k >= 0) & (delta >= 0) + if self.context is not None: + attn_bias = attn_bias & (delta < self.context) + attn_bias = attn_bias[:, None] + else: + attn_bias = None + + x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0) + x = x.transpose(1, 2).reshape(B, T, self.embed_dim) + x = apply_weights_per_step(self.out_projs, self.weights_per_step_schedule, x, offset_cpu) + + if state is not None: + state.offset[:] = torch.where(state.exec_mask, state.offset + T, state.offset) + state.offset_cpu += T + return x + + +# ============================================================================= +# Transformer Layer +# ============================================================================= + + +@dataclass +class LayerState(StreamingState): + offset_cpu: int = 0 + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offset_cpu = 0 + + +class MossAudioTokenizerTransformerLayer(StreamingModule): + """Transformer layer with streaming support.""" + + def __init__( + self, + d_model: int, + num_heads: int, + dim_feedforward: int = 2048, + causal: bool = False, + context: int | None = None, + rope: MossAudioTokenizerRotaryEmbedding | None = None, + norm: str = "layer_norm", + layer_scale: float | None = None, + gating: str = "none", + weights_per_step: int = 0, + weights_per_step_schedule: list[int] | None = None, + activation=F.gelu, + device=None, + dtype=None, + ): + super().__init__() + factory_kwargs = {"device": device, "dtype": dtype} + + self.self_attn = MossAudioTokenizerMultiheadAttention( + embed_dim=d_model, + num_heads=num_heads, + causal=causal, + context=context, + rope=rope, + weights_per_step=weights_per_step, + weights_per_step_schedule=weights_per_step_schedule, + **factory_kwargs, + ) + self.norm1 = create_norm_fn(norm, d_model, **factory_kwargs) + self.norm2 = create_norm_fn(norm, d_model, **factory_kwargs) + + self.weights_per_step = weights_per_step + self.weights_per_step_schedule = weights_per_step_schedule + self.gating: nn.Module | nn.ModuleList | None = None + self.linear1: nn.Module | None = None + self.linear2: nn.Module | None = None + self.activation = activation + + num_weights = 1 + if weights_per_step: + num_weights = max(weights_per_step_schedule) + 1 if weights_per_step_schedule else weights_per_step + + if gating == "none": + self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False, **factory_kwargs) + self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False, **factory_kwargs) + else: + if weights_per_step: + dim_ff_list = [dim_feedforward] * num_weights if isinstance(dim_feedforward, int) else dim_feedforward + self.gating = nn.ModuleList( + [make_gating(gating, d_model, dim, **factory_kwargs) for dim in dim_ff_list] + ) + else: + self.gating = make_gating(gating, d_model, dim_feedforward, **factory_kwargs) + + if layer_scale is None: + self.layer_scale_1 = nn.Identity() + self.layer_scale_2 = nn.Identity() + else: + self.layer_scale_1 = MossAudioTokenizerLayerScale( + channels=d_model, init=layer_scale, channel_last=True, **cast(dict[str, object], factory_kwargs) + ) + self.layer_scale_2 = MossAudioTokenizerLayerScale( + channels=d_model, init=layer_scale, channel_last=True, **cast(dict[str, object], factory_kwargs) + ) + + def _init_streaming_state(self, batch_size: int) -> LayerState: + device = next(iter(self.parameters())).device + return LayerState(batch_size, device, offset_cpu=0) + + def _ff_block(self, x: torch.Tensor) -> torch.Tensor: + state = self._streaming_state + offset = state.offset_cpu if isinstance(state, LayerState) else 0 + + x_orig = x + x = self.norm2(x) + + if self.gating is None: + assert self.linear1 is not None + assert self.linear2 is not None + update = self.linear2(self.activation(self.linear1(x))) + else: + if self.weights_per_step: + assert isinstance(self.gating, nn.ModuleList) + update = apply_weights_per_step(self.gating, self.weights_per_step_schedule, x, offset) + else: + update = self.gating(x) + return x_orig.to(update) + self.layer_scale_2(update) + + def _sa_block(self, x: torch.Tensor): + x_orig = x + x = self.norm1(x) + update = self.self_attn(x, x, x) + return x_orig.to(update) + self.layer_scale_1(update) + + def forward(self, x: torch.Tensor): + x = self._sa_block(x) + x = self._ff_block(x) + state = self._streaming_state + if state is not None: + assert isinstance(state, LayerState) + state.offset_cpu += x.shape[1] + return x + + +# ============================================================================= +# Streaming Transformer +# ============================================================================= + + +@dataclass +class TransformerState(StreamingState): + offsets: torch.Tensor + + def reset(self, reset_mask: torch.Tensor): + super().reset(reset_mask) + self.offsets[:] = torch.where(reset_mask, torch.zeros_like(self.offsets), self.offsets) + + +class MossAudioTokenizerTransformer(StreamingModule): + """Transformer with streaming/causal support.""" + + def __init__( + self, + d_model: int, + num_heads: int, + num_layers: int, + dim_feedforward: int = 2048, + causal: bool = False, + context: int | None = None, + positional_embedding: str = "sin", + max_period: float = 10_000, + positional_scale: float = 1.0, + device=None, + dtype=None, + **kwargs, + ): + super().__init__() + if d_model % num_heads != 0: + raise ValueError(f"d_model must be divisible by num_heads, got d_model={d_model}, num_heads={num_heads}") + + self.positional_embedding = positional_embedding + self.max_period = max_period + self.positional_scale = positional_scale + + self.rope: MossAudioTokenizerRotaryEmbedding | None = None + if positional_embedding in {"rope", "sin_rope"}: + self.rope = MossAudioTokenizerRotaryEmbedding(max_period=max_period) + + self.layers = nn.ModuleList() + for _ in range(num_layers): + self.layers.append( + MossAudioTokenizerTransformerLayer( + d_model=d_model, + num_heads=num_heads, + dim_feedforward=dim_feedforward, + causal=causal, + context=context, + rope=self.rope, + device=device, + dtype=dtype, + **kwargs, + ) + ) + + def _init_streaming_state(self, batch_size: int) -> TransformerState: + device = next(self.parameters()).device + return TransformerState( + batch_size, + device, + offsets=torch.zeros(batch_size, device=device, dtype=torch.long), + ) + + def forward(self, x: torch.Tensor, *args, **kwargs): + B, T, C = x.shape + state = self._streaming_state + offsets = ( + torch.zeros(1, dtype=torch.long, device=x.device) + if state is None + else ( + state.offsets + if isinstance(state, TransformerState) + else torch.zeros(1, dtype=torch.long, device=x.device) + ) + ) + + if self.positional_embedding in {"sin", "sin_rope"}: + positions = torch.arange(T, device=x.device).view(1, -1, 1) + positions = positions + offsets.view(-1, 1, 1) + pos_emb = create_sin_embedding(positions, C, max_period=self.max_period, dtype=x.dtype) + x = x + self.positional_scale * pos_emb + + for layer in self.layers: + x = layer(x, *args, **kwargs) + + if state is not None: + assert isinstance(state, TransformerState) + state.offsets[:] = torch.where(state.exec_mask, state.offsets + T, state.offsets) + return x + + +class MossAudioTokenizerProjectedTransformer(StreamingContainer): + """Transformer with input/output projections.""" + + def __init__( + self, + input_dimension: int, + output_dimension: int, + d_model: int, + *, + conv_layout: bool = False, + module_type: str, + **kwargs, + ): + super().__init__() + self.module_type = module_type + self.downsample_ratio: int = 1 + self.input_dimension = input_dimension + self.output_dimension = output_dimension + + self.input_proj = nn.Linear(input_dimension, d_model, bias=False) + self.transformer = MossAudioTokenizerTransformer(d_model=d_model, **kwargs) + self.conv_layout = conv_layout + self.output_proj = nn.Linear(d_model, output_dimension, bias=False) + + def forward(self, x, input_lengths, *args, **kwargs): + x = self.input_proj(x.transpose(1, 2)) # (B, D, T) -> (B, T, D) + x = self.transformer(x, *args, **kwargs) + x = self.output_proj(x).transpose(1, 2) # (B, T, D) -> (B, D, T) + return x, input_lengths + + +# ============================================================================= +# Patched Pretransform Module +# ============================================================================= + + +class MossAudioTokenizerPatchedPretransform(nn.Module): + """Patching module for downsampling/upsampling.""" + + def __init__(self, patch_size: int, is_downsample: bool, module_type: str, **kwargs): + super().__init__() + self.patch_size = patch_size + self.downsample_ratio: int = patch_size + self.is_downsample = is_downsample + self.module_type = module_type + + def encode(self, x, input_lengths): + b, d, _ = x.shape + h = self.patch_size + x = x.reshape(b, d, -1, h).permute(0, 1, 3, 2).reshape(b, d * h, -1) + # We pad the input waveform to a multiple of `downsample_rate` before applying the encoder. + # Use a ceil division to match that padding and avoid dropping the last (partially padded) frame. + output_lengths = input_lengths // self.patch_size + return x, output_lengths + + def decode(self, x, input_lengths): + b, dh, length = x.shape + h = self.patch_size + d = dh // h + x = x.reshape(b, d, h, length).permute(0, 1, 3, 2).reshape(b, d, length * h) + output_lengths = input_lengths * self.patch_size + return x, output_lengths + + def forward(self, x, input_lengths): + if self.is_downsample: + return self.encode(x, input_lengths) + else: + return self.decode(x, input_lengths) + + +# ============================================================================= +# Vector Quantization +# ============================================================================= + + +def WNConv1d(*args, **kwargs): + """Weight-normalized Conv1d.""" + return nn.utils.parametrizations.weight_norm(nn.Conv1d(*args, **kwargs)) + + +class MossAudioTokenizerVectorQuantize(nn.Module): + """Single codebook vector quantization (inference only).""" + + def __init__( + self, + input_dim: int, + codebook_size: int, + codebook_dim: int, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + if input_dim != codebook_dim: + self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1) + self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1) + else: + self.in_proj = nn.Identity() + self.out_proj = nn.Identity() + + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + @torch.no_grad() + def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Args: + z: Input tensor of shape (B, D, T) + Returns: + z_q: Quantized tensor of shape (B, D, T) + indices: Code indices of shape (B, T) + z_e: Encoded tensor before quantization + """ + z = z.float() + z_e = self.in_proj(z).float() + + encodings = z_e.transpose(1, 2).reshape(-1, z_e.shape[1]) + + codebook_weight = self.codebook.weight + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook_weight.float().t() + + codebook_weight.float().pow(2).sum(1, keepdim=True).t() + ) + + indices = (-dist).max(1)[1] + indices = indices.reshape(z.size(0), -1) + + z_q = self.decode_code(indices) + z_q = self.out_proj(z_q).float() + + return z_q, indices, z_e + + def decode_code(self, embed_id: torch.Tensor) -> torch.Tensor: + """Decode code indices to embeddings.""" + return self.codebook(embed_id).transpose(1, 2).float() + + +class MossAudioTokenizerLFQ(nn.Module): + """LFQ (inference-only) used by ResidualLFQ.""" + + def __init__( + self, + input_dim: int, + codebook_size: int, + codebook_dim: int, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + if self.input_dim != self.codebook_dim: + self.in_proj = WNConv1d(self.input_dim, self.codebook_dim, kernel_size=1) + self.out_proj = WNConv1d(self.codebook_dim, self.input_dim, kernel_size=1) + else: + self.in_proj = nn.Identity() + self.out_proj = nn.Identity() + + self.codebook = nn.Embedding(codebook_size, codebook_dim) + + @torch.no_grad() + def forward(self, z: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize z into codebook vectors.""" + z = z.float() + z_e = self.in_proj(z).float() + z_q, indices = self.decode_latents(z_e) + z_q = (z_e + (z_q - z_e).detach()).float() + z_q = self.out_proj(z_q).float() + return z_q, indices, z_e + + def embed_code(self, embed_id: torch.Tensor) -> torch.Tensor: + return F.embedding(embed_id, self.codebook.weight) + + def decode_code_wo_out_proj(self, embed_id: torch.Tensor) -> torch.Tensor: + return self.embed_code(embed_id).transpose(1, 2) + + def decode_code(self, embed_id: torch.Tensor) -> torch.Tensor: + z_q = self.decode_code_wo_out_proj(embed_id).float() + z_q = self.out_proj(z_q).float() + return z_q + + def decode_latents(self, latents: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Match training LFQ: L2-normalize then argmin squared distance.""" + encodings = latents.transpose(1, 2).reshape(-1, latents.shape[1]).float() + codebook = self.codebook.weight.float() + + encodings = F.normalize(encodings) + codebook = F.normalize(codebook) + + dist = ( + encodings.pow(2).sum(1, keepdim=True) + - 2 * encodings @ codebook.t() + + codebook.pow(2).sum(1, keepdim=True).t() + ) + indices = (-dist).max(1)[1] + indices = indices.reshape(latents.size(0), -1) + z_q = self.decode_code_wo_out_proj(indices).float() + return z_q, indices + + +class MossAudioTokenizerResidualVQ(nn.Module): + """Residual Vector Quantization (inference only).""" + + def __init__( + self, + input_dim: int = 1024, + rvq_dim: int | None = None, + output_dim: int | None = None, + num_quantizers: int = 32, + codebook_size: int = 1024, + codebook_dim: int = 8, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.rvq_dim = rvq_dim or input_dim + self.output_dim = output_dim or input_dim + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + self.input_proj = ( + WNConv1d(input_dim, self.rvq_dim, kernel_size=1) if input_dim != self.rvq_dim else nn.Identity() + ) + self.output_proj = ( + WNConv1d(self.rvq_dim, self.output_dim, kernel_size=1) if self.rvq_dim != self.output_dim else nn.Identity() + ) + + self.quantizers = nn.ModuleList( + [ + MossAudioTokenizerVectorQuantize( + input_dim=self.rvq_dim, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + **kwargs, + ) + for _ in range(num_quantizers) + ] + ) + + @torch.no_grad() + def forward( + self, + z: torch.Tensor, + input_length: torch.Tensor, + n_quantizers: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Args: + z: Input tensor of shape (B, D, T) + input_length: Valid lengths for each sample (B,) + n_quantizers: Number of quantizers to use + Returns: + quantized_out: Quantized output (B, D, T) + all_indices: All code indices (N, B, T) + output_length: Output lengths (B,) + """ + z = self.input_proj(z) + + batch_size, _, max_time = z.shape + mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1) + + quantized_out = torch.zeros_like(z, dtype=torch.float32) + residual = z.clone().float() + all_indices = [] + + n_quantizers = n_quantizers or self.num_quantizers + + for i, quantizer in enumerate(self.quantizers): + if i >= n_quantizers: + break + + masked_residual = residual * mask.unsqueeze(1) + z_q_i, indices_i, _ = quantizer(masked_residual) + + update_mask = mask.unsqueeze(1) + quantized_out = quantized_out + z_q_i * update_mask + residual = residual - z_q_i * update_mask + all_indices.append(indices_i) + + all_indices = torch.stack(all_indices) # (N, B, T) + quantized_out = self.output_proj(quantized_out) + + return quantized_out, all_indices, input_length + + def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: + """Decode codes from multiple quantizers to embeddings.""" + nq, B, T = codes.shape + emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) + + for i, quantizer in enumerate(self.quantizers[:nq]): + quantizer = cast(MossAudioTokenizerVectorQuantize, quantizer) + quantized_i = quantizer.decode_code(codes[i]) + emb += quantized_i + + emb = self.output_proj(emb) + return emb + + +class MossAudioTokenizerResidualLFQ(nn.Module): + """Residual LFQ (inference only).""" + + def __init__( + self, + input_dim: int = 1024, + rvq_dim: int | None = None, + output_dim: int | None = None, + num_quantizers: int = 32, + codebook_size: int = 1024, + codebook_dim: int = 8, + **kwargs, + ): + super().__init__() + self.input_dim = input_dim + self.rvq_dim = rvq_dim or input_dim + self.output_dim = output_dim or input_dim + self.num_quantizers = num_quantizers + self.codebook_size = codebook_size + self.codebook_dim = codebook_dim + + self.input_proj = ( + WNConv1d(input_dim, self.rvq_dim, kernel_size=1) if input_dim != self.rvq_dim else nn.Identity() + ) + self.output_proj = ( + WNConv1d(self.rvq_dim, self.output_dim, kernel_size=1) if self.rvq_dim != self.output_dim else nn.Identity() + ) + + self.quantizers = nn.ModuleList( + [ + MossAudioTokenizerLFQ( + input_dim=self.rvq_dim, + codebook_size=codebook_size, + codebook_dim=codebook_dim, + **kwargs, + ) + for _ in range(num_quantizers) + ] + ) + + @torch.no_grad() + def forward( + self, + z: torch.Tensor, + input_length: torch.Tensor, + n_quantizers: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Inference quantization.""" + z = self.input_proj(z).float() + + batch_size, _, max_time = z.shape + mask = torch.arange(max_time, device=z.device).expand(batch_size, max_time) < input_length.unsqueeze(1) + + quantized_out = torch.zeros_like(z, dtype=torch.float32) + residual = z.clone().float() + all_indices = [] + + n_quantizers = n_quantizers or self.num_quantizers + for i, quantizer in enumerate(self.quantizers): + if i >= n_quantizers: + break + + masked_residual = residual * mask.unsqueeze(1) + z_q_i, indices_i, _ = quantizer(masked_residual) + + update_mask = mask.unsqueeze(1) + quantized_out = quantized_out + z_q_i * update_mask + residual = residual - z_q_i * update_mask + all_indices.append(indices_i) + + all_indices = ( + torch.stack(all_indices) + if all_indices + else torch.empty(0, batch_size, max_time, device=z.device, dtype=torch.long) + ) + quantized_out = self.output_proj(quantized_out) + return quantized_out, all_indices, input_length + + def decode_codes(self, codes: torch.Tensor) -> torch.Tensor: + nq, B, T = codes.shape + emb = torch.zeros(B, self.rvq_dim, T, device=codes.device, dtype=torch.float32) + for i, quantizer in enumerate(self.quantizers[:nq]): + quantizer = cast(MossAudioTokenizerLFQ, quantizer) + emb += quantizer.decode_code(codes[i]).float() + emb = self.output_proj(emb) + return emb + + +# ============================================================================= +# Main Model Classes +# ============================================================================= + + +@auto_docstring +class MossAudioTokenizerPreTrainedModel(PreTrainedAudioTokenizerBase): + """Base class for MossAudioTokenizer models.""" + + config_class = MossAudioTokenizerConfig + base_model_prefix = "" + main_input_name = "input_values" + input_modalities = "audio" + supports_gradient_checkpointing = False + _no_split_modules = [ + "MossAudioTokenizerTransformerLayer", + "MossAudioTokenizerResidualVQ", + "MossAudioTokenizerResidualLFQ", + ] + + +@auto_docstring( + custom_intro=""" + The MossAudioTokenizer neural audio codec model for audio tokenization and synthesis. + """ +) +class MossAudioTokenizerModel(MossAudioTokenizerPreTrainedModel): + """ + MossAudioTokenizer model for audio tokenization and synthesis. + + This model can encode audio waveforms into discrete tokens and decode + tokens back into audio waveforms. + """ + + def __init__(self, config: MossAudioTokenizerConfig): + super().__init__(config) + + self.config = config + _ = config.version + self.sampling_rate = config.sampling_rate + self.downsample_rate = config.downsample_rate + self.number_channels = getattr(config, "number_channels", getattr(config, "num_channels", 1)) + self.enable_channel_interleave = getattr(config, "enable_channel_interleave", True) + self.causal_transformer_context_duration = config.causal_transformer_context_duration + + # Build encoder + channel_interleave_factor = ( + self.number_channels if self.enable_channel_interleave and self.number_channels > 1 else 1 + ) + current_frame_rate: float = float(self.sampling_rate * channel_interleave_factor) + self.encoder = nn.ModuleList() + + for encoder_kwargs_i in config.encoder_kwargs: + encoder_kwargs_i = dict(encoder_kwargs_i) # Make a copy + if encoder_kwargs_i["module_type"] == "PatchedPretransform": + self.encoder.append(MossAudioTokenizerPatchedPretransform(**encoder_kwargs_i, is_downsample=True)) + elif encoder_kwargs_i["module_type"] == "Transformer": + context_duration = float( + encoder_kwargs_i.pop("context_duration", self.causal_transformer_context_duration) + ) + self.encoder.append( + MossAudioTokenizerProjectedTransformer( + **encoder_kwargs_i, + context=int(round(current_frame_rate * context_duration)), + ) + ) + current_frame_rate /= self.encoder[-1].downsample_ratio + + # Build quantizer + quantizer_kwargs = dict(config.quantizer_kwargs) + quantizer_type = quantizer_kwargs.get("quantizer_type", getattr(config, "quantizer_type", "rvq")) + if quantizer_type in {"rvq", "spec_rvq"}: + self.quantizer = MossAudioTokenizerResidualVQ(**quantizer_kwargs) + elif quantizer_type in {"rlfq", "random_prefix_rlfq"}: + self.quantizer = MossAudioTokenizerResidualLFQ(**quantizer_kwargs) + else: + raise ValueError(f"Unsupported quantizer_type: {quantizer_type}") + + # Build decoder + decoder_kwargs_list = copy.deepcopy(config.decoder_kwargs) + self.decoder = nn.ModuleList() + + for decoder_kwargs_i in decoder_kwargs_list: + decoder_kwargs_i = dict(decoder_kwargs_i) + if decoder_kwargs_i["module_type"] == "PatchedPretransform": + self.decoder.append(MossAudioTokenizerPatchedPretransform(**decoder_kwargs_i, is_downsample=False)) + elif decoder_kwargs_i["module_type"] == "Transformer": + context_duration = float( + decoder_kwargs_i.pop("context_duration", self.causal_transformer_context_duration) + ) + self.decoder.append( + MossAudioTokenizerProjectedTransformer( + **decoder_kwargs_i, + context=int(round(current_frame_rate * context_duration)), + ) + ) + current_frame_rate *= self.decoder[-1].downsample_ratio + + expected_output_frame_rate = float(self.sampling_rate * channel_interleave_factor) + if int(round(current_frame_rate)) != int(round(expected_output_frame_rate)): + raise ValueError( + "Decoder stack does not invert the encoder frame rate correctly: " + f"got current_frame_rate={current_frame_rate}, expected={expected_output_frame_rate}." + ) + + self.post_init() + + def _start_streaming(self, batch_size: int): + """Start streaming mode for all modules.""" + + def _start(module): + if isinstance(module, StreamingModule): + module._streaming_state = module._init_streaming_state(batch_size) + + self.apply(_start) + + def _stop_streaming(self): + """Stop streaming mode for all modules.""" + + def _stop(module): + if isinstance(module, StreamingModule): + module._streaming_state = None + + self.apply(_stop) + + def _set_streaming_exec_mask(self, exec_mask: torch.Tensor) -> None: + """Set the active slot mask on every live streaming module.""" + + def _set(module): + if isinstance(module, StreamingModule) and module._streaming_state is not None: + module._streaming_state.set_exec_mask(exec_mask.to(module._streaming_state.device)) + + self.apply(_set) + + @contextmanager + def streaming(self, batch_size: int = 1): + """Context manager for streaming mode.""" + self._start_streaming(batch_size) + try: + yield + finally: + self._stop_streaming() + + def _prepare_waveform_batch(self, wav_list: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + if len(wav_list) == 0: + raise ValueError("`wav_list` must contain at least one waveform.") + + device = wav_list[0].device + dtype = wav_list[0].dtype + batch_size = len(wav_list) + input_lengths = torch.zeros(batch_size, device=device, dtype=torch.long) + + normalized_wavs: list[torch.Tensor] = [] + for i, wav in enumerate(wav_list): + if self.number_channels == 1: + if wav.dim() == 1: + wav_i = wav.unsqueeze(0) + elif wav.dim() == 2 and wav.shape[0] == 1: + wav_i = wav + else: + raise ValueError( + f"Expected wav_list[{i}] to have shape `(T,)` or `(1, T)` for a mono model, " + f"got {tuple(wav.shape)}." + ) + else: + if wav.dim() != 2 or wav.shape[0] != self.number_channels: + raise ValueError( + f"Expected wav_list[{i}] to have shape `({self.number_channels}, T)`, got {tuple(wav.shape)}." + ) + wav_i = wav + + normalized_wavs.append(wav_i) + input_lengths[i] = wav_i.shape[-1] + + max_length = int(input_lengths.max().item()) if batch_size > 0 else 0 + input_values = torch.zeros( + batch_size, + self.number_channels, + max_length, + device=device, + dtype=dtype, + ) + for i, wav_i in enumerate(normalized_wavs): + input_values[i, :, : wav_i.shape[-1]] = wav_i + return input_values, input_lengths + + def _prepare_codes_batch( + self, + codes_list: list[torch.Tensor], + num_quantizers: int | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, int]: + if len(codes_list) == 0: + raise ValueError("`codes_list` must contain at least one code tensor.") + + batch_size = len(codes_list) + device = codes_list[0].device + nqs = [codes.shape[0] for codes in codes_list] + if num_quantizers is None: + num_quantizers = nqs[0] + if any(nq != num_quantizers for nq in nqs): + raise ValueError( + "All elements in `codes_list` must have the same number of quantizers when " + "`num_quantizers` is None. Pass `num_quantizers=...` to decode a common prefix." + ) + else: + min_nq = min(nqs) + if min_nq < num_quantizers: + raise ValueError( + "`num_quantizers` must be <= the number of quantizers for every element in `codes_list`. " + f"Got num_quantizers={num_quantizers}, min(codes.shape[0])={min_nq}." + ) + + max_length = max(codes.shape[-1] for codes in codes_list) + audio_codes = torch.zeros(num_quantizers, batch_size, max_length, device=device, dtype=torch.long) + audio_codes_lengths = torch.zeros(batch_size, device=device, dtype=torch.long) + + for i, codes in enumerate(codes_list): + codes_i = codes[:num_quantizers] + audio_codes[:, i, : codes_i.shape[-1]] = codes_i + audio_codes_lengths[i] = codes_i.shape[-1] + return audio_codes, audio_codes_lengths, num_quantizers + + def _flatten_channels_for_codec( + self, + input_values: torch.Tensor, + input_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if input_values.dim() != 3: + raise ValueError(f"Expected `input_values` with shape `(B, C, T)`, got {tuple(input_values.shape)}.") + if input_values.shape[1] != self.number_channels: + raise ValueError( + f"Expected `input_values.shape[1] == {self.number_channels}`, got {input_values.shape[1]}." + ) + + if input_values.shape[-1] % self.downsample_rate != 0: + pad_length = self.downsample_rate - (input_values.shape[-1] % self.downsample_rate) + input_values = F.pad(input_values, (0, pad_length)) + + if self.number_channels > 1 and self.enable_channel_interleave: + input_values = input_values.transpose(1, 2).contiguous().view(input_values.shape[0], 1, -1) + input_lengths = input_lengths * self.number_channels + return input_values, input_lengths + + def _restore_channels_from_codec( + self, + output_values: torch.Tensor, + output_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.number_channels == 1 or not self.enable_channel_interleave: + return output_values.float(), output_lengths + + output_values = ( + output_values.squeeze(1) + .contiguous() + .view(output_values.shape[0], -1, self.number_channels) + .transpose(1, 2) + .contiguous() + .float() + ) + output_lengths = torch.div(output_lengths, self.number_channels, rounding_mode="floor") + return output_values, output_lengths + + @torch.no_grad() + def batch_encode( + self, wav_list: list[torch.Tensor], num_quantizers: int | None = None + ) -> MossAudioTokenizerEncoderOutput: + """Batch encode a list of audio waveforms. + + Args: + wav_list: List of audio tensors, each of shape `(num_samples,)`. + num_quantizers: Number of quantizers to use. By default, all quantizers are used. + + Returns: + [`MossAudioTokenizerEncoderOutput`] with `audio_codes` and `audio_codes_lengths`. + """ + input_values, input_lengths = self._prepare_waveform_batch(wav_list) + return self._encode_frame(input_values, input_lengths, n_quantizers=num_quantizers) + + @torch.no_grad() + def batch_decode( + self, codes_list: list[torch.Tensor], num_quantizers: int | None = None + ) -> MossAudioTokenizerDecoderOutput: + """Batch decode a list of audio codes. + + Args: + codes_list: List of audio code tensors, each of shape `(num_quantizers, codes_length)`. + num_quantizers: If provided, decode only the first `num_quantizers` quantizers from each element in + `codes_list`. If omitted, all elements in `codes_list` must have the same number of quantizers. + + Returns: + [`MossAudioTokenizerDecoderOutput`] with `audio` and `audio_lengths`. + """ + audio_codes, audio_codes_lengths, _ = self._prepare_codes_batch( + codes_list, + num_quantizers=num_quantizers, + ) + return self._decode_frame(audio_codes, audio_codes_lengths) + + @torch.no_grad() + def _encode_frame( + self, + input_values: torch.Tensor, + input_lengths: torch.Tensor | None = None, + n_quantizers: int | None = None, + ) -> MossAudioTokenizerEncoderOutput: + """Tokenize audio waveform into discrete tokens.""" + # Handle input shape + if input_values.dim() == 1: + input_values = input_values.view(1, 1, -1) + elif input_values.dim() == 2: + if self.number_channels == 1: + input_values = input_values.unsqueeze(1) + else: + input_values = input_values.unsqueeze(0) + + B, _, T = input_values.shape + device = input_values.device + + if input_lengths is None: + input_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + input_values, input_lengths = self._flatten_channels_for_codec(input_values, input_lengths) + + # Encode + e, e_lengths = input_values, input_lengths + for encoder_module in self.encoder: + e, e_lengths = encoder_module(e, e_lengths) + + # Quantize + quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) + zq, audio_codes, audio_codes_lengths = quantizer(e, e_lengths, n_quantizers) + + return MossAudioTokenizerEncoderOutput( + audio_codes=audio_codes, audio_codes_lengths=audio_codes_lengths, encoder_hidden_states=e + ) + + @torch.no_grad() + def _decode_frame( + self, + codes: torch.Tensor, + codes_lengths: torch.Tensor | None = None, + ) -> MossAudioTokenizerDecoderOutput: + """Detokenize discrete tokens into audio waveform.""" + nq, B, T = codes.shape + device = codes.device + + if codes_lengths is None: + codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + # Decode from codes + quantizer = cast(MossAudioTokenizerResidualVQ | MossAudioTokenizerResidualLFQ, self.quantizer) + zq = quantizer.decode_codes(codes) + + d, d_lengths = zq, codes_lengths + for decoder_module in self.decoder: + d, d_lengths = decoder_module(d, d_lengths) + + d, d_lengths = self._restore_channels_from_codec(d, d_lengths) + return MossAudioTokenizerDecoderOutput(audio=d, audio_lengths=d_lengths) + + def encode( # type: ignore[override] + self, + input_values: torch.Tensor, + padding_mask: torch.Tensor | None = None, + num_quantizers: int | None = None, + return_dict: bool | None = None, + chunk_duration: float | None = None, + ) -> MossAudioTokenizerEncoderOutput | tuple[torch.Tensor, torch.Tensor]: + """ + Encodes the input audio waveform into discrete codes. + + Args: + input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`): + Float values of the input audio waveform. + padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to indicate valid audio samples. + num_quantizers (`int`, *optional*): + Number of quantizers to use. By default, all quantizers are used. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + chunk_duration (`float`, *optional*): + If provided, encode the input waveform in successive chunks of `chunk_duration` seconds while keeping a + streaming KV cache for the causal transformers. + + `chunk_duration` must be <= `config.causal_transformer_context_duration`, and + `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. + + Returns: + `MossAudioTokenizerEncoderOutput` or tuple containing audio codes and lengths. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + # Handle input shape + if input_values.dim() == 1: + input_values = input_values.view(1, 1, -1) + elif input_values.dim() == 2: + if self.number_channels == 1: + input_values = input_values.unsqueeze(1) + else: + input_values = input_values.unsqueeze(0) + + B, _, T = input_values.shape + device = input_values.device + + if padding_mask is not None: + input_lengths = padding_mask.sum(dim=-1).long() + else: + input_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + if chunk_duration is None: + encoder_output = self._encode_frame(input_values, input_lengths, num_quantizers) + else: + if chunk_duration <= 0: + raise ValueError("`chunk_duration` must be > 0 when provided.") + if chunk_duration > self.causal_transformer_context_duration: + raise ValueError( + "`chunk_duration` must be <= `config.causal_transformer_context_duration` " + f"({self.causal_transformer_context_duration}), got {chunk_duration}." + ) + if B != 1: + raise ValueError("Streaming encode via `chunk_duration` currently only supports batch_size=1.") + + chunk_length = int(round(chunk_duration * self.sampling_rate)) + if chunk_length <= 0: + raise ValueError("`chunk_duration` is too small and results in chunk_length <= 0.") + if chunk_length % self.downsample_rate != 0: + raise ValueError( + "`chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. " + f"Got chunk_length={chunk_length}, downsample_rate={self.downsample_rate}." + ) + + input_length = int(input_lengths[0].item()) + if input_length <= chunk_length: + encoder_output = self._encode_frame(input_values[..., :input_length], input_lengths, num_quantizers) + else: + codes_chunks: list[torch.Tensor] = [] + hidden_chunks: list[torch.Tensor] = [] + + with ExitStack() as exit_stack: + for encoder_module in self.encoder: + if isinstance(encoder_module, StreamingModule): + exit_stack.enter_context(encoder_module.streaming(batch_size=B)) + + for start_idx in range(0, input_length, chunk_length): + input_length_i = min(chunk_length, input_length - start_idx) + if input_length_i <= 0: + break + + input_lengths_i = torch.tensor([input_length_i], device=device, dtype=torch.long) + input_values_i = input_values[..., start_idx : start_idx + input_length_i] + result_i = self._encode_frame(input_values_i, input_lengths_i, num_quantizers) + + if result_i.audio_codes is None or result_i.audio_codes_lengths is None: + raise RuntimeError("Internal error: `_encode_frame` returned empty audio codes.") + if result_i.encoder_hidden_states is None: + raise RuntimeError("Internal error: `_encode_frame` returned empty encoder hidden states.") + + codes_length_i = result_i.audio_codes_lengths + codes_chunks.append(result_i.audio_codes[:, :, : codes_length_i[0]]) + hidden_chunks.append(result_i.encoder_hidden_states[:, :, : codes_length_i[0]]) + + audio_codes = torch.cat(codes_chunks, dim=-1) + encoder_hidden_states = torch.cat(hidden_chunks, dim=-1) + audio_codes_lengths = torch.tensor([audio_codes.shape[-1]], device=device, dtype=torch.long) + encoder_output = MossAudioTokenizerEncoderOutput( + audio_codes=audio_codes, + audio_codes_lengths=audio_codes_lengths, + encoder_hidden_states=encoder_hidden_states, + ) + + if not return_dict: + assert encoder_output.audio_codes is not None + assert encoder_output.audio_codes_lengths is not None + return ( + cast(torch.Tensor, encoder_output.audio_codes), + cast(torch.Tensor, encoder_output.audio_codes_lengths), + ) + return encoder_output + + def decode( # type: ignore[override] + self, + audio_codes: torch.Tensor, + padding_mask: torch.Tensor | None = None, + return_dict: bool | None = None, + chunk_duration: float | None = None, + num_quantizers: int | None = None, + ) -> MossAudioTokenizerDecoderOutput | tuple[torch.Tensor]: + """ + Decodes the given codes into an output audio waveform. + + Args: + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`): + Discrete code embeddings computed using `model.encode`. + padding_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to indicate valid code positions. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + chunk_duration (`float`, *optional*): + If provided, decode the input codes in successive chunks of `chunk_duration` seconds while keeping a + streaming KV cache for the causal transformers. + + num_quantizers (`int`, *optional*): + Number of quantizers to use. By default, all quantizers in `audio_codes` are used. + + `chunk_duration` must be <= `config.causal_transformer_context_duration`, and + `chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. + + Returns: + `MossAudioTokenizerDecoderOutput` or tuple containing decoded audio. + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + if audio_codes.dim() == 2: + audio_codes = audio_codes.unsqueeze(1) # nq, T -> nq, B=1, T + + if num_quantizers is not None: + if num_quantizers > audio_codes.shape[0]: + raise ValueError( + f"`num_quantizers` ({num_quantizers}) must be <= audio_codes.shape[0] ({audio_codes.shape[0]})." + ) + audio_codes = audio_codes[:num_quantizers] + + _, B, T = audio_codes.shape + device = audio_codes.device + + if padding_mask is not None: + codes_lengths = padding_mask.sum(dim=-1).long() + else: + codes_lengths = torch.full((B,), T, device=device, dtype=torch.long) + + if chunk_duration is None: + decoder_output = self._decode_frame(audio_codes, codes_lengths) + else: + if chunk_duration <= 0: + raise ValueError("`chunk_duration` must be > 0 when provided.") + if chunk_duration > self.causal_transformer_context_duration: + raise ValueError( + "`chunk_duration` must be <= `config.causal_transformer_context_duration` " + f"({self.causal_transformer_context_duration}), got {chunk_duration}." + ) + if B != 1: + raise ValueError("Streaming decode via `chunk_duration` currently only supports batch_size=1.") + + chunk_length = int(round(chunk_duration * self.sampling_rate)) + if chunk_length <= 0: + raise ValueError("`chunk_duration` is too small and results in chunk_length <= 0.") + if chunk_length % self.downsample_rate != 0: + raise ValueError( + "`chunk_duration * config.sampling_rate` must be divisible by `config.downsample_rate`. " + f"Got chunk_length={chunk_length}, downsample_rate={self.downsample_rate}." + ) + + chunk_frame_length = chunk_length // self.downsample_rate + codes_length = int(codes_lengths[0].item()) + if codes_length <= chunk_frame_length: + decoder_output = self._decode_frame(audio_codes[..., :codes_length], codes_lengths) + else: + wav_chunks: list[torch.Tensor] = [] + with ExitStack() as exit_stack: + for decoder_module in self.decoder: + if isinstance(decoder_module, StreamingModule): + exit_stack.enter_context(decoder_module.streaming(batch_size=B)) + + for start_idx in range(0, codes_length, chunk_frame_length): + codes_length_i = min(chunk_frame_length, codes_length - start_idx) + if codes_length_i <= 0: + break + + codes_lengths_i = torch.tensor([codes_length_i], device=device, dtype=torch.long) + codes_i = audio_codes[:, :, start_idx : start_idx + codes_length_i] + result_i = self._decode_frame(codes_i, codes_lengths_i) + + if result_i.audio is None or result_i.audio_lengths is None: + raise RuntimeError("Internal error: `_decode_frame` returned empty audio.") + + wav_chunks.append(result_i.audio[:, :, : result_i.audio_lengths[0]]) + + wav = torch.cat(wav_chunks, dim=-1) + audio_lengths = torch.tensor([wav.shape[-1]], device=device, dtype=torch.long) + decoder_output = MossAudioTokenizerDecoderOutput(audio=wav, audio_lengths=audio_lengths) + + if not return_dict: + assert decoder_output.audio is not None + return (cast(torch.Tensor, decoder_output.audio),) + return decoder_output + + @auto_docstring + def forward( + self, + input_values: torch.FloatTensor | None = None, + padding_mask: torch.BoolTensor | None = None, + audio_codes: torch.Tensor | None = None, + num_quantizers: int | None = None, + return_dict: bool | None = None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None] | MossAudioTokenizerOutput: # type: ignore[override] + r""" + input_values (`torch.FloatTensor` of shape `(batch_size, channels, sequence_length)`, *optional*): + Raw audio input converted to Float. + padding_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid computing on padding token indices. Mask values selected in `[0, 1]`: + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + audio_codes (`torch.LongTensor` of shape `(num_quantizers, batch_size, sequence_length)`, *optional*): + Discrete code embeddings computed using `model.encode`. + num_quantizers (`int`, *optional*): + Number of quantizers (codebooks) to use. By default, all quantizers are used. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + + Examples: + + ```python + >>> import torch + >>> from transformers import MossAudioTokenizerModel + + >>> model = MossAudioTokenizerModel.from_pretrained("moss_audio_tokenizer-model") + + >>> # Create dummy audio input + >>> audio = torch.randn(1, 1, 24000) # 1 second of audio at 24kHz + + >>> outputs = model(input_values=audio) + >>> audio_codes = outputs.audio_codes + >>> audio_values = outputs.audio + ``` + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + output_audio_codes: torch.Tensor | None = None + output_audio_codes_lengths: torch.Tensor | None = None + output_audio: torch.Tensor | None = None + output_audio_lengths: torch.Tensor | None = None + decoded_from_encoded_codes = False + + # Encode if input_values provided + if input_values is not None: + encoder_output = self.encode(input_values, padding_mask, num_quantizers, return_dict=True) + encoder_output = cast(MossAudioTokenizerEncoderOutput, encoder_output) + output_audio_codes = encoder_output.audio_codes + output_audio_codes_lengths = encoder_output.audio_codes_lengths + + # If codes not provided separately, use encoded codes for decoding + if audio_codes is None: + audio_codes = output_audio_codes + decoded_from_encoded_codes = True + + # Decode if codes available + if audio_codes is not None: + # If we're decoding the codes we just produced, use the computed lengths so we don't decode padded garbage. + if decoded_from_encoded_codes and output_audio_codes_lengths is not None: + decoder_output = self._decode_frame(audio_codes, output_audio_codes_lengths) + else: + decoder_output = self.decode( + audio_codes, + padding_mask=padding_mask, + return_dict=True, + num_quantizers=num_quantizers, + ) + decoder_output = cast(MossAudioTokenizerDecoderOutput, decoder_output) + output_audio = decoder_output.audio + output_audio_lengths = decoder_output.audio_lengths + + if not return_dict: + return (output_audio_codes, output_audio, output_audio_lengths) + + return MossAudioTokenizerOutput( + audio=output_audio, + audio_lengths=output_audio_lengths, + audio_codes=output_audio_codes, + audio_codes_lengths=output_audio_codes_lengths, + ) + + +__all__ = ["MossAudioTokenizerModel", "MossAudioTokenizerPreTrainedModel"] diff --git a/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py b/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py new file mode 100644 index 00000000000..c518742a21f --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/configuration_moss_audio_tokenizer_v2.py @@ -0,0 +1,359 @@ +# Copyright 2026 OpenMOSS and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""MossAudioTokenizer model configuration""" + +from typing import Any + +from transformers.configuration_utils import PreTrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + + +class MossAudioTokenizerConfig(PreTrainedConfig): + r""" + This is the configuration class to store the configuration of a [`MossAudioTokenizerModel`]. It is used to + instantiate a MossAudioTokenizer model according to the specified arguments, defining the model architecture. + + Instantiating a configuration with the defaults will yield a similar configuration to that of the + [VoiceAgentGroup/moss_audio_tokenizer](https://huggingface.co/VoiceAgentGroup/moss_audio_tokenizer) architecture. + + Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PreTrainedConfig`] for more information. + + Args: + sampling_rate (`int`, *optional*, defaults to 24000): + The sampling rate at which the audio waveform should be digitalized expressed in hertz (Hz). + downsample_rate (`int`, *optional*, defaults to 1920): + Total downsampling rate from waveform to tokens. + causal_transformer_context_duration (`float`, *optional*, defaults to 10.0): + Context duration in seconds for causal transformer. + encoder_kwargs (`list[dict]`, *optional*): + List of encoder module configurations. Each dict specifies a module type and its parameters. + decoder_kwargs (`list[dict]`, *optional*): + List of decoder module configurations in execution order. + quantizer_type (`str`, *optional*, defaults to `"rvq"`): + Quantizer type. Options include `"rvq"`, `"spec_rvq"`, `"rlfq"`, `"random_prefix_rlfq"`. + quantizer_kwargs (`dict`, *optional*): + Configuration for the quantizer including `input_dim`, `rvq_dim`, `output_dim`, `num_quantizers`, + `codebook_size`, and `codebook_dim`. + + Example: + + ```python + >>> from transformers import MossAudioTokenizerModel, MossAudioTokenizerConfig + + >>> # Initializing a MossAudioTokenizer style configuration + >>> configuration = MossAudioTokenizerConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = MossAudioTokenizerModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "moss-audio-tokenizer" + + # Backward-compatible alias used by some checkpoints. + attribute_map = {"sample_rate": "sampling_rate"} + + sampling_rate: int + downsample_rate: int + causal_transformer_context_duration: float + encoder_kwargs: list[dict[str, Any]] + decoder_kwargs: list[dict[str, Any]] + number_channels: int + enable_channel_interleave: bool + attention_implementation: str + compute_dtype: str + codec_weight_dtype: str + quantizer_type: str + quantizer_kwargs: dict[str, Any] + + def __init__( + self, + version: str | None = None, + sampling_rate: int = 48000, + downsample_rate: int = 3840, + causal_transformer_context_duration: float = 10.0, + encoder_kwargs: list[dict[str, Any]] | None = None, + decoder_kwargs: list[dict[str, Any]] | None = None, + number_channels: int = 2, + enable_channel_interleave: bool = True, + attention_implementation: str = "sdpa", + compute_dtype: str = "fp32", + codec_weight_dtype: str = "fp32", + quantizer_type: str = "rlfq", + quantizer_kwargs: dict[str, Any] | None = None, + **kwargs, + ): + # Some checkpoints might include an incorrect/legacy `model_type` (e.g. "speech_tokenizer"). + # We drop it to avoid overriding the class-level `model_type`. + kwargs.pop("model_type", None) + if "channels_numbers" in kwargs: + number_channels = kwargs.pop("channels_numbers") + if "enable_channel_interleave" in kwargs: + enable_channel_interleave = kwargs.pop("enable_channel_interleave") + if "attention_backend" in kwargs and attention_implementation == "sdpa": + attention_implementation = kwargs.pop("attention_backend") + if "codec_compute_dtype" in kwargs and compute_dtype == "fp32": + compute_dtype = kwargs.pop("codec_compute_dtype") + if "codec_load_dtype" in kwargs and codec_weight_dtype == "fp32": + codec_weight_dtype = kwargs.pop("codec_load_dtype") + + # `version` is accepted for compatibility but not used in modeling. + self.version = version + self.sampling_rate = sampling_rate + self.downsample_rate = downsample_rate + self.causal_transformer_context_duration = causal_transformer_context_duration + self.number_channels = number_channels + self.enable_channel_interleave = enable_channel_interleave + self.attention_implementation = attention_implementation + self.compute_dtype = compute_dtype + self.codec_weight_dtype = codec_weight_dtype + # Default encoder configuration + if encoder_kwargs is None: + encoder_kwargs = [ + { + "module_type": "PatchedPretransform", + "patch_size": 240, + }, + { + "module_type": "Transformer", + "input_dimension": 240, + "output_dimension": 384, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 768, + "output_dimension": 384, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 768, + "output_dimension": 640, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 1280, + "output_dimension": 768, + "d_model": 1280, + "num_heads": 20, + "num_layers": 32, + "dim_feedforward": 5120, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + ] + self.encoder_kwargs = encoder_kwargs + + # Default decoder configuration (execution order) + if decoder_kwargs is None: + decoder_kwargs = [ + { + "module_type": "Transformer", + "input_dimension": 768, + "output_dimension": 1280, + "d_model": 1280, + "num_heads": 20, + "num_layers": 32, + "dim_feedforward": 5120, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 640, + "output_dimension": 768, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 384, + "output_dimension": 768, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 384, + "output_dimension": 768, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 2, + }, + { + "module_type": "Transformer", + "input_dimension": 384, + "output_dimension": 240, + "d_model": 768, + "num_heads": 12, + "num_layers": 12, + "dim_feedforward": 3072, + "causal": True, + "norm": "layer_norm", + "positional_embedding": "rope", + "max_period": 10000, + "gating": "none", + "layer_scale": 0.01, + "conv_layout": True, + }, + { + "module_type": "PatchedPretransform", + "patch_size": 240, + }, + ] + self.decoder_kwargs = decoder_kwargs + + # Default quantizer configuration + if quantizer_kwargs is None: + quantizer_kwargs = { + "input_dim": 768, + "rvq_dim": 512, + "output_dim": 768, + "num_quantizers": 32, + "codebook_size": 1024, + "codebook_dim": 8, + "quantizer_type": "rlfq", + } + + # Handle quantizer_type from kwargs or config + kw_qtype = quantizer_kwargs.get("quantizer_type", None) + if kw_qtype is not None: + self.quantizer_type = kw_qtype + else: + self.quantizer_type = quantizer_type + quantizer_kwargs["quantizer_type"] = quantizer_type + + self.quantizer_kwargs = quantizer_kwargs + + super().__init__(**kwargs) + + @property + def num_quantizers(self) -> int: + """Return the number of quantizers from quantizer_kwargs.""" + return self.quantizer_kwargs.get("num_quantizers", 32) + + @property + def codebook_size(self) -> int: + """Return the codebook size from quantizer_kwargs.""" + return self.quantizer_kwargs.get("codebook_size", 4096) + + @property + def frame_rate(self) -> float: + """Return the frame rate (tokens per second).""" + return self.sampling_rate / self.downsample_rate + + +__all__ = ["MossAudioTokenizerConfig"] diff --git a/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py b/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py index 46147f27c97..42f103f4221 100644 --- a/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py +++ b/vllm_omni/model_executor/models/moss_tts/configuration_moss_tts.py @@ -17,6 +17,7 @@ from transformers import AutoConfig from transformers.configuration_utils import PretrainedConfig +from transformers.models.gpt2.configuration_gpt2 import GPT2Config from transformers.models.qwen3 import Qwen3Config @@ -235,7 +236,80 @@ def get_text_config(self, **_: object) -> Qwen3Config: return self.language_config +class MossTTSLocalConfig(PretrainedConfig): + """Config for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5). + + Like ``MossTTSDelayConfig``, the HF checkpoint stores a nested Qwen3 config + (``qwen3_config``) which we unwrap via ``get_text_config()``. It also stores + a nested GPT2-style config (``gpt2_config``) describing the single-layer + local depth transformer that decodes the ``n_vq`` audio codebooks per frame + (see ``modeling_moss_tts_local_depth.py``). + """ + + model_type = "moss_tts_local" + + def __init__( + self, + qwen3_config: dict | None = None, + gpt2_config: dict | None = None, + n_vq: int = 12, + audio_vocab_size: int = 1024, + audio_pad_token_id: int = 1024, + pad_token_id: int = 151643, + im_start_token_id: int = 151644, + im_end_token_id: int = 151645, + audio_start_token_id: int = 151669, + audio_end_token_id: int = 151670, + audio_user_slot_token_id: int = 151654, + audio_assistant_slot_token_id: int = 151656, + sampling_rate: int = 48000, + audio_tokenizer_name_or_path: str = "OpenMOSS-Team/MOSS-Audio-Tokenizer-v2", + local_text_head_mode: str = "binary", + **kwargs: object, + ) -> None: + if qwen3_config is None: + qwen3_config = {} + if isinstance(qwen3_config, dict): + qwen3_config = dict(qwen3_config) + qwen3_config.pop("model_type", None) + self.qwen3_config = Qwen3Config(**qwen3_config) + else: + self.qwen3_config = qwen3_config + + if isinstance(gpt2_config, dict): + gpt2_config = dict(gpt2_config) + gpt2_config.pop("model_type", None) + self.gpt2_config = GPT2Config(**gpt2_config) + else: + self.gpt2_config = gpt2_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + self.n_vq = n_vq + self.audio_vocab_size = audio_vocab_size + self.audio_pad_token_id = audio_pad_token_id + self.audio_pad_code = audio_pad_token_id + self.im_start_token_id = im_start_token_id + self.im_end_token_id = im_end_token_id + self.audio_start_token_id = audio_start_token_id + self.audio_end_token_id = audio_end_token_id + self.audio_user_slot_token_id = audio_user_slot_token_id + self.audio_assistant_slot_token_id = audio_assistant_slot_token_id + self.sampling_rate = sampling_rate + self.audio_tokenizer_name_or_path = audio_tokenizer_name_or_path + self.local_text_head_mode = local_text_head_mode + + def get_text_config(self, **_: object) -> Qwen3Config: + return self.qwen3_config + + AutoConfig.register("moss_tts_delay", MossTTSDelayConfig) AutoConfig.register("moss_tts_realtime", MossTTSRealtimeConfig) - -__all__ = ["MossTTSDelayConfig", "MossTTSRealtimeConfig", "MossTTSLocalTransformerConfig"] +AutoConfig.register("moss_tts_local", MossTTSLocalConfig) + +__all__ = [ + "MossTTSDelayConfig", + "MossTTSRealtimeConfig", + "MossTTSLocalTransformerConfig", + "MossTTSLocalConfig", +] diff --git a/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py new file mode 100644 index 00000000000..1daa5081e48 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/cuda_graph_streaming_decoder_wrapper.py @@ -0,0 +1,164 @@ +# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +"""CUDA graph wrapper for MOSS-TTS Local streaming codec decode.""" + +from __future__ import annotations + +import time +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.nn as nn +from torch.cuda import CUDAGraph +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +@dataclass +class _CapturedStreamingDecodeGraph: + graph: CUDAGraph + static_codes: torch.Tensor + static_lengths: torch.Tensor + static_audio: torch.Tensor + static_audio_lengths: torch.Tensor + + +class CUDAGraphStreamingDecoderWrapper: + """Replay fixed-size streaming ``_decode_frame`` steps with CUDA graphs. + + The wrapped codec must already be inside ``codec.streaming(batch_size)``. + Graphs are keyed by the streaming step frame count ``T``. The batch width is + fixed to the codec stream slot count. + """ + + def __init__( + self, + codec: nn.Module, + *, + batch_size: int, + num_quantizers: int, + reset_streaming_state: Callable[[], None] | None = None, + ) -> None: + self.codec = codec + self.batch_size = int(batch_size) + self.num_quantizers = int(num_quantizers) + self.reset_streaming_state = reset_streaming_state + self.graphs: dict[int, _CapturedStreamingDecodeGraph] = {} + self._pool = None + self._warmed_up = False + + @property + def is_ready(self) -> bool: + return bool(self.graphs) + + @property + def capture_sizes(self) -> list[int]: + return sorted(self.graphs) + + @torch.no_grad() + def warmup(self, device: torch.device, capture_sizes: list[int]) -> None: + if self._warmed_up: + return + self._warmed_up = True + if device.type != "cuda": + return + sizes = sorted({int(size) for size in capture_sizes if int(size) > 0}, reverse=True) + if not sizes: + return + + logger.info( + "MOSS-TTS streaming decoder CUDA graph warmup: batch_size=%d n_vq=%d capture_sizes=%s", + self.batch_size, + self.num_quantizers, + list(reversed(sizes)), + ) + start_s = time.perf_counter() + with torch.cuda.device(device): + for size in sizes: + try: + self._capture(size, device) + logger.info(" Captured MOSS-TTS streaming decoder CUDA graph for size=%d", size) + except Exception: + self.graphs.pop(size, None) + logger.warning( + " Failed to capture MOSS-TTS streaming decoder CUDA graph for size=%d; " + "this size will use eager decode", + size, + exc_info=True, + ) + logger.info( + "MOSS-TTS streaming decoder CUDA graph warmup complete: %d/%d captured in %.1f ms", + len(self.graphs), + len(sizes), + (time.perf_counter() - start_s) * 1000.0, + ) + + @torch.no_grad() + def _capture(self, size: int, device: torch.device) -> None: + codes = torch.zeros( + self.num_quantizers, + self.batch_size, + size, + dtype=torch.long, + device=device, + ) + lengths = torch.full((self.batch_size,), size, dtype=torch.long, device=device) + exec_mask = torch.ones(self.batch_size, dtype=torch.bool, device=device) + + self.codec._set_streaming_exec_mask(exec_mask) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + _ = self.codec._decode_frame(codes, lengths) + torch.cuda.current_stream().wait_stream(stream) + torch.accelerator.synchronize(device) + + if self.reset_streaming_state is not None: + self.reset_streaming_state() + self.codec._set_streaming_exec_mask(exec_mask) + if self._pool is None: + self._pool = current_platform.get_global_graph_pool() + graph = CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool, capture_error_mode="thread_local"): + output = self.codec._decode_frame(codes, lengths) + + if output.audio is None or output.audio_lengths is None: + raise RuntimeError("MOSS-TTS streaming decoder CUDA graph capture produced empty audio.") + + self.graphs[size] = _CapturedStreamingDecodeGraph( + graph=graph, + static_codes=codes, + static_lengths=lengths, + static_audio=output.audio, + static_audio_lengths=output.audio_lengths, + ) + + @torch.no_grad() + def decode( + self, + codes_step: torch.Tensor, + exec_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + if not codes_step.is_cuda or torch.cuda.is_current_stream_capturing(): + return None + n_vq, batch_size, step_t = codes_step.shape + if int(n_vq) != self.num_quantizers or int(batch_size) != self.batch_size: + return None + # The streaming codec advances its internal KV offsets by tensor T, so + # bucket-padding a smaller step into a larger captured graph would write + # padded frames into state. Capture exact step sizes instead. + entry = self.graphs.get(int(step_t)) + if entry is None: + return None + + self.codec._set_streaming_exec_mask(exec_mask) + entry.static_codes.copy_(codes_step) + entry.graph.replay() + return entry.static_audio, entry.static_audio_lengths + + +__all__ = ["CUDAGraphStreamingDecoderWrapper"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py index b1b6063d47c..8f1a72e9bd6 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_codec.py @@ -5,7 +5,7 @@ from __future__ import annotations -import os +import contextlib from collections.abc import Iterable from typing import Any @@ -20,6 +20,15 @@ MossAudioTokenizerConfig, MossAudioTokenizerModel, ) +from vllm_omni.model_executor.models.moss_tts.audio_tokenizer_v2 import ( + MossAudioTokenizerModel as MossAudioTokenizerV2Model, +) +from vllm_omni.model_executor.models.moss_tts.configuration_moss_audio_tokenizer_v2 import ( + MossAudioTokenizerConfig as MossAudioTokenizerV2Config, +) +from vllm_omni.model_executor.models.moss_tts.cuda_graph_streaming_decoder_wrapper import ( + CUDAGraphStreamingDecoderWrapper, +) from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import ( MossTTSCUDAGraphCodecWrapper, ) @@ -28,6 +37,121 @@ logger = init_logger(__name__) +class _MossCodecStreamSession: + """Persistent streaming decode session for vendored MOSS-Audio-Tokenizer-v2.""" + + def __init__( + self, + codec: nn.Module, + *, + stream_slots: int, + n_vq: int, + cudagraph_capture_sizes: list[int] | None = None, + ) -> None: + self._codec = codec + self._stream_slots = int(stream_slots) + self._batch_size = self._stream_slots + self._n_vq = int(n_vq) + self._device = next(codec.parameters()).device + self._free_stream_slots = list(range(self._stream_slots)) + self._exit_stack = contextlib.ExitStack() + self._closed = False + with torch.no_grad(): + self._exit_stack.enter_context(codec.streaming(self._batch_size)) + self._cudagraph_wrapper: CUDAGraphStreamingDecoderWrapper | None = None + capture_sizes = sorted({int(size) for size in (cudagraph_capture_sizes or []) if int(size) > 0}) + if capture_sizes and self._device.type == "cuda": + self._cudagraph_wrapper = CUDAGraphStreamingDecoderWrapper( + codec, + batch_size=self._batch_size, + num_quantizers=self._n_vq, + reset_streaming_state=lambda: self.reset_slots(list(range(self._batch_size))), + ) + self._cudagraph_wrapper.warmup(self._device, capture_sizes) + self.reset_slots(list(range(self._batch_size))) + if not self._cudagraph_wrapper.is_ready: + self._cudagraph_wrapper = None + + def acquire(self) -> int | None: + if not self._free_stream_slots: + return None + return self._free_stream_slots.pop() + + def release(self, slot: int) -> None: + if self._closed: + return + self.reset_slots([slot]) + self._free_stream_slots.append(slot) + + def reset_slots(self, slots: list[int]) -> None: + if not slots: + return + reset_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) + reset_mask[slots] = True + + def _reset(module: nn.Module) -> None: + state = getattr(module, "_streaming_state", None) + if state is not None: + state.reset(reset_mask.to(state.device)) + + with torch.no_grad(): + self._codec.apply(_reset) + + def close(self) -> None: + if self._closed: + return + with torch.no_grad(): + self._exit_stack.close() + self._closed = True + + @torch.no_grad() + def step(self, slot_codes: dict[int, torch.Tensor]) -> dict[int, torch.Tensor]: + if not slot_codes: + return {} + step_lengths = {int(codes.shape[1]) for codes in slot_codes.values()} + if len(step_lengths) != 1: + raise ValueError(f"MOSS codec streaming step needs uniform T, got {sorted(step_lengths)}") + (step_t,) = step_lengths + codes_step = torch.zeros( + self._n_vq, + self._batch_size, + step_t, + dtype=torch.long, + device=self._device, + ) + codes_lengths = torch.zeros(self._batch_size, dtype=torch.long, device=self._device) + exec_mask = torch.zeros(self._batch_size, dtype=torch.bool, device=self._device) + for slot, codes in slot_codes.items(): + codes_step[:, slot, :] = codes.to(device=self._device, dtype=torch.long) + codes_lengths[slot] = int(codes.shape[1]) + exec_mask[slot] = True + + graph_output: tuple[torch.Tensor, torch.Tensor] | None = None + if self._cudagraph_wrapper is not None: + graph_output = self._cudagraph_wrapper.decode(codes_step, exec_mask) + + used_cudagraph = graph_output is not None + if used_cudagraph: + audio_tensor, lengths_tensor = graph_output + else: + self._codec._set_streaming_exec_mask(exec_mask) + result = self._codec._decode_frame(codes_step, codes_lengths) + if result.audio is None: + return {} + audio_tensor = result.audio + lengths_tensor = result.audio_lengths + + audio = audio_tensor.detach().to("cpu", torch.float32) + lengths = lengths_tensor.detach().to("cpu") if lengths_tensor is not None else None + out: dict[int, torch.Tensor] = {} + for slot in slot_codes: + wav = audio[slot] + if lengths is not None: + wav = wav[..., : int(lengths[slot].item())] + out[slot] = wav.contiguous() + return out + + class MossTTSCodecDecoder(nn.Module): """Stage-1 decoder for all MOSS-TTS variants. @@ -42,10 +166,8 @@ class MossTTSCodecDecoder(nn.Module): ``n_vq=32`` (MOSS-TTS) or ``n_vq=16`` (all other variants) without swapping weights. - The codec checkpoint path defaults to the value stored in - ``vllm_config.model_config.hf_config.codec_model_name_or_path`` but can - be overridden by setting the environment variable - ``MOSS_TTS_CODEC_PATH``. + The codec checkpoint path comes from + ``vllm_config.model_config.hf_config.codec_model_name_or_path``. """ input_modalities = "audio" @@ -64,12 +186,33 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: cfg = vllm_config.model_config.hf_config self._n_vq: int = int(getattr(cfg, "n_vq", getattr(cfg, "rvq", 16))) - self._codec_path: str = str(getattr(cfg, "codec_model_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer")) + self._codec_path: str = str( + getattr( + cfg, + "codec_model_name_or_path", + getattr(cfg, "audio_tokenizer_name_or_path", "OpenMOSS-Team/MOSS-Audio-Tokenizer"), + ) + ) - # Resolved at load_weights() time + # Resolved at load_weights() time, once the codec checkpoint's own + # config (sampling rate, channel count) is known. self._codec: MossAudioTokenizerModel | None = None self._cuda_graph_wrapper: MossTTSCUDAGraphCodecWrapper | None = None + self._n_channels: int = 1 self._sr_tensor = torch.tensor(self._OUTPUT_SAMPLE_RATE, dtype=torch.int32) + self._stream_session: _MossCodecStreamSession | None = None + self._stream_slots: int = self._connector_int("codec_stream_slots", default=0) + self._stream_chunk_frames: int = self._connector_int("codec_chunk_frames", default=0) + self._stream_max_step_frames: int = self._stream_chunk_frames or 100 + self._stream_req_slots: dict[str, int] = {} + self._stream_pending_codes: dict[str, list[torch.Tensor]] = {} + self._stream_starved_reqs: set[str] = set() + self._codec_streaming: bool = self._connector_bool("codec_streaming", default=False) + default_graph_capture_max = self._stream_chunk_frames or min(self._stream_max_step_frames, 32) + default_graph_capture_sizes = list(range(1, max(1, default_graph_capture_max) + 1)) + self._streaming_cudagraph_capture_sizes = self._streaming_cudagraph_capture_sizes_from_compilation_config( + default_graph_capture_sizes + ) # ------------------------------------------------------------------ # vLLM-Omni stubs (codec has no AR loop) @@ -107,9 +250,9 @@ def forward( chunk transfer adapter assigns those to ``request.prompt_token_ids``, which arrives here as ``input_ids`` concatenated across all requests. Per-request slice boundaries are computed from - ``kwargs["num_scheduled_tokens"]`` (a list of token counts, one per - request). ``runtime_additional_information`` carries per-request - metadata such as ``left_context_size``. + ``kwargs["seq_token_counts"]`` (token counts, one per request). + ``runtime_additional_information`` carries per-request metadata such as + ``left_context_size``. Returns ------- @@ -118,8 +261,8 @@ def forward( multimodal_outputs["sr"] — list of scalar int32 tensors """ sr_tensor = self._sr_tensor - empty = torch.zeros((0,), dtype=torch.float32) - info_list: list[dict[str, Any]] = runtime_additional_information or [] + empty = self._empty_audio() + info_list: list[dict[str, Any]] = list(runtime_additional_information or [{}]) num_req = max(len(info_list), 1) if self._codec is None: @@ -132,27 +275,59 @@ def forward( }, ) + if self._codec_streaming and runtime_additional_information is None: + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={ + "model_outputs": [empty] * num_req, + "sr": [sr_tensor] * num_req, + }, + ) + audios: list[torch.Tensor] = [empty] * num_req srs: list[torch.Tensor] = [sr_tensor] * num_req device = next(self._codec.parameters()).device + streaming_work: list[tuple[int, str, torch.Tensor, bool]] = [] if input_ids is None or input_ids.numel() == 0: + for i, wav in self._finish_empty_streaming_requests(info_list).items(): + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, ) - # ``input_ids`` is concatenated across all requests; split by - # query_start_loc-style offsets from ``num_scheduled_tokens`` in - # kwargs if available, else assume one request. + # ``input_ids`` is concatenated across all requests. vLLM-Omni runners + # pass the per-request lengths via the shared code2wav contract + # ``seq_token_counts``. ids_flat = input_ids.reshape(-1).to(dtype=torch.long) - num_scheduled_tokens = kwargs.get("num_scheduled_tokens") - if isinstance(num_scheduled_tokens, list) and len(num_scheduled_tokens) == num_req: - offsets = [0] - for n in num_scheduled_tokens: - offsets.append(offsets[-1] + int(n)) - else: - offsets = [0, int(ids_flat.shape[0])] + token_counts = self._normalize_seq_token_counts(kwargs.get("seq_token_counts")) + if token_counts is None: + raise RuntimeError( + "MossTTS codec requires seq_token_counts; otherwise concatenated " + "codec tokens cannot be split per request." + ) + if sum(token_counts) != int(ids_flat.shape[0]): + raise RuntimeError( + "MossTTS codec seq_token_counts mismatch: " + f"counts={token_counts}, sum={sum(token_counts)}, input_tokens={int(ids_flat.shape[0])}." + ) + + num_req = len(token_counts) + if len(info_list) < num_req: + info_list.extend({} for _ in range(num_req - len(info_list))) + elif len(info_list) > num_req: + info_list = info_list[:num_req] + if len(audios) < num_req: + audios.extend(empty for _ in range(num_req - len(audios))) + srs.extend(sr_tensor for _ in range(num_req - len(srs))) + elif len(audios) > num_req: + audios = audios[:num_req] + srs = srs[:num_req] + + offsets = [0] + for n in token_counts: + offsets.append(offsets[-1] + int(n)) for i, info in enumerate(info_list): if i + 1 >= len(offsets): @@ -160,6 +335,14 @@ def forward( seg = ids_flat[offsets[i] : offsets[i + 1]] if seg.numel() == 0: continue + meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} + finished = bool(meta.get("stream_finished", meta.get("finished", False))) + streaming_enabled = bool(meta.get("codec_streaming", self._codec_streaming)) + code_flat_numel = meta.get("code_flat_numel") + if streaming_enabled and finished and code_flat_numel is not None and int(code_flat_numel) == 0: + for _, wav in self._finish_empty_streaming_requests([info]).items(): + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav + continue if seg.numel() % self._n_vq != 0: logger.warning( "MossTTS codec input length %d not divisible by n_vq %d; skipping.", @@ -176,11 +359,6 @@ def forward( codebook_size = self._codec.config.codebook_size codes_nq_t = codes_nq_t.clamp_(0, int(codebook_size) - 1) - # left_context_size: number of left-context frames prepended to - # this chunk by the stage input processor (overlap-add pattern, - # same as Qwen3-TTS code2wav). We decode them together with the - # real chunk for smooth boundaries, then trim them from the output. - meta = (info.get("meta", {}) if isinstance(info, dict) else {}) or {} left_ctx = meta.get("left_context_size", 0) if isinstance(left_ctx, (list, tuple)): left_ctx = int(left_ctx[0]) if left_ctx else 0 @@ -188,6 +366,12 @@ def forward( left_ctx = int(left_ctx.reshape(-1)[0].item()) if left_ctx.numel() else 0 left_ctx = int(left_ctx) + req_key = self._runtime_request_key(info, meta, i) + + if streaming_enabled: + streaming_work.append((i, req_key, codes_nq_t, finished)) + continue + if self._cuda_graph_wrapper is not None: out = self._cuda_graph_wrapper.decode(codes_nq_t) else: @@ -196,28 +380,292 @@ def forward( if out.audio is None: continue - wav = out.audio.reshape(-1).to(dtype=torch.float32).cpu() + # ``out.audio`` is ``(1, C, T)``; keep the channel axis for + # stereo codecs (Local-v1.5) and flatten to ``(T,)`` for mono + # ones (Delay/Realtime) to preserve their existing output shape. + wav = out.audio[0].to(dtype=torch.float32).cpu() if out.audio_lengths is not None: - wav = wav[: int(out.audio_lengths[0].item())] + wav = wav[..., : int(out.audio_lengths[0].item())] - # Trim left-context samples. + # Trim left-context samples (per-channel sample axis, so the + # trim amount is identical for mono and interleaved-stereo). if left_ctx > 0: - trim = min(left_ctx * self._codec.downsample_rate, wav.shape[0]) + trim = min(left_ctx * self._codec.downsample_rate, wav.shape[-1]) if trim < left_ctx * self._codec.downsample_rate: logger.warning( "left_ctx trim (%d samples) exceeds wav length (%d); returning empty audio.", left_ctx * self._codec.downsample_rate, - wav.shape[0], + wav.shape[-1], ) - wav = wav[trim:] + wav = wav[..., trim:] - audios[i] = wav + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav + + if streaming_work: + for i, wav in self._decode_streaming_batch(streaming_work).items(): + audios[i] = wav.reshape(-1) if wav.ndim == 1 or int(wav.shape[0]) == 1 else wav return OmniOutput( text_hidden_states=None, multimodal_outputs={"model_outputs": audios, "sr": srs}, ) + def _finish_empty_streaming_requests(self, info_list: list[dict[str, Any]]) -> dict[int, torch.Tensor]: + """Release codec stream state for empty finish sentinels. + + Stage-0 can finish on a step that emits no new audio frame. The stage + input processor forwards that as an empty payload with finished=true. + If a request never acquired a stream slot, do not offline-decode its + buffered codes here: streaming requests must only emit client deltas. + """ + session = self._stream_session + outputs: dict[int, torch.Tensor] = {} + if session is None: + return outputs + for i, info in enumerate(info_list): + if not isinstance(info, dict): + continue + meta = (info.get("meta", {}) or {}) if isinstance(info.get("meta", {}), dict) else {} + if not bool(meta.get("codec_streaming", self._codec_streaming)): + continue + finished = bool(meta.get("stream_finished", meta.get("finished", False))) + if not finished: + continue + req_key = self._runtime_request_key(info, meta, i) + slot = self._stream_req_slots.get(req_key) + pending = req_key in self._stream_pending_codes + if slot is not None or pending: + if pending and slot is None: + logger.warning( + "MOSS codec stream request %s finished before a codec stream slot became available; " + "dropping buffered codes instead of offline-decoding a non-delta waveform.", + req_key, + ) + self._finish_stream_request(req_key, session, slot) + return outputs + + @staticmethod + def _normalize_seq_token_counts(value: Any) -> list[int] | None: + if value is None: + return None + if not isinstance(value, (list, tuple)): + raise TypeError( + "MossTTS codec expects seq_token_counts to be a list/tuple of per-request token counts, " + f"got {type(value).__name__}." + ) + counts = [int(item) for item in value] + if not counts: + return None + for count in counts: + if count < 0: + raise ValueError(f"MossTTS codec seq_token_counts must be non-negative, got {counts}.") + return counts + + def _runtime_request_key(self, info: Any, meta: dict[str, Any], index: int) -> str: + for value in ( + meta.get("req_id"), + info.get("request_id") if isinstance(info, dict) else None, + ): + if isinstance(value, (list, tuple)): + value = value[0] if value else None + if value is not None: + return str(value) + return f"moss-codec-stream-{index}" + + def _empty_audio(self) -> torch.Tensor: + if self._n_channels > 1: + return torch.zeros((self._n_channels, 0), dtype=torch.float32) + return torch.zeros((0,), dtype=torch.float32) + + def _ensure_stream_session(self) -> _MossCodecStreamSession | None: + if self._codec is None: + return None + if self._stream_session is not None: + return self._stream_session + slots = self._stream_slots + if slots <= 0: + scheduler_cfg = getattr(self.vllm_config, "scheduler_config", None) + slots = int(getattr(scheduler_cfg, "max_num_seqs", 1) or 1) + self._stream_session = _MossCodecStreamSession( + self._codec, + stream_slots=max(1, slots), + n_vq=self._n_vq, + cudagraph_capture_sizes=self._streaming_cudagraph_capture_sizes, + ) + return self._stream_session + + def _decode_streaming_batch( + self, + items: list[tuple[int, str, torch.Tensor, bool]], + ) -> dict[int, torch.Tensor]: + session = self._ensure_stream_session() + if session is None: + return {} + + outputs: dict[int, torch.Tensor] = {} + grouped: dict[int, list[tuple[int, str, int, torch.Tensor, bool]]] = {} + max_step_frames = max(1, int(self._stream_max_step_frames)) + + for output_index, request_id, codes_nq_t, finished in items: + pending = self._stream_pending_codes.get(request_id) + slot = self._stream_req_slots.get(request_id) + if slot is None: + slot = session.acquire() + if slot is None: + self._append_stream_pending(request_id, codes_nq_t) + if request_id not in self._stream_starved_reqs: + logger.warning( + "MOSS codec streaming slots exhausted; buffering %s until a stream slot is available.", + request_id, + ) + self._stream_starved_reqs.add(request_id) + if finished: + logger.warning( + "MOSS codec stream request %s finished before a codec stream slot became available; " + "dropping buffered codes instead of offline-decoding a non-delta waveform.", + request_id, + ) + self._finish_stream_request(request_id, session, None) + continue + self._stream_req_slots[request_id] = slot + + if pending: + self._append_stream_pending(request_id, codes_nq_t) + replay_codes = self._pop_stream_pending(request_id) + wav = self._decode_stream_slot_sequence(session, slot, replay_codes) + if wav is not None: + outputs[output_index] = wav + if finished: + self._finish_stream_request(request_id, session, slot) + continue + + if int(codes_nq_t.shape[1]) > max_step_frames: + wav = self._decode_stream_slot_sequence(session, slot, codes_nq_t) + if wav is not None: + outputs[output_index] = wav + if finished: + self._finish_stream_request(request_id, session, slot) + continue + + grouped.setdefault(int(codes_nq_t.shape[1]), []).append( + (output_index, request_id, slot, codes_nq_t, finished) + ) + + for group in grouped.values(): + plan = {slot: codes_nq_t for _, _, slot, codes_nq_t, _ in group} + decoded = session.step(plan) + for output_index, request_id, slot, _, finished in group: + wav = decoded.get(slot) + if wav is not None: + outputs[output_index] = wav + if finished: + self._finish_stream_request(request_id, session, slot) + + return outputs + + def _append_stream_pending(self, request_id: str, codes_nq_t: torch.Tensor) -> None: + self._stream_pending_codes.setdefault(request_id, []).append( + codes_nq_t.detach().to("cpu", torch.long).contiguous() + ) + + def _pop_stream_pending(self, request_id: str) -> torch.Tensor: + pending = self._stream_pending_codes.pop(request_id, []) + if not pending: + return torch.empty((self._n_vq, 0), dtype=torch.long) + return torch.cat(pending, dim=1).contiguous() + + def _decode_stream_slot_sequence( + self, + session: _MossCodecStreamSession, + slot: int, + codes_nq_t: torch.Tensor, + ) -> torch.Tensor | None: + if codes_nq_t.numel() == 0: + return None + max_step_frames = max(1, int(self._stream_max_step_frames)) + parts: list[torch.Tensor] = [] + for start in range(0, int(codes_nq_t.shape[1]), max_step_frames): + chunk = codes_nq_t[:, start : start + max_step_frames] + decoded = session.step({slot: chunk}) + wav = decoded.get(slot) + if wav is not None: + parts.append(wav) + if not parts: + return None + return torch.cat(parts, dim=-1) if len(parts) > 1 else parts[0] + + def _finish_stream_request( + self, + request_id: str, + session: _MossCodecStreamSession, + slot: int | None, + ) -> None: + if slot is not None: + session.release(slot) + self._stream_req_slots.pop(request_id, None) + self._stream_pending_codes.pop(request_id, None) + self._stream_starved_reqs.discard(request_id) + + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + """Release codec streaming slots when requests finish outside payload flow. + + Normal streaming completion releases slots from ``_decode_streaming_batch`` + when the Stage-0 payload carries ``finished=True``. Client disconnects + and engine-side aborts can finish a request without delivering that + terminal payload, so the runner calls this hook from its finished-request + path to avoid leaking stream slots and buffered codes. + """ + session = self._stream_session + for req_id in finished_req_ids: + request_id = str(req_id) + slot = self._stream_req_slots.get(request_id) + has_state = ( + slot is not None or request_id in self._stream_pending_codes or request_id in self._stream_starved_reqs + ) + if not has_state: + continue + if session is not None: + self._finish_stream_request(request_id, session, slot) + else: + self._stream_req_slots.pop(request_id, None) + self._stream_pending_codes.pop(request_id, None) + self._stream_starved_reqs.discard(request_id) + + def _connector_int(self, name: str, default: int = 0) -> int: + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) + else: + extra_cfg = getattr(connector_cfg, "extra", None) + if isinstance(extra_cfg, dict) and name in extra_cfg: + return int(extra_cfg[name]) + return default + + def _connector_bool(self, name: str, default: bool = False) -> bool: + model_cfg = getattr(self.vllm_config, "model_config", None) + connector_cfg = getattr(model_cfg, "stage_connector_config", None) + if isinstance(connector_cfg, dict): + extra_cfg: dict | None = connector_cfg.get("extra", connector_cfg) + else: + extra_cfg = getattr(connector_cfg, "extra", None) + if isinstance(extra_cfg, dict) and name in extra_cfg: + return bool(extra_cfg[name]) + return default + + def _streaming_cudagraph_capture_sizes_from_compilation_config( + self, + default: list[int], + ) -> list[int]: + if getattr(self.vllm_config.model_config, "enforce_eager", True): + return [] + compilation_config = getattr(self.vllm_config, "compilation_config", None) + capture_sizes = getattr(compilation_config, "cudagraph_capture_sizes", None) + if capture_sizes: + max_step_frames = max(1, int(self._stream_max_step_frames)) + return sorted({int(size) for size in capture_sizes if 0 < int(size) <= max_step_frames}) + return list(default) + # ------------------------------------------------------------------ # Weight loading # ------------------------------------------------------------------ @@ -234,11 +682,10 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for _ in weights: pass - codec_path = os.environ.get("MOSS_TTS_CODEC_PATH", self._codec_path) + codec_path = self._codec_path logger.info("Loading MOSS Audio Tokenizer from %s", codec_path) - codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path) - codec = MossAudioTokenizerModel(codec_cfg) + codec_cfg, codec = self._build_codec(codec_path) model_loader = DefaultModelLoader(self.vllm_config.load_config) source = DefaultModelLoader.Source( @@ -255,10 +702,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # and the rest stay at their random init, which produces noise that # sounds correct in duration but is structurally garbage. _SUFFIX_REMAP: list[tuple[str, str]] = [ + # v1 (MOSS-Audio-Tokenizer) naming. (".self_attn.in_projs.0.", ".attn.in_proj."), (".self_attn.out_projs.0.", ".attn.out_proj."), (".linear1.", ".ff1."), (".linear2.", ".ff2."), + # v2 checkpoint names use singular in_proj/out_proj and ffn.{0,2}; + # the vendored module keeps the original MOSS layer names. + (".self_attn.in_proj.", ".self_attn.in_projs.0."), + (".self_attn.out_proj.", ".self_attn.out_projs.0."), + (".ffn.0.", ".linear1."), + (".ffn.2.", ".linear2."), (".layer_scale_1.", ".ls1."), (".layer_scale_2.", ".ls2."), (".input_proj.", ".in_proj."), @@ -271,21 +725,39 @@ def _remap(name: str) -> str: return name.replace(src, dst) return name - loaded_total = 0 + loaded_names: set[str] = set() skipped: list[str] = [] + shape_mismatches: list[tuple[str, str, tuple[int, ...], tuple[int, ...]]] = [] for name, tensor in codec_weights: # Try direct name first (e.g. ``quantizer.input_proj.*`` exists # under the same name in both layouts), then the remap (transformer # submodules need ``.linear1.``→``.ff1.`` etc.). tgt = name if name in params_dict else _remap(name) if tgt in params_dict: + expected_shape = tuple(params_dict[tgt].shape) + actual_shape = tuple(tensor.shape) + if expected_shape != actual_shape: + shape_mismatches.append((name, tgt, actual_shape, expected_shape)) + continue default_weight_loader(params_dict[tgt], tensor) - loaded_total += 1 + loaded_names.add(tgt) else: skipped.append(name) + + missing = sorted(set(params_dict) - loaded_names) + if missing or skipped or shape_mismatches: + raise RuntimeError( + "MOSS Audio Tokenizer weights were not fully loaded: " + f"loaded={len(loaded_names)}/{len(params_dict)} " + f"missing={len(missing)} skipped={len(skipped)} " + f"shape_mismatches={len(shape_mismatches)}; " + f"first_missing={missing[:5]} " + f"first_skipped={skipped[:5]} " + f"first_shape_mismatches={shape_mismatches[:3]}" + ) logger.info( "MOSS Audio Tokenizer weights: loaded=%d/%d skipped=%d (first skipped: %s)", - loaded_total, + len(loaded_names), len(params_dict), len(skipped), skipped[:3] if skipped else "none", @@ -295,14 +767,27 @@ def _remap(name: str) -> str: codec.to(device=device, dtype=torch.float32) codec.eval() self._codec = codec + inferred_channels = 2 if "v2" in codec_path.lower() else 1 + self._n_channels = int( + getattr( + codec_cfg, + "number_channels", + getattr(codec_cfg, "num_channels", inferred_channels), + ) + or inferred_channels + ) + self._sr_tensor = torch.tensor(int(codec_cfg.sampling_rate), dtype=torch.int32) logger.info( - "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d", + "MOSS Audio Tokenizer loaded: sampling_rate=%d, n_vq=%d, n_channels=%d", codec_cfg.sampling_rate, codec_cfg.num_quantizers, + self._n_channels, ) - self._maybe_enable_decoder_cudagraph(device) + self._configure_decoder_cudagraph(device) + if self._codec_streaming and self._streaming_cudagraph_capture_sizes: + self._ensure_stream_session() # vLLM's track_weights_loading() compares the returned set against # ``self.named_parameters()``. After ``self._codec = codec`` above, @@ -310,12 +795,45 @@ def _remap(name: str) -> str: # mirror that here. return {f"_codec.{name}" for name, _ in codec.named_parameters()} - def _maybe_enable_decoder_cudagraph(self, device: torch.device) -> None: - """Capture CUDA Graphs for the codec decoder if enforce_eager is False.""" + def _build_codec(self, codec_path: str) -> tuple[Any, nn.Module]: + try: + codec_cfg = MossAudioTokenizerV2Config.from_pretrained(codec_path) + codec = MossAudioTokenizerV2Model(codec_cfg) + logger.info("Using vendored MOSS Audio Tokenizer v2 classes from %s", codec_path) + return codec_cfg, codec + except Exception: + logger.exception( + "Failed to instantiate vendored MOSS Audio Tokenizer v2; falling back to legacy vendored codec." + ) + + codec_cfg = MossAudioTokenizerConfig.from_pretrained(codec_path) + codec = MossAudioTokenizerModel(codec_cfg) + return codec_cfg, codec + + def _configure_decoder_cudagraph(self, device: torch.device) -> None: + """Select the codec CUDA Graph path. + + ``enforce_eager`` is the single graph on/off switch. If graphing is + enabled, ``codec_streaming`` decides whether decode uses the persistent + streaming-state wrapper or the offline full-chunk wrapper. + """ if getattr(self.vllm_config.model_config, "enforce_eager", True): + self._streaming_cudagraph_capture_sizes = [] return if self._codec is None: return + if self._codec_streaming: + logger.info( + "MOSS-TTS codec CUDA Graph selected streaming wrapper: capture_sizes=%s", + self._streaming_cudagraph_capture_sizes, + ) + return + + self._enable_non_streaming_decoder_cudagraph(device) + + def _enable_non_streaming_decoder_cudagraph(self, device: torch.device) -> None: + if self._codec is None: + return # Read capture sizes from the connector's extra config (same convention # as Qwen3-TTS), falling back to a sensible default covering common diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py index bf01a83bf9c..2fd0beefd36 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local.py @@ -126,6 +126,7 @@ def _sample_token( top_k: int, top_p: float, do_sample: bool, + generator: torch.Generator | None = None, ) -> torch.Tensor: """Top-k + top-p sampling (matches upstream's ``sample_token`` for the inference branch). @@ -152,7 +153,7 @@ def _sample_token( probs = F.softmax(logits, dim=-1) flat = probs.reshape(-1, probs.shape[-1]) - sampled = torch.multinomial(flat, num_samples=1).reshape(probs.shape[:-1]) + sampled = torch.multinomial(flat, num_samples=1, generator=generator).reshape(probs.shape[:-1]) return sampled diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py new file mode 100644 index 00000000000..e0d871b46d7 --- /dev/null +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_local_depth.py @@ -0,0 +1,240 @@ +# Copyright 2026 OpenMOSS and the vLLM-Omni team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). +"""Per-frame depth transformer for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5). + +A 1-layer GPT2-style block that decodes the ``n_vq`` audio codebook codes for +one audio frame, run inside the talker's ``make_omni_output`` independent of +vLLM's main scheduler -- mirrors ``MossTTSRealtimeLocalTransformer`` +(``modeling_moss_tts_local.py``) in role, but the algorithm and numerics are +faithful to the official ``gpt2_decoder.py`` / ``modeling_moss_tts.py`` +(GPT2-style LayerNorm + bias, SiLU MLP, **interleaved/GPT-J-style RoPE** -- +not vLLM's neox-style concat-half rotation) rather than Realtime's +Qwen3-style ``CodePredictorBaseModel``. + +Its KV cache and positions reset to 0 every audio frame: there is no +cross-frame state. Submodule names (``h.0.*`` / ``ln_f``) match the +checkpoint 1:1 so ``load_weights()`` needs no remapping. +""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( + _sample_token, +) + + +class _MossTTSLocalAttention(nn.Module): + """GPT2-style fused-QKV self-attention with interleaved RoPE. + + Faithful to the official ``MossTTSNanoGPT2Attention``: ``rotate_half`` + operates on even/odd index pairs (GPT-J style), and ``cos``/``sin`` are + built via ``repeat_interleave(2, dim=-1)`` rather than the neox-style + concat-half construction vLLM uses elsewhere. + """ + + def __init__(self, hidden_size: int, n_head: int, rope_base: float) -> None: + super().__init__() + if hidden_size % n_head != 0: + raise ValueError(f"hidden_size={hidden_size} must be divisible by n_head={n_head}") + self.n_head = n_head + self.head_dim = hidden_size // n_head + self.embed_dim = hidden_size + self.c_attn = nn.Linear(hidden_size, 3 * hidden_size, bias=True) + self.c_proj = nn.Linear(hidden_size, hidden_size, bias=True) + inv_freq = 1.0 / (rope_base ** (torch.arange(0, self.head_dim, 2, dtype=torch.float32) / self.head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def _rope_cos_sin( + self, seq_len: int, device: torch.device, dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + position_ids = torch.arange(seq_len, device=device, dtype=torch.float32) + freqs = torch.einsum("s,d->sd", position_ids, self.inv_freq.to(device=device)) + cos = freqs.cos().repeat_interleave(2, dim=-1).to(dtype) + sin = freqs.sin().repeat_interleave(2, dim=-1).to(dtype) + return cos.view(1, seq_len, 1, self.head_dim), sin.view(1, seq_len, 1, self.head_dim) + + @staticmethod + def _rotate_half(x: torch.Tensor) -> torch.Tensor: + even = x[..., ::2] + odd = x[..., 1::2] + return torch.stack((-odd, even), dim=-1).reshape_as(x) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """``hidden_states``: ``(B, S, H)``. Re-prefills with a fresh causal mask over ``[0, S)`` every call.""" + batch_size, seq_len, _ = hidden_states.shape + qkv = self.c_attn(hidden_states) + query, key, value = qkv.split(self.embed_dim, dim=-1) + query = query.view(batch_size, seq_len, self.n_head, self.head_dim) + key = key.view(batch_size, seq_len, self.n_head, self.head_dim) + value = value.view(batch_size, seq_len, self.n_head, self.head_dim) + + cos, sin = self._rope_cos_sin(seq_len, hidden_states.device, hidden_states.dtype) + query = query * cos + self._rotate_half(query) * sin + key = key * cos + self._rotate_half(key) * sin + + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + attn_output = F.scaled_dot_product_attention(query, key, value, is_causal=True) + attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim) + return self.c_proj(attn_output) + + +class _MossTTSLocalMLP(nn.Module): + def __init__(self, hidden_size: int, inner_size: int) -> None: + super().__init__() + self.fc_in = nn.Linear(hidden_size, inner_size, bias=True) + self.fc_out = nn.Linear(inner_size, hidden_size, bias=True) + self.act = nn.SiLU() + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.fc_out(self.act(self.fc_in(hidden_states))) + + +class _MossTTSLocalBlock(nn.Module): + def __init__(self, hidden_size: int, n_head: int, inner_size: int, rope_base: float, eps: float) -> None: + super().__init__() + self.ln_1 = nn.LayerNorm(hidden_size, eps=eps) + self.attn = _MossTTSLocalAttention(hidden_size, n_head, rope_base) + self.ln_2 = nn.LayerNorm(hidden_size, eps=eps) + self.mlp = _MossTTSLocalMLP(hidden_size, inner_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states + self.attn(self.ln_1(hidden_states)) + hidden_states = hidden_states + self.mlp(self.ln_2(hidden_states)) + return hidden_states + + +class MossTTSLocalDepthTransformer(nn.Module): + """Per-frame depth transformer for MOSS-TTS-Local-Transformer-v1.5. + + Per frame: + - position 0's input is the backbone's last hidden state; its output + feeds BOTH the binary continue/stop head (``local_text_lm_head``) + and codebook-0's head (``audio_lm_heads[0]``) simultaneously. + - codebooks 1..n_vq-1 are sampled sequentially: each sampled code is + re-embedded (``audio_embeddings[c]``) and appended as the next + position, re-prefilling the block over the growing (<=n_vq) sequence + with a fresh causal mask each call -- mathematically identical to + incremental KV-cache decoding since attention is strictly causal and + only the last position is ever read, but avoids any cache plumbing + given the trivially short sequence length. + """ + + def __init__(self, gpt2_config, hidden_size: int | None = None) -> None: + super().__init__() + self.hidden_size = int(hidden_size if hidden_size is not None else gpt2_config.n_embd) + n_head = int(gpt2_config.n_head) + inner_size = int(gpt2_config.n_inner) + eps = float(getattr(gpt2_config, "layer_norm_epsilon", 1e-5)) + rope_base = float(getattr(gpt2_config, "rope_base", 1_000_000.0)) + self.h = nn.ModuleList([_MossTTSLocalBlock(self.hidden_size, n_head, inner_size, rope_base, eps)]) + self.ln_f = nn.LayerNorm(self.hidden_size, eps=eps) + + def _forward_prefix(self, seq_embeds: torch.Tensor) -> torch.Tensor: + hidden_states = seq_embeds + for block in self.h: + hidden_states = block(hidden_states) + return self.ln_f(hidden_states) + + @torch.no_grad() + def generate_frame( + self, + backbone_last_hidden: torch.Tensor, # (B, H) + audio_lm_heads: nn.ModuleList, # n_vq x Linear(H -> audio_vocab_size) + audio_embeddings: nn.ModuleList, # n_vq x Embedding(audio_vocab_size, H) + local_text_lm_head: nn.Module, # Linear(H -> 2): [continue, stop] + *, + n_vq: int, + do_sample: bool = True, + temperature: float = 1.0, + top_k: int = 50, + top_p: float = 1.0, + text_temperature: float = 1.0, + text_top_k: int = 50, + text_top_p: float = 1.0, + repetition_penalty: float = 1.0, + history_per_codebook: list[list[int]] | None = None, + generator: torch.Generator | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Generate one audio frame for batch B. + + Returns ``(should_continue, codes)``: ``should_continue`` is a + ``(B,)`` bool tensor (``True`` iff the binary head picked the + "continue" candidate, i.e. logits index 0); ``codes`` is a + ``(B, n_vq)`` LongTensor of sampled codebook indices. + + ``history_per_codebook[c]`` is a list of recently-emitted token ids + for codebook ``c``; when ``repetition_penalty != 1.0`` those tokens' + logits get scaled down before sampling (mirrors upstream's + ``_apply_repetition_penalty``). + """ + batch_size = backbone_last_hidden.shape[0] + dtype = self.ln_f.weight.dtype + + embeds = backbone_last_hidden.new_zeros((batch_size, n_vq, self.hidden_size), dtype=dtype) + embeds[:, 0, :] = backbone_last_hidden.to(dtype) + + hidden = self._forward_prefix(embeds[:, :1, :]) + local_hidden = hidden[:, 0, :] + + binary_logits = local_text_lm_head(local_hidden).float() + # This is a binary continue/stop gate. The checkpoint expects sampling + # here; greedy argmax is biased toward "continue" and may never stop. + binary_choice = _sample_token( + binary_logits, + text_temperature, + text_top_k, + text_top_p, + do_sample, + generator=generator, + ) + should_continue = binary_choice.eq(0) + import os as _os + + if _os.environ.get("MOSS_TTS_DEBUG_STOP"): + import logging as _logging + + _logging.getLogger("moss_tts_debug").warning( + "binary_logits=%s choice=%s", binary_logits.tolist(), binary_choice.tolist() + ) + + codes = backbone_last_hidden.new_zeros((batch_size, n_vq), dtype=torch.long) + for channel_index in range(n_vq): + channel_logits = audio_lm_heads[channel_index](local_hidden).float() + if ( + repetition_penalty != 1.0 + and history_per_codebook is not None + and channel_index < len(history_per_codebook) + ): + hist = history_per_codebook[channel_index] + if hist: + hist_t = torch.tensor(hist, dtype=torch.long, device=channel_logits.device) + sel = channel_logits.index_select(-1, hist_t) + pos = sel > 0 + sel = torch.where(pos, sel / repetition_penalty, sel * repetition_penalty) + channel_logits.index_copy_(-1, hist_t, sel) + channel_token = _sample_token( + channel_logits, + temperature, + top_k, + top_p, + do_sample, + generator=generator, + ) + codes[:, channel_index] = channel_token + + if channel_index + 1 < n_vq: + embeds[:, channel_index + 1, :] = audio_embeddings[channel_index](channel_token).to(dtype) + hidden = self._forward_prefix(embeds[:, : channel_index + 2, :]) + local_hidden = hidden[:, channel_index + 1, :] + + return should_continue, codes + + +__all__ = ["MossTTSLocalDepthTransformer"] diff --git a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py index e5fcdb2349a..9758c922e8b 100644 --- a/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py +++ b/vllm_omni/model_executor/models/moss_tts/modeling_moss_tts_talker.py @@ -6,7 +6,7 @@ from __future__ import annotations import copy -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any import torch @@ -22,22 +22,58 @@ from vllm_omni.model_executor.models.moss_tts.configuration_moss_tts import ( MossTTSDelayConfig, + MossTTSLocalConfig, MossTTSRealtimeConfig, ) from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local import ( MossTTSRealtimeLocalTransformer, ) +from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import ( + MossTTSLocalDepthTransformer, +) from vllm_omni.model_executor.models.output_templates import OmniOutput logger = init_logger(__name__) -_AUDIO_PAD_ID = 0 # padding id used before delay slot fires - def _maybe_prefix(prefix: str, name: str) -> str: return f"{prefix}.{name}" if prefix else name +def _iter_state_row_spans( + states: Sequence[dict[str, Any]], + spans: Sequence[tuple[int, int]] | None, + num_rows: int, +) -> Iterable[tuple[dict[str, Any], int, int]]: + """Yield state/logit-row ranges without assuming equal request row counts.""" + if num_rows <= 0 or not states: + return + + if spans is not None: + if len(spans) != len(states): + raise RuntimeError(f"request_token_spans has {len(spans)} entries for {len(states)} request states") + if all(row_end <= num_rows for _, row_end in spans): + for state, (row_start, row_end) in zip(states, spans, strict=False): + if row_start < row_end: + yield state, int(row_start), int(row_end) + return + if num_rows != len(states): + raise RuntimeError( + "request_token_spans describe full hidden rows, but current logits " + f"have {num_rows} rows for {len(states)} states and are not one-row-per-request" + ) + + if num_rows == len(states): + for i, state in enumerate(states): + yield state, i, i + 1 + return + + raise RuntimeError( + "MOSS-TTS compute_logits requires request_token_spans or one-row-per-request logits; " + f"got {num_rows} rows for {len(states)} request states" + ) + + # --------------------------------------------------------------------------- # MossTTSDelayTalkerForGeneration # --------------------------------------------------------------------------- @@ -104,9 +140,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.pad_token_id: int = getattr(self.config, "pad_token_id", 151643) self.im_end_token_id: int = getattr(self.config, "im_end_token_id", 151645) - # Qwen3 backbone — weights live under ``language_model.*`` + # Qwen3 backbone — weights live under ``language_model.*``. vLLM's + # Qwen3Model reads num_hidden_layers/hidden_size/etc. from + # vllm_config.model_config.hf_config, but MOSS nests the LLM config under + # ``language_config`` (a Qwen3Config). Rebind the backbone's hf_config to + # it — same pattern as the Qwen3-Omni talker — so the inner model sees a + # real Qwen3Config instead of MossTTSDelayConfig (which lacks + # num_hidden_layers). + lm_vllm_config = vllm_config.with_hf_config(lm_cfg, architectures=["Qwen3ForCausalLM"]) + lm_vllm_config.model_config.hf_text_config = lm_vllm_config.model_config.hf_config self.model = Qwen3Model( - vllm_config=vllm_config, + vllm_config=lm_vllm_config, prefix=_maybe_prefix(prefix, "model"), ) @@ -153,6 +197,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # `make_omni_output`, which runs immediately before the sampler). self._batch_state: list[dict[str, Any]] | None = None + # Stacked weight caches built after load_weights(). Avoids a Python + # loop over n_vq heads/embeddings at every decode step. + # shapes: (n_vq, audio_vocab+1, hidden_size) + self._stacked_audio_head_w: torch.Tensor | None = None + self._stacked_audio_emb_w: torch.Tensor | None = None + # ------------------------------------------------------------------ # vLLM-Omni hooks # ------------------------------------------------------------------ @@ -194,7 +244,7 @@ def compute_logits( hidden_states = hidden_states.text_hidden_states if hidden_states is None: return None - logits = self.logits_processor(self.text_lm_head, hidden_states, sampling_metadata) + logits = self.logits_processor(self.text_lm_head, hidden_states) if logits is None or self._batch_state is None: return logits @@ -205,14 +255,13 @@ def compute_logits( n_vq = self.n_vq device = logits.device vocab_size = logits.shape[-1] - rows_per_state = max(1, logits.shape[0] // max(1, len(states))) - for i, state in enumerate(states): - if not isinstance(state, dict): - continue - row_start = i * rows_per_state - row_end = min(row_start + rows_per_state, logits.shape[0]) - if row_start >= row_end: + for state, row_start, row_end in _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + logits.shape[0], + ): + if state is None: continue row = logits[row_start:row_end] @@ -394,9 +443,15 @@ def preprocess( chunk = ref_codes[ref_offset:end_off] if chunk.numel() > 0 and chunk.shape[0] == span_len: codes = chunk.to(device=device, dtype=torch.long).clamp_(0, self.audio_vocab_size) - audio_embed = torch.zeros_like(embeds) - for i, emb_layer in enumerate(self.audio_embeddings): - audio_embed = audio_embed + emb_layer(codes[:, i]) + if self._stacked_audio_emb_w is not None: + codes_nq = codes.t() # (n_vq, L) + audio_embed = self._stacked_audio_emb_w[ + torch.arange(self.n_vq, device=device)[:, None], codes_nq + ].sum(0) + else: + audio_embed = torch.zeros_like(embeds) + for i, emb_layer in enumerate(self.audio_embeddings): + audio_embed = audio_embed + emb_layer(codes[:, i]) embeds = embeds + audio_embed last_ref_row = codes[-1].detach() @@ -436,10 +491,17 @@ def preprocess( audio_codes_buf = (info_dict.get("audio_codes", {}) or {}).get("current") if isinstance(audio_codes_buf, torch.Tensor) and audio_codes_buf.numel() == self.n_vq: codes = audio_codes_buf.to(device=device, dtype=torch.long) - audio_embed = torch.zeros_like(text_embed) - for i, emb_layer in enumerate(self.audio_embeddings): - code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1) - audio_embed = audio_embed + emb_layer(code_i) + if self._stacked_audio_emb_w is not None: + # Single gather: stacked_emb[i, codes[i]] for each i → (n_vq, H) → sum → (1, H) + n_vq_idx = torch.arange(self.n_vq, device=device) + audio_embed = self._stacked_audio_emb_w[n_vq_idx, codes.clamp(0, self.audio_vocab_size)].sum( + 0, keepdim=True + ) + else: + audio_embed = torch.zeros_like(text_embed) + for i, emb_layer in enumerate(self.audio_embeddings): + code_i = codes[i].clamp(0, self.audio_vocab_size).reshape(1) + audio_embed = audio_embed + emb_layer(code_i) combined = text_embed + audio_embed else: combined = text_embed @@ -501,17 +563,38 @@ def _sample_audio_codes( if not bool(sampling_mask.any()): return codes - # Use deploy YAML defaults. The audio sampler isn't routed through - # vLLM's sampler so we apply the upstream's audio_top_k locally. audio_top_k = 25 audio_temp = 1.7 + if self._stacked_audio_head_w is not None: + # Single batched matmul: (n_vq, V+1, H) @ (H,) → (n_vq, V+1). + # Replaces n_vq serial nn.Linear calls. + all_logits = self._stacked_audio_head_w @ last_h.reshape(-1) # (n_vq, V+1) + all_logits[:, -1] = float("-inf") # invalid sentinel + all_logits[:, self.audio_pad_code] = float("-inf") + + active_idx = sampling_mask.nonzero(as_tuple=True)[0] # (n_active,) + active_logits = all_logits[active_idx] # (n_active, V+1) + active_logits = active_logits / max(audio_temp, 1e-6) + if audio_top_k > 0 and audio_top_k < active_logits.shape[-1]: + top_vals, _ = torch.topk(active_logits, audio_top_k, dim=-1) + kth = top_vals[:, -1:].expand_as(active_logits) + active_logits = torch.where( + active_logits < kth, + torch.full_like(active_logits, float("-inf")), + active_logits, + ) + probs = torch.softmax(active_logits, dim=-1) + codes[active_idx] = torch.multinomial(probs, num_samples=1).squeeze(-1).long() + return codes + + # Fallback: per-head loop (used before load_weights() completes). for i in range(n_vq): if not bool(sampling_mask[i]): continue head = self.audio_heads[i] logits = head(last_h).reshape(-1) # (V,) - logits[..., -1] = float("-inf") # invalid sentinel + logits[..., -1] = float("-inf") logits[..., self.audio_pad_code] = float("-inf") sampled = self._sample_with_top_k(logits, audio_top_k, audio_temp) codes[i] = sampled.long() @@ -531,6 +614,7 @@ def make_omni_output( """ if isinstance(model_outputs, OmniOutput): self._batch_state = None + self._batch_state_spans = None return model_outputs hidden = model_outputs # (S, H) @@ -552,19 +636,37 @@ def make_omni_output( # Stash the per-row state for compute_logits to apply masks. logits # rows align with hidden rows, which align with info_dicts in order. self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] + self._batch_state_spans = kwargs.get("request_token_spans") + + # Per-request (start, end) hidden-row spans from the runner. In mixed + # prefill+decode steps the per-request token counts differ, so an equal + # rows-per-request split would sample codes from the wrong rows. + spans = kwargs.get("request_token_spans") - audio_codes_list: list[torch.Tensor] = [] + # One accumulated-codes tensor per request, in batch order. The generic + # splitter (``to_payload_element``) routes ``per_req_codes[idx]`` to + # request ``idx``, so Stage 1 decodes each request's own codes instead + # of request 0's. Skipped requests get an empty placeholder so indices + # stay aligned (a shorter list would silently fall back to entry 0). + per_req_codes: list[torch.Tensor] = [] + have_codes = False if hidden.numel() > 0: num_rows = hidden.shape[0] - rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) + if spans is None or len(spans) != len(info_dicts): + raise RuntimeError( + "MOSS-TTS make_omni_output requires request_token_spans " + f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}" + ) for i, info in enumerate(info_dicts): if not isinstance(info, dict): + per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) continue - row_start = i * rows_per_req - row_end = min(row_start + rows_per_req, num_rows) + row_start, row_end = spans[i] + row_end = min(int(row_end), num_rows) if row_start >= row_end: + per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) continue # Sample new audio codes from the last hidden state of this request. @@ -582,21 +684,20 @@ def make_omni_output( "current": new_codes, "accumulated": updated_acc, } - audio_codes_list.append(updated_acc) + per_req_codes.append(updated_acc) + have_codes = True - if not audio_codes_list: + if not have_codes: return OmniOutput( text_hidden_states=hidden, multimodal_outputs={}, ) - # Forward the accumulated audio codes (T, NQ) of the first request to - # Stage 1. Multi-request batching for the codec is not yet supported - # by this adapter; the existing code already takes the first stage - # output anyway. + # Emit one accumulated-codes tensor per request as a list; the splitter + # picks ``element[idx]`` per request on the way to Stage 1. return OmniOutput( text_hidden_states=hidden, - multimodal_outputs={"codes": {"audio": audio_codes_list[0]}}, + multimodal_outputs={"codes": {"audio": per_req_codes}}, ) # ------------------------------------------------------------------ @@ -661,6 +762,14 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for n in backbone_loaded: loaded.add(f"model.{n}") + # Build stacked weight caches for vectorised decode ops. + self._stacked_audio_head_w = torch.stack( + [h.weight.detach() for h in self.audio_heads], dim=0 + ) # (n_vq, audio_vocab+1, hidden_size) + self._stacked_audio_emb_w = torch.stack( + [e.weight.detach() for e in self.audio_embeddings], dim=0 + ) # (n_vq, audio_vocab+1, hidden_size) + return loaded @@ -748,6 +857,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.logits_processor = LogitsProcessor(self.text_vocab_size) self._batch_state: list[dict[str, Any]] | None = None + # Stacked weight cache built after load_weights(). + # shape: (n_vq, audio_vocab_size, hidden_size) + self._stacked_audio_emb_w: torch.Tensor | None = None + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { ("audio_codes", "current"), ("audio_codes", "accumulated"), @@ -797,13 +910,12 @@ def compute_logits( logits = hidden_states.new_full((B, V), float("-inf")) states = self._batch_state or [] - rows_per_state = max(1, B // max(1, len(states) or 1)) - for i, state in enumerate(states): - r0 = i * rows_per_state - r1 = min(r0 + rows_per_state, B) - if r0 >= r1: - continue - if not isinstance(state, dict): + for state, r0, r1 in _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + B, + ): + if state is None: logits[r0:r1, self.text_pad_id] = 0.0 continue if state.get("is_stopping"): @@ -837,6 +949,11 @@ def _build_input_embeds( embeds = self.embed_tokens[0](text_ids) if audio_codes is None: return embeds + if self._stacked_audio_emb_w is not None: + # Single gather: (n_vq, vocab, H)[arange[:,None], codes.T] → (n_vq, T, H) → sum → (T, H) + n_vq_idx = torch.arange(self.n_vq, device=audio_codes.device) + codes_t = audio_codes.t().clamp(0, self.audio_vocab_size - 1) # (n_vq, T) + return embeds + self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t].sum(0) for i in range(self.n_vq): col = audio_codes[:, i].clamp_(0, self.audio_vocab_size - 1) embeds = embeds + self.embed_tokens[i + 1](col) @@ -919,6 +1036,7 @@ def make_omni_output( ) -> OmniOutput: if isinstance(model_outputs, OmniOutput): self._batch_state = None + self._batch_state_spans = None return model_outputs hidden = model_outputs # (S, H) @@ -932,16 +1050,29 @@ def make_omni_output( info["audio_state"] = {"is_stopping": False, "step": 0} self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] - - audio_codes_list: list[torch.Tensor] = [] + self._batch_state_spans = kwargs.get("request_token_spans") + + # See delay talker: real per-request row spans from the runner are + # required because mixed prefill+decode steps have unequal row counts. + spans = kwargs.get("request_token_spans") + # Per-request accumulated codes in batch order, pre-filled with empty + # placeholders so every skip path (stopped / eos / bos / pad) keeps + # indices aligned for to_payload_element's per-request routing. A + # shorter list would silently fall back to request 0's codes. + per_req_codes: list[torch.Tensor] = [hidden.new_empty((0, self.n_vq), dtype=torch.long) for _ in info_dicts] + have_codes = False if hidden.numel() > 0 and info_dicts: num_rows = hidden.shape[0] - rows_per_req = max(1, num_rows // max(1, len(info_dicts) or 1)) + if spans is None or len(spans) != len(info_dicts): + raise RuntimeError( + "MOSS-TTS make_omni_output requires request_token_spans " + f"for {len(info_dicts)} request infos, got {0 if spans is None else len(spans)}" + ) for i, info in enumerate(info_dicts): if not isinstance(info, dict): continue - row_start = i * rows_per_req - row_end = min(row_start + rows_per_req, num_rows) + row_start, row_end = spans[i] + row_end = min(int(row_end), num_rows) if row_start >= row_end: continue @@ -972,7 +1103,7 @@ def make_omni_output( history_per_codebook=hist_per_cb, ).squeeze(0) # (n_vq,) if int(state.get("step", 0)) < 5 or int(state.get("step", 0)) % 50 == 0: - logger.info( + logger.debug( "[MossTTSRealtime make_omni] step=%d ch0=%d cursor=%d/%d", int(state.get("step", 0)), int(new_codes[0].item()), @@ -1008,24 +1139,19 @@ def make_omni_output( info["audio_codes"] = {"current": new_codes, "accumulated": updated_acc} state["step"] = int(state.get("step", 0)) + 1 - audio_codes_list.append(updated_acc) + per_req_codes[i] = new_codes.unsqueeze(0) + have_codes = True - if not audio_codes_list: + if not have_codes: return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) - # Forward the first request's accumulated codes (multi-request - # batching for the codec is not yet wired up — same as delay). - # The realtime variant emits raw codes (no delay pattern), so we - # signal that to the chunk processor via a 1-D bool tensor in - # ``meta.finished`` (already a tensor field on the schema). The - # processor reads its truth value as "skip de-delay" only for - # realtime; the delay path doesn't populate ``meta`` at all. - device = audio_codes_list[0].device + # Emit one current-frame tensor per request as a list; the splitter + # routes ``element[idx]`` to request idx, so each Stage-1 request receives + # only its own newly generated raw codec row. return OmniOutput( text_hidden_states=hidden, multimodal_outputs={ - "codes": {"audio": audio_codes_list[0]}, - "meta": {"finished": torch.tensor([True], dtype=torch.bool, device=device)}, + "codes": {"audio": per_req_codes}, }, ) @@ -1084,6 +1210,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for n in backbone_loaded: loaded.add(f"model.{n}") + # Build stacked embedding cache for vectorised decode ops. + # embed_tokens[0] is the text embedding; [1..n_vq] are audio codebooks. + self._stacked_audio_emb_w = torch.stack( + [self.embed_tokens[i + 1].weight.detach() for i in range(self.n_vq)], dim=0 + ) # (n_vq, audio_vocab_size, hidden_size) + logger.info( "[MossTTSRealtime] loaded %d/%d params; skipped=%d (first 5: %s)", len(loaded), @@ -1097,7 +1229,512 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded +# --------------------------------------------------------------------------- +# MossTTSLocalTalkerForGeneration +# --------------------------------------------------------------------------- + + +class MossTTSLocalTalkerForGeneration(nn.Module): + """Stage-0 talker for MossTTSLocalModel (MOSS-TTS-Local-Transformer-v1.5). + + Architecture differs from both existing variants: + + * Backbone (Qwen3; checkpoint prefix ``transformer.*``, unlike Delay/ + Realtime's ``language_model.*``) consumes additive text+audio fusion + the same way Delay does (real text embeddings + masked per-codebook + audio embeddings from ``info["codes"]["ref"]`` / ``info["audio_codes"] + ["current"]``), but there is no delay-pattern countdown: the prompt + already ends right at the audio_start boundary, and every step from + the first decode onward drives exactly one audio frame. + * The checkpoint's own ``text_lm_head`` is loaded (trained for a + teacher-forced text loss) but never consulted at inference: the + official reference's own ``generate()`` hardcodes the next + text-channel input to ``audio_assistant_slot_token_id`` regardless of + what ``text_lm_head`` would predict, because + ``local_text_head_mode == "binary"`` for this checkpoint. + ``compute_logits`` therefore synthesises a one-hot logit row + (Realtime-style), forced to ``audio_assistant_slot_token_id`` while + continuing or ``im_end_token_id`` once the local transformer's binary + head decides to stop. + * Per-step audio generation runs ``MossTTSLocalDepthTransformer`` + (1-layer GPT2-style local transformer, interleaved RoPE, ``n_vq=12`` + codebooks -- structurally different from Realtime's Qwen3-style + ``CodePredictorBaseModel``) in the runner-owned ``talker_mtp`` path. Stop + condition: the local transformer's binary head selects the + "audio_end" candidate (logits index 1). + """ + + have_multimodal_outputs: bool = True + has_preprocess: bool = True + has_postprocess: bool = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.vllm_config = vllm_config + self.config: MossTTSLocalConfig = vllm_config.model_config.hf_config + + qwen3_cfg = self.config.qwen3_config + hidden_size = int(qwen3_cfg.hidden_size) + self.hidden_size = hidden_size + self.n_vq: int = int(self.config.n_vq) + self.audio_vocab_size: int = int(self.config.audio_vocab_size) + self.audio_pad_token_id: int = int(self.config.audio_pad_token_id) + self.text_vocab_size: int = int(qwen3_cfg.vocab_size) + + self.audio_assistant_slot_token_id: int = int(self.config.audio_assistant_slot_token_id) + self.im_end_token_id: int = int(self.config.im_end_token_id) + + # Qwen3 backbone. get_text_config() is sufficient for vLLM to size + # KV cache / heads correctly, as long as AutoConfig resolves to the + # local MossTTSLocalConfig rather than the HF remote-code class. + self.model = Qwen3Model( + vllm_config=vllm_config, + prefix=_maybe_prefix(prefix, "model"), + ) + + # Real text LM head -- loaded from the checkpoint but never consulted + # at inference (see class docstring); kept purely as a destination + # for ``text_lm_head.weight`` so load_weights() loads every tensor. + self.text_lm_head = ParallelLMHead( + self.text_vocab_size, + hidden_size, + bias=False, + prefix=_maybe_prefix(prefix, "text_lm_head"), + ) + self.logits_processor = LogitsProcessor(self.text_vocab_size) + + # Per-codebook audio embeddings/heads. audio_pad_token_id (== + # audio_vocab_size) is outside the embedding table -- masked to zero + # contribution rather than stored as a 1025th row (mirrors upstream + # ``_build_inputs_embeds``'s ``valid_mask`` masking). + self.audio_embeddings = nn.ModuleList( + [nn.Embedding(self.audio_vocab_size, hidden_size) for _ in range(self.n_vq)] + ) + self.audio_lm_heads = nn.ModuleList( + [nn.Linear(hidden_size, self.audio_vocab_size, bias=False) for _ in range(self.n_vq)] + ) + self.local_text_lm_head = nn.Linear(hidden_size, 2, bias=False) + self.local_transformer = MossTTSLocalDepthTransformer(self.config.gpt2_config, hidden_size=hidden_size) + + self._batch_state: list[dict[str, Any]] | None = None + # Stacked embedding cache built after load_weights() for vectorised + # additive-fusion gathers (avoids an n_vq-iteration Python loop). + self._stacked_audio_emb_w: torch.Tensor | None = None + self.mtp_hidden_size = hidden_size + self.talker_mtp_accepts_req_infos = True + + self.gpu_resident_buffer_keys: set[tuple[str, str]] = { + ("audio_codes", "current"), + ("audio_codes", "accumulated"), + ("hidden_states", "last"), + } + + # ------------------------------------------------------------------ + # vLLM hooks + # ------------------------------------------------------------------ + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + return self.model.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **_: Any, + ) -> torch.Tensor | IntermediateTensors: + return self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor | OmniOutput, + sampling_metadata: SamplingMetadata | None = None, + ) -> torch.Tensor | None: + """Synthesise a one-hot text logit row per request. + + The text channel is never freely generated here: the upstream + reference hardcodes it to ``audio_assistant_slot_token_id`` while + audio frames keep coming, and only emits ``im_end_token_id`` once the + local transformer's binary head decides to stop (see class + docstring for why the real ``text_lm_head`` is bypassed). + """ + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None or hidden_states.numel() == 0: + return None + + num_rows = hidden_states.shape[0] + logits = hidden_states.new_full((num_rows, self.text_vocab_size), float("-inf")) + + states = self._batch_state or [] + for state, row_start, row_end in _iter_state_row_spans( + states, + getattr(self, "_batch_state_spans", None), + num_rows, + ): + if state is not None and state.get("is_stopping"): + forced_im_end_count = int(state.get("_forced_im_end_count", 0)) + if forced_im_end_count == 1: + logger.warning( + "MOSS-TTS Local request is still running after " + "im_end_token_id=%d was forced once; verify stage-0 " + "stop_token_ids/eos includes this id.", + self.im_end_token_id, + ) + state["_forced_im_end_count"] = forced_im_end_count + 1 + logits[row_start:row_end, self.im_end_token_id] = 0.0 + else: + logits[row_start:row_end, self.audio_assistant_slot_token_id] = 0.0 + if not states: + logits[:, self.audio_assistant_slot_token_id] = 0.0 + return logits + + # ------------------------------------------------------------------ + # Embedding (text + audio codebooks, additive) + # ------------------------------------------------------------------ + + def _audio_embed(self, codes: torch.Tensor) -> torch.Tensor: + """``codes``: ``(T, n_vq)`` long, possibly containing ``audio_pad_token_id`` + (outside the valid embedding range). Returns the ``(T, H)`` additive + audio embedding, masking out pad positions per codebook. + """ + device = codes.device + valid_mask = codes.ne(self.audio_pad_token_id) # (T, n_vq) + safe_codes = codes.masked_fill(~valid_mask, 0).clamp(0, self.audio_vocab_size - 1) + if self._stacked_audio_emb_w is not None: + n_vq_idx = torch.arange(self.n_vq, device=device) + codes_t = safe_codes.t() # (n_vq, T) + gathered = self._stacked_audio_emb_w[n_vq_idx.unsqueeze(1), codes_t] # (n_vq, T, H) + gathered = gathered * valid_mask.t().unsqueeze(-1) + return gathered.sum(0) + embeds: torch.Tensor | None = None + for i, emb_layer in enumerate(self.audio_embeddings): + col = emb_layer(safe_codes[:, i]) * valid_mask[:, i].unsqueeze(-1) + embeds = col if embeds is None else embeds + col + return embeds + + def preprocess( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Any, + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """Build per-step input embeddings (text + audio additive fusion). + + Prefill: initialise the per-request ``{is_stopping, step}`` state. + Decode: combine the forced text token (from ``compute_logits``) with + the audio codes the local transformer produced last step. + """ + device = input_ids.device + span_len = int(input_ids.shape[0]) + audio_state = info_dict.get("audio_state") + is_first_call = not isinstance(audio_state, dict) + + if span_len > 1 or is_first_call: + embeds = self.model.embed_tokens(input_ids) + + ref_codes = (info_dict.get("codes", {}) or {}).get("ref") + ref_offset = int(info_dict.get("ref_offset", 0)) + last_ref_row = None + if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: + if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: + ref_codes = ref_codes.view(-1, self.n_vq) + if isinstance(ref_codes, torch.Tensor) and ref_codes.dim() == 2: + end_off = ref_offset + span_len + chunk = ref_codes[ref_offset:end_off] + if chunk.numel() > 0 and chunk.shape[0] == span_len: + codes = chunk.to(device=device, dtype=torch.long) + embeds = embeds + self._audio_embed(codes) + last_ref_row = codes[-1].detach() + + current_codes = ( + last_ref_row + if last_ref_row is not None + else torch.full((self.n_vq,), self.audio_pad_token_id, dtype=torch.long, device=device) + ) + + max_new_frames = info_dict.get("max_new_frames", [-1])[0] + + info_update: dict[str, Any] = { + "audio_state": { + "is_stopping": False, + "step": 0, + "max_new_frames": max_new_frames, + }, + "audio_codes": {"current": current_codes, "emit": False}, + "ref_offset": ref_offset + span_len, + } + return input_ids, embeds, info_update + + # Decode step: return the plain forced text-token embedding. The local + # depth transformer runs in OmniGPUModelRunner.talker_mtp before the + # backbone forward; it adds the freshly sampled audio-code embedding + # and stores the generated code frame in model_intermediate_buffer. + text_embed = self.model.embed_tokens(input_ids.reshape(1)) # (1, H) + hs = info_dict.get("hidden_states", {}) or {} + last_hidden = hs.get("last") + if isinstance(last_hidden, torch.Tensor) and last_hidden.numel() > 0: + mtp_hidden = last_hidden.to(device=device, dtype=text_embed.dtype).reshape(1, -1) + else: + mtp_hidden = torch.zeros((1, self.hidden_size), device=device, dtype=text_embed.dtype) + return input_ids, text_embed, {"mtp_inputs": (mtp_hidden, torch.zeros_like(mtp_hidden))} + + def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]: + if hidden_states.numel() == 0: + return {} + return {"hidden_states": {"last": hidden_states[-1].detach()}} + + @torch.inference_mode() + def talker_mtp( + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor, + last_talker_hidden: torch.Tensor, + text_step: torch.Tensor, + do_sample: bool | None = None, + temperature: float | None = None, + top_k: int | None = None, + top_p: float | None = None, + generator: torch.Generator | None = None, + req_infos: list[dict[str, Any]] | None = None, + **_: Any, + ) -> tuple[torch.Tensor, None]: + bsz = int(input_embeds.shape[0]) + dev = input_embeds.device + input_embeds_out = input_embeds.reshape(bsz, -1).clone() + req_infos = req_infos or [{} for _ in range(bsz)] + + def write_update(info: dict[str, Any], update: dict[str, Any]) -> None: + for key, value in update.items(): + if isinstance(value, dict): + target = info.setdefault(key, {}) + target.update(value) + else: + info[key] = value + + # Match MOSS-TTS Local v1.5's model-card/SGLang defaults. The generic + # stage SamplingParams are for the backbone token loop and should not + # silently retune the local depth transformer's audio-channel sampler. + audio_temperature = 1.7 + audio_top_k = 25 + audio_top_p = 0.8 + do_sample_val = True if do_sample is None else bool(do_sample) + + for i in range(bsz): + info = req_infos[i] if i < len(req_infos) and isinstance(req_infos[i], dict) else {} + state = dict(info.get("audio_state", {}) or {}) + step = int(state.get("step", 0)) + max_new_frames = int(state.get("max_new_frames", -1)) + acc = (info.get("audio_codes", {}) or {}).get("accumulated") + + if state.get("is_stopping"): + write_update( + info, + { + "audio_state": state, + "audio_codes": { + "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long), + "emit": False, + }, + }, + ) + continue + + rep_window = 50 + if isinstance(acc, torch.Tensor) and acc.numel() > 0: + tail = acc[-rep_window:].long().cpu().tolist() + hist_per_cb = [[row[cb] for row in tail] for cb in range(self.n_vq)] + else: + hist_per_cb = [[] for _ in range(self.n_vq)] + + should_continue_t, new_codes_b = self.local_transformer.generate_frame( + last_talker_hidden[i : i + 1], + self.audio_lm_heads, + self.audio_embeddings, + self.local_text_lm_head, + n_vq=self.n_vq, + do_sample=do_sample_val, + temperature=audio_temperature, + top_k=audio_top_k, + top_p=audio_top_p, + repetition_penalty=1.0, + history_per_codebook=hist_per_cb, + generator=generator, + ) + should_continue = bool(should_continue_t[0].item()) + new_codes = new_codes_b.squeeze(0).to(device=dev, dtype=torch.long) + + if not should_continue or (0 <= max_new_frames <= step): + state["is_stopping"] = True + state["step"] = step + 1 + write_update( + info, + { + "audio_state": state, + "audio_codes": { + "current": input_embeds.new_empty((0, self.n_vq), dtype=torch.long), + "accumulated": acc, + "emit": False, + }, + }, + ) + continue + + if isinstance(acc, torch.Tensor) and acc.numel() > 0: + updated_acc = torch.cat([acc.to(dev), new_codes.unsqueeze(0)], dim=0) + else: + updated_acc = new_codes.unsqueeze(0) + + state["step"] = step + 1 + current = new_codes.unsqueeze(0) + input_embeds_out[i : i + 1] = input_embeds_out[i : i + 1] + self._audio_embed(current).to( + dtype=input_embeds_out.dtype + ) + write_update( + info, + { + "audio_state": state, + "audio_codes": { + "current": current, + "accumulated": updated_acc, + "emit": True, + }, + }, + ) + + return input_embeds_out, None + + # ------------------------------------------------------------------ + # Package runner-generated audio frames + # ------------------------------------------------------------------ + + def make_omni_output( + self, + model_outputs: torch.Tensor | OmniOutput, + **kwargs: Any, + ) -> OmniOutput: + if isinstance(model_outputs, OmniOutput): + self._batch_state = None + self._batch_state_spans = None + return model_outputs + + hidden = model_outputs # (S, H) + info_dicts: list[dict[str, Any]] = ( + kwargs.get("model_intermediate_buffer") or kwargs.get("runtime_additional_information") or [] + ) + + for info in info_dicts: + if isinstance(info, dict) and not isinstance(info.get("audio_state"), dict): + info["audio_state"] = {"is_stopping": False, "step": 0, "max_new_frames": -1} + + self._batch_state = [(info["audio_state"] if isinstance(info, dict) else {}) for info in info_dicts] + self._batch_state_spans = kwargs.get("request_token_spans") + + per_req_codes: list[torch.Tensor] = [] + have_codes = False + for info in info_dicts: + current = (info.get("audio_codes", {}) or {}).get("current") if isinstance(info, dict) else None + emit = bool((info.get("audio_codes", {}) or {}).get("emit")) if isinstance(info, dict) else False + if emit and isinstance(current, torch.Tensor) and current.numel() > 0: + codes = current.to(device=hidden.device, dtype=torch.long) + if codes.dim() == 1: + codes = codes.unsqueeze(0) + per_req_codes.append(codes) + have_codes = True + else: + per_req_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) + + if not have_codes: + return OmniOutput(text_hidden_states=hidden, multimodal_outputs={}) + + # Local-v1.5 emits the current raw code frame (no delay pattern); the + # pipeline routes those rows through ``talker2codec_raw_async_chunk``. + per_req_ref_codes: list[torch.Tensor] = [] + for info in info_dicts: + ref_codes = (info.get("codes", {}) or {}).get("ref") if isinstance(info, dict) else None + if isinstance(ref_codes, torch.Tensor) and ref_codes.numel() > 0: + ref_codes = ref_codes.detach() + if ref_codes.dim() == 1 and ref_codes.numel() % self.n_vq == 0: + ref_codes = ref_codes.view(-1, self.n_vq) + if ref_codes.dim() == 2 and ref_codes.shape[1] == self.n_vq: + per_req_ref_codes.append(ref_codes.to(device=hidden.device, dtype=torch.long)) + continue + per_req_ref_codes.append(hidden.new_empty((0, self.n_vq), dtype=torch.long)) + + return OmniOutput( + text_hidden_states=hidden, + multimodal_outputs={ + "codes": {"audio": per_req_codes, "ref": per_req_ref_codes}, + }, + ) + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Map checkpoint weight names to vLLM-Omni module names. + + Checkpoint layout (prefix ``transformer.*`` for the backbone -- + unlike Delay/Realtime's ``language_model.*``): + transformer.* → model.* (Qwen3 backbone) + text_lm_head.weight → text_lm_head.weight + audio_embeddings.{i}.weight → audio_embeddings.{i}.weight + audio_lm_heads.{i}.weight → audio_lm_heads.{i}.weight + local_text_lm_head.weight → local_text_lm_head.weight + local_transformer.h.0.* → local_transformer.h.0.* + local_transformer.ln_f.* → local_transformer.ln_f.* + + Every "tied" pair in the upstream config is stored as a full, + separate tensor in this checkpoint (confirmed via the raw + safetensors header: 438 keys, no dedup), so aside from the backbone + prefix this is pure 1:1 name-based loading. + """ + loaded: set[str] = set() + params_dict = dict(self.named_parameters()) + + backbone_weights: list[tuple[str, torch.Tensor]] = [] + skipped: list[str] = [] + + for name, tensor in weights: + if name.startswith("transformer."): + backbone_weights.append((name[len("transformer.") :], tensor)) + continue + if name in params_dict: + default_weight_loader(params_dict[name], tensor) + loaded.add(name) + else: + skipped.append(name) + + backbone_loaded = self.model.load_weights(iter(backbone_weights)) + for n in backbone_loaded: + loaded.add(f"model.{n}") + + self._stacked_audio_emb_w = torch.stack( + [e.weight.detach() for e in self.audio_embeddings], dim=0 + ) # (n_vq, audio_vocab_size, hidden_size) + + logger.info( + "[MossTTSLocal] loaded %d/%d params; skipped=%d (first 5: %s)", + len(loaded), + len(params_dict), + len(skipped), + skipped[:5], + ) + not_loaded = [n for n in params_dict if n not in loaded] + if not_loaded: + logger.warning("[MossTTSLocal] %d params NOT loaded (first 5: %s)", len(not_loaded), not_loaded[:5]) + return loaded + + __all__ = [ "MossTTSDelayTalkerForGeneration", "MossTTSRealtimeTalkerForGeneration", + "MossTTSLocalTalkerForGeneration", ] diff --git a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py index ed03c018e91..87f2550728d 100644 --- a/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py +++ b/vllm_omni/model_executor/models/moss_tts/moss_codec_cudagraph.py @@ -1,6 +1,6 @@ """CUDA Graph acceleration for the MOSS Audio Tokenizer codec decoder. -Captures MossAudioTokenizerModel._decode for a set of fixed frame-count +Captures MossAudioTokenizerModel._decode_frame for a set of fixed frame-count bucket sizes, then replays the captured graph at inference time to eliminate kernel-launch overhead. Inputs that exceed all captured sizes fall back to eager execution transparently. @@ -24,7 +24,7 @@ class MossTTSCUDAGraphCodecWrapper: - """CUDA Graph wrapper for MossAudioTokenizerModel._decode. + """CUDA Graph wrapper for MossAudioTokenizerModel._decode_frame. Graphs are keyed by padded_T (int). On each call the actual T is bucket-matched to the smallest pre-captured size >= T. The static code @@ -99,7 +99,7 @@ def warmup(self, device: torch.device) -> None: dummy_codes = torch.zeros(nq, 1, size, dtype=torch.long, device=device) dummy_lengths = torch.tensor([size], dtype=torch.long, device=device) with torch.no_grad(): - _ = self.model._decode(dummy_codes, dummy_lengths) + _ = self.model._decode_frame(dummy_codes, dummy_lengths) torch.accelerator.synchronize(device) @@ -132,13 +132,13 @@ def _capture(self, size: int, device: torch.device) -> None: # Extra eager warmup inside capture to ensure all kernels are compiled. with torch.no_grad(): - _ = self.model._decode(static_codes, static_lengths) + _ = self.model._decode_frame(static_codes, static_lengths) torch.accelerator.synchronize(device) graph = CUDAGraph() with torch.no_grad(): with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): - static_out = self.model._decode(static_codes, static_lengths) + static_out = self.model._decode_frame(static_codes, static_lengths) self.graphs[size] = graph self.static_codes[size] = static_codes diff --git a/vllm_omni/model_executor/models/moss_tts/pipeline.py b/vllm_omni/model_executor/models/moss_tts/pipeline.py index fbbdd8897ff..254bc18fe42 100644 --- a/vllm_omni/model_executor/models/moss_tts/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts/pipeline.py @@ -22,6 +22,7 @@ MOSS_TTS_PIPELINE = PipelineConfig( model_type="moss_tts", + default_deploy_config_name="moss_tts.yaml", model_arch="MossTTSDelayModel", # HF architectures string stages=( StagePipelineConfig( @@ -31,7 +32,7 @@ input_sources=(), owns_tokenizer=True, engine_output_type="latent", - async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"), + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_delay_async_chunk"), sampling_constraints={ "detokenize": False, }, @@ -53,6 +54,7 @@ MOSS_TTS_REALTIME_PIPELINE = PipelineConfig( model_type="moss_tts_realtime", + default_deploy_config_name="moss_tts_realtime.yaml", model_arch="MossTTSRealtime", # different talker class from the delay variant stages=( StagePipelineConfig( @@ -62,7 +64,7 @@ input_sources=(), owns_tokenizer=True, engine_output_type="latent", - async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_async_chunk"), + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"), sampling_constraints={"detokenize": False}, ), StagePipelineConfig( @@ -80,9 +82,43 @@ ), ) -# The pipeline config is otherwise the same for all 5 variants; the per-variant +MOSS_TTS_LOCAL_PIPELINE = PipelineConfig( + model_type="moss_tts_local", + default_deploy_config_name="moss_tts_local.yaml", + model_arch="MossTTSLocalModel", # different talker class: GPT2-style local depth transformer + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="moss_tts_local", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + owns_tokenizer=True, + engine_output_type="latent", + async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2codec_raw_async_chunk"), + sampling_constraints={ + "detokenize": False, + "stop_token_ids": [151645], + }, + ), + StagePipelineConfig( + stage_id=1, + model_stage="moss_tts_local_codec", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="MossTTSCodecDecoder", + sync_process_input_func=f"{_PROC}.talker2codec", + sampling_constraints={"detokenize": True}, + ), + ), +) + +# The pipeline config is otherwise the same for all variants; the per-variant # differences (n_vq, backbone size, generation strategy) are encoded in the -# HF config.json and the deploy YAML. Realtime is split out because it has a -# different talker architecture (MossTTSRealtime vs MossTTSDelayModel). +# HF config.json and the deploy YAML. Realtime and Local are split out because +# they have different talker architectures from the delay variant +# (MossTTSRealtime / MossTTSLocalModel vs MossTTSDelayModel). -__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE"] +__all__ = ["MOSS_TTS_PIPELINE", "MOSS_TTS_REALTIME_PIPELINE", "MOSS_TTS_LOCAL_PIPELINE"] diff --git a/vllm_omni/model_executor/models/moss_tts_nano/__init__.py b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py index 896ead57d9b..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/__init__.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/__init__.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .modeling_moss_tts_nano import MossTTSNanoForGeneration -__all__ = ["MossTTSNanoForGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py deleted file mode 100644 index a5a9f25acd3..00000000000 --- a/vllm_omni/model_executor/models/moss_tts_nano/configuration_moss_tts_nano.py +++ /dev/null @@ -1,60 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Configuration for MOSS-TTS-Nano in vLLM-Omni single-stage pipeline.""" - -from transformers.configuration_utils import PretrainedConfig - - -class MossTTSNanoConfig(PretrainedConfig): - """Config for MOSS-TTS-Nano (OpenMOSS-Team/MOSS-TTS-Nano). - - The model is a Global-Local GPT2-based AR LM (0.1B) paired with the - MOSS-Audio-Tokenizer-Nano codec. Both components run in a single - vLLM-Omni generation stage. - - Relevant fields from config.json: - gpt2_config – backbone config (n_layer=12, n_embd=768, n_head=12, ...) - n_vq – number of RVQ codebooks (16) - audio_vocab_size – per-codebook vocabulary size (1024) - audio_tokenizer_pretrained_name_or_path – HF hub path for codec model - """ - - model_type = "moss_tts_nano" - - def __init__(self, **kwargs): - gpt2_cfg = kwargs.pop("gpt2_config", None) or {} - if hasattr(gpt2_cfg, "to_dict"): - gpt2_cfg = gpt2_cfg.to_dict() - - super().__init__(**kwargs) - - # --- GPT2 backbone parameters (exposed at top level for vLLM) --- - self.hidden_size = gpt2_cfg.get("n_embd", 768) - self.num_hidden_layers = gpt2_cfg.get("n_layer", 12) - self.num_attention_heads = gpt2_cfg.get("n_head", 12) - self.num_key_value_heads = self.num_attention_heads # no GQA in GPT2 - self.head_dim = self.hidden_size // self.num_attention_heads - self.vocab_size = gpt2_cfg.get("vocab_size", 16384) - self.max_position_embeddings = gpt2_cfg.get("n_positions", 32768) - self.intermediate_size = gpt2_cfg.get("n_inner", self.hidden_size * 4) - - # --- Audio codec parameters --- - self.n_vq: int = getattr(self, "n_vq", 16) - self.audio_vocab_size: int = getattr(self, "audio_vocab_size", 1024) - self.audio_start_token_id: int = getattr(self, "audio_start_token_id", 6) - self.audio_end_token_id: int = getattr(self, "audio_end_token_id", 7) - self.audio_user_slot_token_id: int = getattr(self, "audio_user_slot_token_id", 8) - self.audio_assistant_slot_token_id: int = getattr(self, "audio_assistant_slot_token_id", 9) - self.audio_tokenizer_pretrained_name_or_path: str = getattr( - self, - "audio_tokenizer_pretrained_name_or_path", - "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano", - ) - self.audio_tokenizer_sample_rate: int = getattr(self, "audio_tokenizer_sample_rate", 48000) - - # vLLM requires speculative_config to be absent or None - self.speculative_config = None - - def get_text_config(self, **kwargs): - """Return self so vLLM uses our top-level config.""" - return self diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index 87b93288684..a0b31bdf2db 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -38,6 +38,7 @@ reinit_rotary_inv_freq, transformers_keys_to_ignore_compat, ) +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -172,7 +173,6 @@ class MossTTSNanoForGeneration(nn.Module): has_preprocess = False has_postprocess = False enable_update_additional_information = True - inject_omni_request_id_into_runtime_info = True def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() @@ -183,7 +183,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # Eager construction (not in load_weights) -- see module docstring. _patch_torchaudio_load() - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() self._device: torch.device = device if device.type == "cuda" and torch.cuda.is_bf16_supported(): @@ -194,7 +194,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: tts_dtype = torch.float32 logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype) - from transformers import AutoModelForCausalLM + from transformers import AutoModelForCausalLM, PreTrainedModel + + # Shim _tp_plan to prevent TypeError in transformers when loading remote-code models + if getattr(PreTrainedModel, "_tp_plan", None) is None: + PreTrainedModel._tp_plan = {} # Two orthogonal trust_remote_code transformers fixes: # - _hf_load_without_tp_warmup: tp_plan=None warmup crash (transformers 5.8.x) diff --git a/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py b/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py index 749c4f98029..c089ea90eef 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/pipeline.py @@ -18,6 +18,7 @@ MOSS_TTS_NANO_PIPELINE = PipelineConfig( model_type="moss_tts_nano", + default_deploy_config_name="moss_tts_nano.yaml", model_arch="MossTTSNanoForCausalLM", stages=( StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice.py b/vllm_omni/model_executor/models/omnivoice/omnivoice.py index c5026086e23..c167aa1c0a1 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice.py @@ -38,6 +38,7 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.platforms import current_omni_platform from vllm_omni.transformers_utils.configs.omnivoice import OmniVoiceConfig logger = init_logger(__name__) @@ -379,13 +380,12 @@ def _forward_generator(self, input_ids: torch.Tensor, kwargs: dict) -> OmniOutpu target_ids = torch.full((num_codebooks, target_len), mask_id, dtype=torch.long, device=device) # Conditional: [text] [ref_audio?] [target_mask] + cond_audio_start = text_ids.shape[1] if ref_audio_tokens is not None: ref_tokens = ref_audio_tokens.to(device) # [8, T_ref] cond_ids = torch.cat([text_ids, ref_tokens, target_ids], dim=1) - cond_audio_start = text_ids.shape[1] else: cond_ids = torch.cat([text_ids, target_ids], dim=1) - cond_audio_start = text_ids.shape[1] cond_len = cond_ids.shape[1] @@ -495,7 +495,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: try: device = next(self.parameters()).device except StopIteration: - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = current_omni_platform.get_torch_device() model_dir = self._resolve_model_dir() diff --git a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py index fc5befeb959..7eacfb9bb92 100644 --- a/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py +++ b/vllm_omni/model_executor/models/omnivoice/omnivoice_generator.py @@ -6,8 +6,9 @@ Generates 8-codebook audio tokens from text via 32-step non-autoregressive iterative masked prediction with classifier-free guidance. -Uses vLLM-Omni's DiffusionAttention for optimized full (bidirectional) attention -via FlashAttention/SageAttention/SDPA backends. +Uses full bidirectional attention computed directly with PyTorch SDPA +(torch.nn.functional.scaled_dot_product_attention); no auto-selected +FlashAttention/SageAttention/DiffusionAttention backend is used. """ from __future__ import annotations @@ -237,7 +238,7 @@ def _gumbel_sample(logits: torch.Tensor, temperature: float, generator: torch.Ge # --------------------------------------------------------------------------- -# Qwen3-style transformer blocks using DiffusionAttention +# Qwen3-style transformer blocks using PyTorch SDPA # --------------------------------------------------------------------------- @@ -256,7 +257,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class OmniVoiceAttention(nn.Module): - """Qwen3-style GQA attention using DiffusionAttention backend.""" + """Qwen3-style GQA attention using PyTorch SDPA (full bidirectional).""" def __init__(self, config: OmniVoiceConfig): super().__init__() @@ -315,18 +316,17 @@ def forward( v = v.permute(0, 2, 1, 3) # Convert [B, 1, S, S] bool mask to float mask for SDPA + # (0.0 where attend, -inf where masked) sdpa_mask = None if attention_mask is not None: - sdpa_mask = attention_mask.to(dtype=q.dtype) - sdpa_mask = sdpa_mask.masked_fill(~attention_mask, float("-inf")) - sdpa_mask = sdpa_mask.masked_fill(attention_mask, 0.0) + sdpa_mask = torch.zeros_like(attention_mask, dtype=q.dtype).masked_fill_(~attention_mask, float("-inf")) out = F.scaled_dot_product_attention( q, k, v, attn_mask=sdpa_mask, - scale=1.0 / math.sqrt(self.head_dim), + scale=self.scale, ) # Back to (batch, seq, heads * head_dim) @@ -351,7 +351,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class OmniVoiceTransformerBlock(nn.Module): - """Single Qwen3 transformer block with DiffusionAttention.""" + """Single Qwen3 transformer block with PyTorch SDPA attention.""" def __init__(self, config: OmniVoiceConfig): super().__init__() @@ -643,10 +643,9 @@ class OmniVoiceGenerator(nn.Module): - 32-step iterative unmasking with classifier-free guidance Optimizations: - - DiffusionAttention (FlashAttn/SageAttn/SDPA auto-selected) - - TeaCache / Cache-DiT compatible (hook-based, non-intrusive) + - Full bidirectional attention via PyTorch SDPA (no auto-selected + FlashAttn/SageAttn/DiffusionAttention backend) - regionally_compile() compatible for torch.compile on repeated blocks - - Sequence parallelism via SP hooks for multi-GPU """ # For regionally_compile() support diff --git a/vllm_omni/model_executor/models/omnivoice/pipeline.py b/vllm_omni/model_executor/models/omnivoice/pipeline.py new file mode 100644 index 00000000000..760d5314fb1 --- /dev/null +++ b/vllm_omni/model_executor/models/omnivoice/pipeline.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""OmniVoice pipeline topology (frozen). + +Single-stage diffusion model for text-to-speech. +Stage 0: DiT — text → iterative unmasking → codebook tokens → DAC decode → 24kHz audio +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +OMNIVOICE_PIPELINE = PipelineConfig( + model_type="omnivoice", + default_deploy_config_name="omnivoice.yaml", + model_arch="OmniVoicePipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="audio", + owns_tokenizer=False, + requires_multimodal_data=False, + engine_output_type="audio", + ), + ), +) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py index afd0a92a531..4ce5ef20748 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/pipeline.py @@ -17,6 +17,7 @@ QWEN2_5_OMNI_PIPELINE = PipelineConfig( model_type="qwen2_5_omni", + default_deploy_config_name="qwen2_5_omni.yaml", model_arch="Qwen2_5OmniForConditionalGeneration", stages=( StagePipelineConfig( @@ -38,7 +39,6 @@ execution_type=StageExecutionType.LLM_AR, input_sources=(0,), engine_output_type="latent", - custom_process_input_func=f"{_PROC}.thinker2talker", sync_process_input_func=f"{_PROC}.thinker2talker_token_only", custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", sampling_constraints={ @@ -54,7 +54,6 @@ final_output=True, final_output_type="audio", engine_output_type="audio", - custom_process_input_func=f"{_PROC}.talker2code2wav", sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", sampling_constraints={"detokenize": True}, ), diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 8d5341f3ef4..7b2c2558ceb 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -17,7 +17,6 @@ from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, - Qwen2_5OmniThinkerProcessingInfo, ) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix from vllm.model_executor.models.vision import ( @@ -39,6 +38,7 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges from vllm_omni.platforms import current_omni_platform @@ -81,6 +81,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # Initialize thinker model (multimodal processing) self.thinker = init_vllm_registered_model( vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "thinker"), hf_config=thinker_config, # Use registry architecture key architectures=["Qwen2_5OmniThinkerModel"], diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 5d58a63d48c..03edcfa10c0 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -1,9 +1,11 @@ """Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker with minimal overrides.""" -from collections.abc import Iterable, Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping, Sequence +from functools import partial from typing import Any import numpy as np +import PIL.Image as PILImage import torch from torch import nn from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( @@ -26,16 +28,22 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniAudioFeatureInputs, Qwen2_5OmniThinkerDummyInputsBuilder, - Qwen2_5OmniThinkerProcessingInfo, check_interleaved_audio_video, + create_qwen2_5_omni_thinker_field_factory, merge_interleaved_embeddings, ) from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase, ) +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerMultiModalDataParser as _Qwen2_5OmniThinkerMultiModalDataParserBase, +) from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerMultiModalProcessor as _Qwen2_5OmniThinkerMultiModalProcessorBase, ) +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerProcessingInfo as _Qwen2_5OmniThinkerProcessingInfoBase, +) from vllm.model_executor.models.qwen2_5_vl import ( Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs, @@ -45,6 +53,9 @@ Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs, ) +from vllm.model_executor.models.qwen2_audio import ( + _get_feat_extract_output_lengths, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -54,14 +65,22 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFeatureSpec, + MultiModalFieldConfig, MultiModalKwargsItems, ) -from vllm.multimodal.parse import MultiModalDataItems +from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems, VideoProcessorItems +from vllm.multimodal.processing.context import TimingContext +from vllm.multimodal.processing.inputs import ProcessorInputs from vllm.multimodal.processing.processor import ( + MultiModalProcessingInfo, MultiModalPromptUpdates, PlaceholderFeaturesInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, ) from vllm.sequence import IntermediateTensors +from vllm.utils.collection_utils import is_list_of from vllm_omni.quantization.component_config import ( resolve_encoder_quant_config, @@ -73,12 +92,602 @@ flash_attn = None logger = init_logger(__name__) +_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY = "_vllm_omni_per_video_use_audio_in_video" + + +def _normalize_use_audio_in_video( + value: object, + num_videos: int, +) -> list[bool]: + if isinstance(value, torch.Tensor): + if value.numel() == 1: + return [bool(value.item())] * num_videos + values = value.flatten().tolist() + elif isinstance(value, np.ndarray): + if value.size == 1: + return [bool(value.item())] * num_videos + values = value.flatten().tolist() + elif isinstance(value, bool): + return [value] * num_videos + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + values = list(value) + else: + return [bool(value)] * num_videos + + if len(values) != num_videos: + raise ValueError( + "use_audio_in_video must be a boolean or contain one boolean per " + f"video, but found {len(values)} values for {num_videos} videos." + ) + return [bool(v) for v in values] + + +def _normalize_per_video_use_audio_in_video(value: object) -> list[bool]: + if isinstance(value, torch.Tensor): + values = value.flatten().tolist() + elif isinstance(value, np.ndarray): + values = value.flatten().tolist() + elif isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + values = list(value) + else: + values = [value] + return [bool(v) for v in values] + + +def _is_per_video_use_audio_in_video(value: object) -> bool: + if isinstance(value, bool): + return False + if isinstance(value, torch.Tensor): + return value.numel() != 1 + if isinstance(value, np.ndarray): + return value.size != 1 + return isinstance(value, Sequence) and not isinstance(value, (str, bytes)) + + +def _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs: Mapping[str, object], + num_videos: int, +) -> list[bool]: + value = hf_processor_mm_kwargs.get( + _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY, + hf_processor_mm_kwargs.get("use_audio_in_video", False), + ) + return _normalize_use_audio_in_video(value, num_videos) + + +def _prompt_update_has_audio_token( + mm_prompt_updates: MultiModalPromptUpdates, + item_idx: int, + audio_token_id: int, +) -> bool: + updates = mm_prompt_updates.get("video", []) + if item_idx >= len(updates): + return False + return any( + audio_token_id in update.content.full for update in updates[item_idx] if isinstance(update.content.full, list) + ) + + +def _get_video_second_per_grid_t( + out_mm_data: Mapping[str, object], + hf_processor_mm_kwargs: Mapping[str, object], + item_idx: int, + default: float, +) -> float: + second_per_grid_ts = out_mm_data.get("second_per_grid_ts") + if second_per_grid_ts is None: + second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) + if second_per_grid_ts is None: + return default + + if isinstance(second_per_grid_ts, torch.Tensor): + if second_per_grid_ts.numel() == 1: + return float(second_per_grid_ts.item()) + return float(second_per_grid_ts.flatten()[item_idx].item()) + + if isinstance(second_per_grid_ts, np.ndarray): + if second_per_grid_ts.size == 1: + return float(second_per_grid_ts.item()) + return float(second_per_grid_ts.flatten()[item_idx].item()) + + if isinstance(second_per_grid_ts, Sequence) and not isinstance(second_per_grid_ts, (str, bytes)): + return float(second_per_grid_ts[item_idx]) + + return float(second_per_grid_ts) + + +def _filter_video_use_audio_in_video_for_uncached_items( + hf_processor_mm_kwargs: Mapping[str, object], + video_is_cached: Sequence[bool] | None, +) -> Mapping[str, object]: + use_audio_in_video = hf_processor_mm_kwargs.get( + _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY, + hf_processor_mm_kwargs.get("use_audio_in_video"), + ) + if video_is_cached is None or not _is_per_video_use_audio_in_video(use_audio_in_video): + return hf_processor_mm_kwargs + + mask = _normalize_use_audio_in_video(use_audio_in_video, len(video_is_cached)) + missing_mask = [use_audio for use_audio, is_cached in zip(mask, video_is_cached) if not is_cached] + + filtered_kwargs = dict(hf_processor_mm_kwargs) + filtered_kwargs["use_audio_in_video"] = missing_mask + filtered_kwargs[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] = missing_mask + second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts") + if ( + isinstance(second_per_grid_ts, Sequence) + and not isinstance(second_per_grid_ts, (str, bytes)) + and len(second_per_grid_ts) == len(video_is_cached) + ): + filtered_kwargs["second_per_grid_ts"] = [ + second_per_grid_t + for second_per_grid_t, is_cached in zip(second_per_grid_ts, video_is_cached) + if not is_cached + ] + return filtered_kwargs + + +def _coerce_use_audio_in_video_for_hf_processor( + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], +) -> Mapping[str, object]: + use_audio_in_video = mm_kwargs.get("use_audio_in_video") + if not _is_per_video_use_audio_in_video(use_audio_in_video): + return mm_kwargs + + video_use_audio_in_video = _normalize_per_video_use_audio_in_video(use_audio_in_video) + + hf_mm_kwargs = dict(mm_kwargs) + # HF processors only support a global bool. For per-video masks, vLLM + # consumes the list while building prompt updates and placeholders. + hf_mm_kwargs["use_audio_in_video"] = False + hf_mm_kwargs[_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY] = video_use_audio_in_video + return hf_mm_kwargs + + +def _presampled_videos_hf_kwargs( + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], +) -> Mapping[str, object]: + """Adjust HF video kwargs for videos pre-sampled by vLLM's video loader. + + When ``video_metadata`` is present, the frames were already sampled + according to ``media_io_kwargs``. The HF processor should not resample + them and should use the sampled fps when computing temporal metadata. + """ + video_metadata = mm_data.get("video_metadata") + if not video_metadata: + return mm_kwargs + + mm_kwargs = dict(mm_kwargs) + videos_kwargs = dict(mm_kwargs.get("videos_kwargs") or {}) + videos_kwargs["do_sample_frames"] = False + + def _compute_sampled_video_fps(metadata) -> float | None: + duration = getattr(metadata, "duration", None) + indices = getattr(metadata, "frames_indices", None) + if not duration or not indices or duration <= 0: + return None + return len(indices) / float(duration) + + if "fps" not in videos_kwargs: + fps_values = [_compute_sampled_video_fps(m) for m in video_metadata] + known_fps = [fps for fps in fps_values if fps is not None] + unique_fps = set(known_fps) + if len(unique_fps) == 1: + videos_kwargs["fps"] = known_fps[0] + elif len(unique_fps) > 1: + logger.warning( + "Mixed sampled FPS %s in one request; HF accepts a single fps, using %s.", + sorted(unique_fps), + known_fps[0], + ) + videos_kwargs["fps"] = known_fps[0] + + mm_kwargs["videos_kwargs"] = videos_kwargs + return mm_kwargs + + +class Qwen2_5OmniVideoProcessorItems(VideoProcessorItems): + """Video items that carry the loader's ``(frames, metadata)`` tuples.""" + + def get_processor_data(self) -> Mapping[str, object]: + from transformers.video_utils import VideoMetadata + + videos = [frames for frames, _ in self.data] + video_metadata = [ + VideoMetadata(**{k: v for k, v in metadata.items() if k != "do_sample_frames"}) for _, metadata in self.data + ] + return {"videos": videos, "video_metadata": video_metadata} + + +class Qwen2_5OmniThinkerMultiModalDataParser(_Qwen2_5OmniThinkerMultiModalDataParserBase): + def _parse_video_data(self, data): + if data is None or isinstance(data, dict) or self.is_embeddings(data): + return super()._parse_video_data(data) + + if (is_list_of(data, PILImage.Image) and len(data) > 0) or ( + isinstance(data, (np.ndarray, torch.Tensor)) and data.ndim == 4 + ): + data_items = [data] + elif isinstance(data, (np.ndarray, torch.Tensor)): + data_items = [elem for elem in data] + elif isinstance(data, tuple) and len(data) == 2: + data_items = [data] + else: + data_items = data + + videos_with_metadata = [self._get_video_with_metadata(item) for item in data_items] + if any(metadata is None for _, metadata in videos_with_metadata): + return super()._parse_video_data(data) + + return Qwen2_5OmniVideoProcessorItems( + videos_with_metadata, + metadata=[metadata for _, metadata in videos_with_metadata], + ) + + +class Qwen2_5OmniThinkerProcessingInfo(_Qwen2_5OmniThinkerProcessingInfoBase): + def get_data_parser(self): + feature_extractor = self.get_feature_extractor() + + return Qwen2_5OmniThinkerMultiModalDataParser( + spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size, + target_sr=feature_extractor.sampling_rate, + target_channels=self.get_target_channels(), + expected_hidden_size=self._get_expected_hidden_size(), + ) + class Qwen2_5OmniThinkerMultiModalProcessor( _Qwen2_5OmniThinkerMultiModalProcessorBase, ): """Override to fix use_audio_in_video detection when mm cache returns None.""" + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ): + mm_kwargs = _presampled_videos_hf_kwargs(mm_data, mm_kwargs) + mm_kwargs = _coerce_use_audio_in_video_for_hf_processor(mm_data, mm_kwargs) + hf_inputs = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + per_video_mask = mm_kwargs.get(_PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY) + if per_video_mask is not None: + hf_inputs["use_audio_in_video"] = torch.tensor(per_video_mask) + return hf_inputs + + def _get_mm_fields_config( + self, + hf_inputs, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + fields_config = dict( + create_qwen2_5_omni_thinker_field_factory(self.info.get_hf_config().vision_config.spatial_merge_size)( + hf_inputs + ) + ) + use_audio_in_video = hf_inputs.get("use_audio_in_video") + is_batched_mask = isinstance(use_audio_in_video, torch.Tensor) and use_audio_in_video.numel() > 1 + if _PER_VIDEO_USE_AUDIO_IN_VIDEO_KEY in hf_processor_mm_kwargs or is_batched_mask: + fields_config["use_audio_in_video"] = MultiModalFieldConfig.batched("video") + return fields_config + + def _get_video_use_audio_in_video( + self, + mm_kwargs: MultiModalKwargsItems, + mm_prompt_updates: MultiModalPromptUpdates, + ) -> list[bool]: + video_kwargs = mm_kwargs.get("video", []) + if not video_kwargs: + return [] + + audio_token_id = self.info.get_hf_config().audio_token_id + video_use_audio_in_video = [] + has_use_audio_in_video = any(item is not None and "use_audio_in_video" in item for item in video_kwargs) + for item_idx, item in enumerate(video_kwargs): + if has_use_audio_in_video: + if item is None or "use_audio_in_video" not in item: + video_use_audio_in_video.append( + _prompt_update_has_audio_token( + mm_prompt_updates, + item_idx, + audio_token_id, + ) + ) + continue + use_audio_in_video_tensor = item["use_audio_in_video"].data + if use_audio_in_video_tensor.numel() > 0: + video_use_audio_in_video.append(bool(use_audio_in_video_tensor.item())) + continue + video_use_audio_in_video.append(False) + continue + + video_use_audio_in_video.append( + _prompt_update_has_audio_token( + mm_prompt_updates, + item_idx, + audio_token_id, + ) + ) + + return video_use_audio_in_video + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) + vocab = tokenizer.get_vocab() + + audio_token = processor.audio_token + image_token = processor.image_token + video_token = processor.video_token + audio_token_id = vocab[audio_token] + image_token_id = vocab[image_token] + video_token_id = vocab[video_token] + + out_mm_data = out_mm_kwargs.get_data() + audio_feature_lengths = out_mm_data.get("audio_feature_lengths") + feature_attention_mask = out_mm_data.get("feature_attention_mask") + if audio_feature_lengths is None and feature_attention_mask is None: + audio_output_lengths = [] + elif audio_feature_lengths is not None: + _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths) + audio_output_lengths = audio_output_lens.tolist() + elif feature_attention_mask is not None: + assert isinstance(feature_attention_mask, torch.Tensor) + _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1)) + audio_output_lengths = audio_output_lens.tolist() + + audio_in_video_item_idx = 0 + + def get_replacement_qwen2_audio(item_idx: int): + item_idx += audio_in_video_item_idx + + num_features = audio_output_lengths[item_idx] + if num_features == 0: + audios = mm_items.get_items("audio", AudioProcessorItems) + audio = audios.get(item_idx) + raise ValueError( + f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model" + ) + + return [audio_token_id] * num_features + + def get_replacement_qwen2_vision(item_idx: int, modality: str): + grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx] + assert isinstance(grid_thw, torch.Tensor) + merge_length = image_processor.merge_size**2 + + token_id = image_token_id if modality == "image" else video_token_id + return [token_id] * (int(grid_thw.prod()) // merge_length) + + num_videos = mm_items.get_all_counts().get("video", 0) + video_use_audio_in_video = _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs, + num_videos, + ) + thinker_config = self.info.get_hf_config() + + def get_replacement_qwen2_video(item_idx: int): + nonlocal audio_in_video_item_idx + + if not video_use_audio_in_video[item_idx]: + return get_replacement_qwen2_vision(item_idx, modality="video") + + audio_num_features = audio_output_lengths[audio_in_video_item_idx] + video_grid_thw = out_mm_data["video_grid_thw"][item_idx] + + audio_in_video_item_idx += 1 + + video_second_per_grid_t = _get_video_second_per_grid_t( + out_mm_data, + hf_processor_mm_kwargs, + item_idx, + default=1.0, + ) + + updates = self.omni_get_updates_use_audio_in_video( + thinker_config=thinker_config, + audio_len=audio_num_features, + video_grid_thw=video_grid_thw, + video_second_per_grid_t=video_second_per_grid_t, + ) + + return PromptUpdateDetails.select_token_id( + seq=updates, + embed_token_id=video_token_id, + ) + + return [ + PromptReplacement( + modality="audio", + target=audio_token, + replacement=get_replacement_qwen2_audio, + ), + PromptReplacement( + modality="image", + target=image_token, + replacement=partial(get_replacement_qwen2_vision, modality="image"), + ), + PromptReplacement( + modality="video", + target=video_token, + replacement=get_replacement_qwen2_video, + ), + ] + + def _apply_hf_processor_mm_only( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ): + mm_counts = mm_items.get_all_counts() + + if "video" in mm_counts: + video_use_audio_in_video = _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs, + mm_counts["video"], + ) + if any(video_use_audio_in_video): + assert "audio" in mm_counts + mm_counts["audio"] -= sum(video_use_audio_in_video) + + _, mm_processed_data, _ = self._apply_hf_processor_text_mm( + prompt_text=self.dummy_inputs.get_dummy_text(mm_counts), + mm_items=mm_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=tokenization_kwargs, + ) + + return mm_processed_data + + def _cached_apply_hf_processor( + self, + inputs: ProcessorInputs, + timing_ctx: TimingContext, + ) -> tuple[list[int], MultiModalProcessingInfo, bool]: + cache = self.cache + + _, passthrough_data = self._get_hf_mm_data(inputs.mm_data_items) + if cache is None or passthrough_data: + return self._apply_hf_processor(inputs, timing_ctx) + + with timing_ctx.record("get_mm_hashes"): + mm_hashes = inputs.get_mm_hashes(self.info.model_id) + + with timing_ctx.record("get_cache_missing_items"): + mm_is_cached, mm_missing_data_items = self._get_cache_missing_items( + cache=cache, + mm_data_items=inputs.mm_data_items, + mm_hashes=mm_hashes, + ) + + hf_processor_mm_kwargs = _filter_video_use_audio_in_video_for_uncached_items( + inputs.hf_processor_mm_kwargs, + mm_is_cached.get("video"), + ) + + with timing_ctx.record("apply_hf_processor"): + ( + prompt_ids, + mm_missing_processed_data, + is_update_applied, + ) = self._apply_hf_processor_main( + prompt=inputs.prompt, + mm_items=mm_missing_data_items, + hf_processor_mm_kwargs=hf_processor_mm_kwargs, + tokenization_kwargs=inputs.tokenization_kwargs, + enable_hf_prompt_update=False, + ) + + mm_missing_kwargs = MultiModalKwargsItems.from_hf_inputs( + mm_missing_processed_data, + self._get_mm_fields_config( + mm_missing_processed_data, + hf_processor_mm_kwargs, + ), + ) + + mm_missing_prompt_updates = self._get_mm_prompt_updates( + mm_missing_data_items, + hf_processor_mm_kwargs, + mm_missing_kwargs, + ) + + with timing_ctx.record("merge_mm_kwargs"): + mm_kwargs, mm_prompt_updates = self._merge_mm_kwargs( + cache, + mm_hashes=mm_hashes, + mm_is_cached=mm_is_cached, + mm_missing_kwargs=mm_missing_kwargs, + mm_missing_prompt_updates=mm_missing_prompt_updates, + ) + + mm_info = MultiModalProcessingInfo( + kwargs=mm_kwargs, + hashes=mm_hashes, + prompt_updates=mm_prompt_updates, + ) + + return prompt_ids, mm_info, is_update_applied + + def _derive_audio_from_video_placeholders( + self, + placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], + mm_prompt_updates: MultiModalPromptUpdates, + video_use_audio_in_video: Sequence[bool] | None = None, + ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: + if "video" not in placeholders: + return placeholders + + num_videos = len(placeholders["video"]) + if video_use_audio_in_video is None: + video_use_audio_in_video = [True] * num_videos + elif len(video_use_audio_in_video) != num_videos: + raise ValueError( + "use_audio_in_video must contain one boolean per video, " + f"but found {len(video_use_audio_in_video)} values for " + f"{num_videos} videos." + ) + + num_audio_in_video = sum(video_use_audio_in_video) + num_audios = len(mm_prompt_updates.get("audio", [])) + if num_audios != num_audio_in_video: + raise ValueError( + "use_audio_in_video requires equal number of audio and video " + f"items using audio, got {num_audios=}, {num_audio_in_video=}" + ) + + tokenizer = self.info.get_tokenizer() + processor = self.info.get_hf_processor() + audio_token_id = tokenizer.get_vocab()[processor.audio_token] + + result_placeholders = dict(placeholders) + audio_placeholders = [] + video_placeholders = [] + + audio_idx = 0 + for video_idx, video_placeholder in enumerate(placeholders["video"]): + audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id + + if video_use_audio_in_video[video_idx]: + audio_placeholder = PlaceholderFeaturesInfo( + modality="audio", + item_idx=audio_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=audio_is_embed, + ) + audio_placeholders.append(audio_placeholder) + audio_idx += 1 + + video_placeholder_with_mask = PlaceholderFeaturesInfo( + modality="video", + item_idx=video_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=~audio_is_embed, + ) + video_placeholders.append(video_placeholder_with_mask) + + result_placeholders["audio"] = audio_placeholders + result_placeholders["video"] = video_placeholders + return result_placeholders + def _maybe_apply_prompt_updates( self, mm_items: MultiModalDataItems, @@ -91,24 +700,8 @@ def _maybe_apply_prompt_updates( self._validate_mm_kwargs(mm_kwargs, mm_item_counts) self._validate_mm_updates(mm_prompt_updates, mm_item_counts) - # Detect use_audio_in_video from mm_kwargs - use_audio_in_video = False - if "video" in mm_kwargs: - for item in mm_kwargs["video"]: - if item and item.get("use_audio_in_video"): - use_audio_in_video_tensor = item["use_audio_in_video"].data - if use_audio_in_video_tensor.numel() > 0: - use_audio_in_video = bool(use_audio_in_video_tensor.item()) - break - # for mutilmodality cache - if any(item is None for item in mm_kwargs["video"]): - video_token_id = self.info.get_hf_config().video_token_id - audio_token_id = self.info.get_hf_config().audio_token_id - video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids) - audio_updates_num = len(mm_prompt_updates.get("audio", [])) - video_updates_num = len(mm_prompt_updates.get("video", [])) - if video_audio_item_num != video_updates_num + audio_updates_num: - use_audio_in_video = True + video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) + use_audio_in_video = any(video_use_audio_in_video) if is_update_applied: mm_placeholders = self._find_mm_placeholders( @@ -126,7 +719,11 @@ def _maybe_apply_prompt_updates( prompt_ids, filtered_updates, ) - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) + mm_placeholders = self._derive_audio_from_video_placeholders( + mm_placeholders, + mm_prompt_updates, + video_use_audio_in_video, + ) else: prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index 27eb25720bd..91eafd2a876 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -339,14 +339,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue - if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)): - # Loading kv cache quantization scales - param = params_dict[scale_name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0] - weight_loader(param, loaded_weight) - loaded_params.add(scale_name) - continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue diff --git a/vllm_omni/model_executor/models/qwen3_omni/__init__.py b/vllm_omni/model_executor/models/qwen3_omni/__init__.py index a469704f614..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/__init__.py +++ b/vllm_omni/model_executor/models/qwen3_omni/__init__.py @@ -1,3 +1,8 @@ -from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -__all__ = ["Qwen3OmniMoeForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py index 043c61393b8..9baecdd3ede 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_omni/pipeline.py @@ -7,17 +7,21 @@ Stage 2: Code2Wav — RVQ codes → audio waveform """ +from transformers import Qwen3OmniMoeConfig + from vllm_omni.config.endpoint_policy import EndpointRestriction, OmniServingCapability from vllm_omni.config.stage_config import ( PipelineConfig, StageExecutionType, StagePipelineConfig, + pipeline_cfg_resolver, ) _PROC = "vllm_omni.model_executor.stage_input_processors.qwen3_omni" QWEN3_OMNI_PIPELINE = PipelineConfig( model_type="qwen3_omni_moe", + default_deploy_config_name="qwen3_omni_moe.yaml", model_arch="Qwen3OmniMoeForConditionalGeneration", endpoint_restrictions=( EndpointRestriction( @@ -48,7 +52,6 @@ input_sources=(0,), hf_config_name="talker_config", engine_output_type="latent", - custom_process_input_func=f"{_PROC}.thinker2talker", sync_process_input_func=f"{_PROC}.thinker2talker_token_only", custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"), async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"), @@ -66,7 +69,6 @@ final_output_type="audio", hf_config_name="thinker_config", engine_output_type="audio", - custom_process_input_func=f"{_PROC}.talker2code2wav", sampling_constraints={"detokenize": True}, ), ), @@ -91,3 +93,18 @@ ), ), ) + + +@pipeline_cfg_resolver(config_type=Qwen3OmniMoeConfig) +def resolve_qwen3_omni_pipeline( + hf_config: Qwen3OmniMoeConfig, +) -> PipelineConfig: + """Select the right pipeline variant based on the HF config, since some variants, + e.g., Qwen3-Omni-30B-A3B-Captioner, are thinker only. + + By default, we load the full pipeline, as this is the common case. + """ + # If we have a config and it explicitly disabled audio input, load thinker only + if not hf_config.enable_audio_output: + return QWEN3_OMNI_THINKER_ONLY_PIPELINE + return QWEN3_OMNI_PIPELINE diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py deleted file mode 100644 index 95ae62627e2..00000000000 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py +++ /dev/null @@ -1,173 +0,0 @@ -from __future__ import annotations - -import torch -import torch.nn.functional as F -from torch import nn -from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.layer import MoERunner -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_moe import ( - Qwen3MoeDecoderLayer, - Qwen3MoeMLP, - Qwen3MoeModel, # as _BaseQwen3MoeModel, -) -from vllm.model_executor.models.qwen3_moe import ( - Qwen3MoeForCausalLM as _BaseQwen3MoeForCausalLM, -) -from vllm.model_executor.models.utils import ( - PPMissingLayer, - maybe_prefix, -) - -logger = init_logger(__name__) - - -# Individual expert MoE block using Qwen3MoeMLP instead of FusedMoE -class Qwen3OmniMoeSparseMoeBlock(nn.Module): - """Sparse MoE block using individual Qwen3MoeMLP experts instead of FusedMoE.""" - - def __init__( - self, - vllm_config: VllmConfig, - prefix: str = "", - ): - super().__init__() - - config = vllm_config.model_config.hf_text_config - quant_config = vllm_config.quant_config - - self.num_experts = config.num_experts - self.top_k = config.num_experts_per_tok - self.norm_topk_prob = config.norm_topk_prob - self.hidden_size = config.hidden_size - - # Create individual expert MLPs - self.experts = nn.ModuleList( - [ - Qwen3MoeMLP( - hidden_size=config.hidden_size, - intermediate_size=config.moe_intermediate_size, - hidden_act=config.hidden_act, - quant_config=quant_config, - reduce_results=False, - prefix=f"{prefix}.experts.{i}", - ) - for i in range(self.num_experts) - ] - ) - - # Router for expert selection - from vllm.model_executor.layers.linear import ReplicatedLinear - - self.gate = ReplicatedLinear( - config.hidden_size, config.num_experts, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate" - ) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - """Forward pass using individual experts.""" - # Handle 3D inputs (batch, seq_len, hidden_size) by reshaping to 2D - orig_shape = hidden_states.shape - if hidden_states.dim() == 3: - batch_size, seq_len, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - elif hidden_states.dim() == 2: - num_tokens, hidden_dim = hidden_states.shape - elif hidden_states.dim() == 1: - hidden_states = hidden_states.unsqueeze(0) - num_tokens, hidden_dim = hidden_states.shape - else: - raise ValueError( - f"Qwen3OmniMoeSparseMoeBlock only supports 1D, 2D, or 3D inputs, got {hidden_states.dim()}D" - ) - - is_input_1d = len(orig_shape) == 1 - hidden_states = hidden_states.view(-1, hidden_dim) - - # Get router logits and select experts (matching transformers) - router_logits, _ = self.gate(hidden_states) - selected_experts, routing_weights = self._route_tokens(router_logits) - - # Forward through individual experts - final_hidden_states = self._forward_experts(hidden_states, selected_experts, routing_weights) - - # Reshape back to original shape - if is_input_1d: - return final_hidden_states.squeeze(0) - elif len(orig_shape) == 3: - # Reshape back to 3D (batch, seq_len, hidden_dim) - return final_hidden_states.view(orig_shape) - else: - return final_hidden_states - - def _route_tokens(self, router_logits: torch.Tensor): - """Route tokens to experts using top-k selection (matching transformers).""" - routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float) - routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) - if self.norm_topk_prob: - routing_weights /= routing_weights.sum(dim=-1, keepdim=True) - routing_weights = routing_weights.to(router_logits.dtype) - return selected_experts, routing_weights - - def _forward_experts( - self, hidden_states: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor - ): - """Forward through individual experts (matching transformers implementation).""" - final_hidden_states = torch.zeros_like(hidden_states) - expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) - - expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - for expert_idx in expert_hit: - idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) - current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1]) - current_hidden_states = self.experts[expert_idx](current_state) * routing_weights[top_x, idx, None] - final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) - - return final_hidden_states - - -class Qwen3MoeForCausalLM(_BaseQwen3MoeForCausalLM): - """Thin wrapper to swap in the patched `Qwen3MoeModel`.""" - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - # Don't call super().__init__() to avoid duplicate layer registration. - nn.Module.__init__(self) - config = vllm_config.model_config.hf_text_config - quant_config = vllm_config.quant_config - self.config = config - self.quant_config = quant_config - self.model = Qwen3MoeModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) - self.lm_head = ParallelLMHead( - config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head") - ) - if self.config.tie_word_embeddings: - self.lm_head.weight = self.model.embed_tokens.weight - self.logits_processor = LogitsProcessor(config.vocab_size) - self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors - - # Set MoE hyperparameters for individual experts - self.expert_weights = [] - - self.moe_layers: list[MoERunner] = [] - example_layer = None - for layer in self.model.layers: - if isinstance(layer, PPMissingLayer): - continue - - assert isinstance(layer, Qwen3MoeDecoderLayer) - if isinstance(layer.mlp, MoERunner): - example_layer = layer.mlp - self.moe_layers.append(layer.mlp) - - if example_layer is None: - raise RuntimeError("No Qwen3OmniMoe layer found in the model.layers.") - - self.num_moe_layers = len(self.moe_layers) - self.num_expert_groups = 1 - self.num_shared_experts = 0 - self.num_logical_experts = example_layer.n_logical_experts - self.num_physical_experts = example_layer.n_physical_experts - self.num_local_physical_experts = example_layer.n_local_physical_experts - self.num_routed_experts = example_layer.n_routed_experts - self.num_redundant_experts = example_layer.n_redundant_experts diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 4a144bff027..b65cfdad018 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -46,6 +46,7 @@ Qwen3OmniMoeThinkerProcessingInfo, ) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, safe_tensor_reshape +from vllm_omni.platforms import current_omni_platform # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -105,6 +106,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.have_multimodal_outputs = True self.has_preprocess = False self.has_postprocess = False + self.use_async_omni_output = False config: Qwen3OmniMoeConfig = vllm_config.model_config.hf_config multimodal_config = vllm_config.model_config.multimodal_config @@ -129,6 +131,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.model_stage = vllm_config.model_config.model_stage if self.model_stage == "thinker": + self.use_async_omni_output = True # Initialize thinker model (multimodal processing + text generation) # Create a new vllm_config with thinker_config as the hf_config thinker_vllm_config = vllm_config.with_hf_config( @@ -152,6 +155,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): multimodal_config.skip_mm_profiling = True self.has_preprocess = True self.has_postprocess = True + # Talker only ships codec codes to code2wav; skip latent hidden D2H. + # Build Omni output asynchronously like thinker, but run lightweight + # postprocess eagerly so hidden_states.last stays on GPU before the + # next decode step. + self.use_async_omni_output = True + self.eager_omni_postprocess_before_async_output = True + self.omni_pooler_payload_include_hidden = False self.set_custom_preprocess(self.talker_preprocess) self.set_custom_postprocess(self.talker_postprocess) self.thinker = None @@ -1213,7 +1223,7 @@ def _talker_to_code_predictor( else ( talker_hidden_states.device if isinstance(talker_hidden_states, torch.Tensor) - else torch.device("cuda" if torch.cuda.is_available() else "cpu") + else current_omni_platform.get_torch_device() ) ) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 2b1a7ccf221..3a85599d616 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -314,10 +314,20 @@ def chunked_decode_streaming( # Fallback: assume all batch elements share the same sequence length. code_seq_lens = [codes.shape[-1]] * codes.shape[0] for idx, code_seq_len in enumerate(code_seq_lens): - # Remove context from output (left_context_size * total_upsample samples) - wav_chunk = batch_wav[ - idx, :, left_context_size[idx] * self.total_upsample : code_seq_len * self.total_upsample - ] + # The eager decoder (self(codes)) emits a fixed `tail` fewer samples than + # code_seq_len*total_upsample: the causal conv stack trims the sequence end + # because it lacks right context. The original slice assumed full length, so + # under async_chunk every streaming chunk silently dropped its last `tail` + # samples -> a ~23ms gap at every chunk boundary + cumulative ~1.2% time + # compression. Shift the start back by `tail` so this chunk refills the gap the + # previous chunk left (those frames are re-decoded here with proper context). + # `tail` is measured per call, so this is a no-op only when the decoder returns + # the nominal full length (tail==0, e.g. a C==0 decoder like Qwen3-TTS); for + # Qwen3-Omni Code2Wav both the eager forward AND the CUDA-graph wrapper + # (_trim_replay_output -> actual*upsample - 555) are short, so it applies to both. + tail = max(0, int(code_seq_len * self.total_upsample) - batch_wav.shape[-1]) + start = max(0, left_context_size[idx] * self.total_upsample - tail) + wav_chunk = batch_wav[idx, :, start : code_seq_len * self.total_upsample] wavs.append(wav_chunk) return wavs diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index 44aaaf4d376..65ec102a12e 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -74,7 +74,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniAudioFeatureInputs, Qwen2_5OmniThinkerDummyInputsBuilder, - Qwen2_5OmniThinkerMultiModalProcessor, check_interleaved_audio_video, merge_interleaved_embeddings, ) @@ -120,6 +119,11 @@ from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, + Qwen2_5OmniThinkerMultiModalDataParser, + Qwen2_5OmniThinkerMultiModalProcessor, + _get_request_video_use_audio_in_video, + _get_video_second_per_grid_t, + _presampled_videos_hf_kwargs, ) from vllm_omni.quantization.component_config import ( PRE_QUANTIZED_METHODS, @@ -133,6 +137,14 @@ logger = init_logger(__name__) +_THINKER_ARCHITECTURE = "Qwen3OmniMoeThinkerForConditionalGeneration" + + +def _ensure_thinker_architecture(config: PretrainedConfig) -> None: + """Give the nested Thinker config the architecture required by MoE LoRA.""" + if not config.architectures: + config.architectures = [_THINKER_ARCHITECTURE] + class Qwen3Omni_VisionTransformer(_Qwen3Omni_VisionTransformer): """Subclass that fixes Qwen2_5_VisionAttention.forward() compatibility. @@ -577,7 +589,20 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.quant_config = quant_config self.model = Qwen3MoeLLMModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) - self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, quant_config=quant_config) + # Pass the prefix so ModelOpt's prefix-based exclude-list can match this + # head. vLLM 0.23.0 made ParallelLMHead quantizable by NVFP4 (0.22's + # ModelOpt.get_quant_method only matched LinearBase, so the head was + # never a quant candidate and the prefix was irrelevant). On 0.23.0 the + # head is kept BF16 only via the exclude-list (the W4A4 ckpt ignores + # "*.lm_head"); with prefix="" that match fails, the head is wrongly + # NVFP4-quantized (FP4-packed -> width 1024), and the BF16 [vocab, 2048] + # weight fails to load. Mirrors vLLM's stock Qwen3MoeForCausalLM. + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) if self.config.tie_word_embeddings: self.lm_head.weight = self.model.embed_tokens.weight self.logits_processor = LogitsProcessor(config.vocab_size) @@ -608,6 +633,16 @@ def get_feature_extractor(self, **kwargs: object): assert isinstance(feature_extractor, WhisperFeatureExtractor) return feature_extractor + def get_data_parser(self): + feature_extractor = self.get_feature_extractor() + + return Qwen2_5OmniThinkerMultiModalDataParser( + spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size, + target_sr=feature_extractor.sampling_rate, + target_channels=1, + expected_hidden_size=self._get_expected_hidden_size(), + ) + def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"audio": None, "image": None, "video": None} @@ -658,6 +693,7 @@ def _call_hf_processor( tok_kwargs: Mapping[str, object], ) -> BatchFeature: mm_data = dict(mm_data) + mm_kwargs = dict(_presampled_videos_hf_kwargs(mm_data, mm_kwargs)) audios = mm_data.pop("audios", []) def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray: @@ -756,22 +792,8 @@ def _maybe_apply_prompt_updates( mm_item_counts = mm_items.get_all_counts() self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - use_audio_in_video = False - if "video" in mm_kwargs: - for item in mm_kwargs["video"]: - if item and item["use_audio_in_video"].data: - use_audio_in_video = True - else: - use_audio_in_video = False - # for mutilmodality cache - if any(item is None for item in mm_kwargs["video"]): - video_token_id = self.info.get_hf_config().video_token_id - audio_token_id = self.info.get_hf_config().audio_token_id - video_audio_item_num = sum(id in (video_token_id, audio_token_id) for id in prompt_ids) - audio_updates_num = len(mm_prompt_updates.get("audio", [])) - video_updates_num = len(mm_prompt_updates.get("video", [])) - if video_audio_item_num != video_updates_num + audio_updates_num: - use_audio_in_video = True + video_use_audio_in_video = self._get_video_use_audio_in_video(mm_kwargs, mm_prompt_updates) + use_audio_in_video = any(video_use_audio_in_video) # normal case with `use_audio_in_video=False` if is_update_applied: @@ -790,7 +812,11 @@ def _maybe_apply_prompt_updates( prompt_ids, filtered_updates, ) - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) + mm_placeholders = self._derive_audio_from_video_placeholders( + mm_placeholders, + mm_prompt_updates, + video_use_audio_in_video, + ) else: prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, @@ -904,21 +930,30 @@ def get_replacement_qwen2_vision(item_idx: int, modality: str): token_id = image_token_id if modality == "image" else video_token_id return [token_id] * (int(grid_thw.prod()) // merge_length) - use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) + num_videos = len(out_mm_data.get("video_grid_thw", [])) + video_use_audio_in_video = _get_request_video_use_audio_in_video( + hf_processor_mm_kwargs, + num_videos, + ) thinker_config = self.info.get_hf_config() - def get_replacement_qwen2_use_audio_in_video(item_idx: int): + def get_replacement_qwen2_video(item_idx: int): nonlocal audio_in_video_item_idx + + if not video_use_audio_in_video[item_idx]: + return get_replacement_qwen2_vision(item_idx, modality="video") + audio_num_features = audio_output_lengths[audio_in_video_item_idx] video_grid_thw = out_mm_data["video_grid_thw"][item_idx] audio_in_video_item_idx += 1 - second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) - if second_per_grid_ts: - video_second_per_grid_t = second_per_grid_ts[item_idx] - else: - video_second_per_grid_t = 2.0 + video_second_per_grid_t = _get_video_second_per_grid_t( + out_mm_data, + hf_processor_mm_kwargs, + item_idx, + default=2.0, + ) placeholder = self.get_updates_use_audio_in_video( thinker_config=thinker_config, @@ -928,12 +963,6 @@ def get_replacement_qwen2_use_audio_in_video(item_idx: int): ) return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id) - video_replacement_fn = ( - get_replacement_qwen2_use_audio_in_video - if use_audio_in_video - else partial(get_replacement_qwen2_vision, modality="video") - ) - return [ PromptReplacement( modality="audio", @@ -948,7 +977,7 @@ def get_replacement_qwen2_use_audio_in_video(item_idx: int): PromptReplacement( modality="video", target=video_token, - replacement=video_replacement_fn, + replacement=get_replacement_qwen2_video, ), ] @@ -956,6 +985,7 @@ def _derive_audio_from_video_placeholders( self, placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], mm_prompt_updates: MultiModalPromptUpdates, + video_use_audio_in_video: Sequence[bool] | None = None, ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: """ Helper to derive audio placeholders from video placeholders when @@ -965,34 +995,62 @@ def _derive_audio_from_video_placeholders( return placeholders num_videos = len(placeholders["video"]) + if video_use_audio_in_video is None: + video_use_audio_in_video = [True] * num_videos + elif len(video_use_audio_in_video) != num_videos: + raise ValueError( + "use_audio_in_video must contain one boolean per video, " + f"but found {len(video_use_audio_in_video)} values for " + f"{num_videos} videos." + ) + + num_audio_in_video = sum(video_use_audio_in_video) num_audios = len(mm_prompt_updates.get("audio", [])) - if num_audios != num_videos: + if num_audios != num_audio_in_video: raise ValueError( - f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}" + "use_audio_in_video requires equal number of audio and video " + f"items using audio, got {num_audios=}, {num_audio_in_video=}" ) tokenizer = self.info.get_tokenizer() processor = self.info.get_hf_processor() audio_token_id = tokenizer.get_vocab()[processor.audio_token] + video_token_id = tokenizer.get_vocab()[processor.video_token] result_placeholders = dict(placeholders) audio_placeholders = [] + video_placeholders = [] - # Each video is paired with one audio + # Each video using audio is paired with one audio + audio_idx = 0 for video_idx, video_placeholder in enumerate(placeholders["video"]): # Create is_embed mask selecting only audio tokens - audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id + placeholder_tokens = torch.tensor(video_placeholder.tokens) + audio_is_embed = placeholder_tokens == audio_token_id + video_is_embed = placeholder_tokens == video_token_id + + if video_use_audio_in_video[video_idx]: + audio_placeholder = PlaceholderFeaturesInfo( + modality="audio", + item_idx=audio_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=audio_is_embed, + ) + audio_placeholders.append(audio_placeholder) + audio_idx += 1 - audio_placeholder = PlaceholderFeaturesInfo( - modality="audio", + video_placeholder_with_mask = PlaceholderFeaturesInfo( + modality="video", item_idx=video_idx, start_idx=video_placeholder.start_idx, tokens=video_placeholder.tokens, - is_embed=audio_is_embed, + is_embed=video_is_embed, ) - audio_placeholders.append(audio_placeholder) + video_placeholders.append(video_placeholder_with_mask) result_placeholders["audio"] = audio_placeholders + result_placeholders["video"] = video_placeholders return result_placeholders def _get_raw_input_ids( @@ -1071,6 +1129,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( Qwen3OmniMoeConditionalGenerationMixin, SupportsTranscription, ): + # PEFT stores the expert LoRA matrices with an explicit expert dimension. + is_3d_moe_weight: bool = True + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "thinker.lm_head.": "language_model.lm_head.", @@ -1114,6 +1175,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vllm_config = vllm_config # needed for torch compile forward context hf_config = vllm_config.model_config.hf_config thinker_config: Qwen3OmniMoeThinkerConfig = getattr(hf_config, "thinker_config", None) or hf_config + _ensure_thinker_architecture(thinker_config) quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config self.config = thinker_config @@ -1494,10 +1556,12 @@ def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator Iterate over multimodal features sorted by position offset. Yields: (offset, modality, feature_data) where feature_data contains: - - image: {"grid_t", "grid_h", "grid_w", "t_factor"} + - image: {"grid_t", "grid_h", "grid_w", "t_factor", + "placeholder_len"} - video: {"grid_t", "grid_h", "grid_w", "t_factor", - "use_audio_in_video", "audio_feature_length"} - - audio: {"audio_feature_length"} + "use_audio_in_video", "audio_feature_length", + "placeholder_len"} + - audio: {"audio_feature_length", "placeholder_len"} """ config = self.config spatial_merge_size = config.vision_config.spatial_merge_size @@ -1520,6 +1584,7 @@ def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator "grid_h": h // spatial_merge_size, "grid_w": w // spatial_merge_size, "t_factor": position_id_per_seconds, + "placeholder_len": mm_feature.mm_position.length, }, ) elif modality == "video": @@ -1541,12 +1606,20 @@ def iter_mm_features(self, mm_features: list[MultiModalFeatureSpec]) -> Iterator "t_factor": second_per_grid_ts * position_id_per_seconds, "use_audio_in_video": use_audio_in_video, "audio_feature_length": audio_for_video.get(offset), + "placeholder_len": mm_feature.mm_position.length, }, ) elif modality == "audio": if offset not in paired_audio_offsets: audio_len = mm_feature.data["audio_feature_lengths"].data.item() - yield offset, "audio", {"audio_feature_length": audio_len} + yield ( + offset, + "audio", + { + "audio_feature_length": audio_len, + "placeholder_len": mm_feature.mm_position.length, + }, + ) def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) -> tuple[np.ndarray, int]: """ @@ -1559,9 +1632,13 @@ def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) - grid_h = data["grid_h"] grid_w = data["grid_w"] t_factor = data["t_factor"] - audio_feature_length = data["audio_feature_length"] - - audio_len = self._compute_audio_token_count(audio_feature_length) + num_video = grid_t * grid_h * grid_w + placeholder_len = data.get("placeholder_len") + if placeholder_len is None: + audio_feature_length = data["audio_feature_length"] + audio_len = self._compute_audio_token_count(audio_feature_length) + else: + audio_len = max(0, int(placeholder_len) - num_video - 2) h_index = np.tile(np.arange(grid_h).reshape(1, -1, 1), (grid_t, 1, grid_w)).flatten() w_index = np.tile(np.arange(grid_w).reshape(1, 1, -1), (grid_t, grid_h, 1)).flatten() @@ -1577,8 +1654,6 @@ def _compute_interleaved_positions(self, start_idx: int, data: dict[str, Any]) - pos_ids_list: list[np.ndarray] = [] video_idx, audio_idx = 0, 0 - num_video = grid_t * grid_h * grid_w - while video_idx < num_video and audio_idx < audio_len: if video_t_values[video_idx] <= audio_t_values[audio_idx]: pos_ids_list.append(video_pos[:, video_idx : video_idx + 1]) @@ -1668,9 +1743,18 @@ def get_mrope_input_positions( seq_len = len(input_tokens) llm_pos_ids_list: list[np.ndarray] = [] + feature_summaries: list[dict[str, Any]] = [] st = 0 for offset, modality, data in self.iter_mm_features(mm_features): + feature_summaries.append( + { + "offset": offset, + "modality": modality, + "placeholder_len": data.get("placeholder_len"), + "use_audio_in_video": data.get("use_audio_in_video"), + } + ) text_len = offset - st st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 @@ -1678,19 +1762,11 @@ def get_mrope_input_positions( llm_pos_ids_list.append(np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx) st_idx += text_len - bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(bos_pos) - st_idx += 1 - if modality == "audio": audio_tokens = self._compute_audio_token_count(data["audio_feature_length"]) audio_pos = np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx llm_pos_ids_list.append(audio_pos) - st_idx = int(audio_pos.max()) + 1 - - eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(eos_pos) - st = offset + 1 + audio_tokens + 1 + st = offset + int(data.get("placeholder_len", audio_tokens)) elif modality == "image": grid_t = data["grid_t"] @@ -1704,11 +1780,7 @@ def get_mrope_input_positions( llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx) image_len = grid_t * grid_h * grid_w - st_idx = int(llm_pos_ids_list[-1].max()) + 1 - - eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(eos_pos) - st = offset + 1 + image_len + 1 + st = offset + int(data.get("placeholder_len", image_len)) elif modality == "video": grid_t = data["grid_t"] @@ -1723,26 +1795,20 @@ def get_mrope_input_positions( llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx) video_len = grid_t * grid_h * grid_w - st_idx = int(llm_pos_ids_list[-1].max()) + 1 - - eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) - llm_pos_ids_list.append(eos_pos) - st = offset + 1 + video_len + 1 + st = offset + int(data.get("placeholder_len", video_len)) else: - audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1)) + audio_bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(audio_bos_pos) + st_idx += 1 - pos_ids, _ = self._compute_interleaved_positions(st_idx, data) + pos_ids, token_count = self._compute_interleaved_positions(st_idx, data) llm_pos_ids_list.append(pos_ids) st_idx = int(pos_ids.max()) + 1 eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1)) llm_pos_ids_list.append(eos_pos) - llm_pos_ids_list.append(eos_pos) - video_len = grid_t * grid_h * grid_w - audio_len = self._compute_audio_token_count(data["audio_feature_length"]) - st = offset + 2 + video_len + audio_len + 2 + st = offset + int(data.get("placeholder_len", token_count + 2)) if st < seq_len: st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0 @@ -1751,7 +1817,11 @@ def get_mrope_input_positions( llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1) if llm_positions.shape[1] != seq_len: - raise RuntimeError("Position ids length mismatch with input ids length") + raise RuntimeError( + "Position ids length mismatch with input ids length: " + f"positions={llm_positions.shape[1]}, input_ids={seq_len}, " + f"features={feature_summaries}" + ) mrope_position_delta = int(llm_positions.max()) + 1 - seq_len return torch.from_numpy(llm_positions), mrope_position_delta diff --git a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py index 8bb921a1eae..28ca4f36289 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py +++ b/vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py @@ -27,7 +27,7 @@ class Qwen3TTSSpeakerEncoderConfig(PretrainedConfig): Args: mel_dim (`int`, *optional*, defaults to 128): The dimension of the input mel-spectrogram. - enc_dim (`int`, *optional*, defaults to 192): + enc_dim (`int`, *optional*, defaults to 1024): The dimension of the final speaker embedding. enc_channels (`list[int]`, *optional*, defaults to `[512, 512, 512, 512, 1536]`): A list of output channels for each TDNN/SERes2Net layer in the encoder. @@ -80,11 +80,11 @@ class Qwen3TTSTalkerCodePredictorConfig(PretrainedConfig): Args: - vocab_size (`int`, *optional*, defaults to 151936): + vocab_size (`int`, *optional*, defaults to 2048): Vocabulary size of the Qwen3TTSTalkerCodePredictor model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`Qwen3TTSTalkerCodePredictorModel`] - hidden_size (`int`, *optional*, defaults to 4096): + hidden_size (`int`, *optional*, defaults to 1024): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 22016): Dimension of the MLP representations. @@ -271,11 +271,11 @@ class Qwen3TTSTalkerConfig(PretrainedConfig): Args: - vocab_size (`int`, *optional*, defaults to 151936): + vocab_size (`int`, *optional*, defaults to 3072): Vocabulary size of the Qwen3TTSTalker model. Defines the number of different tokens that can be represented by the `inputs_ids` passed when calling [`Qwen3TTSTalkerModel`] - hidden_size (`int`, *optional*, defaults to 2048): + hidden_size (`int`, *optional*, defaults to 1024): Dimension of the hidden representations. intermediate_size (`int`, *optional*, defaults to 6144): Dimension of the MLP representations. @@ -436,7 +436,6 @@ def __init__( self.rope_scaling["rope_type"] = self.rope_scaling["type"] if code_predictor_config is None: - code_predictor_config = {} self.code_predictor_config = Qwen3TTSTalkerCodePredictorConfig() logger.info("code_predictor_config is None. Initializing code_predictor model with default values") elif isinstance(code_predictor_config, Qwen3TTSTalkerCodePredictorConfig): diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py index 604946a35c0..abeaf243def 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py +++ b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py @@ -44,7 +44,13 @@ def _batched_chunked_decode( left_context_size: int = 25, max_batch_size: int = 0, ) -> torch.Tensor: - """Decode a padded batch by grouping same-round chunks across requests.""" + """Decode a padded batch by grouping same-round chunks across requests. + + Assumes ``decode_fn`` returns exactly ``chunk_len * total_upsample`` samples per + chunk (causal-conv trim constant C = 0, e.g. Qwen3-TTS) and raises otherwise. It + is not used by short-output decoders such as Qwen3-Omni Code2Wav (C = 555), which + go through the per-request ``chunked_decode`` / ``chunked_decode_streaming`` paths. + """ if codes.dim() < 3: raise ValueError(f"Expected codes with shape [B, Q, F], got {tuple(codes.shape)}") if chunk_size <= 0: @@ -511,8 +517,7 @@ def _decode(self, codes: torch.Tensor, *, clone_graph_output: bool) -> torch.Ten static_input.zero_() static_input[:, :, :actual_size] = codes self._compiled_graphs[compile_key].replay() - actual_out_len = actual_size * self.decoder.total_upsample - output = self._compiled_static_outputs[compile_key][..., :actual_out_len] + output = self._trim_replay_output(self._compiled_static_outputs[compile_key], actual_size, compiled_size) if clone_graph_output: return output.clone() return output @@ -542,12 +547,26 @@ def _decode(self, codes: torch.Tensor, *, clone_graph_output: bool) -> torch.Ten static_input[:, :, :actual_size] = codes self.graphs[graph_key].replay() - actual_out_len = actual_size * self.decoder.total_upsample - output = self.static_outputs[graph_key][..., :actual_out_len] + output = self._trim_replay_output(self.static_outputs[graph_key], actual_size, padded_size) if clone_graph_output: return output.clone() return output + def _trim_replay_output(self, static_output: torch.Tensor, actual_size: int, padded_size: int) -> torch.Tensor: + """Trim a graph/compiled replay output to the eager-equivalent length. + + The captured ``static_output`` already reflects the decoder's TRUE output + length for ``padded_size`` frames, which for causal decoders (e.g. + Qwen3-Omni Code2Wav) is shorter than ``padded_size * total_upsample`` by a + fixed amount. Each zero-padded frame beyond ``actual_size`` contributes + ``total_upsample`` trailing samples that are stale buffer content, so trim + relative to the captured length instead of the nominal length. This is a + no-op for decoders whose output equals the nominal length (e.g. Qwen3-TTS). + """ + drop = (padded_size - actual_size) * self.decoder.total_upsample + actual_out_len = max(0, static_output.shape[-1] - drop) + return static_output[..., :actual_out_len] + def decode(self, codes: torch.Tensor) -> torch.Tensor: return self._decode(codes, clone_graph_output=True) diff --git a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py index 7f50931ddf7..8a68e63cd22 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/pipeline.py +++ b/vllm_omni/model_executor/models/qwen3_tts/pipeline.py @@ -15,6 +15,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig( model_type="qwen3_tts", + default_deploy_config_name="qwen3_tts.yaml", # Pipeline-level default; the code2wav stage overrides per-stage below. model_arch="Qwen3TTSTalkerForConditionalGeneration", stages=( diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index b0294b72583..0d8f032949e 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -1320,9 +1320,7 @@ def _get_ref_audio() -> tuple[np.ndarray, int]: # Keep it at least 1D; embedding on a 0-d tensor can return 1D. spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long) spk_embed = codec_embed(spk_tensor) - if spk_embed.ndim == 1: - spk_embed = spk_embed.view(1, 1, -1) - elif spk_embed.ndim == 2: + if spk_embed.ndim in (1, 2): spk_embed = spk_embed.view(1, 1, -1) speaker_embed = spk_embed codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1) @@ -1523,9 +1521,7 @@ def _first(x: object, default: object) -> object: ref_code_len: int | None = None if isinstance(ref_code, list): - if ref_code and isinstance(ref_code[0], list): - ref_code_len = len(ref_code) - elif ref_code: + if ref_code: ref_code_len = len(ref_code) elif hasattr(ref_code, "shape"): try: diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 8d2f0686ae0..9d00cc2bae9 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -13,6 +13,7 @@ CodePredictorWrapper, CodePredictorWrapperConfig, ) +from vllm_omni.platforms import current_omni_platform from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig @@ -21,7 +22,7 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper): - """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection).""" + """Qwen3-TTS code predictor (per-call sampling, projection).""" def __init__( self, @@ -35,7 +36,7 @@ def __init__( vllm_config=vllm_config, cp_config=config, wrapper_config=CodePredictorWrapperConfig( - use_cuda_graphs=True, + use_cuda_graphs=current_omni_platform.is_npu(), use_parallel_embedding=False, use_projection=(config.hidden_size != talker_config.hidden_size), return_proj_buf=False, diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py index 64eb3d25320..3daeb4cf285 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py @@ -1,7 +1,7 @@ from __future__ import annotations import copy -from collections.abc import Iterable, Mapping +from collections.abc import Iterable, Mapping, Sequence from typing import Any import numpy as np @@ -30,7 +30,7 @@ ) from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig -from .prompt_embeds_builder import Qwen3TTSPromptEmbedsBuilder +from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder @@ -40,6 +40,19 @@ _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64 +def _has_tts_text_conditioning(info_dict: dict[str, Any], hidden_states: Any | None = None) -> bool: + text_list = info_dict.get("text") + if isinstance(text_list, list) and bool(text_list) and bool(text_list[0]): + return True + if PRECOMPUTED_TEXT_IDS_KEY in info_dict: + return True + if isinstance(hidden_states, dict): + tail = hidden_states.get("trailing_text") + if isinstance(tail, torch.Tensor): + return True + return False + + # --------------------------------------------------------------------------- # Components ported from the HuggingFace Qwen3-TTS reference implementation. # Only the classes actually needed by the vLLM AR Talker are kept here. @@ -286,8 +299,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) self._codec_eos_token_id = int(getattr(self.talker_config, "codec_eos_token_id", -1)) - self._eos_logit_bias: float = 0.0 - self.have_multimodal_outputs = True self.has_preprocess = True self.has_postprocess = True @@ -298,6 +309,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # tensor read; postprocess receives the tail-only slice instead, which # avoids ~18 ms merge + ~6 ms write per step (Sy0307 profile, #3665). self.requires_full_prefix_cached_hidden_states = False + # Stage 1 (code2wav) consumes runtime audio codes, not hidden states, + # so the inter-stage pooler payload never needs a CPU hidden-states + # view. Opting out lets the runner skip the per-step blocking + # ``hidden_states[:n].to("cpu")`` (measured at 24% of stage-0 on-CPU + # samples at c=64 on H20). + self.omni_pooler_payload_include_hidden = False # ``codes.audio`` is only needed for future prefix-hit reconstruction # after a request has produced codec rows. Keep per-step rows on GPU and # materialize the CPU OmniTensorPrefixCache entry once at completion. @@ -307,7 +324,33 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # OmniGPUModelRunner will store talker_mtp output under this key in # per-request additional_information. self.talker_mtp_output_key = ("codes", "audio") - + self.talker_mtp_graph_safe = True + # talker_mtp samples with per-row generators, so explicitly-seeded + # requests stay batched instead of one scalar forward per row (#4883). + # This is only valid while talker_mtp receives the *unpadded* active + # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp + # (padded batches + a single captured RNG stream) whenever full + # cudagraphs are enabled, so disable per-row generators in that case. + # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is + # not reproducible (matching Qwen3-Omni, which has no per-row seeding); + # eager mode keeps the per-row generator fast-path. + # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an + # upstream (runner) concern, and the runner already re-derives the same + # wrap decision in ``_init_talker_mtp``. Once all TTS models declare + # ``talker_mtp_graph_safe`` we should move this gating into the model + # runner (models only expose the static capability flag, and users can + # then configure ``cudagraph_mode`` freely). This also means mtp seeding + # is silently dropped when the decode batch > 1 under full cudagraphs; + # the follow-up should refactor the per-row generator in the runner to + # be cudagraph-safe so seeded batches stay reproducible. + cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None) + talker_mtp_graph_wrapped = ( + self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs() + ) + self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped + self.use_async_omni_output = True + self.eager_omni_postprocess_before_async_output = True + self.omni_pooler_payload_include_hidden = False self.model = Qwen3Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) if get_pp_group().is_last_rank: @@ -517,11 +560,6 @@ def compute_logits( # Mask out invalid codec ids using the pre-built constant buffer. logits = logits.masked_fill(self._codec_disallowed_mask, float("-inf")) - if self._eos_logit_bias != 0.0: - eos_id = self._codec_eos_token_id - if 0 <= eos_id < logits.shape[-1]: - logits[:, eos_id] = logits[:, eos_id] + self._eos_logit_bias - return logits # -------------------- Omni multimodal output plumbing -------------------- @@ -586,7 +624,6 @@ def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: A audio_codes = torch.cat(audio_codes_list, dim=0) span_len = int(audio_codes.shape[0]) - hidden = hidden[:span_len] mm: OmniPayload = {"codes": {"audio": audio_codes}} if ref_code_len_list: mm.setdefault("meta", {})["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len] @@ -632,9 +669,8 @@ def preprocess( except Exception: is_prefill = span_len > 1 - text_list = info_dict.get("text") - if not isinstance(text_list, list) or not text_list or not text_list[0]: - raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.") + if not _has_tts_text_conditioning(info_dict, hs): + raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.") task_type = (info_dict.get("task_type") or ["CustomVoice"])[0] codec_streaming_val = meta.get("codec_streaming") @@ -657,7 +693,6 @@ def preprocess( if is_prefill: # Prefill (prompt embeddings) prompt_embeds_cpu = embed.get("prefill") - # First prefill round: prompt_embeds_cpu is not yet populated. # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round. is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2 @@ -714,6 +749,14 @@ def preprocess( info_update.setdefault("codes", {})["audio"] = zeros return input_ids_out, prompt_embeds, info_update + if span_len > 1: + inputs_embeds_out = ( + self.embed_input_ids(input_ids.reshape(-1, 1).to(torch.long)) + .to(device=input_ids.device, dtype=dtype) + .reshape(span_len, -1) + ) + return input_ids, inputs_embeds_out, {"meta": {"codec_streaming": codec_streaming}} + # Decode: span_len == 1 # Pop one text-step vector from tailing_text_hidden queue. # ``tts_pad_embed`` was materialized above from :attr:`_tts_pad_embed` @@ -815,9 +858,8 @@ def preprocess_decode_batch( hs = payload.get("hidden_states", {}) meta = payload.get("meta", {}) - text_list = info_dict.get("text") - if not isinstance(text_list, list) or not text_list or not text_list[0]: - raise ValueError("Missing additional_information.text for Qwen3-TTS AR talker.") + if not _has_tts_text_conditioning(info_dict, hs): + raise ValueError("Missing Qwen3-TTS text conditioning: provide `text` or precomputed text token ids.") task_type = (info_dict.get("task_type") or ["CustomVoice"])[0] codec_streaming_val = meta.get("codec_streaming") @@ -1046,6 +1088,7 @@ def talker_mtp( top_k: int | None = None, top_p: float | None = None, generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor]: """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1). @@ -1085,6 +1128,7 @@ def talker_mtp( top_k=top_k, top_p=top_p, generator=generator, + generators=generators, ) # [B, Q] # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes. diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 6eb78eeae0d..447823abb11 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -912,11 +912,6 @@ def enable_cudagraph( self._cudagraph_wrapper.compile_shapes, ) - def disable_cudagraph(self): - self._cudagraph_enabled = False - self._cudagraph_wrapper = None - logger.info("CUDA Graph disabled for decoder") - def forward(self, codes): if codes.shape[1] != self.config.num_quantizers: raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}") diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py index 5609abb394c..4014014533d 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/core_vq.py @@ -33,7 +33,6 @@ import random import typing as tp -from math import ceil import numpy as np import torch @@ -45,10 +44,6 @@ logger = init_logger(__name__) -def round_up_multiple(num, mult): - return ceil(num / mult) * mult - - def default(val: tp.Any, d: tp.Any) -> tp.Any: return val if val is not None else d diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 9906fa4e0b4..6749d5cfc63 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -294,6 +294,12 @@ "modeling_moss_tts_talker", "MossTTSRealtimeTalkerForGeneration", ), + # MossTTSLocalModel: MOSS-TTS-Local-Transformer-v1.5 + "MossTTSLocalModel": ( + "moss_tts", + "modeling_moss_tts_talker", + "MossTTSLocalTalkerForGeneration", + ), # Stage-1 codec decoder (shared by all 5 variants) "MossTTSCodecDecoder": ( "moss_tts", @@ -305,6 +311,17 @@ "dynin_omni", "DyninOmniForConditionalGeneration", ), + ## IndexTTS2 + "IndexTTS2TalkerForConditionalGeneration": ( + "indextts2", + "indextts2_talker", + "IndexTTS2TalkerForConditionalGeneration", + ), + "IndexTTS2S2MelDecoder": ( + "indextts2", + "indextts2_s2mel_decoder", + "IndexTTS2S2MelDecoder", + ), ## Ming-flash-omni-2.0 "MingFlashOmniForConditionalGeneration": ( "ming_flash_omni", @@ -343,6 +360,11 @@ "minicpmo_4_5_omni_tts", "MiniCPMO45OmniTTSForConditionalGeneration", ), + "AuraQwen3VLForConditionalGeneration": ( + "aura_omni", + "qwen3_vl", + "AuraQwen3VLForConditionalGeneration", + ), } diff --git a/vllm_omni/model_executor/models/soulx_singer/__init__.py b/vllm_omni/model_executor/models/soulx_singer/__init__.py new file mode 100644 index 00000000000..1ab85a8d3ce --- /dev/null +++ b/vllm_omni/model_executor/models/soulx_singer/__init__.py @@ -0,0 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +# Keep this package lightweight. The pipeline registry imports pipeline.py +# directly, while diffusion model classes are loaded lazily from their registry. diff --git a/vllm_omni/model_executor/models/soulx_singer/pipeline.py b/vllm_omni/model_executor/models/soulx_singer/pipeline.py new file mode 100644 index 00000000000..4d24205288c --- /dev/null +++ b/vllm_omni/model_executor/models/soulx_singer/pipeline.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""SoulX-Singer pipeline topologies.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +SOULXSINGER_SVS_PIPELINE = PipelineConfig( + model_type="soulxsinger_svs", + default_deploy_config_name="soulxsinger_svs.yaml", + model_arch="SoulXSingerPipeline", + hf_architectures=("SoulXSingerPipeline",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="audio", + model_arch="SoulXSingerPipeline", + ), + ), +) + +SOULXSINGER_SVC_PIPELINE = PipelineConfig( + model_type="soulxsinger_svc", + default_deploy_config_name="soulxsinger_svc.yaml", + model_arch="SoulXSingerSVCPipeline", + hf_architectures=("SoulXSingerSVCPipeline",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="audio", + model_arch="SoulXSingerSVCPipeline", + ), + ), +) diff --git a/vllm_omni/model_executor/models/step_audio2/__init__.py b/vllm_omni/model_executor/models/step_audio2/__init__.py index 0b6f07e2807..e69de29bb2d 100644 --- a/vllm_omni/model_executor/models/step_audio2/__init__.py +++ b/vllm_omni/model_executor/models/step_audio2/__init__.py @@ -1,3 +0,0 @@ -from .step_audio2 import StepAudio2ForConditionalGeneration - -__all__ = ["StepAudio2ForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/step_audio2/pipeline.py b/vllm_omni/model_executor/models/step_audio2/pipeline.py new file mode 100644 index 00000000000..7ba6d9ee15b --- /dev/null +++ b/vllm_omni/model_executor/models/step_audio2/pipeline.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Step-Audio2 pipeline topologies (frozen).""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.step_audio2" +_HF_ARCHITECTURES = ( + "StepAudio2ForCausalLM", + "StepAudio2ForConditionalGeneration", +) + + +STEP_AUDIO2_PIPELINE = PipelineConfig( + model_type="step_audio_2", + model_arch="StepAudio2ForConditionalGeneration", + hf_architectures=_HF_ARCHITECTURES, + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="step_audio2_thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="latent", + model_arch="StepAudio2ThinkerForConditionalGeneration", + async_chunk_process_next_stage_input_func=f"{_PROC}.thinker2token2wav_async_chunk", + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="token2wav", + execution_type=StageExecutionType.LLM_GENERATION, + input_sources=(0,), + final_output=True, + final_output_type="audio", + engine_output_type="audio", + model_arch="StepAudio2Token2WavModel", + sync_process_input_func=f"{_PROC}.thinker2token2wav", + sampling_constraints={"detokenize": False}, + ), + ), +) + + +STEP_AUDIO2_ASR_PIPELINE = PipelineConfig( + model_type="step_audio_2_asr", + model_arch="StepAudio2ForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="step_audio2_thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + engine_output_type="text", + model_arch="StepAudio2ThinkerForConditionalGeneration", + sampling_constraints={"detokenize": True}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2.py b/vllm_omni/model_executor/models/step_audio2/step_audio2.py index 82e5609ff40..0a94ab9defb 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2.py @@ -26,17 +26,17 @@ class StepAudio2ForConditionalGeneration(nn.Module, SupportsMultiModal, Supports Step-Audio2 Main Controller Manages two-stage inference pipeline: - - Stage 1 (Thinker): Audio understanding and token generation - - Stage 2 (Token2Wav): Audio token to waveform synthesis + - Stage 0 (Thinker): Audio understanding and token generation + - Stage 1 (Token2Wav): Audio token to waveform synthesis Usage: - # Stage 1: Thinker + # Stage 0: Thinker model = StepAudio2ForConditionalGeneration( vllm_config=config, model_stage="thinker" ) - # Stage 2: Token2Wav + # Stage 1: Token2Wav model = StepAudio2ForConditionalGeneration( vllm_config=config, model_stage="token2wav" diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py index 802cd63bef8..2f2918a1be3 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_constants.py @@ -42,39 +42,6 @@ def audio_end(self) -> int: return self.audio_start + self.audio_vocab_size - 1 -@dataclass(frozen=True) -class StepAudio2EncoderConfig: - """ - Step-Audio2 audio encoder architecture configuration. - - These values define the encoder structure and should match your model weights. - """ - - # Mel spectrogram parameters - n_mels: int = 128 - """Number of mel frequency bins""" - - # Encoder architecture - n_audio_ctx: int = 1500 - """Audio context length (max sequence length)""" - - n_audio_state: int = 512 - """Audio encoder hidden state dimension""" - - n_audio_head: int = 8 - """Number of attention heads in encoder""" - - n_audio_layer: int = 6 - """Number of encoder layers""" - - # Adapter parameters - kernel_size: int = 3 - """Convolution kernel size for adapter""" - - adapter_stride: int = 2 - """Stride for adapter convolution (downsampling factor)""" - - @dataclass(frozen=True) class StepAudio2Token2WavConfig: """Token2Wav synthesis configuration.""" @@ -141,7 +108,6 @@ class StepAudio2ModelConfig: # Default configuration instances DEFAULT_TOKEN_CONFIG = StepAudio2TokenConfig() -DEFAULT_ENCODER_CONFIG = StepAudio2EncoderConfig() DEFAULT_TOKEN2WAV_CONFIG = StepAudio2Token2WavConfig() DEFAULT_STREAM_CONFIG = StepAudio2StreamConfig() DEFAULT_MODEL_CONFIG = StepAudio2ModelConfig() @@ -162,12 +128,10 @@ class StepAudio2ModelConfig: __all__ = [ "StepAudio2TokenConfig", - "StepAudio2EncoderConfig", "StepAudio2Token2WavConfig", "StepAudio2StreamConfig", "StepAudio2ModelConfig", "DEFAULT_TOKEN_CONFIG", - "DEFAULT_ENCODER_CONFIG", "DEFAULT_TOKEN2WAV_CONFIG", "DEFAULT_STREAM_CONFIG", "DEFAULT_MODEL_CONFIG", diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index 7fd0bf4fbd7..c92f2397813 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -609,7 +609,7 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S Architecture: AudioEncoder (6 layers, 512 hidden) → Adaptor (512 → LLM dim) → - Qwen2 LLM (vocab=64012) + Qwen2 LLM (vocab=158720) """ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): @@ -868,31 +868,3 @@ def separate_tokens( def has_audio_output(token_ids: list[int], audio_start: int = DEFAULT_TOKEN_CONFIG.audio_start) -> bool: """Check if generated tokens contain audio tokens""" return any(tid >= audio_start for tid in token_ids) - - -class StepAudio2OutputProcessor: - """Helper class to process Step-Audio2 outputs""" - - @staticmethod - def process_output(output_ids: torch.Tensor, tokenizer, remove_audio_padding: bool = True) -> dict: - """Process model output and separate text/audio tokens""" - if isinstance(output_ids, torch.Tensor): - output_ids = output_ids.squeeze().tolist() - - # Separate tokens - text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(output_ids) - - # Remove audio padding if requested - if remove_audio_padding and audio_tokens: - audio_tokens = [t for t in audio_tokens if t < DEFAULT_TOKEN_CONFIG.audio_eos] - - # Decode text - text = tokenizer.decode(text_tokens, skip_special_tokens=False) - - return { - "text": text, - "text_tokens": text_tokens, - "audio_tokens": audio_tokens, - "has_audio": len(audio_tokens) > 0, - "all_tokens": output_ids, - } diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index a59ba3fef19..1ebad86cef6 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -94,8 +94,22 @@ def __init__( self.mel_cache_len = DEFAULT_STREAM_CONFIG.mel_cache_len self.source_cache_len = STREAM_SOURCE_CACHE_LEN self.speech_window: torch.Tensor | None = None # created lazily on device - self.stream_cache: dict | None = None # flow model causal cache - self.hift_cache_dict: dict = {} # HiFT vocoder overlap cache + + # On Ascend the talker worker never routes through + # ``NPUOmniPlatform.set_device``, so gate on ``current_omni_platform`` + # and apply the Token2Wav NPU patch here (idempotent) before + # ``_ensure_models_loaded`` runs. Covers every + # construction site (Step-Audio2 wrapper + MiniCPM-o facade): patched + # ``_ensure_models_loaded`` offloads HiFT to CPU and patched ``forward`` + # expands CosyVoice DiT masks + forces MATH SDPA (avoids FA 161001). + from vllm_omni.platforms import current_omni_platform + + if current_omni_platform.is_npu(): + from vllm_omni.platforms.npu.models.step_audio2_token2wav import ( + apply_step_audio2_token2wav_npu_patch, + ) + + apply_step_audio2_token2wav_npu_patch() def _ensure_models_loaded(self): """Lazy load models on first use""" @@ -253,125 +267,6 @@ def forward( else: return wav - # ------------------------------------------------------------------ - # Streaming (chunked) inference - # ------------------------------------------------------------------ - - def setup_stream(self, prompt_wav: str) -> None: - """Initialize flow + HiFT caches for chunked streaming inference. - - Must be called once before the first ``stream_chunk()`` call for a - given utterance. - """ - if prompt_wav not in self.cache: - self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) - ( - prompt_speech_tokens, - _, - spk_emb, - prompt_mels, - _, - ) = self.cache[prompt_wav] - - # Lazily create the Hamming window on the correct device - if self.speech_window is None or self.speech_window.device != self.device: - self.speech_window = torch.from_numpy(np.hamming(2 * self.source_cache_len)).to( - device=self.device, dtype=torch.float32 - ) - - # Flow model: setup causal cache with prompt tokens + lookahead - pre_lookahead = DEFAULT_STREAM_CONFIG.pre_lookahead_len - self.stream_cache = self.flow.setup_cache( - torch.cat([prompt_speech_tokens, prompt_speech_tokens[:, :pre_lookahead]], dim=1), - prompt_mels, - spk_emb, - n_timesteps=self.n_timesteps, - ) - - # HiFT vocoder: empty overlap buffers - self.hift_cache_dict = { - "mel": torch.zeros(1, prompt_mels.shape[2], 0, device=self.device), - "source": torch.zeros(1, 1, 0, device=self.device), - "speech": torch.zeros(1, 0, device=self.device), - } - - def stream_chunk( - self, - audio_tokens: list[int], - prompt_wav: str, - last_chunk: bool = False, - ) -> torch.Tensor: - """Process one chunk of audio tokens and return a waveform segment. - - Args: - audio_tokens: Token IDs for this chunk (includes lookahead for - non-last chunks). - prompt_wav: Path to the prompt wav (for cached speaker embeddings). - last_chunk: Whether this is the final chunk of the utterance. - - Returns: - 1-D float tensor of audio samples (24 kHz). - """ - if self.stream_cache is None: - raise ValueError("stream_cache not initialised – call setup_stream() first") - - if prompt_wav not in self.cache: - self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) - _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] - - token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) - - with torch.amp.autocast( - str(self.device.type), - dtype=torch.float16 if self.float16 else torch.float32, - ): - chunk_mel, self.stream_cache = self.flow.inference_chunk( - token=token_tensor, - spk=spk_emb, - cache=self.stream_cache, - last_chunk=last_chunk, - n_timesteps=self.n_timesteps, - ) - - # Trim estimator attention cache to avoid unbounded growth - keep = DEFAULT_STREAM_CONFIG.estimator_cache_keep - est_att = self.stream_cache.get("estimator_att_cache") - if est_att is not None and est_att.shape[4] > (prompt_mels.shape[1] + keep): - self.stream_cache["estimator_att_cache"] = torch.cat( - [est_att[:, :, :, :, : prompt_mels.shape[1]], est_att[:, :, :, :, -keep:]], - dim=4, - ) - - # ---- HiFT vocoder with overlap-add ---- - hift_cache_mel = self.hift_cache_dict["mel"] - hift_cache_source = self.hift_cache_dict["source"] - hift_cache_speech = self.hift_cache_dict["speech"] - - mel = torch.cat([hift_cache_mel, chunk_mel], dim=2) - speech, source = self.hift(mel, hift_cache_source) - - # Cross-fade overlap region for smooth concatenation - if hift_cache_speech.shape[-1] > 0: - speech = fade_in_out(speech, hift_cache_speech, self.speech_window) - - # Update vocoder cache for the next chunk - self.hift_cache_dict = { - "mel": mel[..., -self.mel_cache_len :].clone().detach(), - "source": source[:, :, -self.source_cache_len :].clone().detach(), - "speech": speech[:, -self.source_cache_len :].clone().detach(), - } - - # For non-last chunks, trim the tail that will be re-synthesised - if not last_chunk: - speech = speech[:, : -self.source_cache_len] - - return speech.squeeze(0) # [samples] - - def reset_stream(self) -> None: - """Clear all streaming state after an utterance is complete.""" - self.stream_cache = None - self.hift_cache_dict = {} - # ------------------------------------------------------------------ # Per-request streaming (does NOT mutate self – uses external state) # ------------------------------------------------------------------ diff --git a/vllm_omni/model_executor/models/voxcpm2/__init__.py b/vllm_omni/model_executor/models/voxcpm2/__init__.py index 77bd8dfb518..fb6738a2967 100644 --- a/vllm_omni/model_executor/models/voxcpm2/__init__.py +++ b/vllm_omni/model_executor/models/voxcpm2/__init__.py @@ -1,5 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from .voxcpm2_talker import VoxCPM2TalkerForConditionalGeneration -__all__ = ["VoxCPM2TalkerForConditionalGeneration"] +# NOTE: Do not import model classes in this file. Importing any +# submodule in this package triggers __init__.py execution, and +# both the model registry and pipeline registry import submodules +# directly — heavy imports here would be loaded as a side effect +# even though nothing depends on these re-exports. diff --git a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py index b87ec5aafef..13f1d0937d3 100644 --- a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py +++ b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py @@ -65,7 +65,6 @@ def __init__( prefix: str = "", ) -> None: super().__init__() - self.layer_idx = layer_idx self.hidden_size = hidden_size self.num_heads = num_attention_heads self.head_dim = kv_channels if kv_channels else hidden_size // num_attention_heads @@ -318,18 +317,12 @@ def precompute_fused_qkv(self) -> None: ).detach() def compile_selective(self) -> list[str]: - """Compile the full model forward as one graph. - - Earlier versions compiled ``layer.mlp`` + ``layer.self_attn.o_proj`` - (PR #2690) and then the whole ``layer`` (perf/voxcpm2-streaming-vae). - Both still paid one Dynamo dispatch per layer per decode step. - V3 profiling showed 1,332 per-layer dispatches (~28 layers × ~47 - decode steps) costing ~726 ms of CPU self-time for a long prompt. - - Compiling ``forward`` at the model level lets Dynamo unroll the - 28-layer Python loop inside the graph. Graph breaks at - PagedAttention produce sub-graphs but Dynamo memoises the whole - trace once, so the per-step dispatch drops from 28 to just a few. + """Compile the full model ``forward`` as one graph. + + ``torch.compile`` is applied at the model level so Dynamo unrolls the + per-layer Python loop inside the graph. Graph breaks at PagedAttention + produce sub-graphs, but Dynamo memoises the whole trace once, so the + per-decode-step dispatch drops to just a few instead of one per layer. """ if self._compiled_layers: return [] diff --git a/vllm_omni/model_executor/models/voxcpm2/pipeline.py b/vllm_omni/model_executor/models/voxcpm2/pipeline.py index 347fce17707..c92cf61ba92 100644 --- a/vllm_omni/model_executor/models/voxcpm2/pipeline.py +++ b/vllm_omni/model_executor/models/voxcpm2/pipeline.py @@ -15,6 +15,7 @@ VOXCPM2_PIPELINE = PipelineConfig( model_type="voxcpm2", + default_deploy_config_name="voxcpm2.yaml", model_arch="VoxCPM2TalkerForConditionalGeneration", stages=( StagePipelineConfig( @@ -26,6 +27,7 @@ final_output_type="audio", owns_tokenizer=True, engine_output_type="audio", + scheduler_cls="vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler", sampling_constraints={ "detokenize": False, "stop_token_ids": [1], diff --git a/vllm_omni/model_executor/models/voxcpm2/runtime_config.py b/vllm_omni/model_executor/models/voxcpm2/runtime_config.py new file mode 100644 index 00000000000..1446729e7c2 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm2/runtime_config.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import dataclasses +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclasses.dataclass(frozen=True) +class _VoxCPM2RuntimeConfig: + enable_profiling: bool = False + enable_nvtx_profile: bool = False + enable_loc_dit_layer_nvtx: bool = False + enable_loc_dit_fused_qkv: bool = True + enable_loc_dit_fused_mlp: bool = True + enable_loc_dit_zero_dt_cache: bool = True + enable_loc_dit_fast_rope: bool = False + enable_loc_dit_skip_qkv_contig: bool = True + enable_loc_dit_reduce_overhead_no_cg: bool = False + enable_loc_dit_fullgraph_no_cg: bool = False + cfg_cutoff_ratio: float = 1.0 + decode_graph_capture_policy: str = "all" + enable_vae_cuda_graph: bool = False + enable_cfm_cuda_graph: bool = False + enable_cfm_prealloc_output: bool = False + enable_batched_cfm: bool = True + deterministic_cfm_noise: bool = False + deterministic_cfm_seed: int = 20260601 + audio_emit_every: int = 1 + vae_decode_every: int = 3 + enable_delayed_audio_copy: bool = False + delayed_audio_copy_use_events: bool = False + coalesce_audio_d2h: bool = True + enable_batched_vae_decode: bool = True + enable_batched_fsq_fusion: bool = False + batched_fsq_fusion_max_batch: int = 32 + enable_batched_prefill_tail: bool = False + enable_unified_decode_graph: bool = True + unified_decode_graph_max_batch_size: int = 64 + + @classmethod + def from_vllm_config(cls, vllm_config: Any) -> _VoxCPM2RuntimeConfig: + model_config = vllm_config.model_config + raw = getattr(model_config, "voxcpm2_runtime_config", None) + if raw is None: + raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None) + if raw is None: + return cls() + if hasattr(raw, "to_dict"): + raw = raw.to_dict() + elif not isinstance(raw, dict) and hasattr(raw, "__dict__"): + raw = vars(raw) + if not isinstance(raw, dict): + logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw) + return cls() + + fields = {field.name: field for field in dataclasses.fields(cls)} + values: dict[str, Any] = {} + for key, value in raw.items(): + if key not in fields: + logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key) + continue + default = getattr(cls(), key) + values[key] = cls._coerce_value(key, value, default) + return cls(**values)._normalized() + + def _normalized(self) -> _VoxCPM2RuntimeConfig: + cfg = self + if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph: + logger.warning( + "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; " + "disabling CFM CUDA Graph and keeping batched CFM." + ) + cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False) + if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph: + logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.") + cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False) + if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy: + logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.") + cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False) + if cfg.enable_batched_vae_decode and ( + cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph + ): + logger.warning( + "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, " + "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode." + ) + cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False) + if cfg.enable_unified_decode_graph: + if cfg.enable_cfm_cuda_graph: + logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.") + cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False) + if not cfg.enable_batched_cfm: + logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.") + cfg = dataclasses.replace(cfg, enable_batched_cfm=True) + return cfg + + def unified_decode_graph_available(self, *, use_cuda_graph: bool) -> bool: + return bool(use_cuda_graph and self.enable_unified_decode_graph and not self.deterministic_cfm_noise) + + @staticmethod + def _coerce_value(key: str, value: Any, default: Any) -> Any: + if isinstance(default, bool): + if isinstance(value, str): + return value.strip().lower() in ("1", "true", "yes", "on") + return bool(value) + if isinstance(default, int) and not isinstance(default, bool): + value = int(value) + if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}: + return max(1, value) + return value + if isinstance(default, float): + if key == "cfg_cutoff_ratio": + return min(1.0, max(0.0, float(value))) + return float(value) + if isinstance(default, str): + return str(value) + return value diff --git a/vllm_omni/model_executor/models/voxcpm2/scheduler.py b/vllm_omni/model_executor/models/voxcpm2/scheduler.py new file mode 100644 index 00000000000..a780c695556 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm2/scheduler.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from vllm.v1.request import RequestStatus + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler +from vllm_omni.platforms import current_omni_platform + +from .runtime_config import _VoxCPM2RuntimeConfig + + +class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler): + """VoxCPM2 scheduler variant for full unified decode graph serving. + + VoxCPM2's full unified decode graph only applies to pure decode batches. + When a decode-ready request is already running, this scheduler defers new + waiting admissions for the current tick so the decode batch can stay on the + VoxCPM2 graph path. This is a model-local serving policy, not a generic AR + scheduler rule. + """ + + def _unified_decode_graph_enabled(self) -> bool: + runtime_config = _VoxCPM2RuntimeConfig.from_vllm_config(self.vllm_config) + return runtime_config.unified_decode_graph_available(use_cuda_graph=current_omni_platform.is_cuda()) + + def _should_defer_waiting_for_unified_decode_graph(self) -> bool: + if not self._unified_decode_graph_enabled(): + return False + if not self.waiting or not self.running: + return False + + for request in self.running: + if getattr(request, "status", None) != RequestStatus.RUNNING or request.is_finished(): + continue + if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens: + return True + return False + + def _should_defer_waiting_admission(self) -> bool: + return self._should_defer_waiting_for_unified_decode_graph() diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 3fcc683f3aa..94ac454d87e 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -16,12 +16,13 @@ import logging import math import time -from collections.abc import Iterable +from collections.abc import Callable, Iterable, Sequence from types import MethodType -from typing import Any +from typing import Any, NamedTuple, Protocol, TypedDict import torch import torch.nn as nn +from typing_extensions import Unpack from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, override_forward_context from vllm.inputs import tokens_input @@ -42,8 +43,10 @@ iter_custom_voice_profiles, load_validated_profile_tensors, ) +from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM +from .runtime_config import _VoxCPM2RuntimeConfig from .voxcpm2_import_utils import import_voxcpm2_core logger = init_logger(__name__) @@ -56,103 +59,50 @@ _ACTIVE_STATE_LEAK_WARN_MIN = 512 -@dataclasses.dataclass(frozen=True) -class _VoxCPM2RuntimeConfig: - enable_profiling: bool = False - enable_nvtx_profile: bool = False - enable_loc_dit_layer_nvtx: bool = False - enable_loc_dit_fused_qkv: bool = True - enable_loc_dit_fused_mlp: bool = True - enable_loc_dit_zero_dt_cache: bool = True - enable_loc_dit_fast_rope: bool = False - enable_loc_dit_skip_qkv_contig: bool = True - enable_loc_dit_reduce_overhead_no_cg: bool = False - enable_loc_dit_fullgraph_no_cg: bool = False - cfg_cutoff_ratio: float = 1.0 - decode_graph_capture_policy: str = "all" - enable_vae_cuda_graph: bool = False - enable_cfm_cuda_graph: bool = False - enable_cfm_prealloc_output: bool = False - enable_batched_cfm: bool = True - deterministic_cfm_noise: bool = False - deterministic_cfm_seed: int = 20260601 - audio_emit_every: int = 1 - vae_decode_every: int = 3 - enable_delayed_audio_copy: bool = False - delayed_audio_copy_use_events: bool = False - coalesce_audio_d2h: bool = True - enable_batched_vae_decode: bool = True - enable_batched_fsq_fusion: bool = False - batched_fsq_fusion_max_batch: int = 32 - - @classmethod - def from_vllm_config(cls, vllm_config: VllmConfig) -> _VoxCPM2RuntimeConfig: - model_config = vllm_config.model_config - raw = getattr(model_config, "voxcpm2_runtime_config", None) - if raw is None: - raw = getattr(getattr(model_config, "hf_config", None), "voxcpm2_runtime_config", None) - if raw is None: - return cls() - if hasattr(raw, "to_dict"): - raw = raw.to_dict() - elif not isinstance(raw, dict) and hasattr(raw, "__dict__"): - raw = vars(raw) - if not isinstance(raw, dict): - logger.warning("Ignoring invalid voxcpm2_runtime_config=%r; expected a dict.", raw) - return cls() - - fields = {field.name: field for field in dataclasses.fields(cls)} - values: dict[str, Any] = {} - for key, value in raw.items(): - if key not in fields: - logger.warning("Ignoring unknown VoxCPM2 runtime config key: %s", key) - continue - default = getattr(cls(), key) - values[key] = cls._coerce_value(key, value, default) - return cls(**values)._normalized() - - def _normalized(self) -> _VoxCPM2RuntimeConfig: - cfg = self - if cfg.enable_batched_cfm and cfg.enable_cfm_cuda_graph: - logger.warning( - "VoxCPM2 batched CFM and CFM CUDA Graph are mutually exclusive; " - "disabling CFM CUDA Graph and keeping batched CFM." - ) - cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False) - if cfg.enable_cfm_prealloc_output and not cfg.enable_cfm_cuda_graph: - logger.warning("VoxCPM2 CFM preallocated output requires CFM CUDA Graph; disabling it.") - cfg = dataclasses.replace(cfg, enable_cfm_prealloc_output=False) - if cfg.delayed_audio_copy_use_events and not cfg.enable_delayed_audio_copy: - logger.warning("VoxCPM2 delayed audio copy events require delayed audio copy; disabling event polling.") - cfg = dataclasses.replace(cfg, delayed_audio_copy_use_events=False) - if cfg.enable_batched_vae_decode and ( - cfg.enable_delayed_audio_copy or cfg.audio_emit_every > 1 or cfg.enable_vae_cuda_graph - ): - logger.warning( - "VoxCPM2 batched VAE decode is incompatible with delayed audio copy, " - "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode." - ) - cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False) - return cfg +class _VoxCPM2PromptConfigLike(Protocol): + audio_vae_config: dict[str, Any] + patch_size: int + + +class _AttentionMetadataLike(Protocol): + scheduler_metadata: Any - @staticmethod - def _coerce_value(key: str, value: Any, default: Any) -> Any: - if isinstance(default, bool): - if isinstance(value, str): - return value.strip().lower() in ("1", "true", "yes", "on") - return bool(value) - if isinstance(default, int) and not isinstance(default, bool): - value = int(value) - if key in {"audio_emit_every", "vae_decode_every", "batched_fsq_fusion_max_batch"}: - return max(1, value) - return value - if isinstance(default, float): - if key == "cfg_cutoff_ratio": - return min(1.0, max(0.0, float(value))) - return float(value) - if isinstance(default, str): - return str(value) - return value + +class _ForwardContextLike(Protocol): + attn_metadata: dict[str, _AttentionMetadataLike] | Any + + +class VoxCPM2PreprocessInput(TypedDict, total=False): + additional_information: dict[str, Any] + request_id: str + text_token_ids: list[list[int]] + reference_audio: object + ref_audio: object + prompt_audio: object + prompt_text: str | list[str] | None + voice_profile: dict[str, Any] | list[dict[str, Any]] | None + voice_name: str | list[str] | None + voice_created_at: int | str | None + + +class VoxCPM2PostprocessInput(TypedDict, total=False): + request_id: str + + +class _PrefillInputs(NamedTuple): + text_token: torch.Tensor + audio_feat: torch.Tensor + text_mask: torch.Tensor + audio_mask: torch.Tensor + + +class _PrefillResidualMeta(NamedTuple): + lm_hidden: torch.Tensor + prefix_feat_cond: torch.Tensor + + +class _DecodeResidualMeta(NamedTuple): + new_lm_hidden: torch.Tensor def is_cjk_char(c: str) -> bool: @@ -195,11 +145,11 @@ def split_multichar_chinese(token_ids: list[int], split_map: dict[int, list[int] def build_voxcpm2_prompt( - hf_config: Any, + hf_config: _VoxCPM2PromptConfigLike, tokenizer: Any, split_map: dict[int, list[int]], text: str, - ref_audio: Any | None = None, + ref_audio: Sequence[float] | torch.Tensor | None = None, ref_sr: int | None = None, ref_text: str | None = None, voice_profile: dict[str, Any] | None = None, @@ -234,6 +184,8 @@ def build_voxcpm2_prompt( prefill_len += len(prompt_ids) prefill_len += audio_feat_len elif ref_audio is not None: + if ref_sr is None: + raise ValueError("VoxCPM2 ref_sr is required when ref_audio is provided.") vae = hf_config.audio_vae_config patch_samples = hf_config.patch_size * math.prod(vae["encoder_rates"]) ref_len = math.ceil(math.ceil(len(ref_audio) * vae["sample_rate"] / ref_sr) / patch_samples) @@ -312,7 +264,6 @@ class _RequestState: decode_step_count: int = 0 request_start_time: float = 0.0 prefill_completed: bool = False - prefill_text: str = "" prompt_cache: dict | None = None prefill_masks: tuple | None = None is_stopping: bool = False @@ -344,6 +295,46 @@ class _CapturedCFMGraph: buffers: _CFMBufferManager +@dataclasses.dataclass +class _CapturedUnifiedDecodeGraph: + graph: torch.cuda.CUDAGraph + batch_size: int + # Static input buffers (copied before replay) + input_embeds: torch.Tensor + positions: torch.Tensor + prev_feat_embed: torch.Tensor + prefix_feat_cond: torch.Tensor + # Static output buffers (read after replay) + next_feat_embed: torch.Tensor + cfm_output: torch.Tensor + lm_hidden: torch.Tensor + # Internal buffers for CFM solver + cfm_buffers: _CFMBufferManager + cfm_noise: torch.Tensor + + +@dataclasses.dataclass +class _UnifiedDecodeGraphStats: + captures: int = 0 + replays: int = 0 + skips: dict[str, int] = dataclasses.field(default_factory=dict) + real_batch_sizes: dict[int, int] = dataclasses.field(default_factory=dict) + graph_bucket_sizes: dict[int, int] = dataclasses.field(default_factory=dict) + logged_replays: int = 0 + logged_skips: int = 0 + + def record_skip(self, reason: str) -> None: + self.skips[reason] = self.skips.get(reason, 0) + 1 + + def record_replay_batch(self, *, real_batch_size: int, graph_bucket_size: int) -> None: + self.real_batch_sizes[real_batch_size] = self.real_batch_sizes.get(real_batch_size, 0) + 1 + self.graph_bucket_sizes[graph_bucket_size] = self.graph_bucket_sizes.get(graph_bucket_size, 0) + 1 + + @property + def total_skips(self) -> int: + return sum(self.skips.values()) + + @dataclasses.dataclass class _PendingAudioCopy: host: torch.Tensor @@ -928,6 +919,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self._max_cached_graphs = self._max_batch_size self._cuda_graph_warmup_steps = 0 self._cuda_graph_warmup_threshold = 3 + self._unified_graphs: dict[int, _CapturedUnifiedDecodeGraph] = {} + self._enable_unified_decode_graph = self._runtime_config.unified_decode_graph_available( + use_cuda_graph=use_cuda_graph + ) + self._unified_decode_graph_max_batch_size = self._runtime_config.unified_decode_graph_max_batch_size + self._unified_graph_bucket_sizes: frozenset[int] = self._build_unified_graph_bucket_sizes( + min(self._max_batch_size, self._unified_decode_graph_max_batch_size) + ) + self._unified_graph_stats = _UnifiedDecodeGraphStats() + self._runner_assisted_unified_decode_graph_active = False + self._runner_assisted_unified_decode_graph_batch_size = 0 self._multichar_zh_split: dict[int, list[int]] | None = None @@ -1036,10 +1038,10 @@ def _flush_deferred_cleanup(self) -> None: def _build_prompt_cache( self, - ref_audio: Any = None, - prompt_audio: Any = None, + ref_audio: object = None, + prompt_audio: object = None, prompt_text: str | None = None, - ) -> dict | None: + ) -> dict[str, Any] | None: """Build prompt cache, handling both file paths and raw audio data. The OpenAI speech API sends decoded audio as [samples_list, sr] @@ -1107,6 +1109,84 @@ def _setup_cfm_buffers(self) -> None: max_batch_size=self._max_batch_size, ) + @staticmethod + def _voxcpm2_unwrap_torch_compile(module: nn.Module | Callable) -> nn.Module | Callable: + """Return the eager module behind torch.compile wrappers.""" + seen: set[int] = set() + while hasattr(module, "_orig_mod") and id(module) not in seen: + seen.add(id(module)) + module = getattr(module, "_orig_mod") + return module + + @staticmethod + def _voxcpm2_compile_without_inductor_cudagraphs( + module: nn.Module | Callable, + *, + mode: str | None, + fullgraph: bool, + ) -> Callable: + """Compile VoxCPM2 LocDiT for capture inside VoxCPM2-owned CUDA graphs.""" + kwargs: dict[str, Any] = { + "fullgraph": fullgraph, + "options": { + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + } + if mode is not None: + kwargs["mode"] = mode + return torch.compile(module, **kwargs) + + @staticmethod + def _raise_if_cuda_runtime_error(error: BaseException) -> None: + if isinstance(error, torch.cuda.OutOfMemoryError): + raise error + if isinstance(error, RuntimeError): + message = str(error).lower() + cuda_markers = ( + "cuda error", + "device-side assert", + "illegal memory access", + "out of memory", + "cublas", + "cudnn", + "cufft", + ) + if any(marker in message for marker in cuda_markers): + raise error + + def _voxcpm2_compile_unified_capture_estimator(self, module: nn.Module | Callable) -> Callable: + eager_module = self._voxcpm2_unwrap_torch_compile(module) + if getattr(self, "_estimator_for_unified_capture_source", None) is eager_module and hasattr( + self, "_estimator_for_unified_capture" + ): + return self._estimator_for_unified_capture + compiled = self._voxcpm2_compile_without_inductor_cudagraphs( + eager_module, + mode=None, + fullgraph=False, + ) + compiled._compiled = True + self._estimator_for_unified_capture_source = eager_module + self._estimator_for_unified_capture = compiled + return compiled + + def _voxcpm2_compile_unified_capture_feat_encoder(self, module: nn.Module | Callable) -> Callable: + eager_module = self._voxcpm2_unwrap_torch_compile(module) + if getattr(self, "_feat_encoder_for_unified_capture_source", None) is eager_module and hasattr( + self, "_feat_encoder_for_unified_capture" + ): + return self._feat_encoder_for_unified_capture + compiled = self._voxcpm2_compile_without_inductor_cudagraphs( + eager_module, + mode=None, + fullgraph=False, + ) + compiled._compiled = True + self._feat_encoder_for_unified_capture_source = eager_module + self._feat_encoder_for_unified_capture = compiled + return compiled + def _setup_torch_compile(self) -> None: if not self._enable_torch_compile: return @@ -1134,30 +1214,42 @@ def _setup_torch_compile(self) -> None: if _install_locdit_zero_dt_cache(estimator): targets.append("LocDiT zero-dt embedding cache") + external_cfm_capture = self._enable_cfm_cuda_graph or self._enable_unified_decode_graph if cfg.enable_loc_dit_layer_nvtx: patched = _install_locdit_layer_nvtx(estimator) if patched: targets.append(f"LocDiT layer NVTX ({patched} layers, compile skipped)") estimator._compiled = True - elif self._enable_cfm_cuda_graph: + elif self._enable_unified_decode_graph: + try: + compiled_ro = torch.compile(estimator, mode="reduce-overhead", fullgraph=False) + compiled_ro._compiled = True + tts.feat_decoder.estimator = compiled_ro + targets.append("LocDiT (reduce-overhead serving + lazy no-cg unified capture)") + except Exception as e: + self._raise_if_cuda_runtime_error(e) + logger.warning("torch.compile LocDiT dual-mode failed: %s", e) + try: + tts.feat_decoder.estimator = self._voxcpm2_compile_without_inductor_cudagraphs( + estimator, + mode="reduce-overhead", + fullgraph=False, + ) + tts.feat_decoder.estimator._compiled = True + self._estimator_for_unified_capture = tts.feat_decoder.estimator + targets.append("LocDiT (no-cg only, unified capture)") + except Exception as inner_e: + self._raise_if_cuda_runtime_error(inner_e) + logger.warning("torch.compile LocDiT failed completely: %s", inner_e) + elif external_cfm_capture: try: if cfg.enable_loc_dit_reduce_overhead_no_cg: - import torch._inductor.config as inductor_config - - old_cudagraphs = inductor_config.triton.cudagraphs - old_cudagraph_trees = inductor_config.triton.cudagraph_trees - inductor_config.triton.cudagraphs = False - inductor_config.triton.cudagraph_trees = False - try: - tts.feat_decoder.estimator = torch.compile( - estimator, - mode="reduce-overhead", - fullgraph=False, - ) - finally: - inductor_config.triton.cudagraphs = old_cudagraphs - inductor_config.triton.cudagraph_trees = old_cudagraph_trees - targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + CUDA Graph)") + tts.feat_decoder.estimator = self._voxcpm2_compile_without_inductor_cudagraphs( + estimator, + mode="reduce-overhead", + fullgraph=False, + ) + targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + external CUDA Graph)") else: tts.feat_decoder.estimator = torch.compile( estimator, @@ -1167,25 +1259,11 @@ def _setup_torch_compile(self) -> None: "triton.cudagraph_trees": False, }, ) - graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned" - targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)") + targets.append("LocDiT (no-inductor-cudagraph + external CUDA Graph)") tts.feat_decoder.estimator._compiled = True except Exception as e: - logger.warning("torch.compile LocDiT reduce-overhead for CUDA Graph failed: %s", e) - try: - tts.feat_decoder.estimator = torch.compile( - estimator, - fullgraph=cfg.enable_loc_dit_fullgraph_no_cg, - options={ - "triton.cudagraphs": False, - "triton.cudagraph_trees": False, - }, - ) - tts.feat_decoder.estimator._compiled = True - graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned" - targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)") - except Exception as inner_e: - logger.warning("torch.compile LocDiT for CUDA Graph failed: %s", inner_e) + self._raise_if_cuda_runtime_error(e) + logger.warning("torch.compile LocDiT for external CUDA Graph failed: %s", e) else: try: tts.feat_decoder.estimator = torch.compile( @@ -1196,14 +1274,34 @@ def _setup_torch_compile(self) -> None: tts.feat_decoder.estimator._compiled = True targets.append("LocDiT") except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("torch.compile LocDiT failed: %s", e) try: if not hasattr(tts.feat_encoder, "_compiled"): - tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False) - tts.feat_encoder._compiled = True - targets.append("feat_encoder") + if self._enable_unified_decode_graph: + feat_encoder = tts.feat_encoder + compiled_ro = torch.compile(feat_encoder, mode="reduce-overhead", fullgraph=False) + compiled_ro._compiled = True + tts.feat_encoder = compiled_ro + targets.append("feat_encoder (reduce-overhead serving + lazy no-cg unified capture)") + elif external_cfm_capture: + tts.feat_encoder = torch.compile( + tts.feat_encoder, + fullgraph=False, + options={ + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + ) + tts.feat_encoder._compiled = True + targets.append("feat_encoder (no-inductor-cudagraph)") + else: + tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False) + tts.feat_encoder._compiled = True + targets.append("feat_encoder") except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("torch.compile feat_encoder failed: %s", e) if self._compile_vae and not self._enable_vae_cuda_graph: @@ -1213,6 +1311,7 @@ def _setup_torch_compile(self) -> None: tts.audio_vae._compiled = True targets.append("AudioVAE") except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("torch.compile AudioVAE failed: %s", e) if not self._enable_cuda_graph: @@ -1221,6 +1320,7 @@ def _setup_torch_compile(self) -> None: targets.extend(f"scaffold.{t}" for t in self.model.compile_selective()) self.model._selective_compiled = True except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("scaffold compile failed: %s", e) if not getattr(self.residual_model, "_selective_compiled", False): @@ -1228,6 +1328,7 @@ def _setup_torch_compile(self) -> None: targets.extend(f"residual.{t}" for t in self.residual_model.compile_selective()) self.residual_model._selective_compiled = True except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("residual compile failed: %s", e) else: self.model.precompute_fused_qkv() @@ -1241,6 +1342,7 @@ def _setup_torch_compile(self) -> None: self._projections_compiled = True targets.append("projections") except Exception as e: + self._raise_if_cuda_runtime_error(e) self._compiled_dit_proj = self._compiled_stop_fn = None logger.warning("projections compile failed: %s", e) @@ -1256,7 +1358,7 @@ def _stop_fn(self, lm_h: torch.Tensor) -> torch.Tensor: return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h))) @staticmethod - def _nullify_volatile_metadata(ctx: Any) -> Any: + def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike: """Set ``scheduler_metadata`` to None on all attention layers. This is the only tensor FA3 reallocates each step (variable shape). @@ -1352,6 +1454,31 @@ def _should_use_decode_graph(self, batch_size: int) -> bool: ) return True + @staticmethod + def _build_unified_graph_bucket_sizes(max_size: int) -> frozenset[int]: + if max_size <= 0: + return frozenset() + sizes = {1} + size = 2 + while size <= min(max_size, 16): + sizes.add(size) + size <<= 1 + bucket = 24 + while bucket <= max_size: + sizes.add(bucket) + bucket += 8 + sizes.add(max_size) + return frozenset(size for size in sizes if 0 < size <= max_size) + + def _select_unified_graph_bucket_size(self, batch_size: int) -> int | None: + if ( + batch_size <= 0 + or batch_size > self._unified_decode_graph_max_batch_size + or batch_size > self._max_cached_graphs + ): + return None + return min((size for size in self._unified_graph_bucket_sizes if size >= batch_size), default=None) + def _capture_vae_graph(self, feat: torch.Tensor) -> _CapturedVAEGraph: batch_size = feat.shape[0] num_frames = feat.shape[-1] @@ -1536,8 +1663,297 @@ def _run_cfm_cuda_graph_to_state_buffer( state.cfm_output_gpu.copy_(graph.output) return state.cfm_output_gpu.transpose(1, 2) + # -------------------- unified decode graph -------------------- + + def _capture_unified_decode_graph(self, batch_size: int) -> _CapturedUnifiedDecodeGraph: + H = self.config.hidden_size + D = self._feat_dim + P = self._patch_size + dev = torch.device(self._device) + dtype = self._side_dtype + + self._setup_cfm_buffers() + if self._enable_torch_compile: + self._setup_torch_compile() + self.model.precompute_fused_qkv() + self.residual_model.precompute_fused_qkv() + + tts = self.tts + dit_hidden = tts.lm_to_dit_proj.out_features + tts.res_to_dit_proj.out_features + cfm_buffers = _CFMBufferManager( + device=dev, + dtype=dtype, + feat_dim=D, + patch_size=P, + dit_hidden_size=dit_hidden, + max_batch_size=batch_size, + ) + + g = _CapturedUnifiedDecodeGraph( + graph=torch.cuda.CUDAGraph(), + batch_size=batch_size, + input_embeds=torch.zeros(batch_size, H, device=dev, dtype=dtype), + positions=torch.zeros(batch_size, device=dev, dtype=torch.long), + prev_feat_embed=torch.zeros(batch_size, H, device=dev, dtype=dtype), + prefix_feat_cond=torch.zeros(batch_size, P, D, device=dev, dtype=dtype), + next_feat_embed=torch.zeros(batch_size, H, device=dev, dtype=dtype), + cfm_output=torch.zeros(batch_size, D, P, device=dev, dtype=dtype), + lm_hidden=torch.zeros(batch_size, H, device=dev, dtype=dtype), + cfm_buffers=cfm_buffers, + cfm_noise=torch.zeros(batch_size, D, P, device=dev, dtype=dtype), + ) + + dit_proj = self._dit_proj_fn + + def unified_fwd(): + scaffold_h = self.model(input_ids=None, positions=g.positions, inputs_embeds=g.input_embeds) + if isinstance(scaffold_h, tuple): + scaffold_h = scaffold_h[0] + lm_h = tts.fsq_layer(scaffold_h) + res_input = tts.fusion_concat_proj(torch.cat([lm_h, g.prev_feat_embed], dim=-1)) + res_h = self.residual_model(positions=g.positions, inputs_embeds=res_input) + dit_h = dit_proj(lm_h, res_h) + cond = g.prefix_feat_cond.transpose(1, 2).contiguous() + cfm_out = _optimized_solve_euler_with_noise( + tts.feat_decoder, + dit_h, + P, + cond, + g.cfm_noise, + self._inference_timesteps, + self._cfg_value, + cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + ) + feat_enc = tts.feat_encoder(cfm_out.transpose(1, 2).unsqueeze(1)).squeeze(1) + next_embed = tts.enc_to_lm_proj(feat_enc) + return next_embed, cfm_out, lm_h + + ctx = get_forward_context() + capture_context = override_forward_context(self._nullify_volatile_metadata(ctx)) + + original_estimator = tts.feat_decoder.estimator + capture_estimator = self._voxcpm2_compile_unified_capture_estimator(original_estimator) + tts.feat_decoder.estimator = capture_estimator + + original_feat_encoder = tts.feat_encoder + capture_feat_encoder = self._voxcpm2_compile_unified_capture_feat_encoder(original_feat_encoder) + tts.feat_encoder = capture_feat_encoder + + try: + with capture_context: + with torch.no_grad(): + for _ in range(3): + g.cfm_noise.normal_() + unified_fwd() + + g.cfm_noise.normal_() + with torch.cuda.graph(g.graph, pool=current_platform.get_global_graph_pool()): + g.next_feat_embed, g.cfm_output, g.lm_hidden = unified_fwd() + finally: + tts.feat_decoder.estimator = original_estimator + tts.feat_encoder = original_feat_encoder + self._unified_graph_stats.captures += 1 + logger.info( + "CUDA Graph captured for unified decode (batch_size=%d, captures=%d)", + batch_size, + self._unified_graph_stats.captures, + ) + return g + + def _unified_decode_graph_skip_reason( + self, + *, + can_use_graph: bool, + is_all_decode: bool, + num_reqs: int, + ) -> str | None: + if not self._enable_unified_decode_graph: + return "disabled" + if not can_use_graph: + return "graph_not_ready" + if not is_all_decode: + return "not_all_decode" + if num_reqs > self._unified_decode_graph_max_batch_size: + return "batch_exceeds_unified_max" + if num_reqs > self._max_cached_graphs: + return "batch_too_large" + graph_size = self._select_unified_graph_bucket_size(num_reqs) + if graph_size is None: + return "capture_policy" + if num_reqs > 1 and not self._runner_assisted_unified_decode_graph_active: + return "runner_full_metadata_missing" + for req_id, _, _, _ in self._pending_requests: + state = self._active_states[req_id] + if not state.prefill_completed: + return "prefill_incomplete" + if state.prev_feat_embed is None or state.curr_prefix_feat_cond is None: + return "state_not_ready" + return None + + def _maybe_log_unified_graph_stats(self) -> None: + if not self._enable_profiling: + return + stats = self._unified_graph_stats + total_skips = stats.total_skips + should_log = (stats.replays - stats.logged_replays >= 100) or (total_skips - stats.logged_skips >= 50) + if not should_log: + return + stats.logged_replays = stats.replays + stats.logged_skips = total_skips + logger.info( + "VoxCPM2 unified graph stats: captures=%d replays=%d skips=%s real_batch_sizes=%s graph_bucket_sizes=%s", + stats.captures, + stats.replays, + stats.skips, + dict(sorted(stats.real_batch_sizes.items())), + dict(sorted(stats.graph_bucket_sizes.items())), + ) + + def _forward_unified_decode( + self, + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + num_reqs: int, + ) -> torch.Tensor: + graph_size = self._select_unified_graph_bucket_size(num_reqs) + if graph_size is None: + raise RuntimeError(f"No unified decode graph bucket for batch size {num_reqs}") + g = self._unified_graphs.get(graph_size) + if g is None: + g = self._capture_unified_decode_graph(graph_size) + self._unified_graphs[graph_size] = g + + self._unified_graph_stats.replays += 1 + if self._enable_profiling: + self._unified_graph_stats.record_replay_batch(real_batch_size=num_reqs, graph_bucket_size=graph_size) + self._maybe_log_unified_graph_stats() + states: list[_RequestState] = [] + commit_mask: list[bool] = [] + for req_id, _is_prefill, _embeds, _n in self._pending_requests: + state = self._active_states[req_id] + states.append(state) + already_stopping = state.is_stopping + commit_mask.append(not already_stopping) + if not already_stopping: + state.decode_step_count += 1 + if state.decode_step_count >= self._max_decode_steps: + state.is_stopping = True + + self._perf.start("unified.copy_inputs") + for i, state in enumerate(states): + pfe = state.prev_feat_embed + g.prev_feat_embed[i].copy_(pfe.squeeze(0) if pfe.ndim > 1 else pfe) + pfc = state.curr_prefix_feat_cond + if pfc.ndim == 2: + g.prefix_feat_cond[i].copy_(pfc) + else: + g.prefix_feat_cond[i].copy_(pfc.squeeze(0)) + + g.input_embeds[:num_reqs].copy_(inputs_embeds[:num_reqs]) + g.positions[:num_reqs].copy_(positions[:num_reqs]) + if graph_size > num_reqs: + # Runner-assisted FULL metadata pads attention to graph_size. Keep + # padded rows as valid duplicates instead of position-0 zero rows; + # some attention backends may still touch padded slots before + # honoring seq_lens=0. + last = num_reqs - 1 + g.input_embeds[num_reqs:graph_size].copy_(g.input_embeds[last : last + 1].expand(graph_size - num_reqs, -1)) + g.positions[num_reqs:graph_size].copy_(g.positions[last : last + 1].expand(graph_size - num_reqs)) + g.prev_feat_embed[num_reqs:graph_size].copy_( + g.prev_feat_embed[last : last + 1].expand(graph_size - num_reqs, -1) + ) + g.prefix_feat_cond[num_reqs:graph_size].copy_( + g.prefix_feat_cond[last : last + 1].expand(graph_size - num_reqs, -1, -1) + ) + g.cfm_noise.normal_() + self._perf.stop("unified.copy_inputs") + + self._perf.start("unified.replay") + g.graph.replay() + self._perf.stop("unified.replay") + + self._perf.start("unified.commit") + with torch.no_grad(): + all_stop_logits = self._stop_fn(g.lm_hidden[:num_reqs]) + for i, state in enumerate(states): + stop_logits_i = all_stop_logits[i : i + 1] + if not commit_mask[i]: + state.precomputed_stop_logits = stop_logits_i + continue + next_embed_i = g.next_feat_embed[i : i + 1].clone() + cfm_out_i = g.cfm_output[i : i + 1].transpose(1, 2) + self._commit_decode_state(state, stop_logits_i, next_embed_i, cfm_out_i) + self._perf.stop("unified.commit") + + self._perf.start("unified.audio") + self._precompute_stop_flags_for_audio_collect(states) + ready_audio = self._drain_ready_audio_copies_for_states(states) + audio_by_req = self._collect_audio_batch(states, initial_delayed_chunks_by_req=ready_audio) + for state in states: + self._results_queue.append((state.request_id, state.precomputed_stop_logits)) + self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) + + self._pending_requests.clear() + self._flush_deferred_cleanup() + self._perf.stop("unified.audio") + return g.next_feat_embed[:num_reqs] + # -------------------- vllm hooks -------------------- + def get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_scheduled_tokens: Sequence[int], + num_computed_tokens: Sequence[int], + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest | None: + if not self._enable_unified_decode_graph or not self._enable_cuda_graph: + return None + if num_reqs <= 1 or num_reqs > self._unified_decode_graph_max_batch_size: + return None + if num_reqs > self._max_cached_graphs or self._select_unified_graph_bucket_size(num_reqs) is None: + return None + if max_num_scheduled_tokens != 1 or len(num_scheduled_tokens) != num_reqs: + return None + if any(int(n) != 1 for n in num_scheduled_tokens): + return None + if len(req_ids) != num_reqs or len(num_computed_tokens) != num_reqs: + return None + if any(int(n) <= 0 for n in num_computed_tokens): + return None + tts_compiled = getattr(self.tts.feat_decoder.estimator, "_compiled", False) if self._tts is not None else False + if not tts_compiled or self._cuda_graph_warmup_steps < self._cuda_graph_warmup_threshold: + return None + for req_id in req_ids: + state = self._active_states.get(req_id) + if ( + state is None + or not state.prefill_completed + or state.prev_feat_embed is None + or state.curr_prefix_feat_cond is None + ): + return None + bucket_size = self._select_unified_graph_bucket_size(num_reqs) + if bucket_size is None: + return None + return RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=bucket_size, + for_cudagraph_capture=bucket_size not in self._unified_graphs, + ) + + def set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> None: + enabled = bool(enabled) + self._runner_assisted_unified_decode_graph_active = enabled + self._runner_assisted_unified_decode_graph_batch_size = int(num_reqs) if enabled else 0 + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -1566,12 +1982,30 @@ def forward( self._enable_cuda_graph and graph_ready and intermediate_tensors is None and inputs_embeds is not None ) - if can_use_graph and is_all_decode and num_reqs <= self._max_cached_graphs: - use_scaffold_graph = self._should_use_decode_graph(num_reqs) - else: - use_scaffold_graph = False + unified_skip_reason = self._unified_decode_graph_skip_reason( + can_use_graph=can_use_graph, + is_all_decode=is_all_decode, + num_reqs=num_reqs, + ) + if unified_skip_reason is None: + self._perf.start("unified_decode") + result = self._forward_unified_decode(inputs_embeds, positions, num_reqs) + self._perf.stop("unified_decode") + self._perf.stop("forward_total") + return result + if self._enable_unified_decode_graph and num_reqs > 0 and unified_skip_reason is not None: + self._unified_graph_stats.record_skip(unified_skip_reason) + self._maybe_log_unified_graph_stats() + + use_segmented_decode_graph = ( + can_use_graph + and is_all_decode + and num_reqs <= self._max_cached_graphs + and (not self._enable_unified_decode_graph or unified_skip_reason is not None) + and self._should_use_decode_graph(num_reqs) + ) - if use_scaffold_graph: + if use_segmented_decode_graph: self._perf.start("scaffold_fwd") if num_reqs not in self._scaffold_graphs: self._scaffold_graphs[num_reqs] = self._capture_graph(self.model, num_reqs, "scaffold") @@ -1591,7 +2025,7 @@ def forward( token_offset = 0 residual_inputs: list[torch.Tensor] = [] residual_positions: list[torch.Tensor] = [] - req_metas: list[tuple[_RequestState, bool, dict]] = [] + req_metas: list[tuple[_RequestState, bool, _PrefillResidualMeta | _DecodeResidualMeta]] = [] pending_decode_fsq: list[tuple[_RequestState, torch.Tensor, torch.Tensor]] = [] def flush_decode_fsq_batch() -> None: @@ -1657,6 +2091,7 @@ def flush_decode_fsq_batch() -> None: residual_batch_size = batch_in.shape[0] use_residual_graph = ( self._enable_cuda_graph + and use_segmented_decode_graph and is_all_decode and graph_ready and residual_batch_size == num_reqs # 1 token per request @@ -1680,12 +2115,12 @@ def flush_decode_fsq_batch() -> None: batch_out = self.residual_model(batch_pos, batch_in) self._perf.stop("residual_fwd") - can_batch_decode_tail = ( + can_finish_decode_batch = ( bool(req_metas) and all(not is_prefill for _, is_prefill, _ in req_metas) and all(x.shape[0] == 1 for x in residual_inputs) ) - if can_batch_decode_tail: + if can_finish_decode_batch: self._finish_decode_batch(req_metas, batch_out) self._precompute_stop_flags_for_audio_collect([state for state, _, _ in req_metas]) ready_audio_by_req = self._drain_ready_audio_copies_for_states([state for state, _, _ in req_metas]) @@ -1699,17 +2134,24 @@ def flush_decode_fsq_batch() -> None: else: offset = 0 decoded_states: list[_RequestState] = [] + prefill_batch: list[tuple[_RequestState, _PrefillResidualMeta, torch.Tensor]] = [] for idx, (state, is_prefill, meta) in enumerate(req_metas): n = residual_inputs[idx].shape[0] res_out = batch_out[offset : offset + n] offset += n if is_prefill: - self._finish_prefill(state, meta, res_out, dev) + prefill_batch.append((state, meta, res_out)) else: - self._finish_decode(state, meta, res_out, dev) + self._finish_decode(state, meta, res_out) decoded_states.append(state) + if prefill_batch and self._runtime_config.enable_batched_prefill_tail: + self._finish_prefill_batch(prefill_batch) + else: + for state, meta, res_out in prefill_batch: + self._finish_prefill(state, meta, res_out, dev) + collect_states = [state for state, _, _ in req_metas] self._precompute_stop_flags_for_audio_collect(collect_states) ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) @@ -1732,7 +2174,12 @@ def flush_decode_fsq_batch() -> None: # -------------------- prefill / decode helpers -------------------- - def _prepare_residual_prefill(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any): + def _prepare_residual_prefill( + self, + state: _RequestState, + base_lm_out: torch.Tensor, + dev: torch.device, + ) -> tuple[torch.Tensor, _PrefillResidualMeta]: tts = self.tts text_mask, feat_mask, feat, feat_embed = state.prefill_masks state.prefill_masks = None @@ -1755,10 +2202,15 @@ def _prepare_residual_prefill(self, state: _RequestState, base_lm_out: torch.Ten lm_hidden = enc_outputs[:, -1, :] residual_input = tts.fusion_concat_proj(torch.cat([enc_outputs, feat_mask.unsqueeze(-1) * feat_embed], dim=-1)) - meta = {"lm_hidden": lm_hidden, "prefix_feat_cond": prefix_feat_cond} + meta = _PrefillResidualMeta(lm_hidden=lm_hidden, prefix_feat_cond=prefix_feat_cond) return residual_input.squeeze(0), meta - def _prepare_residual_decode(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any): + def _prepare_residual_decode( + self, + state: _RequestState, + base_lm_out: torch.Tensor, + dev: torch.device, + ) -> tuple[torch.Tensor, _DecodeResidualMeta]: tts = self.tts state.decode_step_count += 1 @@ -1775,14 +2227,14 @@ def _prepare_residual_decode(self, state: _RequestState, base_lm_out: torch.Tens if prev.ndim == 1: prev = prev.unsqueeze(0) res_input = tts.fusion_concat_proj(torch.cat([lm_h, prev], dim=-1)) - return res_input, {"new_lm_hidden": lm_h} + return res_input, _DecodeResidualMeta(new_lm_hidden=lm_h) def _prepare_residual_decode_batch( self, states: list[_RequestState], base_lm_outs: list[torch.Tensor], - dev: Any, - ) -> list[tuple[torch.Tensor, dict]]: + dev: torch.device, + ) -> list[tuple[torch.Tensor, _DecodeResidualMeta]]: tts = self.tts for state in states: @@ -1812,7 +2264,7 @@ def _prepare_residual_decode_batch( return [ ( res_batch[i : i + 1], - {"new_lm_hidden": lm_h_batch[i : i + 1]}, + _DecodeResidualMeta(new_lm_hidden=lm_h_batch[i : i + 1]), ) for i in range(len(states)) ] @@ -1907,14 +2359,96 @@ def _fill_deterministic_cfm_noise(self, state: _RequestState, out: torch.Tensor) out.normal_(generator=gen) state.cfm_noise_step += 1 - def _finish_prefill(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any): + def _finish_prefill_batch(self, batch: list[tuple[_RequestState, _PrefillResidualMeta, torch.Tensor]]) -> None: + if len(batch) == 1: + state, meta, res_out = batch[0] + self._finish_prefill(state, meta, res_out, None) + return + + self._perf.start("prefill_tail_batch") + tts = self.tts + self._setup_cfm_buffers() + if self._enable_torch_compile: + self._setup_torch_compile() + + dit_hs: list[torch.Tensor] = [] + conds: list[torch.Tensor] = [] + stop_logits_list: list[torch.Tensor] = [] + + for state, meta, res_out in batch: + lm_hidden = meta.lm_hidden + prefix_feat_cond = meta.prefix_feat_cond + residual_hidden = res_out[-1:, :] + stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach() + dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1) + cond = prefix_feat_cond.transpose(1, 2).contiguous() + dit_hs.append(dit_h) + conds.append(cond) + stop_logits_list.append(stop_logits) + + batched_dit_h = torch.cat(dit_hs, dim=0) + batched_cond = torch.cat(conds, dim=0) + b = batched_dit_h.shape[0] + + if self._deterministic_cfm_noise and self._cfm_buffers is not None: + noise = self._cfm_buffers.noise[:b] + for i, (state, _, _) in enumerate(batch): + self._fill_deterministic_cfm_noise(state, noise[i : i + 1]) + else: + noise = torch.randn( + b, + self._feat_dim, + self._patch_size, + device=batched_dit_h.device, + dtype=batched_dit_h.dtype, + ) + pred_feats = _optimized_solve_euler_with_noise( + tts.feat_decoder, + batched_dit_h, + self._patch_size, + batched_cond, + noise, + self._inference_timesteps, + self._cfg_value, + self._cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + ).transpose(1, 2) + + with torch.no_grad(): + curr_embeds = tts.enc_to_lm_proj(tts.feat_encoder(pred_feats.unsqueeze(1))).squeeze(1) + + for i, (state, meta, res_out) in enumerate(batch): + pred_feat_i = pred_feats[i : i + 1] + curr_embed_i = curr_embeds[i : i + 1] + self._commit_decode_state(state, stop_logits_list[i], curr_embed_i, pred_feat_i) + state.decode_step_count = 0 + state.request_start_time = time.perf_counter() + state.prefill_completed = True + + self._perf.stop("prefill_tail_batch") + if self._enable_profiling: + logger.info("PREFILL_BATCH[%d requests] tail breakdown:\n%s", len(batch), self._perf.breakdown()) + self._perf.reset() + + def _finish_prefill( + self, + state: _RequestState, + meta: _PrefillResidualMeta, + res_out: torch.Tensor, + dev: torch.device | None, + ) -> None: + self._perf.start("prefill_tail") tts = self.tts - lm_hidden = meta["lm_hidden"] - prefix_feat_cond = meta["prefix_feat_cond"] + lm_hidden = meta.lm_hidden + prefix_feat_cond = meta.prefix_feat_cond residual_hidden = res_out[-1:, :] + self._perf.start("prefill.stop_fn") stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach() + self._perf.stop("prefill.stop_fn") + self._perf.start("prefill.dit_proj") dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1) + self._perf.stop("prefill.dit_proj") self._setup_cfm_buffers() if self._enable_torch_compile: @@ -1922,24 +2456,29 @@ def _finish_prefill(self, state: _RequestState, meta: dict, res_out: torch.Tenso pred_feat = self._run_cfm_for_state(state, dit_h, prefix_feat_cond.transpose(1, 2).contiguous()) + self._perf.start("prefill.feat_encoder") with torch.no_grad(): curr_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1) + self._perf.stop("prefill.feat_encoder") self._commit_decode_state(state, stop_logits, curr_embed, pred_feat) state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True + self._perf.stop("prefill_tail") if logger.isEnabledFor(logging.DEBUG): # Only compute the norm (which forces a GPU->CPU sync) if we will log it. logger.debug("PREFILL[%s]: patch norm=%.4f", state.request_id, pred_feat.norm().item()) + if self._enable_profiling: + logger.info("PREFILL[%s] tail breakdown:\n%s", state.request_id, self._perf.breakdown()) self._perf.reset() - def _finish_decode(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any): + def _finish_decode(self, state: _RequestState, meta: _DecodeResidualMeta, res_out: torch.Tensor) -> None: self._perf.start("decode_step") tts = self.tts - lm_h = meta["new_lm_hidden"] + lm_h = meta.new_lm_hidden res_h = res_out.unsqueeze(0) if res_out.ndim == 1 else res_out dit_proj = getattr(self, "_compiled_dit_proj", None) or self._dit_proj_fn @@ -1960,7 +2499,11 @@ def _finish_decode(self, state: _RequestState, meta: dict, res_out: torch.Tensor if self._enable_profiling and state.decode_step_count % 20 == 0: logger.info("Step %d[%s]:\n%s", state.decode_step_count, state.request_id, self._perf.breakdown()) - def _finish_decode_batch(self, req_metas: list[tuple], batch_out: torch.Tensor): + def _finish_decode_batch( + self, + req_metas: list[tuple[_RequestState, bool, _DecodeResidualMeta]], + batch_out: torch.Tensor, + ) -> None: self._perf.start("decode_step") tts = self.tts @@ -1969,7 +2512,7 @@ def _finish_decode_batch(self, req_metas: list[tuple], batch_out: torch.Tensor): if self._enable_batched_cfm and not self._enable_cfm_cuda_graph: states = [state for state, _, _ in req_metas] - lm_h = torch.cat([meta["new_lm_hidden"] for _, _, meta in req_metas], dim=0) + lm_h = torch.cat([meta.new_lm_hidden for _, _, meta in req_metas], dim=0) pfc = torch.cat( [ state.curr_prefix_feat_cond.to(self._side_dtype).unsqueeze(0) @@ -2017,7 +2560,7 @@ def _finish_decode_batch(self, req_metas: list[tuple], batch_out: torch.Tensor): return for i, (state, _, meta) in enumerate(req_metas): - lm_h = meta["new_lm_hidden"] + lm_h = meta.new_lm_hidden res_h = batch_out[i : i + 1] pfc = state.curr_prefix_feat_cond.to(self._side_dtype) if pfc.ndim == 2: @@ -2475,7 +3018,10 @@ def _get_multichar_zh_split(self) -> dict[int, list[int]]: # -------------------- preprocess / postprocess -------------------- def preprocess( - self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Unpack[VoxCPM2PreprocessInput], ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: additional = info_dict.get("additional_information") if isinstance(additional, dict): @@ -2510,7 +3056,6 @@ def preprocess( token_ids = token_ids[1:] state = self._get_or_create_state(req_id) - state.prefill_text = "" state.decode_pad = None state.prefill_completed = False state.decode_step_count = 0 @@ -2572,37 +3117,28 @@ def preprocess( ) if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)): - if state.prompt_cache is None: - try: - state.prompt_cache = self._build_prompt_cache( - ref_audio=ref_audio, - prompt_audio=prompt_audio, - prompt_text=prompt_text, - ) - if ( - voice_name - and state.prompt_cache is not None - and state.prompt_cache.get("mode") == "reference" - and "ref_audio_feat" in state.prompt_cache - ): - _key = self._speaker_cache.make_cache_key( - voice_name, model_type="voxcpm2", created_at=_created_at - ) - self._speaker_cache.put( - _key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()} - ) - logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name) - except Exception as e: - logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e) - state.prompt_cache = None + state.prompt_cache = self._build_prompt_cache( + ref_audio=ref_audio, + prompt_audio=prompt_audio, + prompt_text=prompt_text, + ) + if ( + voice_name + and state.prompt_cache is not None + and state.prompt_cache.get("mode") == "reference" + and "ref_audio_feat" in state.prompt_cache + ): + _key = self._speaker_cache.make_cache_key(voice_name, model_type="voxcpm2", created_at=_created_at) + self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}) + logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name) inputs = self._build_prefill_inputs(token_ids, dev, req_id) tts = self.tts - feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs["audio_feat"])) - text_embed = self.model.embed_input_ids(inputs["text_token"].to(dev)) - text_mask, feat_mask = inputs["text_mask"], inputs["audio_mask"] + feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs.audio_feat)) + text_embed = self.model.embed_input_ids(inputs.text_token.to(dev)) + text_mask, feat_mask = inputs.text_mask, inputs.audio_mask embeds = (text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed).squeeze(0) - state.prefill_masks = (text_mask, feat_mask, inputs["audio_feat"], feat_embed) + state.prefill_masks = (text_mask, feat_mask, inputs.audio_feat, feat_embed) else: state = self._active_states.get(req_id) curr = state.curr_embed_for_next if state else None @@ -2614,23 +3150,31 @@ def preprocess( self._pending_requests.append((req_id, is_prefill, embeds, span_len)) return input_ids, embeds, {} - def postprocess(self, hidden_states: torch.Tensor, **info: Any) -> dict[str, Any]: + def postprocess(self, hidden_states: torch.Tensor, **info: Unpack[VoxCPM2PostprocessInput]) -> dict[str, Any]: req_id = info.get("request_id", self._current_request_id or "default") if self._enable_profiling: state = self._active_states.get(req_id) if state and state.decode_step_count > 0: logger.info( - "REQUEST DONE[%s]: %d steps, %.2fs\n%s", + "REQUEST DONE[%s]: %d steps, %.2fs\n%s\nUnified graph: captures=%d replays=%d skips=%s", req_id, state.decode_step_count, time.perf_counter() - state.request_start_time, self._perf.breakdown(), + self._unified_graph_stats.captures, + self._unified_graph_stats.replays, + self._unified_graph_stats.skips, ) return {} # -------------------- build prefill inputs -------------------- - def _build_prefill_inputs(self, token_ids: list[int], dev: Any, req_id: str = "default") -> dict: + def _build_prefill_inputs( + self, + token_ids: list[int], + dev: torch.device, + req_id: str = "default", + ) -> _PrefillInputs: tts = self.tts dtype = self._side_dtype state = self._active_states.get(req_id) @@ -2678,12 +3222,12 @@ def _build_prefill_inputs(self, token_ids: list[int], dev: Any, req_id: str = "d text_mask = torch.cat([rtm.cpu(), ones_t, zeros_p]) audio_mask = torch.cat([ram.cpu(), zeros_t, ones_p]) - return { - "text_token": text_token.unsqueeze(0).to(dev), - "audio_feat": audio_feat.unsqueeze(0).to(dev).to(dtype), - "text_mask": text_mask.unsqueeze(0).to(dev), - "audio_mask": audio_mask.unsqueeze(0).to(dev), - } + return _PrefillInputs( + text_token=text_token.unsqueeze(0).to(dev), + audio_feat=audio_feat.unsqueeze(0).to(dev).to(dtype), + text_mask=text_mask.unsqueeze(0).to(dev), + audio_mask=audio_mask.unsqueeze(0).to(dev), + ) # -------------------- weight loading -------------------- diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py index 491f185a4b8..5898136c997 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py +++ b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py @@ -262,8 +262,8 @@ def __call__( fake_eos = self.static_fake_eos[padded_size][:actual_size].clone() audio_codes = self.static_audio_codes[padded_size][:actual_size].clone() - # Package into expected format: (fake_eos, {"audio": [list of tensors]}) + # Package into expected format: (fake_eos, {"codes": {"audio": [list of tensors]}}) audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) - mm_tokens = {"audio": audio_list} + mm_tokens = {"codes": {"audio": audio_list}} return fake_eos, mm_tokens diff --git a/vllm_omni/model_executor/models/voxtral_tts/pipeline.py b/vllm_omni/model_executor/models/voxtral_tts/pipeline.py index 3bc2743c6cd..e6bc712fe08 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/pipeline.py +++ b/vllm_omni/model_executor/models/voxtral_tts/pipeline.py @@ -16,6 +16,7 @@ VOXTRAL_TTS_PIPELINE = PipelineConfig( model_type="voxtral_tts", + default_deploy_config_name="voxtral_tts.yaml", model_arch="VoxtralTTSForConditionalGeneration", stages=( StagePipelineConfig( diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index d12034c243b..40d4ec62069 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -197,17 +197,33 @@ def sampler(self): return self.model.sampler return Sampler() + @staticmethod + def _is_prefill_step(info_dict: Mapping[str, object], input_ids: torch.Tensor) -> bool: + is_prefill_raw = info_dict.get("_omni_is_prefill") + if isinstance(is_prefill_raw, bool): + return is_prefill_raw + try: + return int(info_dict["_omni_num_computed_tokens"]) < int(info_dict["_omni_prompt_len"]) + except (KeyError, TypeError, ValueError): + return input_ids.shape[0] > 1 + def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict | None): self.post_process_idx = 0 - audio_tokens = info_dict.pop("audio", None) + codes = info_dict.get("codes") + audio_tokens = codes.get("audio") if isinstance(codes, Mapping) else None + if audio_tokens is None: + # Backward compatible with request state written before #4527 moved + # Voxtral TTS feedback frames under ``codes.audio``. + audio_tokens = info_dict.pop("audio", None) if audio_tokens is not None: kwargs = {"audio_tokens": audio_tokens.to(input_ids.device)} multimodal_embeddings = self.model.embed_multimodal(**kwargs) if input_ids[0] == self._audio_token_id: input_embeds = multimodal_embeddings[0] return input_ids, input_embeds, info_dict - voice = info_dict.pop("voice", None) - if voice is not None: + voice = info_dict.get("voice") + if voice is not None and self._is_prefill_step(info_dict, input_ids): + voice = info_dict.pop("voice") if isinstance(voice, list): voice = voice[0] multimodal_embeddings = self.voice_to_embedding[voice].to(input_ids.device).clone().detach() @@ -220,12 +236,15 @@ def tts_preprocess(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, ** def tts_postprocess(self, hidden_states: torch.Tensor, multimodal_outputs: object, **info_dict: object | None): update_dict = {} - if isinstance(multimodal_outputs, Mapping) and "audio" in multimodal_outputs: - assert self.post_process_idx < len(multimodal_outputs["audio"]), ( - f"Expect {self.post_process_idx=} < {len(multimodal_outputs['audio'])=}" - ) - update_dict["audio"] = multimodal_outputs["audio"][self.post_process_idx] - self.post_process_idx += 1 + if isinstance(multimodal_outputs, Mapping): + codes = multimodal_outputs.get("codes") + if isinstance(codes, Mapping) and "audio" in codes: + audio_frames = codes["audio"] + assert self.post_process_idx < len(audio_frames), ( + f"Expect {self.post_process_idx=} < {len(audio_frames)=}" + ) + update_dict["codes"] = {"audio": audio_frames[self.post_process_idx]} + self.post_process_idx += 1 return update_dict def embed_input_ids( @@ -244,9 +263,6 @@ def embed_multimodal(self, **kwargs): # Delegate to generation model for multimodal processing return self.model.embed_multimodal(**kwargs) - def last_index_of(self, list, value): - return len(list) - 1 - list[::-1].index(value) - def forward( self, input_ids: torch.Tensor, diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 2e50e34f32d..370639fb092 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -246,7 +246,6 @@ def __init__( self.n_local_heads: int = args.n_heads self.n_local_kv_heads: int = args.n_kv_heads - self.repeats = self.n_local_heads self.layer_id = layer_id self.head_dim = args.head_dim @@ -276,7 +275,6 @@ def __init__( bias=args.use_biases, ) - self.softmax_scale: float = self.args.head_dim**-0.5 self.repeats = self.n_local_heads // self.n_local_kv_heads def _native_attention( @@ -420,13 +418,6 @@ def __init__( self.acoustic_transformer_args = args assert isinstance(self.acoustic_transformer_args, AcousticTransformerArgs) - # currently assuming always 1 semantic codebook + N acoustic codebooks - self.num_non_acoustic_embeddings = 1 - self.num_acoustic_codebooks = len(self.model_args.get_codebook_sizes()) - self.num_non_acoustic_embeddings - - # flow matching utils - self.sigma = args.sigma - # codebook sizes acoustic_codebook_sizes = self.model_args.get_codebook_sizes( pad_to_multiple=None, include_special_tokens=False @@ -715,7 +706,6 @@ def __call__( ) text_tokens_for_audio = list[torch.Tensor]() - assert audios is not None audio_tokens_pt = list[torch.Tensor]() for audio_token_array in audio_tokens: assert isinstance(audio_token_array, np.ndarray) @@ -1008,7 +998,7 @@ def _parse_and_validate_audio_arrays( if audio_arrays is not None: if not isinstance(audio_arrays, (torch.Tensor, list)): - raise ValueError(f"Incorrect type of images. Got type: {type(audio_arrays)}") + raise ValueError(f"Incorrect type of audio_arrays. Got type: {type(audio_arrays)}") if isinstance(audio_arrays, torch.Tensor) and audio_arrays.dim() == 3: audio_arrays = flatten_bn(audio_arrays) if isinstance(audio_arrays, torch.Tensor): @@ -1053,24 +1043,11 @@ def fake_logits_for_audio_tokens( fake_logits[~is_eos, self.audio_tok_id] = 1.0 return fake_logits - # TODO(chenyo): Remove this - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - text_logits = self.language_model.compute_logits( - hidden_states, - ) - assert text_logits is not None - return text_logits - def compute_mm_logits( self, hidden_states: torch.Tensor, cfg_alpha: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor | None]: - audio_codes = None - mm_tokens = None audio_codes = self.acoustic_transformer( llm_hidden=hidden_states, cfg_alpha=cfg_alpha, @@ -1083,12 +1060,11 @@ def compute_mm_logits( # BxC -> Bx1xC since this is per-step # Make it a list for vllm-omni processing audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) - mm_tokens = {"audio": audio_list} + mm_tokens = {"codes": {"audio": audio_list}} return fake_eos, mm_tokens def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - # fmt: on remapping_rules = [ (r"^acoustic_transformer\.(.*)$", r"\1"), # noqa: E501 (r"^audio_tokenizer\.(.*)$", r"\1"), # noqa: E501 diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py index fc753a58f37..57a86a5e91c 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py @@ -652,7 +652,7 @@ def __init__( def layer_id(self) -> int: return self._layer_id - def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + def forward(self, x: torch.Tensor) -> torch.Tensor: r = self.attention.forward(self.attention_norm(x)) if self.post_attention_norm is not None: r = self.post_attention_norm(r) @@ -757,7 +757,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # transformer layer_args = deepcopy(args) layer_args.attn_sliding_window_size = cur_window_size - assert layer_args.qk_norm, "qk_norm must be True for decoder" + assert layer_args.qk_norm, "qk_norm must be True for encoder" encoder_transformer = Transformer( args=layer_args, n_layers=n_layers, diff --git a/vllm_omni/model_executor/models/wan2_2/__init__.py b/vllm_omni/model_executor/models/wan2_2/__init__.py new file mode 100644 index 00000000000..208f01a7cb5 --- /dev/null +++ b/vllm_omni/model_executor/models/wan2_2/__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/wan2_2/pipeline.py b/vllm_omni/model_executor/models/wan2_2/pipeline.py new file mode 100644 index 00000000000..7059845faaf --- /dev/null +++ b/vllm_omni/model_executor/models/wan2_2/pipeline.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Wan2.2 TI2V pipeline topology.""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +WAN2_2_TI2V_PIPELINE = PipelineConfig( + model_type="wan2_2_ti2v", + model_arch="WanPipeline", + diffusers_class_name="WanPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="video", + model_arch="WanPipeline", + ), + ), +) diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml deleted file mode 100644 index 097afc020c6..00000000000 --- a/vllm_omni/model_executor/stage_configs/hunyuan_video_15_dit_fp8.yaml +++ /dev/null @@ -1,38 +0,0 @@ -# Stage config for running HunyuanVideo-1.5 DiT with ModelOpt FP8 auto-detect. -# Single GPU. Bump `tensor_parallel_size` and `devices` for multi-GPU TP. -# -# Use with a ModelOpt FP8 checkpoint (e.g. produced by -# examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py). - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - model_stage: dit - model_class_name: HunyuanVideo15Pipeline - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - # HunyuanVideo-1.5 peak VRAM is dominated by the single-shot VAE decode of - # the full video (the DiT contributes little), so tiling the decode is the - # main memory lever — far more than FP8 (e.g. 480x832x33: ~117 -> ~47 GiB). - # Default it on here; override with --vae-use-tiling/false if you have headroom. - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 1 - - final_output: true - final_output_type: video - is_comprehension: false - default_sampling_params: - seed: 42 - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml deleted file mode 100644 index f4f523ac599..00000000000 --- a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml +++ /dev/null @@ -1,42 +0,0 @@ -# Stage config for running MammothModa2 with full AR→DiT pipeline (text-to-image). -# Stage 0: AR (multimodal understanding + latent generation) -# Stage 1: DiT (latent → image) -# -# For text/image understanding tasks (text output only), use mammoth_moda2_ar.yaml. - -stage_args: - - stage_id: 0 - runtime: - devices: "0" - engine_args: - model_stage: ar - max_num_seqs: 100 - model_arch: MammothModa2ForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 8192 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - final_output: false - - - stage_id: 1 - runtime: - devices: "0" - engine_args: - model_stage: dit - max_num_seqs: 1 - model_arch: MammothModa2ForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - gpu_memory_utilization: 0.3 - enforce_eager: true - trust_remote_code: true - engine_output_type: image - enable_prefix_caching: false - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit - final_output: true - final_output_type: image diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml deleted file mode 100644 index d455c7973e3..00000000000 --- a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml +++ /dev/null @@ -1,24 +0,0 @@ -# Stage config for running MammothModa2 with single AR stage (understanding / summarization). -# Stage 0: AR (multimodal understanding + text generation) -# -# Use this config for image/video understanding tasks where only text output is needed. -# For text-to-image generation (AR → DiT), use mammoth_moda2.yaml instead. - -stage_args: - - stage_id: 0 - runtime: - devices: "0" - engine_args: - model_stage: ar - max_num_seqs: 16 - model_arch: MammothModa2ForConditionalGeneration - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - max_model_len: 8192 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - engine_output_type: text - enable_prefix_caching: false - final_output: true - final_output_type: text diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml deleted file mode 100644 index 493c05f004d..00000000000 --- a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml +++ /dev/null @@ -1,61 +0,0 @@ - -# MiMo Audio deploy: 2-stage thinker+talker → code2wav. -# -# Default mode is async-chunk streaming on a single GPU (both stages on -# device 0) through SharedMemoryConnector. For the legacy 2-GPU sync -# pipeline (stage 1 on a second card), pass ``--no-async-chunk -# --stage-1-devices 1 --stage-1-max-model-len 18192 -# --stage-1-max-num-batched-tokens 18192`` to the serve CLI. -async_chunk: true -dtype: bfloat16 - -connectors: - connector_of_shared_memory: - name: SharedMemoryConnector - extra: - shm_threshold_bytes: 365536 - codec_streaming: true - connector_get_sleep_s: 0.001 - connector_get_max_wait_first_chunk: 3000 - connector_get_max_wait: 300 - codec_chunk_frames: 30 - codec_left_context_frames: 40 - -stages: - - stage_id: 0 - max_num_seqs: 1 - gpu_memory_utilization: 0.5 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: 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.6 - top_p: 0.95 - top_k: 50 - max_tokens: 18192 - seed: 42 - repetition_penalty: 1.1 - - - stage_id: 1 - max_num_seqs: 1 - gpu_memory_utilization: 0.4 - enforce_eager: true - trust_remote_code: 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: 18192 - seed: 42 diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2.yaml deleted file mode 100644 index b560fea9b50..00000000000 --- a/vllm_omni/model_executor/stage_configs/step_audio_2.yaml +++ /dev/null @@ -1,122 +0,0 @@ -# Stage config for Step-Audio2 with vLLM-Omni -# -# Two-stage pipeline: -# Stage 0 (Thinker): Audio understanding → Text + Audio tokens -# Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz) -# -# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both) - -# Model configuration reference (for documentation) -# Actual values are set in hf_overrides below -# -# Token configuration: -# text_max: 151688 # Maximum text token ID (inclusive) -# audio_start: 151696 # First audio token ID in absolute vocabulary -# audio_vocab_size: 6562 # Total number of audio tokens -# audio_eos: 6561 # Audio EOS token ID (relative to audio_start) -# audio_patch_token_id: 151690 # <audio_patch> placeholder token ID -# -# Encoder configuration: -# n_mels: 128 # Number of mel frequency bins -# n_audio_ctx: 1500 # Audio context length -# n_audio_state: 512 # Audio encoder hidden state dimension -# n_audio_head: 8 # Number of attention heads -# n_audio_layer: 6 # Number of encoder layers -# kernel_size: 3 # Convolution kernel size for adapter -# adapter_stride: 2 # Stride for adapter convolution - -stage_args: - # Stage 0: Thinker (Audio Understanding + Token Generation) - - stage_id: 0 - runtime: - process: true # Run in separate process - devices: "0,1,2,3" # Use 3 GPUs for tensor parallelism - max_batch_size: 1 # Single sample at a time - engine_args: - model_stage: step_audio2_thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - tensor_parallel_size: 4 # Split model across 3 GPUs - # NOTE: All default values are defined in step_audio2_constants.py (DEFAULT_MODEL_CONFIG) - # Only override here if your model differs from the defaults - # - # Default values (for reference, from constants.py): - # hidden_size: 4096 # Qwen2-7B (change to 1536 for 1.5B, 896 for 0.5B) - # text_max: 151688 - # audio_start: 151696 - # audio_vocab_size: 6562 - # audio_eos: 6561 - # audio_patch_token_id: 151690 - # n_mels: 128 - # n_audio_ctx: 1500 - # n_audio_state: 512 - # n_audio_head: 8 - # n_audio_layer: 6 - # kernel_size: 3 - # adapter_stride: 2 - # - # Uncomment and modify only if needed: - # hidden_size: 4096 - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.8 - enforce_eager: true # Required for now - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 8192 - engine_output_type: text # Outputs text + audio tokens - is_comprehension: true # This stage does comprehension - final_output: true # Can be final output (for ASR mode) - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 1024 - seed: 42 - detokenize: true - repetition_penalty: 1.05 - - # Stage 1: Token2Wav (Audio Synthesis) - - stage_id: 1 - runtime: - process: true - devices: "3" # Use GPU 3 (separate from Stage 0's 0-2) - max_batch_size: 1 - engine_args: - model_stage: token2wav - model_arch: StepAudio2Token2WavModel - hf_overrides: - architectures: ["StepAudio2Token2WavForConditionalGeneration"] - worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - - gpu_memory_utilization: 0.3 # Token2Wav needs less memory - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - max_num_batched_tokens: 4096 - engine_output_type: audio # Outputs waveform - engine_input_source: [0] # Takes input from Stage 0 - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav - final_output: true # Final output stage - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: false - -# Runtime configuration -runtime: - enabled: true - defaults: - window_size: -1 # Wait for full input before triggering - max_inflight: 1 # Process one at a time - - edges: - - from: 0 # Thinker → Token2Wav - to: 1 - window_size: -1 # Wait for complete output from Thinker diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml deleted file mode 100644 index 1cecfddc556..00000000000 --- a/vllm_omni/model_executor/stage_configs/step_audio_2_asr.yaml +++ /dev/null @@ -1,35 +0,0 @@ -# Stage config for Step-Audio2 ASR (speech-to-text only) -# -# Single-stage pipeline: only Thinker, no Token2Wav. -# Use this for audio_to_text tasks where no audio output is needed. - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1" - max_batch_size: 1 - engine_args: - model_stage: step_audio2_thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.8 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 2048 - tensor_parallel_size: 2 - engine_output_type: text - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 1024 - seed: 42 - detokenize: true - repetition_penalty: 1.05 diff --git a/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml deleted file mode 100644 index 98c4840721d..00000000000 --- a/vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml +++ /dev/null @@ -1,78 +0,0 @@ -# Stage config for Step-Audio2 with async_chunk support -# -# Two-stage pipeline with streaming between stages: -# Stage 0 (Thinker): Audio understanding → Text + Audio tokens (streamed in chunks) -# Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz) -# -# Key difference from step_audio_2.yaml: -# - async_chunk: true enables streaming from Thinker to Token2Wav -# - Thinker uses custom_process_next_stage_input_func to send audio token chunks -# - Token2Wav receives chunks incrementally instead of waiting for full output -# -# Hardware requirement: 1-2 GPUs (can run on single GPU with devices: "0" for both) - -async_chunk: true -stage_args: - # Stage 0: Thinker (Audio Understanding + Token Generation) - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1" - max_batch_size: 1 - engine_args: - model_stage: step_audio2_thinker - model_arch: StepAudio2ThinkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.8 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 1024 - tensor_parallel_size: 2 - engine_output_type: latent - custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav_async_chunk - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.7 - top_p: 0.9 - top_k: -1 - max_tokens: 1024 - seed: 42 - detokenize: true - repetition_penalty: 1.05 - - # Stage 1: Token2Wav (Audio Synthesis) - - stage_id: 1 - stage_type: llm - runtime: - process: true - devices: "3" - max_batch_size: 1 - engine_args: - model_stage: token2wav - model_arch: StepAudio2Token2WavModel - worker_type: generation - hf_overrides: - architectures: ["StepAudio2Token2WavForConditionalGeneration"] - scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler - gpu_memory_utilization: 0.3 - enforce_eager: true - trust_remote_code: true - async_scheduling: false - enable_prefix_caching: false - max_num_batched_tokens: 2048 - engine_output_type: audio - engine_input_source: [0] - final_output: true - final_output_type: audio - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 1 - seed: 42 - detokenize: false diff --git a/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml b/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml deleted file mode 100644 index 1b1c40537a3..00000000000 --- a/vllm_omni/model_executor/stage_configs/wan2_2_ti2v_dit_fp8.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# Stage config for running Wan2.2 TI2V-5B DiT with ModelOpt FP8 auto-detect. -# Single GPU (TI2V-5B fits 80GB BF16; FP8 drops by ~half). -# For the A14B MoE variants, bump `tensor_parallel_size` and `devices`. -# -# Use with a ModelOpt FP8 checkpoint (e.g. produced by -# examples/quantization/quantize_wan2_2_modelopt_fp8.py). - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - devices: "0" - max_batch_size: 1 - engine_args: - model_stage: dit - # Wan2.2 TI2V-5B is served by the base Wan22Pipeline, which auto-detects - # expand_timesteps (TI2V-5B mode) from model_index.json. (The standalone - # Wan22TI2VPipeline was consolidated into Wan22Pipeline upstream.) - model_class_name: Wan22Pipeline - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - # Tile the VAE decode by default. Wan2.2's 480p peak is DiT-dominated (so - # FP8 is the main lever here), but tiling adds headroom for higher - # resolutions / longer clips at negligible cost. Override with --vae-use-tiling/false. - vae_use_tiling: true - parallel_config: - tensor_parallel_size: 1 - - final_output: true - final_output_type: video - is_comprehension: false - default_sampling_params: - seed: 42 - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_input_processors/aura_omni.py b/vllm_omni/model_executor/stage_input_processors/aura_omni.py new file mode 100644 index 00000000000..789ed2c6416 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/aura_omni.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage processors for the AURA Omni pipeline.""" + +from __future__ import annotations + +import math +import os +from pathlib import Path +from typing import Any + +import soundfile as sf + +from vllm_omni.inputs.data import OmniTokensPrompt +from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import ( + PRECOMPUTED_TEXT_IDS_KEY, +) + +DEFAULT_AURA_SYSTEM_PROMPT = ( + "You are receiving a live video stream where the final frame is the present moment. " + "Respond only when a response is needed based on the user's message or the visual context. " + "Otherwise, output '<|silent|>' to signify silence. Respond in Chinese." +) + +SILENT_TEXT = "<|silent|>" +QWEN_IM_START_ID = 151644 +QWEN_IM_END_ID = 151645 +QWEN_ASSISTANT_ID = 77091 +QWEN_NEWLINE_ID = 198 +QWEN_ASSISTANT_PREFIX_IDS = [QWEN_IM_START_ID, QWEN_ASSISTANT_ID, QWEN_NEWLINE_ID] +QWEN_ASSISTANT_SUFFIX_IDS = [ + QWEN_IM_END_ID, + QWEN_NEWLINE_ID, + QWEN_IM_START_ID, + QWEN_ASSISTANT_ID, + QWEN_NEWLINE_ID, +] +DEFAULT_QWEN3_TTS_REF_AUDIO = "vllm-omni/tests/assets/qwen3_tts/clone_2.wav" +DEFAULT_QWEN3_TTS_REF_TEXT = ( + "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you." +) + + +def default_qwen3_tts_ref_audio_path() -> str: + """Return absolute path to the bundled ``clone_2.wav`` reference asset.""" + bundled = Path(__file__).resolve().parents[3] / "tests" / "assets" / "qwen3_tts" / "clone_2.wav" + if bundled.is_file(): + return str(bundled) + return DEFAULT_QWEN3_TTS_REF_AUDIO + + +def _as_list(value: Any) -> list[Any]: + if value is None: + return [] + return value if isinstance(value, list) else [value] + + +def _as_prompt_dict(prompt_item: Any) -> dict[str, Any]: + return prompt_item if isinstance(prompt_item, dict) else {} + + +def _first_value(value: Any, default: Any = None) -> Any: + if isinstance(value, list): + return value[0] if value else default + return default if value is None else value + + +def _first_bool(value: Any, default: bool = False) -> bool: + value = _first_value(value, default) + if isinstance(value, str): + return value.strip().lower() in {"1", "true", "yes", "on"} + return bool(value) + + +def _normalize_qwen3_tts_speaker(speaker: Any) -> Any: + if not isinstance(speaker, str): + return speaker + speaker = speaker.strip() + if not speaker: + return speaker + if "_" in speaker: + return speaker + return speaker[0].upper() + speaker[1:].lower() + + +def _extract_output(source_output: Any) -> Any: + outputs = getattr(source_output, "outputs", None) + if isinstance(outputs, list) and outputs: + return outputs[0] + return source_output + + +def _extract_text(source_output: Any) -> str: + output = _extract_output(source_output) + cumulative_text = getattr(output, "cumulative_text", None) + if isinstance(cumulative_text, str) and cumulative_text: + return cumulative_text + text = getattr(output, "text", None) + if isinstance(text, str): + return text + mm = getattr(output, "multimodal_output", None) + if isinstance(mm, dict): + for key in ("text", "transcript", "asr_text"): + value = mm.get(key) + if isinstance(value, str): + return value + if isinstance(value, list) and value and isinstance(value[0], str): + return value[0] + return "" + + +def _extract_token_ids(source_output: Any) -> list[int]: + output = _extract_output(source_output) + token_ids = getattr(output, "cumulative_token_ids", None) + if isinstance(token_ids, list): + return [int(token_id) for token_id in token_ids if isinstance(token_id, int)] + return [] + + +def _trim_aura_response_token_ids(token_ids: list[int]) -> list[int]: + ids = list(token_ids) + if ids[: len(QWEN_ASSISTANT_PREFIX_IDS)] == QWEN_ASSISTANT_PREFIX_IDS: + ids = ids[len(QWEN_ASSISTANT_PREFIX_IDS) :] + if QWEN_IM_END_ID in ids: + ids = ids[: ids.index(QWEN_IM_END_ID)] + while ids and ids[-1] in {QWEN_IM_START_ID, QWEN_IM_END_ID, QWEN_NEWLINE_ID}: + ids.pop() + return ids + + +def _qwen3_tts_assistant_token_ids_from_aura(source_output: Any) -> list[int]: + content_ids = _trim_aura_response_token_ids(_extract_token_ids(source_output)) + if not content_ids: + return [] + return QWEN_ASSISTANT_PREFIX_IDS + content_ids + QWEN_ASSISTANT_SUFFIX_IDS + + +def _source_prompt_by_request_id(source_outputs: list[Any], prompt: Any) -> dict[str, dict[str, Any]]: + prompts = _as_list(prompt) + return { + str(getattr(source_output, "request_id", idx)): _as_prompt_dict(prompt_item) + for idx, (source_output, prompt_item) in enumerate(zip(source_outputs, prompts)) + } + + +def _vision_placeholder(multi_modal_data: dict[str, Any]) -> str: + if "video" in multi_modal_data: + return "<|vision_start|><|video_pad|><|vision_end|>" + if "image" in multi_modal_data: + return "<|vision_start|><|image_pad|><|vision_end|>" + return "" + + +def _vision_multimodal_data(multi_modal_data: dict[str, Any]) -> dict[str, Any]: + return {key: value for key, value in multi_modal_data.items() if key in {"image", "video"}} + + +def _aura_prompt(system_prompt: str, transcript: str, multi_modal_data: dict[str, Any]) -> str: + vision = _vision_placeholder(multi_modal_data) + query = transcript.strip() + user_body = f"{vision}{query}" if query else vision + return ( + f"<|im_start|>system\n{system_prompt}<|im_end|>\n" + f"<|im_start|>user\n{user_body}<|im_end|>\n" + "<|im_start|>assistant\n" + ) + + +def asr2aura( + source_outputs: list[Any], + prompt: Any = None, + requires_multimodal_data: bool = True, +) -> list[dict[str, Any]]: + """Build AURA Qwen3-VL prompts from ASR transcripts and original video payloads.""" + prompt_by_request_id = _source_prompt_by_request_id(source_outputs, prompt) + next_inputs: list[dict[str, Any]] = [] + for idx, source_output in enumerate(source_outputs): + src_prompt = prompt_by_request_id.get(str(getattr(source_output, "request_id", idx)), {}) + additional_info = src_prompt.get("additional_information") or {} + system_prompt = _first_value(additional_info.get("aura_system_prompt"), DEFAULT_AURA_SYSTEM_PROMPT) + transcript = _extract_text(source_output) + multi_modal_data = {} + source_multi_modal_data = src_prompt.get("multi_modal_data") or {} + if isinstance(source_multi_modal_data, dict): + multi_modal_data.update(source_multi_modal_data) + deferred_multi_modal_data = additional_info.get("deferred_multi_modal_data") or {} + if isinstance(deferred_multi_modal_data, dict): + multi_modal_data.update(deferred_multi_modal_data) + multi_modal_data = _vision_multimodal_data(multi_modal_data) + + next_input: dict[str, Any] = { + "prompt": _aura_prompt(str(system_prompt), transcript, multi_modal_data), + } + if requires_multimodal_data: + next_input["multi_modal_data"] = multi_modal_data + if src_prompt.get("mm_processor_kwargs") is not None: + next_input["mm_processor_kwargs"] = src_prompt.get("mm_processor_kwargs") + next_inputs.append(next_input) + return next_inputs + + +def _estimate_ref_code_len_from_ref_audio(ref_audio: Any) -> int | None: + """Estimate Qwen3-TTS ref_code length from a ref-audio payload. + + For Qwen3-TTS 12Hz models, code length is approximately: + ceil(duration_seconds * 12.5) + i.e. one codec frame per 1920 samples at 24kHz. + """ + + codec_frame_rate = 24000.0 / 1920.0 + + # Unwrap common list wrappers. + item = ref_audio + while isinstance(item, list) and item: + item = item[0] + + # Accept tuple/list like (wav, sr). + if isinstance(item, (tuple, list)) and len(item) == 2 and isinstance(item[1], (int, float)): + wav, sr = item + sr_i = int(sr) + if sr_i <= 0: + return None + if hasattr(wav, "__len__"): + n_samples = len(wav) + elif hasattr(wav, "shape"): + shape = getattr(wav, "shape", None) + if not shape: + return None + n_samples = shape[-1] if len(shape) > 1 else shape[0] + else: + return None + if n_samples <= 0: + return None + return max(1, int(math.ceil((float(n_samples) / float(sr_i)) * codec_frame_rate))) + + # Accept file path (wav only). + if isinstance(item, str) and item: + audio_path = item + if not os.path.isfile(audio_path) or not audio_path.lower().endswith(".wav"): + return None + try: + info = sf.info(audio_path) + n_frames = int(info.frames) + sr = int(info.samplerate) + if n_frames <= 0 or sr <= 0: + return None + return max(1, int(math.ceil((float(n_frames) / float(sr)) * codec_frame_rate))) + except Exception: + return None + + return None + + +def _estimate_tts_prompt_len_from_token_ids( + token_ids: list[int], + *, + task_type: str = "Base", + language: str = "Chinese", + instruct: str = "", + x_vector_only_mode: bool = False, + non_streaming_mode: bool | None = None, + ref_code_len: int | None = None, +) -> int: + """Estimate Talker prefill length from prompt structure. + + This mirrors Qwen3-TTS prompt assembly at length level: + prompt_len = instruct_len + role_len + codec_prefix_len + text/icl term + """ + + # Official defaults: Base -> streaming, others -> non-streaming. + if non_streaming_mode is None: + non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign") + + # We do not have tokenizer here; use char length as a monotonic proxy. + instruct_len = len(instruct.strip()) if isinstance(instruct, str) else 0 + assistant_len = max(0, len(token_ids)) + + # role_len = 3; codec_prefix_len = (prefill_len + speaker_len + 2) - 1 + # prefill_len = 4 when language_id exists else 3. Use non-auto language as + # the language-id-present proxy. + has_language_id = isinstance(language, str) and language.strip().lower() != "auto" + prefill_len = 4 if has_language_id else 3 + speaker_len = 1 if task_type in ("CustomVoice", "Base") else 0 + base_len = instruct_len + 3 + (prefill_len + speaker_len + 2 - 1) + if task_type in ("CustomVoice", "VoiceDesign"): + if non_streaming_mode: + prompt_len = base_len + max(0, assistant_len - 6) + else: + prompt_len = base_len + 1 + return int(prompt_len) + + if task_type == "Base": + in_context_mode = not bool(x_vector_only_mode) + if in_context_mode and ref_code_len is not None: + codec_lens = 1 + int(ref_code_len) + if non_streaming_mode: + # Exact non-streaming ICL needs ref_ids token length; unavailable + # in this processor. Keep a conservative upper estimate. + prompt_len = base_len + codec_lens + max(0, assistant_len - 8) + 1 + else: + # Streaming ICL exact length term: 1 + ref_code_len + prompt_len = base_len + codec_lens + else: + # Base x-vector-only (or missing ref_code length) follows CV shape. + if non_streaming_mode: + prompt_len = base_len + max(0, assistant_len - 6) + else: + prompt_len = base_len + 1 + return int(prompt_len) + + # Defensive fallback for unknown task types. + return int(base_len + max(assistant_len, 1)) + + +def aura2tts( + source_outputs: list[Any], + prompt: Any = None, + requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + """Convert AURA text output into Qwen3-TTS Talker requests.""" + del requires_multimodal_data + prompt_by_request_id = _source_prompt_by_request_id(source_outputs, prompt) + next_inputs: list[OmniTokensPrompt] = [] + for idx, source_output in enumerate(source_outputs): + text = _extract_text(source_output).strip() + if not text or text == SILENT_TEXT: + continue + + src_prompt = prompt_by_request_id.get(str(getattr(source_output, "request_id", idx)), {}) + additional_info = src_prompt.get("additional_information") or {} + task_type = _first_value(additional_info.get("tts_task_type"), "Base") + language = _first_value(additional_info.get("tts_language"), "English") + instruct = _first_value(additional_info.get("tts_instruct"), "") + x_vector_only_mode = _first_bool(additional_info.get("tts_x_vector_only_mode"), False) + non_streaming_mode_raw = _first_value(additional_info.get("tts_non_streaming_mode"), None) + non_streaming_mode = non_streaming_mode_raw if isinstance(non_streaming_mode_raw, bool) else None + ref_code_len_raw = _first_value(additional_info.get("tts_ref_code_length"), None) + ref_code_len = int(ref_code_len_raw) if isinstance(ref_code_len_raw, int) else None + ref_audio = None + ref_text = None + if task_type == "Base" and not x_vector_only_mode and ref_code_len is None: + ref_audio = _first_value(additional_info.get("tts_ref_audio"), None) + ref_code_len = _estimate_ref_code_len_from_ref_audio(ref_audio) + + assistant_token_ids_for_len = _qwen3_tts_assistant_token_ids_from_aura(source_output) + pass_token_ids = _first_bool(additional_info.get("tts_pass_token_ids"), False) + tts_info = { + "task_type": [task_type], + "language": [language], + "instruct": [instruct], + "max_new_tokens": [int(_first_value(additional_info.get("tts_max_new_tokens"), 2048))], + } + if pass_token_ids and assistant_token_ids_for_len: + tts_info[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_token_ids_for_len] + else: + tts_info["text"] = [text] + if ref_code_len is not None: + tts_info["ref_code_length"] = [int(ref_code_len)] + prompt_len = _estimate_tts_prompt_len_from_token_ids( + assistant_token_ids_for_len if assistant_token_ids_for_len else [0] * max(0, len(text)), + task_type=str(task_type), + language=str(language), + instruct=str(instruct), + x_vector_only_mode=x_vector_only_mode, + non_streaming_mode=non_streaming_mode, + ref_code_len=ref_code_len, + ) + + if task_type == "Base": + ref_audio = ref_audio or _first_value(additional_info.get("tts_ref_audio"), None) + ref_text = _first_value(additional_info.get("tts_ref_text"), None) + if not ref_audio or not ref_text: + raise ValueError("AURA Base TTS requires tts_ref_audio and tts_ref_text.") + x_vector_only_mode = _first_bool(additional_info.get("tts_x_vector_only_mode"), False) + tts_info["ref_audio"] = [ref_audio] + tts_info["ref_text"] = [ref_text] + tts_info["x_vector_only_mode"] = [x_vector_only_mode] + elif task_type == "CustomVoice": + tts_info["speaker"] = [ + _normalize_qwen3_tts_speaker(_first_value(additional_info.get("tts_speaker"), "Vivian")) + ] + next_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * prompt_len, + additional_information=tts_info, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return next_inputs diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index 072fc9cf6a0..2956be94f25 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py @@ -21,8 +21,6 @@ logger = init_logger(__name__) -_COSYVOICE3_SPEECH_TOKEN_SIZE = 6561 - def _build_prompt_embed_struct(prompt_payload: dict[str, Any]) -> EmbeddingsStruct | None: """Wrap prompt_payload's flat speech_token/speech_feat/embedding tensors into EmbeddingsStruct.""" @@ -85,20 +83,6 @@ def _prompt_speech_token_ids(multi_modal_data: dict[str, Any]) -> list[int]: return _to_token_id_list(speech_token) -def _has_speech_stop_token(output_ids: list[Any]) -> bool: - return any(token_id >= _COSYVOICE3_SPEECH_TOKEN_SIZE for token_id in _to_token_id_list(output_ids)) - - -def _set_non_stream_prompt_trim(additional_info: dict[str, Any], prompt_speech_len: int) -> None: - if prompt_speech_len <= 0: - return - meta = additional_info.get("meta") - if not isinstance(meta, dict): - meta = {} - additional_info["meta"] = meta - meta["talker_prefill_offset"] = prompt_speech_len - - def _to_cpu_tensor(x: Any) -> torch.Tensor | None: if isinstance(x, list): if not x: @@ -135,32 +119,6 @@ def _decode_additional_information(raw_info: Any) -> dict[str, Any]: return decoded -def text2flow( - source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt = None, - _requires_multimodal_data: bool = True, -): - """Build stage-1 inputs by prefixing stage-0 prompt ids to its outputs.""" - engine_inputs: list[OmniTokensPrompt] = [] - for source_output in source_outputs: - output = source_output.outputs[0] - multi_modal_data = output.multimodal_output - if multi_modal_data is None: - raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}") - - prefix_ids = _ensure_list(source_output.prompt_token_ids) - raw_output_ids = _ensure_list(output.cumulative_token_ids) - prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data) - output_ids = _strip_prompt_prefix(raw_output_ids, prefix_ids) - output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids) - additional_info = dict(multi_modal_data) - if _has_speech_stop_token(raw_output_ids): - _set_non_stream_prompt_trim(additional_info, len(prompt_speech_ids)) - additional_info.setdefault("ids", {})["prompt"] = prefix_ids - engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info)) - return engine_inputs - - def talker2code2wav_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, @@ -389,8 +347,6 @@ def text2flow_token_only( prompt_speech_ids = _prompt_speech_token_ids(multi_modal_data) output_ids = _strip_prompt_prefix(output_ids, prompt_speech_ids) additional_info: dict[str, Any] = dict(multi_modal_data) - if _has_speech_stop_token(raw_output_ids): - _set_non_stream_prompt_trim(additional_info, len(prompt_speech_ids)) additional_info.setdefault("ids", {})["prompt"] = prefix_ids engine_inputs.append( OmniTokensPrompt( diff --git a/vllm_omni/model_executor/stage_input_processors/covo_audio.py b/vllm_omni/model_executor/stage_input_processors/covo_audio.py index 7b5ed8c266d..a15cca012c6 100644 --- a/vllm_omni/model_executor/stage_input_processors/covo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/covo_audio.py @@ -23,33 +23,6 @@ def _filter_audio_codes(token_ids: list[int]) -> list[int]: return audio_codes -def llm2code2wav( - source_outputs: list[Any], - prompt: Any = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """Legacy orchestrator-path builder (retained for async_chunk + back-compat). - - The non-async-chunk path now goes through ``llm2code2wav_token_only`` + - worker connector + ``llm2code2wav_full_payload``. - """ - talker_outputs = source_outputs - code2wav_inputs = [] - - for talker_output in talker_outputs: - output = talker_output.outputs[0] - audio_codes = _filter_audio_codes(list(output.token_ids)) - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=audio_codes, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return code2wav_inputs - - def llm2code2wav_token_only( source_outputs: list[Any], prompt: Any = None, diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index 668ff2ed485..f428c34a454 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -114,37 +114,6 @@ def _extract_last_frame(multimodal_output: dict[str, Any]) -> torch.Tensor | Non raise ValueError(f"Invalid audio_codes shape for Fish Speech async_chunk: {tuple(audio_codes.shape)}") -def slow_ar_to_dac_decoder( - source_outputs: list[Any], - _prompt: Any = None, - _requires_multimodal_data: bool = False, -) -> list[Any]: - """Non-async processor: wait for Slow AR to finish, then pass all codes to DAC decoder.""" - from vllm_omni.inputs.data import OmniTokensPrompt - - slow_ar_outputs = source_outputs - dac_inputs: list[OmniTokensPrompt] = [] - - for output in slow_ar_outputs: - out = output.outputs[0] - # audio_codes shape: [num_frames, num_codebooks] - audio_codes = out.multimodal_output["audio_codes"].to(torch.long) - # Filter zero-padded frames. - valid_mask = audio_codes.any(dim=1) - audio_codes = audio_codes[valid_mask] - # Codebook-major flat: [num_codebooks * num_frames] - codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist() - dac_inputs.append( - OmniTokensPrompt( - prompt_token_ids=codec_codes, - multi_modal_data=None, - mm_processor_kwargs=None, - additional_information=None, - ) - ) - return dac_inputs - - def slow_ar_to_dac_decoder_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, diff --git a/vllm_omni/model_executor/stage_input_processors/glm_image.py b/vllm_omni/model_executor/stage_input_processors/glm_image.py index bf644abd0e3..d31d8c44169 100644 --- a/vllm_omni/model_executor/stage_input_processors/glm_image.py +++ b/vllm_omni/model_executor/stage_input_processors/glm_image.py @@ -217,192 +217,192 @@ def ar2diffusion( prompt: OmniTokensPrompt | TextPrompt | list | None = None, requires_multimodal_data: bool = False, streaming_context: Any | None = None, -) -> list[dict[str, Any]]: +) -> dict[str, Any] | None: """Process AR stage outputs to create Diffusion stage inputs. - This processor accepts the stage-pool transition interface: - ``ar2diffusion(source_outputs, prompt, requires_multimodal_data)``. + GLM-Image only produces one downstream diffusion request per AR request. + ``source_outputs`` may still include CFG companion outputs, but only the + first AR output is used to build the diffusion payload. """ del streaming_context _t_total = time.perf_counter() - ar_outputs = source_outputs - diffusion_inputs = [] - - # Normalize prompt to list - if not isinstance(prompt, list): - prompt = [prompt] if prompt is not None else [{}] - - for i, ar_output in enumerate(ar_outputs): - _t_req = time.perf_counter() - output = ar_output.outputs[0] - generated_token_ids = output.cumulative_token_ids - - # Get original prompt info - original_prompt = prompt[i] if i < len(prompt) else {} - if isinstance(original_prompt, dict): - pass - elif hasattr(original_prompt, "_asdict"): - original_prompt = original_prompt._asdict() - elif hasattr(original_prompt, "__dict__"): - original_prompt = vars(original_prompt) - else: - original_prompt = {} - - mm_processor_kwargs = original_prompt.get("mm_processor_kwargs") - - def _coerce_dim(v: Any, default: int) -> int: - try: - iv = int(v) - return iv if iv > 0 else default - except (TypeError, ValueError): - return default - - # Prefer GLM-Image target size from mm_processor_kwargs (set by serving layer), - # then fall back to top-level fields for backward compatibility. - height = _coerce_dim( - mm_processor_kwargs.get("target_h") if isinstance(mm_processor_kwargs, dict) else None, - _coerce_dim(original_prompt.get("height"), 1024), - ) - width = _coerce_dim( - mm_processor_kwargs.get("target_w") if isinstance(mm_processor_kwargs, dict) else None, - _coerce_dim(original_prompt.get("width"), 1024), - ) - text_prompt = original_prompt.get("prompt", "") - - # Detect i2i mode. - # Prefer normalized prompt multi_modal_data source-image presence, with - # multimodal output as secondary signal. - _t_mode = time.perf_counter() - is_i2i = False + if not source_outputs: + return None - prompt_modalities = original_prompt.get("modalities") - if isinstance(prompt_modalities, list) and "img2img" in prompt_modalities: - is_i2i = True + ar_output = source_outputs[0] + _t_req = time.perf_counter() + output = ar_output.outputs[0] + generated_token_ids = output.cumulative_token_ids - prompt_mm_data = original_prompt.get("multi_modal_data") - if _has_source_image(prompt_mm_data): - is_i2i = True + if isinstance(prompt, list): + original_prompt = prompt[0] if prompt else {} + elif prompt is not None: + original_prompt = prompt + else: + original_prompt = {} + + if isinstance(original_prompt, dict): + pass + elif hasattr(original_prompt, "_asdict"): + original_prompt = original_prompt._asdict() + elif hasattr(original_prompt, "__dict__"): + original_prompt = vars(original_prompt) + else: + original_prompt = {} - if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: - mm_output = ar_output.multimodal_output - if isinstance(mm_output, Mapping) and mm_output.get("ids", {}).get("prior_image") is not None: - is_i2i = True - _dt_mode = (time.perf_counter() - _t_mode) * 1000 + mm_processor_kwargs = original_prompt.get("mm_processor_kwargs") - # Parse and upsample prior tokens - _t_parse = time.perf_counter() + def _coerce_dim(v: Any, default: int) -> int: try: - prior_token_ids, pixel_h, pixel_w = _parse_generated_tokens( - generated_token_ids, - height, - width, - is_i2i=is_i2i, - ) - except ValueError as e: - logger.warning( - "[ar2diffusion] Request %s: skip due to token parse failure: %s " - "(target=%sx%s, mode=%s, raw_tokens=%s, tail=%s)", - i, - e, - height, - width, - "i2i" if is_i2i else "t2i", - len(generated_token_ids), - generated_token_ids[-8:] if len(generated_token_ids) >= 8 else generated_token_ids, - ) - continue - _dt_parse = (time.perf_counter() - _t_parse) * 1000 - - # Get prior_token_image_ids from AR model output (for i2i mode) - # This contains VQ-VAE tokens from input image, used for KV cache conditioning - # NOTE: multimodal_output is attached to ar_output (RequestOutput), NOT output (CompletionOutput) - _t_prior_img = time.perf_counter() - prior_token_image_ids = None - - # Check ar_output (RequestOutput) for multimodal_output - this is the correct location - if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: - mm_output = ar_output.multimodal_output + iv = int(v) + return iv if iv > 0 else default + except (TypeError, ValueError): + return default + + # Prefer GLM-Image target size from mm_processor_kwargs (set by serving layer), + # then fall back to top-level fields for backward compatibility. + height = _coerce_dim( + mm_processor_kwargs.get("target_h") if isinstance(mm_processor_kwargs, dict) else None, + _coerce_dim(original_prompt.get("height"), 1024), + ) + width = _coerce_dim( + mm_processor_kwargs.get("target_w") if isinstance(mm_processor_kwargs, dict) else None, + _coerce_dim(original_prompt.get("width"), 1024), + ) + text_prompt = original_prompt.get("prompt", "") + + # Detect i2i mode. + # Prefer normalized prompt multi_modal_data source-image presence, with + # multimodal output as secondary signal. + _t_mode = time.perf_counter() + is_i2i = False + + prompt_modalities = original_prompt.get("modalities") + if isinstance(prompt_modalities, list) and "img2img" in prompt_modalities: + is_i2i = True + + prompt_mm_data = original_prompt.get("multi_modal_data") + if _has_source_image(prompt_mm_data): + is_i2i = True + + if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: + mm_output = ar_output.multimodal_output + if isinstance(mm_output, Mapping) and mm_output.get("ids", {}).get("prior_image") is not None: + is_i2i = True + _dt_mode = (time.perf_counter() - _t_mode) * 1000 + + # Parse and upsample prior tokens + _t_parse = time.perf_counter() + try: + prior_token_ids, pixel_h, pixel_w = _parse_generated_tokens( + generated_token_ids, + height, + width, + is_i2i=is_i2i, + ) + except ValueError as e: + logger.warning( + "[ar2diffusion] Request %s: skip due to token parse failure: %s " + "(target=%sx%s, mode=%s, raw_tokens=%s, tail=%s)", + 0, + e, + height, + width, + "i2i" if is_i2i else "t2i", + len(generated_token_ids), + generated_token_ids[-8:] if len(generated_token_ids) >= 8 else generated_token_ids, + ) + return None + _dt_parse = (time.perf_counter() - _t_parse) * 1000 + + # Get prior_token_image_ids from AR model output (for i2i mode) + # This contains VQ-VAE tokens from input image, used for KV cache conditioning + # NOTE: multimodal_output is attached to ar_output (RequestOutput), NOT output (CompletionOutput) + _t_prior_img = time.perf_counter() + prior_token_image_ids = None + + # Check ar_output (RequestOutput) for multimodal_output - this is the correct location + if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: + mm_output = ar_output.multimodal_output + if isinstance(mm_output, Mapping): + raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image") + if raw_prior_image_ids is not None: + # Handle different formats: + # 1. Single tensor -> wrap in list + # 2. List of tensors -> use as-is + # 3. List of Python lists (from serialization) -> convert to tensors + if isinstance(raw_prior_image_ids, torch.Tensor): + prior_token_image_ids = [raw_prior_image_ids] + elif isinstance(raw_prior_image_ids, list): + # Check if elements are tensors or Python lists + if raw_prior_image_ids and isinstance(raw_prior_image_ids[0], torch.Tensor): + prior_token_image_ids = raw_prior_image_ids + elif raw_prior_image_ids and isinstance(raw_prior_image_ids[0], list): + # Convert Python lists back to tensors + prior_token_image_ids = [torch.tensor(ids, dtype=torch.long) for ids in raw_prior_image_ids] + else: + logger.warning( + f"[ar2diffusion] Request 0: unexpected prior_token_image_ids format: " + f"{type(raw_prior_image_ids[0]) if raw_prior_image_ids else 'empty'}" + ) + else: + # Fallback: also check output (CompletionOutput) in case of different vLLM versions + if hasattr(output, "multimodal_output") and output.multimodal_output: + mm_output = output.multimodal_output + logger.debug("[ar2diffusion] Request 0: found multimodal_output on CompletionOutput (fallback)") if isinstance(mm_output, Mapping): raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image") if raw_prior_image_ids is not None: - # Handle different formats: - # 1. Single tensor -> wrap in list - # 2. List of tensors -> use as-is - # 3. List of Python lists (from serialization) -> convert to tensors if isinstance(raw_prior_image_ids, torch.Tensor): prior_token_image_ids = [raw_prior_image_ids] elif isinstance(raw_prior_image_ids, list): - # Check if elements are tensors or Python lists - if raw_prior_image_ids and isinstance(raw_prior_image_ids[0], torch.Tensor): - prior_token_image_ids = raw_prior_image_ids - elif raw_prior_image_ids and isinstance(raw_prior_image_ids[0], list): - # Convert Python lists back to tensors - prior_token_image_ids = [torch.tensor(ids, dtype=torch.long) for ids in raw_prior_image_ids] - else: - logger.warning( - f"[ar2diffusion] Request {i}: unexpected prior_token_image_ids format: " - f"{type(raw_prior_image_ids[0]) if raw_prior_image_ids else 'empty'}" - ) - else: - # Fallback: also check output (CompletionOutput) in case of different vLLM versions - if hasattr(output, "multimodal_output") and output.multimodal_output: - mm_output = output.multimodal_output - logger.debug(f"[ar2diffusion] Request {i}: found multimodal_output on CompletionOutput (fallback)") - if isinstance(mm_output, Mapping): - raw_prior_image_ids = mm_output.get("ids", {}).get("prior_image") - if raw_prior_image_ids is not None: - if isinstance(raw_prior_image_ids, torch.Tensor): - prior_token_image_ids = [raw_prior_image_ids] - elif isinstance(raw_prior_image_ids, list): - prior_token_image_ids = raw_prior_image_ids - _dt_prior_img = (time.perf_counter() - _t_prior_img) * 1000 - - diffusion_input = { - "prompt": text_prompt, - "height": pixel_h, - "width": pixel_w, - "extra": { - "prior_token_ids": prior_token_ids, - "prior_token_image_ids": prior_token_image_ids, - }, - } - - if requires_multimodal_data: - mm_data = original_prompt.get("multi_modal_data") - if mm_data: - pil_image = _first_source_image(mm_data) - diffusion_input["pil_image"] = pil_image - - for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]: - if key in original_prompt: - diffusion_input[key] = original_prompt[key] - - _dt_req = (time.perf_counter() - _t_req) * 1000 - logger.info( - "[ar2diffusion] req=%d mode=%s target=%dx%d " - "raw_tokens=%d prior_tokens=%d prior_image_ids=%s " - "timing: mode_detect=%.3fms parse+upsample=%.3fms " - "prior_image_ids_extract=%.3fms req_total=%.3fms", - i, - "i2i" if is_i2i else "t2i", - pixel_h, - pixel_w, - len(generated_token_ids), - len(prior_token_ids), - "yes" if prior_token_image_ids is not None else "no", - _dt_mode, - _dt_parse, - _dt_prior_img, - _dt_req, - ) - diffusion_inputs.append(diffusion_input) + prior_token_image_ids = raw_prior_image_ids + _dt_prior_img = (time.perf_counter() - _t_prior_img) * 1000 + + diffusion_input = { + "prompt": text_prompt, + "height": pixel_h, + "width": pixel_w, + "extra": { + "prior_token_ids": prior_token_ids, + "prior_token_image_ids": prior_token_image_ids, + }, + } + + if requires_multimodal_data: + mm_data = original_prompt.get("multi_modal_data") + if mm_data: + pil_image = _first_source_image(mm_data) + diffusion_input["pil_image"] = pil_image + + for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]: + if key in original_prompt: + diffusion_input[key] = original_prompt[key] + + _dt_req = (time.perf_counter() - _t_req) * 1000 + logger.info( + "[ar2diffusion] req=%d mode=%s target=%dx%d " + "raw_tokens=%d prior_tokens=%d prior_image_ids=%s " + "timing: mode_detect=%.3fms parse+upsample=%.3fms " + "prior_image_ids_extract=%.3fms req_total=%.3fms", + 0, + "i2i" if is_i2i else "t2i", + pixel_h, + pixel_w, + len(generated_token_ids), + len(prior_token_ids), + "yes" if prior_token_image_ids is not None else "no", + _dt_mode, + _dt_parse, + _dt_prior_img, + _dt_req, + ) _dt_total = (time.perf_counter() - _t_total) * 1000 logger.info( - "[ar2diffusion] batch done: %d reqs, total=%.3fms", - len(diffusion_inputs), + "[ar2diffusion] request done: 1 req, total=%.3fms", _dt_total, ) - return diffusion_inputs + return diffusion_input diff --git a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py index 9a81a8cc3ab..963adb9c2a0 100644 --- a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py +++ b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py @@ -95,128 +95,130 @@ def ar2diffusion( source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | list | None = None, requires_multimodal_data: bool = False, -) -> list[dict[str, Any]]: +) -> dict[str, Any] | None: """Process AR stage outputs to create Diffusion stage inputs. + HunyuanImage3 produces one downstream diffusion request per parent AR + request. ``source_outputs`` may include companion outputs, but the parent + output is expected first: the orchestrator bundles diffusion inputs as + ``[parent_output, *cfg_companion_outputs]`` before invoking this bridge. + Args: prompt: Original user prompt (may contain multimodal data). requires_multimodal_data: Whether to forward multimodal data. Returns: - List of dicts, each consumable by the HunyuanImage3 diffusion pipeline. + One prompt dict consumable by the HunyuanImage3 diffusion pipeline, or + ``None`` when no parent output is available. """ - ar_outputs = source_outputs - diffusion_inputs = [] - - # Normalize prompt to list - if not isinstance(prompt, list): - prompt = [prompt] if prompt is not None else [{}] - - for i, ar_output in enumerate(ar_outputs): - output = ar_output.outputs[0] - generated_token_ids = output.cumulative_token_ids - # Prefer cumulative_text, fallback to text if aggregation dropped it - generated_text = getattr(output, "cumulative_text", None) or getattr(output, "text", "") or "" - - # Get original prompt info - original_prompt = prompt[i] if i < len(prompt) else {} - if isinstance(original_prompt, dict): - pass - elif hasattr(original_prompt, "_asdict"): - original_prompt = original_prompt._asdict() - elif hasattr(original_prompt, "__dict__"): - original_prompt = vars(original_prompt) - else: - original_prompt = {} - - height = original_prompt.get("height", 1024) - width = original_prompt.get("width", 1024) - text_prompt = original_prompt.get("prompt", "") - use_system_prompt = original_prompt.get("use_system_prompt") - custom_system_prompt = original_prompt.get("system_prompt") - - # Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>` - # tail emitted by `HunyuanImage3ForCausalMM.sample` under the - # ratio-restriction logits processor) over the carried-through - # height/width, which the serving layer fills with the first - # reference image's bucket and so collapses non-square targets to - # square in the multi-image / mismatched-aspect case. Mirrors the - # official upstream where `reso_group[ratio_index]` is the - # canonical source of the diffusion target shape. - ratio_idx = _extract_ratio_index(generated_token_ids) - ar_predicted = False - if ratio_idx is not None: - base_size = int(original_prompt.get("image_base_size", 1024)) - reso_group: ResolutionGroup = get_cached_resolution_group(base_size=base_size) - try: - reso: Resolution = reso_group[ratio_idx] - height = reso.height - width = reso.width - ar_predicted = True - except IndexError: - logger.warning( - "[ar2diffusion] Request %d: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d", - i, - ratio_idx, - len(reso_group), - height, - width, - ) - - cot_text_for_dit = _truncate_at_cot_end(generated_text) - - logger.info( - "[ar2diffusion] Request %d: AR generated %d tokens, text length=%d, " - "cot_text length=%d, target size=%dx%d (%s)", - i, - len(generated_token_ids), - len(generated_text), - len(cot_text_for_dit), - height, - width, - f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)", - ) - - diffusion_input: dict[str, Any] = { - "prompt": text_prompt, - "height": height, - "width": width, - "extra": { - "ar_generated_text": cot_text_for_dit, - }, - } - - # Forward use_system_prompt so the DiT can build the same system prefix. - # Also forward the custom system prompt body when sys_type=custom so - # DiT's `get_system_prompt(use, "image", body)` doesn't fall back to - # an empty prefix and silently diverge from AR. - if use_system_prompt is not None: - diffusion_input["use_system_prompt"] = use_system_prompt - if custom_system_prompt is not None: - diffusion_input["system_prompt"] = custom_system_prompt - - # Forward multimodal data (original image for IT2I conditioning). - # The diffusion pre_process_func reads multi_modal_data["image"], which - # matches vLLM's standard prompt schema, so we only need to pass it once. - mm_data = original_prompt.get("multi_modal_data") - if mm_data: - prompt_images = mm_data.get("image") - if prompt_images is None: - prompt_images = mm_data.get("images") - if prompt_images is not None: - diffusion_input["multi_modal_data"] = {"image": prompt_images} - - # Forward multimodal output from AR (if any) - if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: - mm_output = ar_output.multimodal_output - if isinstance(mm_output, Mapping): - diffusion_input["extra"]["ar_multimodal_output"] = mm_output - - # Forward sampling params - for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]: - if key in original_prompt: - diffusion_input[key] = original_prompt[key] - - diffusion_inputs.append(diffusion_input) - - return diffusion_inputs + if not source_outputs: + return None + + # The orchestrator constructs this list as [parent, *cfg_companions]. + ar_output = source_outputs[0] + output = ar_output.outputs[0] + generated_token_ids = output.cumulative_token_ids + # Prefer cumulative_text, fallback to text if aggregation dropped it + generated_text = getattr(output, "cumulative_text", None) or getattr(output, "text", "") or "" + + if isinstance(prompt, list): + original_prompt = prompt[0] if prompt else {} + elif prompt is not None: + original_prompt = prompt + else: + original_prompt = {} + if isinstance(original_prompt, dict): + pass + elif hasattr(original_prompt, "_asdict"): + original_prompt = original_prompt._asdict() + elif hasattr(original_prompt, "__dict__"): + original_prompt = vars(original_prompt) + else: + original_prompt = {} + + height = original_prompt.get("height", 1024) + width = original_prompt.get("width", 1024) + text_prompt = original_prompt.get("prompt", "") + use_system_prompt = original_prompt.get("use_system_prompt") + custom_system_prompt = original_prompt.get("system_prompt") + + # Prefer the AR's predicted output aspect (`<img_size_*><img_ratio_*>` + # tail emitted by `HunyuanImage3ForCausalMM.sample` under the + # ratio-restriction logits processor) over the carried-through + # height/width, which the serving layer fills with the first + # reference image's bucket and so collapses non-square targets to + # square in the multi-image / mismatched-aspect case. Mirrors the + # official upstream where `reso_group[ratio_index]` is the + # canonical source of the diffusion target shape. + ratio_idx = _extract_ratio_index(generated_token_ids) + ar_predicted = False + if ratio_idx is not None: + base_size = int(original_prompt.get("image_base_size", 1024)) + reso_group: ResolutionGroup = get_cached_resolution_group(base_size=base_size) + try: + reso: Resolution = reso_group[ratio_idx] + height = reso.height + width = reso.width + ar_predicted = True + except IndexError: + logger.warning( + "[ar2diffusion] Request 0: ratio_index=%d out of range [0,%d), keeping prompt size %dx%d", + ratio_idx, + len(reso_group), + height, + width, + ) + + cot_text_for_dit = _truncate_at_cot_end(generated_text) + + logger.info( + "[ar2diffusion] Request 0: AR generated %d tokens, text length=%d, cot_text length=%d, target size=%dx%d (%s)", + len(generated_token_ids), + len(generated_text), + len(cot_text_for_dit), + height, + width, + f"AR ratio_idx={ratio_idx}" if ar_predicted else "from prompt (no AR ratio token)", + ) + + diffusion_input: dict[str, Any] = { + "prompt": text_prompt, + "height": height, + "width": width, + "extra": { + "ar_generated_text": cot_text_for_dit, + }, + } + + # Forward use_system_prompt so the DiT can build the same system prefix. + # Also forward the custom system prompt body when sys_type=custom so + # DiT's `get_system_prompt(use, "image", body)` doesn't fall back to + # an empty prefix and silently diverge from AR. + if use_system_prompt is not None: + diffusion_input["use_system_prompt"] = use_system_prompt + if custom_system_prompt is not None: + diffusion_input["system_prompt"] = custom_system_prompt + + # Forward multimodal data (original image for IT2I conditioning). + # The diffusion pre_process_func reads multi_modal_data["image"], which + # matches vLLM's standard prompt schema, so we only need to pass it once. + mm_data = original_prompt.get("multi_modal_data") + if mm_data: + prompt_images = mm_data.get("image") + if prompt_images is None: + prompt_images = mm_data.get("images") + if prompt_images is not None: + diffusion_input["multi_modal_data"] = {"image": prompt_images} + + # Forward multimodal output from AR (if any) + if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: + mm_output = ar_output.multimodal_output + if isinstance(mm_output, Mapping): + diffusion_input["extra"]["ar_multimodal_output"] = mm_output + + # Forward sampling params + for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]: + if key in original_prompt: + diffusion_input[key] = original_prompt[key] + + return diffusion_input diff --git a/vllm_omni/model_executor/stage_input_processors/indextts2.py b/vllm_omni/model_executor/stage_input_processors/indextts2.py new file mode 100644 index 00000000000..39d9d011cbc --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/indextts2.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage input processor for IndexTTS2: Talker (GPT AR) → S2Mel decoder.""" + +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +STOP_MEL_TOKEN = 8193 + + +def _cpu_view(tensor: torch.Tensor) -> torch.Tensor: + """Return a contiguous CPU tensor suitable for connector serialization. + + Always returns an owning copy so downstream consumers cannot mutate the + source tensors that may still be referenced by the producer stage. + """ + out = tensor.detach() + if out.device.type != "cpu": + return out.cpu().contiguous() + return out.clone().contiguous() + + +def _strip_stop_token( + codes: torch.Tensor, + latent: torch.Tensor, + stop_mel_token: int = STOP_MEL_TOKEN, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Strip at the first stop token, matching official IndexTTS2 v2. + + Returns: (codes [B, T'], latent [B, T', D], code_lens [B]). + """ + if codes.ndim == 1: + codes = codes.unsqueeze(0) + if latent.ndim == 2: + latent = latent.unsqueeze(0) + + if codes.device.type != "cpu": + codes = codes.detach().cpu() + if latent.device.type != "cpu": + latent = latent.detach().cpu() + + device = codes.device + code_lens = [] + codes_out = [] + latent_out = [] + + for i in range(codes.shape[0]): + code = codes[i] + lat = latent[i] + + stop_mask = (code == stop_mel_token).nonzero(as_tuple=False) + if stop_mask.numel() > 0: + valid_len = int(stop_mask[0].item()) + else: + valid_len = int(code.shape[0]) + code_lens.append(valid_len) + codes_out.append(code[:valid_len]) + latent_out.append(lat[:valid_len]) + + max_len = max(code_lens) if code_lens else 0 + if max_len == 0: + return ( + torch.zeros(codes.shape[0], 0, dtype=torch.long, device=device), + torch.zeros(codes.shape[0], 0, latent.shape[-1], device=device, dtype=latent.dtype), + torch.zeros(codes.shape[0], dtype=torch.long, device=device), + ) + + padded_codes = torch.zeros((len(codes_out), max_len), dtype=torch.long, device=device) + padded_latent = torch.zeros( + len(latent_out), + max_len, + latent_out[0].shape[-1], + device=device, + dtype=latent_out[0].dtype, + ) + for i, (c, lat) in enumerate(zip(codes_out, latent_out)): + padded_codes[i, : c.shape[0]] = c + padded_latent[i, : lat.shape[0]] = lat + + return padded_codes, padded_latent, torch.tensor(code_lens, dtype=torch.long, device=device) + + +def _normalize_mel_sequence(mel_codes: torch.Tensor) -> torch.Tensor: + mel_codes = mel_codes.to(torch.long) + if mel_codes.ndim <= 1: + return mel_codes.reshape(-1).contiguous() if mel_codes.ndim == 0 else mel_codes.contiguous() + if mel_codes.ndim == 2 and (mel_codes.shape[0] == 1 or mel_codes.shape[1] == 1): + return mel_codes.reshape(-1).contiguous() + return mel_codes.reshape(-1).contiguous() + + +def _normalize_latent_sequence(latent: torch.Tensor) -> torch.Tensor: + if latent.ndim <= 1: + return latent.reshape(1, -1).contiguous() + if latent.ndim == 2: + return latent.contiguous() + if latent.ndim == 3 and latent.shape[0] == 1: + return latent[0].contiguous() + return latent.reshape(-1, latent.shape[-1]).contiguous() + + +def _build_s2mel_additional_information( + mel_codes: torch.Tensor, + latent: torch.Tensor, + meta: dict[str, Any], + *, + context: str, +) -> dict[str, Any]: + """Build the Stage-1 S2Mel tensor contract shared by legacy and connector paths.""" + mel_codes_clean, latent_clean, code_lens = _strip_stop_token(mel_codes, latent) + + additional_information = { + "latent": _cpu_view(latent_clean), + "mel_codes": _cpu_view(mel_codes_clean), + "code_lens": _cpu_view(code_lens), + } + + for key in ("S_ref", "ref_mel", "style"): + val = meta.get(key) + if isinstance(val, torch.Tensor): + additional_information[key] = _cpu_view(val) + else: + logger.warning("[%s] %s MISSING — Stage 1 will use fallback", context, key) + + return additional_information + + +def _get_payload_value(pooling_output: dict[str, Any], dotted_key: str, nested_parent: str, nested_key: str) -> Any: + value = pooling_output.get(dotted_key) + if value is not None: + return value + nested = pooling_output.get(nested_parent) + if isinstance(nested, dict): + return nested.get(nested_key) + return None + + +def _request_id(request: Any) -> str: + return str(getattr(request, "external_req_id", None) or getattr(request, "request_id", "?")) + + +def talker2s2mel_token_only( + source_outputs: list[Any], + prompt: Any = None, + _requires_multimodal_data: bool = False, +) -> list[Any]: + """Sync-side placeholder for Stage 1; tensors arrive via full-payload connector.""" + from vllm_omni.inputs.data import OmniTokensPrompt + + del prompt + s2mel_inputs: list[OmniTokensPrompt] = [] + for talker_output in source_outputs: + if not talker_output.finished: + continue + s2mel_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0], + additional_information=None, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + return s2mel_inputs + + +def talker2s2mel_full_payload( + transfer_manager: Any, + pooling_output: dict[str, Any], + request: Any, + **_: Any, +) -> dict[str, Any] | None: + """Build the complete S2Mel input from accumulated per-step talker deltas.""" + del transfer_manager + rid = _request_id(request) + if not isinstance(pooling_output, dict): + logger.warning("talker2s2mel_full_payload: pooling_output not a dict for req=%s", rid) + return None + + mel_codes = _get_payload_value(pooling_output, "codes.mel", "codes", "mel") + if not isinstance(mel_codes, torch.Tensor) or mel_codes.numel() == 0: + logger.warning("talker2s2mel_full_payload: missing codes.mel for req=%s", rid) + return None + + latent = _get_payload_value(pooling_output, "hidden_states.latent", "hidden_states", "latent") + if not isinstance(latent, torch.Tensor) or latent.numel() == 0: + logger.warning("talker2s2mel_full_payload: missing hidden_states.latent for req=%s", rid) + return None + + mel_seq = _normalize_mel_sequence(mel_codes) + latent_seq = _normalize_latent_sequence(latent) + if mel_seq.numel() == 0 or latent_seq.numel() == 0: + logger.warning("talker2s2mel_full_payload: empty normalized mel/latent for req=%s", rid) + return None + + common_len = min(int(mel_seq.shape[0]), int(latent_seq.shape[0])) + if common_len <= 0: + logger.warning("talker2s2mel_full_payload: no common mel/latent length for req=%s", rid) + return None + + mel_seq = mel_seq[:common_len] + latent_seq = latent_seq[:common_len] + + meta = { + "S_ref": _get_payload_value(pooling_output, "meta.S_ref", "meta", "S_ref"), + "ref_mel": _get_payload_value(pooling_output, "meta.ref_mel", "meta", "ref_mel"), + "style": _get_payload_value(pooling_output, "meta.style", "meta", "style"), + } + return _build_s2mel_additional_information(mel_seq, latent_seq, meta, context="full_payload") diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py index 7c821c243b5..7b23f8e4986 100644 --- a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py +++ b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py @@ -3,31 +3,69 @@ from collections.abc import Mapping from typing import Any -import torch from vllm.inputs import TextPrompt from vllm_omni.inputs.data import OmniTokensPrompt +def _as_dict(prompt: Any) -> dict[str, Any]: + """Coerce an original-stage prompt to a dict. + + It may arrive as a dict, a NamedTuple/object, or a bare string depending on + the calling flow (the shared text_to_image example vs the bespoke script). + """ + if isinstance(prompt, dict): + return prompt + if hasattr(prompt, "_asdict"): + return prompt._asdict() + if hasattr(prompt, "__dict__"): + return vars(prompt) + return {} + + +def _coerce_dim(value: Any, default: int) -> int: + try: + iv = int(value) + except (TypeError, ValueError): + return default + return iv if iv > 0 else default + + def ar2dit( source_outputs: list[Any], - prompts: OmniTokensPrompt | TextPrompt | None = None, + prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None, _requires_multimodal_data: bool = False, ) -> list[OmniTokensPrompt]: """Convert AR stage outputs to DiT stage inputs.""" ar_outputs = source_outputs + # The shared text_to_image example forwards a single prompt (not a list); normalize + # so a lone dict isn't iterated as its keys. Mirrors glm_image.ar2diffusion. + if not isinstance(prompts, list): + prompts = [prompts] if prompts is not None else [{}] + dit_inputs: list[OmniTokensPrompt] = [] - for ar_output, prompt in zip(ar_outputs, prompts): - addi_info = prompt["additional_information"] - image_height = addi_info["image_height"][0] - image_width = addi_info["image_width"][0] - text_guidance_scale = addi_info["text_guidance_scale"][0] - cfg_range = addi_info["cfg_range"] - num_inference_steps = addi_info["num_inference_steps"][0] - gen_vocab_start_index = addi_info["visual_token_start_id"][0] - # ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"] - visual_ids = addi_info["visual_ids"] + for i, ar_output in enumerate(ar_outputs): + prompt_dict = _as_dict(prompts[i] if i < len(prompts) else {}) + addi_info = prompt_dict.get("additional_information") or {} + mm_kwargs = prompt_dict.get("mm_processor_kwargs") or {} + + # Image size: prefer mm_processor_kwargs target_h/target_w (set by the serving + # layer), fall back to additional_information, then a 1024 default. + image_height = _coerce_dim( + mm_kwargs.get("target_h"), + _coerce_dim((addi_info.get("image_height") or [None])[0], 1024), + ) + image_width = _coerce_dim( + mm_kwargs.get("target_w"), + _coerce_dim((addi_info.get("image_width") or [None])[0], 1024), + ) + + # Sampling knobs arrive on the DiT stage via extra_body -> extra_args; these are + # defensive fallbacks (defaults mirror the former bespoke script's argparse). + text_guidance_scale = (addi_info.get("text_guidance_scale") or [9.0])[0] + cfg_range = addi_info.get("cfg_range") or [0.0, 1.0] + num_inference_steps = (addi_info.get("num_inference_steps") or [50])[0] prompt_token_ids = ar_output.prompt_token_ids # exclude the last token because it has no corresponding hidden state @@ -48,36 +86,16 @@ def ar2dit( f"Hidden states length mismatch: expected {len(prompt_token_ids) + len(gen_token_ids)}, got {hidden_total}" ) - mask_device = full_hidden_states.device - full_token_ids_t = torch.tensor(full_token_ids, dtype=torch.long, device=mask_device) - attention_mask = torch.ones_like(full_token_ids_t, dtype=torch.bool) - - pos = torch.arange(full_token_ids_t.shape[0], device=mask_device) - answer_start_index = len(prompt_token_ids) - questions_mask = pos < answer_start_index - answers_mask = ~questions_mask - - gen_token_mask = full_token_ids_t >= gen_vocab_start_index - - visual_token_mask = torch.isin( - full_token_ids_t, - torch.tensor(visual_ids, dtype=torch.long, device=mask_device), - ) - - text_condition_token_mask = questions_mask & ~(visual_token_mask | gen_token_mask) & attention_mask - image_condition_token_mask = answers_mask & gen_token_mask & attention_mask - - text_condition = full_hidden_states[text_condition_token_mask] - image_condition = full_hidden_states[image_condition_token_mask] - - text_prompt_embeds = text_condition.to(dtype=torch.float32).contiguous() - image_prompt_embeds = image_condition.to(dtype=torch.float32).contiguous() - + # The text/image condition split is performed in the DiT pipeline, which sources + # the distinguishing token ids (gen_vocab_start_index, vision placeholder ids) + # from the model config. Pass through the raw AR hidden states + token ids and + # the question/answer boundary so the pipeline can reconstruct the masks. additional_information = { - "text_prompt_embeds": text_prompt_embeds, - "text_prompt_embeds_shape": list(text_prompt_embeds.shape), - "image_prompt_embeds": image_prompt_embeds, - "image_prompt_embeds_shape": list(image_prompt_embeds.shape), + # float32 so the tensor crosses the stage boundary (the serializer uses + # numpy, which has no bf16); the DiT re-casts to the model dtype. + "full_hidden_states": full_hidden_states.float().contiguous(), + "full_token_ids": full_token_ids, + "answer_start_index": [len(prompt_token_ids)], "image_height": [int(image_height)], "image_width": [int(image_width)], "text_guidance_scale": [float(text_guidance_scale)], diff --git a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py index 7c0db4131ff..cd9ca7bf211 100644 --- a/vllm_omni/model_executor/stage_input_processors/mimo_audio.py +++ b/vllm_omni/model_executor/stage_input_processors/mimo_audio.py @@ -1,7 +1,6 @@ from typing import Any import torch -from vllm.inputs import TextPrompt from vllm.logger import init_logger from vllm_omni.data_entry_keys import ( @@ -10,7 +9,6 @@ OmniPayload, OmniPayloadStruct, ) -from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.mimo_audio.config_mimo_audio import TALKER_CODEC_PAD_TOKEN_ID logger = init_logger(__name__) @@ -19,8 +17,8 @@ # sequence fed to stage-1 must not exceed this, otherwise gpu_input_batch # add_request will fail with a broadcast error when copying prompt_token_ids # into token_ids_cpu. Keep in sync with the stage-1 ``max_model_len`` in -# ``vllm_omni/model_executor/stage_configs/mimo_audio.yaml`` and the offline -# example ``examples/offline_inference/mimo_audio/end2end.py``. +# ``vllm_omni/deploy/mimo_audio.yaml`` and the offline example +# ``examples/offline_inference/mimo_audio/end2end.py``. MAX_CODE2WAV_TOKENS = 18192 # Minimum safe values for codec streaming parameters. @@ -230,88 +228,6 @@ def llm2code2wav_async_chunk( ) -def llm2code2wav( - source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, -) -> list[OmniTokensPrompt]: - """ - Process talker outputs to create code2wav inputs. - - Workflow: - 1. Extract talker's codec code outputs (8-layer RVQ codes) - 2. Flatten codes for code2wav input - 3. Package for code2wav stage - - Args: - prompt: Original prompt data - requires_multimodal_data: Whether multimodal data is required - - Returns: - List of OmniTokensPrompt for code2wav stage - """ - talker_outputs = source_outputs - code2wav_inputs = [] - - # Process each talker output - for i, talker_output in enumerate(talker_outputs): - output = talker_output.outputs[0] - - # Extract codec codes from talker output - # Expected shape: [8, seq_len] (8-layer RVQ codes) - mm = output.multimodal_output - mm_codes = mm.get("codes", {}) - mm_hs = mm.get("hidden_states", {}) - if "audio" in mm_codes: - codec_codes = mm_codes["audio"].to(torch.long) # [seq_batch_size, 1, 8, 4] - is_all_zero = (codec_codes == 0).all(dim=(1, 2, 3)) - non_zero_indices = (~is_all_zero).nonzero(as_tuple=True)[0] - if len(non_zero_indices) == 0: - # All codec codes are zero - skip this request with a warning - request_id = getattr(talker_output, "request_id", f"unknown_{i}") - logger.warning( - "Skipping request %s: all codec codes are zero (empty output from Stage-0). " - "This may indicate the model failed to generate valid audio codes.", - request_id, - ) - else: - if len(non_zero_indices) < codec_codes.shape[0]: - codec_codes = codec_codes[non_zero_indices] - elif "output" in mm_hs and "audio" not in mm_codes: - codec_codes = torch.zeros(1, 1, 8, 4, dtype=torch.long) - else: - raise ValueError(f"Invalid multimodal_output: {output.multimodal_output}") - - pad_vec = torch.tensor([TALKER_CODEC_PAD_TOKEN_ID] * 4) - - code_final = prepend_and_flatten_colmajor(codec_codes, pad_vec) - code_final = code_final.tolist() - - # Guard against flattened sequences longer than code2wav's max_model_len. - # Without this, add_request raises ``could not broadcast input array - # from shape (N,) into shape (max_model_len,)`` and kills the engine - # core (see issue #2683). Mirrors the offline end2end.py safeguard. - if len(code_final) > MAX_CODE2WAV_TOKENS: - request_id = getattr(talker_output, "request_id", f"unknown_{i}") - logger.warning( - "Request %s: code_final len=%d > MAX_CODE2WAV_TOKENS=%d, truncating.", - request_id, - len(code_final), - MAX_CODE2WAV_TOKENS, - ) - code_final = code_final[:MAX_CODE2WAV_TOKENS] - - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=code_final, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return code2wav_inputs - - # ============================================================================ # Worker-connector data plane (non-async-chunk path). # AR runner's `flatten_payload` converts the model emit @@ -329,8 +245,9 @@ def llm2code2wav( def _filter_zero_codec_rows(codec_codes: torch.Tensor) -> torch.Tensor: """Drop zero-padded codec rows from a 4-D `[N, 1, 8, 4]` tensor. - Mirrors the zero-row filter in the orchestrator-path `llm2code2wav` - body (see this file's ``llm2code2wav`` around line 224). + Shared by the sync placeholder builder (``llm2code2wav_token_only``) and + the full-payload producer (``llm2code2wav_full_payload``) so both size the + downstream codec sequence off the same non-zero frames. """ if codec_codes.ndim != 4 or codec_codes.numel() == 0: return codec_codes diff --git a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py index 8f0df7269bd..653e3477789 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_flash_omni.py @@ -104,10 +104,9 @@ def expand_cfg_prompts( ) -> list[_CfgExpandedPrompt]: """Expand a text-to-image request into one CFG-text companion (opt-in). - Triggers only when a non-empty - ``sampling_params.extra_args["image_gen"]["negative_prompt"]`` is set on - the stage-0 params; otherwise returns ``[]`` and the pipeline falls back - to zero negative (Ming's default behavior). + Triggers only when a non-empty `negative_prompt` is set flat on the stage-0 params + (`sampling_params.extra_args`); otherwise returns an empty list + and the pipeline falls back to zero negative (Ming's default behavior). """ if not isinstance(prompt, dict): return [] @@ -115,8 +114,7 @@ def expand_cfg_prompts( return [] extra_args = getattr(sampling_params, "extra_args", None) or {} - image_gen_args = extra_args.get("image_gen") or {} - negative = image_gen_args.get("negative_prompt") + negative = extra_args.get("negative_prompt") if not isinstance(negative, str) or not negative.strip(): return [] @@ -328,7 +326,7 @@ def _resolve_token_ids_from_stage_or_defaults( def _extract_byte5_from_sampling_params(sampling_params: Any) -> list[str] | None: """Read ``byte5_text`` from the diffusion-stage sampling_params. - Looks up ``sampling_params.extra_args["image_gen"]["byte5_text"]`` (the + Looks up `sampling_params.extra_args["byte5_text"]` key (the explicit API surface for ByT5 glyph text). Returns ``None`` if absent or malformed, so callers can fall back to other sources. """ @@ -337,10 +335,9 @@ def _extract_byte5_from_sampling_params(sampling_params: Any) -> list[str] | Non extra = getattr(sampling_params, "extra_args", None) if not isinstance(extra, dict): return None - image_gen = extra.get("image_gen") - if not isinstance(image_gen, dict): - return None - texts = image_gen.get("byte5_text") + texts = extra.get("byte5_text") + if isinstance(texts, str): + texts = [texts] if isinstance(texts, list) and texts: return [t for t in texts if isinstance(t, str)] return None @@ -361,10 +358,9 @@ def thinker2imagegen( CFG negative conditioning. Unknown-suffix outputs are skipped. ``sampling_params`` is the diffusion stage's own SamplingParams, supplied - by the orchestrator. ByT5 explicit ``byte5_text`` is read from - ``sampling_params.extra_args.image_gen.byte5_text`` (preferred); falls - back to ``prompt.image_gen_extra_args.byte5_text``, then to auto-extraction - from quoted prompt text. + by the orchestrator. ByT5 explicit ``byte5_text`` is read from the flat + `sampling_params.extra_args.byte5_text` key (preferred); otherwise it is + auto-extracted from quoted prompt text. """ thinker_outputs = source_outputs image_patch_token_id, image_end_token_id, num_query_tokens = _resolve_token_ids_from_stage_or_defaults(stage=None) @@ -417,17 +413,10 @@ def thinker2imagegen( if ref_image is not None: extra["reference_image"] = ref_image - # ByT5 glyph text: prefer sampling_params (stage-1 explicit API), then - # prompt.image_gen_extra_args (serving_chat path), then auto-extract - # from quoted prompt text. + # ByT5 glyph text: prefer the flat byte5_text on sampling_params + # (stage-1 explicit API), else auto-extract from quoted prompt text. prompt_text = prompt.get("prompt", "") - sp_byte5 = _extract_byte5_from_sampling_params(sampling_params) - if sp_byte5: - byte5_texts: list[str] | None = sp_byte5 - else: - ig_extra = prompt.get("image_gen_extra_args") or {} - cand = ig_extra.get("byte5_text") - byte5_texts = cand if isinstance(cand, list) and cand else None + byte5_texts: list[str] | None = _extract_byte5_from_sampling_params(sampling_params) if byte5_texts: extra["byte5_text"] = [ @@ -511,25 +500,6 @@ def _build_talker_inputs( return talker_inputs -def thinker2talker( - source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | None = None, - _requires_multimodal_data: bool = False, - _streaming_context: Any | None = None, -) -> list[OmniTokensPrompt]: - """Build talker stage inputs from thinker stage outputs.""" - return _build_talker_inputs(source_outputs, prompt) - - -# ming_flash_omni is not in ``_OMNI_CONNECTOR_INIT_ARCHS`` or -# ``_FULL_PAYLOAD_INPUT_STAGES``, so the worker connector is not -# initialised for this arch and the consumer never waits on a connector -# payload. Data flows through ``additional_information`` written by -# ``thinker2talker_token_only`` (wired as ``sync_process_input_func`` -# in the pipeline) or ``thinker2talker`` (wired as -# ``custom_process_input_func``). - - def thinker2talker_token_only( source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | None = None, @@ -542,29 +512,9 @@ def thinker2talker_token_only( thinker2talker_token_only._is_sync_input = True -def thinker2talker_full_payload( - transfer_manager, - pooling_output, - request, -): - """Producer-side payload builder — no-op. - - ming_flash_omni's thinker emits no heavy tensor to ship via the - worker connector (the bridge passes text only, and speaker metadata - arrives through the USER request's additional_information). - ming_flash_omni is not in ``_OMNI_CONNECTOR_INIT_ARCHS`` so this - function is never invoked at runtime; it is retained for forward - compatibility with the connector path. - """ - del transfer_manager, pooling_output, request - return None - - __all__ = [ "CFG_TEXT_SUFFIX", "expand_cfg_prompts", "thinker2imagegen", - "thinker2talker", - "thinker2talker_full_payload", "thinker2talker_token_only", ] diff --git a/vllm_omni/model_executor/stage_input_processors/ming_tts.py b/vllm_omni/model_executor/stage_input_processors/ming_tts.py index 0f2effdcb07..058cda7880f 100644 --- a/vllm_omni/model_executor/stage_input_processors/ming_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/ming_tts.py @@ -12,6 +12,7 @@ from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.ming_tts.config_ming_tts import ( + INITIAL_LATENT_CHUNK_SIZE, KEY_CHUNK_ID, KEY_REQUEST_ID, LATENT_CHUNK_SIZE, @@ -116,7 +117,7 @@ def _decode_stop_reason(value: Any) -> str | None: return MING_STOP_REASON_BY_CODE.get(int(value)) -def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]: +def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int, int]: connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, Mapping) else {} @@ -131,9 +132,19 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]: LATENT_LEFT_CONTEXT, ) chunk_size = int(cfg.get("latent_chunk_size", LATENT_CHUNK_SIZE)) + initial_chunk_size = int(cfg.get("initial_latent_chunk_size", INITIAL_LATENT_CHUNK_SIZE)) left_context = int(cfg.get("latent_left_context", LATENT_LEFT_CONTEXT)) if chunk_size <= 0: raise ValueError(f"Invalid Ming latent_chunk_size={chunk_size}") + if initial_chunk_size < 0: + raise ValueError(f"Invalid Ming initial_latent_chunk_size={initial_chunk_size}") + if initial_chunk_size > chunk_size: + logger.warning( + "Ming initial_latent_chunk_size=%d > latent_chunk_size=%d, clamping to latent_chunk_size.", + initial_chunk_size, + chunk_size, + ) + initial_chunk_size = chunk_size # Stage-2 VAE caches past_key_values and stream_state by request_id. # Replaying left-context latents would double-feed cached decoder state. if left_context != 0: @@ -142,7 +153,7 @@ def _get_async_chunk_config(transfer_manager: Any) -> tuple[int, int]: "Ming boundary continuity is handled by per-request decoder state cache, not " f"latent replay. Got latent_left_context={left_context}." ) - return chunk_size, left_context + return chunk_size, initial_chunk_size, left_context def _build_chunk_observability( @@ -196,7 +207,7 @@ def llm2audio_vae_async_chunk( if patch is not None: transfer_manager.code_prompt_token_ids[request_id].append(patch) - chunk_size, _ = _get_async_chunk_config(transfer_manager) + chunk_size, initial_chunk_size, _ = _get_async_chunk_config(transfer_manager) patches = transfer_manager.code_prompt_token_ids[request_id] seen_patch_len = int(state.get("seen_patch_len", 0)) @@ -225,11 +236,17 @@ def llm2audio_vae_async_chunk( ) return None - chunk_length = length % chunk_size - if chunk_length != 0 and not finished: - return None - - emit_count = chunk_length if chunk_length != 0 else chunk_size + use_first_chunk = 0 < initial_chunk_size < chunk_size and seen_patch_len == 0 + if finished: + emit_count = length + elif use_first_chunk: + if length < initial_chunk_size: + return None + emit_count = initial_chunk_size + else: + if length < chunk_size: + return None + emit_count = chunk_size emit_patches = list(new_patches[:emit_count]) state["seen_patch_len"] = seen_patch_len + len(emit_patches) latent_patches = torch.stack(emit_patches, dim=0) diff --git a/vllm_omni/model_executor/stage_input_processors/moss_tts.py b/vllm_omni/model_executor/stage_input_processors/moss_tts.py index 48c0f98cb2d..3dfe2a6032a 100644 --- a/vllm_omni/model_executor/stage_input_processors/moss_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/moss_tts.py @@ -5,6 +5,8 @@ from __future__ import annotations +from collections import defaultdict +from collections.abc import Mapping from typing import Any import torch @@ -15,6 +17,8 @@ logger = init_logger(__name__) +_MOSS_AUDIO_PAD_CODE = 1024 + # --------------------------------------------------------------------------- # Helpers @@ -90,13 +94,13 @@ def talker2codec( # --------------------------------------------------------------------------- -def talker2codec_async_chunk( +def talker2codec_delay_async_chunk( transfer_manager: Any, multimodal_output: dict[str, Any] | None, request: Any, is_finished: bool = False, ) -> OmniPayloadStruct | None: - """Emit accumulated audio codes to Stage 1 as they arrive from Stage 0. + """Emit delay-patterned accumulated audio codes to Stage 1. State is maintained in ``transfer_manager`` keyed by request ID. A chunk is forwarded to Stage 1 when either: @@ -138,14 +142,6 @@ def talker2codec_async_chunk( req_state["accumulated"] = new_rows else: req_state["accumulated"] = torch.cat([req_state["accumulated"], new_rows], dim=0) - # Realtime variant emits raw (un-delay-patterned) codes; record that - # so the emit step below skips the apply_de_delay_pattern transform. - # Realtime sets ``meta.finished`` to a 1-D bool tensor; the delay - # variant leaves ``meta`` empty. - meta_in = pooling_output.get("meta", {}) or {} - flag = meta_in.get("finished") - if isinstance(flag, torch.Tensor) and flag.numel() >= 1 and bool(flag.reshape(-1)[0].item()): - req_state["skip_dedelay"] = True acc = req_state["accumulated"] if acc is None or acc.numel() == 0: @@ -181,15 +177,12 @@ def talker2codec_async_chunk( # 1. de-delay ``(T+nq-1, nq)`` → ``(T, nq)`` # 2. drop rows that are entirely pad (separators between audio segments, # and the leading text-mode rows that precede the first audio_start). - # The realtime talker emits raw codes (no delay), so steps 1 is skipped. chunk_codes_long = chunk_codes.to(torch.long).cpu().contiguous() # (T_chunk, NQ) nq = int(chunk_codes_long.shape[1]) t_chunk = int(chunk_codes_long.shape[0]) audio_pad_code = 1024 # MOSS-TTS audio_pad_code; same value across variants. - if req_state.get("skip_dedelay"): - de_delayed = chunk_codes_long - elif t_chunk > nq: + if t_chunk > nq: de_delayed = chunk_codes_long.new_zeros((t_chunk - nq + 1, nq)) for i in range(nq): de_delayed[:, i] = chunk_codes_long[i : i + de_delayed.shape[0], i] @@ -224,9 +217,133 @@ def talker2codec_async_chunk( codes=CodesStruct(audio=codec_flat), meta=MetaStruct( left_context_size=left_context, - finished=torch.tensor(is_finished, dtype=torch.bool), + finished=torch.tensor(bool(is_finished), dtype=torch.bool), + ), + ) + + +def talker2codec_raw_async_chunk( + transfer_manager: Any, + multimodal_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> OmniPayloadStruct | None: + """Async processor for MOSS-TTS Local/Realtime raw codec rows. + + Stage 0 emits newly generated raw codec rows shaped ``[T, n_vq]`` (normally + ``[1, n_vq]`` per decode step). This processor buffers those new rows until + a codec chunk is ready, then forwards the chunk to Stage 1. No delay-pattern + de-delay is applied on this path. + """ + external_req_id = getattr(request, "external_req_id", None) + req_id = str(external_req_id if external_req_id is not None else getattr(request, "request_id", id(request))) + + if not hasattr(transfer_manager, "code_prompt_token_ids"): + transfer_manager.code_prompt_token_ids = defaultdict(list) + if not hasattr(transfer_manager, "request_payload"): + transfer_manager.request_payload = {} + if not hasattr(transfer_manager, "put_req_chunk"): + transfer_manager.put_req_chunk = defaultdict(int) + + pending_frames = transfer_manager.code_prompt_token_ids[req_id] + + if isinstance(multimodal_output, Mapping): + codes_dict = multimodal_output.get("codes", {}) or {} + new_frames = codes_dict.get("audio") + if isinstance(new_frames, torch.Tensor) and new_frames.numel() > 0: + frames_cpu = new_frames.detach().to("cpu", torch.long).contiguous() + if frames_cpu.ndim == 1: + frames_cpu = frames_cpu.reshape(1, -1) + if frames_cpu.ndim != 2: + raise ValueError(f"MOSS raw codec frames must be 2-D, got {tuple(frames_cpu.shape)}") + valid_rows = frames_cpu.ne(_MOSS_AUDIO_PAD_CODE).any(dim=1) + for frame in frames_cpu[valid_rows]: + pending_frames.append(frame.clone()) + # Raw/local streaming should mirror the non-streaming path: the codec + # decodes only generated audio rows. Reference audio conditions the + # talker, but feeding its codes into the codec streaming state adds a + # long first-packet prime step and changes the decoder state relative + # to non-streaming output. + + connector = getattr(transfer_manager, "connector", None) + raw_cfg = getattr(connector, "config", {}) or {} + cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} + cfg = cfg if isinstance(cfg, dict) else {} + chunk_frames = int(cfg.get("codec_chunk_frames", 15) or 15) + initial_chunk_frames = int(cfg.get("initial_codec_chunk_frames") or 0) + if chunk_frames <= 0: + raise ValueError(f"codec_chunk_frames must be positive for MOSS raw streaming, got {chunk_frames}") + if initial_chunk_frames < 0: + raise ValueError(f"initial_codec_chunk_frames must be non-negative, got {initial_chunk_frames}") + if initial_chunk_frames > chunk_frames: + logger.warning( + "initial_codec_chunk_frames=%d > codec_chunk_frames=%d, clamping.", + initial_chunk_frames, + chunk_frames, + ) + initial_chunk_frames = chunk_frames + + pending = len(pending_frames) + emitted_any = int(transfer_manager.put_req_chunk.get(req_id, 0)) > 0 + threshold = initial_chunk_frames if initial_chunk_frames > 0 and not emitted_any else chunk_frames + if pending <= 0: + if is_finished: + transfer_manager.code_prompt_token_ids.pop(req_id, None) + transfer_manager.request_payload.pop(req_id, None) + return OmniPayloadStruct( + # A non-empty sentinel is required so Stage-1 is scheduled. + # ``code_flat_numel=0`` tells the codec this is a control-only + # finish packet, not an audio code. + codes=CodesStruct(audio=torch.tensor([0], dtype=torch.long)), + meta=MetaStruct( + req_id=[req_id], + left_context_size=0, + codec_streaming=True, + codec_chunk_frames=0, + codec_left_context_frames=0, + code_flat_numel=0, + stream_finished=torch.tensor(True, dtype=torch.bool), + finished=torch.tensor(True, dtype=torch.bool), + ), + request_id=req_id, + ) + return None + if not is_finished and pending < threshold: + return None + + emit_frames = pending if is_finished else threshold + chunk_rows = pending_frames[:emit_frames] + del pending_frames[:emit_frames] + chunk_codes = torch.stack( + [row.to(torch.long).cpu() for row in chunk_rows], + dim=0, + ).contiguous() + finished = bool(is_finished and len(pending_frames) == 0) + + codec_flat = chunk_codes.transpose(0, 1).contiguous().reshape(-1).to(torch.long) + + if finished: + transfer_manager.code_prompt_token_ids.pop(req_id, None) + transfer_manager.request_payload.pop(req_id, None) + + return OmniPayloadStruct( + codes=CodesStruct(audio=codec_flat), + meta=MetaStruct( + req_id=[req_id], + left_context_size=0, + codec_streaming=True, + codec_chunk_frames=int(chunk_codes.shape[0]), + codec_left_context_frames=0, + code_flat_numel=int(codec_flat.numel()), + stream_finished=torch.tensor(finished, dtype=torch.bool), + finished=torch.tensor(finished, dtype=torch.bool), ), + request_id=req_id, ) -__all__ = ["talker2codec", "talker2codec_async_chunk"] +__all__ = [ + "talker2codec", + "talker2codec_delay_async_chunk", + "talker2codec_raw_async_chunk", +] diff --git a/vllm_omni/model_executor/stage_input_processors/omnivoice.py b/vllm_omni/model_executor/stage_input_processors/omnivoice.py deleted file mode 100644 index ce863c13568..00000000000 --- a/vllm_omni/model_executor/stage_input_processors/omnivoice.py +++ /dev/null @@ -1,34 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inter-stage processor for OmniVoice: Generator → Decoder.""" - -from typing import Any - -from vllm.inputs import TextPrompt - -from vllm_omni.inputs.data import OmniTokensPrompt - - -def tokens2audio( - source_outputs: list[Any], - _prompt: OmniTokensPrompt | TextPrompt = None, - _requires_multimodal_data: bool = True, -): - """Build stage-1 (decoder) inputs from stage-0 (generator) outputs. - - Takes the 8-codebook audio tokens from the generator and packages - them for the HiggsAudioV2 decoder. - """ - source_output = source_outputs[0] - output = source_output.outputs[0] - - multi_modal_data = output.multimodal_output - if multi_modal_data is None: - raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}") - - # Pass audio_tokens from generator to decoder - engine_input = OmniTokensPrompt( - prompt_token_ids=output.cumulative_token_ids, - additional_information=multi_modal_data, - ) - return [engine_input] diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 225674360c2..42377faa8b7 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -20,78 +20,6 @@ TALKER_CODEC_END_TOKEN_ID = 8294 -def thinker2talker( - source_outputs, - prompt: OmniTokensPrompt | TextPrompt = None, - requires_multimodal_data: bool = False, -): - thinker_outputs = source_outputs - talker_inputs = [] - if not isinstance(prompt, list): - prompt = [prompt] - multi_modal_data = { - thinker_output.request_id: p.get("multi_modal_data", None) for thinker_output, p in zip(thinker_outputs, prompt) - } - - for i, thinker_output in enumerate(thinker_outputs): - output = thinker_output.outputs[0] - prompt_token_ids = thinker_output.prompt_token_ids - thinker_output_ids = output.cumulative_token_ids - prompt_token_ids_len = len(prompt_token_ids) - mm: OmniPayload = output.multimodal_output - latent = mm["latent"] - thinker_hidden_states = latent.clone().detach().to(latent.device) - decode_hidden = thinker_hidden_states[prompt_token_ids_len:].to(torch.float32) - prefill_hidden = thinker_hidden_states[:prompt_token_ids_len].to(torch.float32) - additional_information = to_dict( - OmniPayloadStruct( - hidden_states=HiddenStatesStruct(output=decode_hidden, output_shape=list(decode_hidden.shape)), - embed=EmbeddingsStruct(prefill=prefill_hidden, prefill_shape=list(prefill_hidden.shape)), - ids=IdsStruct(prompt=list(prompt_token_ids), output=list(thinker_output_ids)), - ) - ) - talker_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID] - + [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids)) - + [TALKER_CODEC_END_TOKEN_ID], - additional_information=additional_information, - multi_modal_data=( - multi_modal_data[thinker_output.request_id] - if requires_multimodal_data and multi_modal_data is not None - else None - ), - mm_processor_kwargs=None, - ) - ) - return talker_inputs - - -def talker2code2wav( - source_outputs, - _prompt: OmniTokensPrompt | TextPrompt = None, - _requires_multimodal_data: bool = False, -): - code2wav_inputs = [] - for talker_output in source_outputs: - output = talker_output.outputs[0] - token_ids = list(output.cumulative_token_ids) - if token_ids and token_ids[0] == TALKER_CODEC_START_TOKEN_ID: - token_ids = token_ids[1:] - if token_ids and token_ids[-1] == TALKER_CODEC_END_TOKEN_ID: - token_ids = token_ids[:-1] - if not token_ids: - continue - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=token_ids, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - return code2wav_inputs - - # ============================================================================ # Worker-connector data plane (non-async-chunk path). # Both transitions ship payloads via the worker connector @@ -101,9 +29,7 @@ def talker2code2wav( # packs an OmniPayload-shaped dict (embed.prefill / # hidden_states.output / ids.prompt / ids.output) for the talker, which # the talker's ``talker_preprocess`` reads from -# ``model_intermediate_buffer``. The shape matches what legacy -# ``thinker2talker`` writes into ``additional_information`` as a debug -# fallback; ``thinker2talker_token_only`` only allocates prompt slots. +# ``model_intermediate_buffer``. # - talker->code2wav strips TALKER_CODEC_{START,END} boundary tokens # and ships the codec token ids. # ============================================================================ @@ -212,9 +138,7 @@ def talker2code2wav_full_payload( # prefill+decode trajectory and packs an OmniPayload-shaped dict that # the talker's ``talker_preprocess`` reads from # ``model_intermediate_buffer``. ``thinker2talker_token_only`` only -# allocates the talker's codec prompt slots; legacy -# ``thinker2talker`` above remains as a debug fallback that bundles the -# same shape into ``additional_information``. +# allocates the talker's codec prompt slots. # ============================================================================ diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index d624b4a49a1..edde6972d08 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -10,7 +10,6 @@ import torch from vllm.inputs import TextPrompt -from vllm.platforms import current_platform from vllm_omni.data_entry_keys import ( CodesStruct, @@ -569,25 +568,31 @@ def thinker2talker_full_payload( output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or []) all_token_ids = list(prompt_token_ids) + list(output_token_ids) - # Trim the trailing stop-token row from the accumulated thinker output. - # The accumulator captures one hidden-state row per executed thinker - # forward (prefill + every decode step including the one that emitted - # the stop_token), so for a finished request thinker_emb has exactly one - # row more than the rows the talker should consume. async_chunk's - # chunk-0 path naturally captures only the prefill / non-stop portion, - # which is why the [async_chunk] parametrization passes while [default] - # over-generates one codec frame on short outputs (e.g. - # test_one_word_prompt_001[default]: audio extends "London" with - # spurious phonemes). - if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0: + # Drop the terminal stop-token row only when more than one row was + # accumulated; trimming a single row would ship 0 conditioning tensors + # while ids still has tokens and break talker prefill alignment. + if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 1: thinker_emb_prefill = thinker_emb[:-1] else: thinker_emb_prefill = thinker_emb - if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0: + if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 1: thinker_hid_prefill = thinker_hid[:-1] else: thinker_hid_prefill = thinker_hid + emb_rows = int(thinker_emb_prefill.shape[0]) if isinstance(thinker_emb_prefill, torch.Tensor) else 0 + hid_rows = int(thinker_hid_prefill.shape[0]) if isinstance(thinker_hid_prefill, torch.Tensor) else 0 + if len(all_token_ids) > 0 and (emb_rows == 0 or hid_rows == 0): + logger.warning( + "thinker2talker_full_payload: empty thinker conditioning for req=%s " + "(ids_len=%s embed_rows=%s hidden_rows=%s); withholding payload.", + rid, + len(all_token_ids), + emb_rows, + hid_rows, + ) + return None + payload: OmniPayload = { "embed": { "prefill": thinker_emb_prefill.detach().cpu(), @@ -599,7 +604,6 @@ def thinker2talker_full_payload( "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)}, "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, } - payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu") speaker = extract_speaker_from_request(request) if speaker is not None: payload["speaker"] = speaker @@ -609,154 +613,30 @@ def thinker2talker_full_payload( return payload -def thinker2talker( +def thinker2talker_token_only( source_outputs: list[Any], prompt: OmniTokensPrompt | TextPrompt | None = None, requires_multimodal_data: bool = False, streaming_context: Any | None = None, ) -> list[OmniTokensPrompt]: - """ - Process thinker outputs to create talker inputs. - - Workflow: - 1. Extract thinker's text generation outputs (token IDs + hidden states) - 2. Split hidden states into: prompt embeddings + generated embeddings - 3. Package for talker with additional information - - In PD disaggregation mode, merges prefill-stage prompt embeddings with - decode-stage generated embeddings before handing off to the talker. - - Args: - prompt: Original prompt data - requires_multimodal_data: Whether multimodal data is required - - Returns: - List of OmniTokensPrompt for talker stage - """ - thinker_outputs = source_outputs - talker_inputs: list[OmniTokensPrompt] = [] - - device = torch.device(current_platform.device_type) - - # Process each thinker output - for i, thinker_output in enumerate(thinker_outputs): - output = thinker_output.outputs[0] - req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) - prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids) - output_ids = _ensure_list(output.cumulative_token_ids) - is_streaming_session = bool(getattr(streaming_context, "enabled", False)) - if is_streaming_session: - prompt_token_ids, output_ids = _get_streaming_talker_tokens( - req_id, - prompt_token_ids, - output_ids, - getattr(streaming_context, "new_prompt_len_snapshot", None), - streaming_context, - clear_state=bool(getattr(thinker_output, "finished", False)), - ) - thinker_sequences = prompt_token_ids + output_ids - thinker_input_ids = prompt_token_ids - new_seq_length = len(prompt_token_ids + output_ids) - 1 - thinker_mm_raw = getattr(output, "multimodal_output", None) - if not isinstance(thinker_mm_raw, Mapping): - logger.debug("thinker2talker: skip req=%s due to empty multimodal_output", req_id) - continue - thinker_mm: OmniPayload = thinker_mm_raw - mm_hs = thinker_mm.get("hidden_states", {}) - mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {} - emb_layer = _layer_tensor(mm_layers, _EMBED_LAYER_KEY) - hid_layer = _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) - if emb_layer is None or hid_layer is None: - logger.debug("thinker2talker: skip req=%s due to missing hidden-state layers", req_id) - continue - thinker_emb = emb_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:] - thinker_hid = hid_layer.detach().to(device=device, dtype=torch.float)[-new_seq_length:] - - prefill_mm: dict[str, Any] | None = None - prefill_mm = _get_prefill_multimodal_output(req_id, streaming_context) - - if prefill_mm is not None: - expected_total = len(prompt_token_ids) + len(output_ids) - try: - thinker_emb, thinker_hid = _merge_pd_embeddings( - thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total - ) - except Exception as exc: - logger.warning("[PD] Could not merge prefill embeddings: %s", exc) - - payload = OmniPayloadStruct( - embed=EmbeddingsStruct( - prefill=thinker_emb, - tts_bos=_resolve_tts_token_embedding( - "tts_bos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device - ), - tts_eos=_resolve_tts_token_embedding( - "tts_eos", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device - ), - tts_pad=_resolve_tts_token_embedding( - "tts_pad", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device - ), - ), - hidden_states=HiddenStatesStruct(output=thinker_hid), - ids=IdsStruct(all=thinker_sequences, prompt=thinker_input_ids), - speaker=extract_speaker_from_prompt(prompt, index=i), - language=extract_language_from_prompt(prompt, index=i), - ) - info = to_dict(payload) - prompt_len = _compute_talker_prompt_ids_length(info, device=device) - - talker_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[0] * prompt_len, - additional_information=info, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) + """Orchestrator-side placeholder builder for Stage-1 (Talker) when + ``async_chunk=False``. - return talker_inputs + After the communication-layer refactor, this function only allocates a + placeholder ``prompt_token_ids`` of the correct length so the scheduler can + reserve KV-cache slots. It does **not** forward bulk tensors. + Bulk talker conditioning is sent through the connector. Speaker and + language are also copied from the original prompt so they survive when + Stage-0 request metadata is unavailable to the connector payload. -def thinker2talker_token_only( - source_outputs: list[Any], - prompt: OmniTokensPrompt | TextPrompt | None = None, - requires_multimodal_data: bool = False, - streaming_context: Any | None = None, -) -> list[OmniTokensPrompt]: - """Non-async-chunk Stage-1 input builder for the connector data plane. - - The worker connector (Stage-0 ``thinker2talker_full_payload`` → - ``_sync_local_stage_payloads``) supplies the bulk talker conditioning - tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``. - The orchestrator only needs to ship a placeholder prefill prompt of the - correct length so the scheduler can allocate KV-cache slots. - - Small per-request voice metadata (``speaker`` / ``language``) is forwarded - here from the user prompt so the worker's line-408 buffer seed picks it - up. The connector-side ``extract_speaker_from_request`` reads the - strongly-typed ``request.additional_information.entries["speaker"]`` which - currently does not always round-trip the user-supplied voice; until that - plumbing is normalized, providing the small fields directly preserves - voice selection (regression discovered on Buildkite 9668: - ``test_speaker_002[default]`` lost the preset voice). + ``prompt`` / ``requires_multimodal_data`` are kept for call-site signature + compatibility with other orchestrator input processors; they are unused. """ talker_inputs: list[OmniTokensPrompt] = [] for i, thinker_output in enumerate(source_outputs): output = thinker_output.outputs[0] req_id = str(getattr(thinker_output, "request_id", f"idx-{i}")) - # Skip-on-missing parity with thinker2talker_full_payload: if the - # connector builder would drop this request (no MM dict or missing - # hidden-state layers), do the same here so the worker buffer - # presence agrees with the orchestrator's scheduling decision. - thinker_mm_raw = getattr(output, "multimodal_output", None) - if not isinstance(thinker_mm_raw, Mapping): - logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id) - continue - mm_hs = thinker_mm_raw.get("hidden_states", {}) - mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, Mapping) else {} - if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None: - logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id) - continue prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids) output_ids = _ensure_list(output.cumulative_token_ids) is_streaming_session = bool(getattr(streaming_context, "enabled", False)) @@ -773,21 +653,17 @@ def thinker2talker_token_only( thinker_input_ids = prompt_token_ids info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}} prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu") - - # Forward only small voice metadata; bulk tensors come from the - # connector path via _sync_local_stage_payloads. - small_info: dict[str, Any] = {} - speaker = extract_speaker_from_prompt(prompt, index=i) - if speaker is not None: - small_info["speaker"] = speaker - language = extract_language_from_prompt(prompt, index=i) - if language is not None: - small_info["language"] = language - + # Keep this fallback until the connector reliably preserves voice metadata. + additional_information = to_dict( + OmniPayloadStruct( + speaker=extract_speaker_from_prompt(prompt, index=i), + language=extract_language_from_prompt(prompt, index=i), + ) + ) talker_inputs.append( OmniTokensPrompt( prompt_token_ids=[0] * prompt_len, - additional_information=(small_info if small_info else None), + additional_information=additional_information or None, multi_modal_data=None, mm_processor_kwargs=None, ) @@ -949,59 +825,5 @@ def talker2code2wav_full_payload( ) return { "codes": {"audio": codec_codes}, - "code_predictor_codes": codec_codes, "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, } - - -def talker2code2wav( - source_outputs: list[Any], - _prompt: OmniTokensPrompt | TextPrompt | None = None, - _requires_multimodal_data: bool = False, - streaming_context: Any | None = None, -) -> list[OmniTokensPrompt]: - """ - Process talker outputs to create code2wav inputs. - - Workflow: - 1. Extract talker's codec code outputs (8-layer RVQ codes) - 2. Flatten codes for code2wav input - 3. Package for code2wav stage - - Args: - Returns: - List of OmniTokensPrompt for code2wav stage - """ - talker_outputs = source_outputs - code2wav_inputs: list[OmniTokensPrompt] = [] - # Process each talker output - for i, talker_output in enumerate(talker_outputs): - output = talker_output.outputs[0] - req_id = str(getattr(talker_output, "request_id", f"idx-{i}")) - cur_seq_len = len(output.cumulative_token_ids) - 1 - seq_len = cur_seq_len - is_streaming_session = bool(getattr(streaming_context, "enabled", False)) - if is_streaming_session: - seq_len = _get_streaming_codec_delta_len(cur_seq_len, req_id, talker_output, streaming_context) - mm_raw = getattr(output, "multimodal_output", None) - if not isinstance(mm_raw, Mapping): - logger.debug("talker2code2wav: skip req=%s due to empty multimodal_output", req_id) - continue - mm: OmniPayload = mm_raw - if "codes" not in mm or not isinstance(mm.get("codes"), dict) or "audio" not in mm["codes"]: - logger.debug("talker2code2wav: skip req=%s due to missing codes.audio", req_id) - continue - # Extract codec codes from talker output - # Expected shape: [8, seq_len] (8-layer RVQ codes) - codec_codes = ( - mm["codes"]["audio"][-seq_len:].to(torch.long).transpose(0, 1).cpu().to(torch.long).reshape(-1).tolist() - ) # 16, seq_len - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=codec_codes, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return code2wav_inputs diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index c67e4a1af3c..597fc4db53e 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -27,101 +27,15 @@ logger = init_logger(__name__) -def talker2code2wav( - source_outputs: list[Any], - prompt: Any = None, - _requires_multimodal_data: bool = False, -) -> list[Any]: - """Non-async: collect all talker codes, then pass to code2wav at once.""" - from vllm_omni.inputs.data import OmniTokensPrompt - - talker_outputs = source_outputs - code2wav_inputs: list[OmniTokensPrompt] = [] - for i, talker_output in enumerate(talker_outputs): - if not talker_output.finished: - # Non-async decode should only run once, after talker has - # accumulated the final code sequence. - continue - output = talker_output.outputs[0] - mm = output.multimodal_output - mm_codes = mm.get("codes", {}) - - # audio_codes shape: [num_frames, Q] where Q=num_quantizers (16) - audio_codes = mm_codes["audio"].to(torch.long) - token_ids = output.cumulative_token_ids - - # token_ids provides an upper bound on the newly generated codec span. - # audio_codes may still contain zero-padded / invalid rows, so trim only - # after filtering valid frames instead of trying to align EOS indices. - seq_len = max(len(token_ids) - 1, 0) - # Filter invalid frames: zero-padded (EOS), out-of-range values (e.g. - # stop_token_id=2150 exceeds codebook_size=2048), and negative - # sentinels (e.g. -1 padding). - _CODEBOOK_SIZE = 2048 - valid_mask = ( - (audio_codes >= 0).all(dim=1) & audio_codes.any(dim=1) & (audio_codes.max(dim=1).values < _CODEBOOK_SIZE) - ) - audio_codes = audio_codes[valid_mask] - if seq_len > 0 and audio_codes.ndim == 2 and int(audio_codes.shape[0]) > seq_len: - audio_codes = audio_codes[-seq_len:] - ref_code = mm_codes.get("ref") - ref_code_len = mm.get("meta", {}).get("ref_code_len") - if isinstance(ref_code_len, torch.Tensor): - ref_code_len = int(ref_code_len.reshape(-1)[-1].item()) if ref_code_len.numel() > 0 else 0 - elif ref_code_len is None: - ref_code_len = 0 - else: - ref_code_len = int(ref_code_len) - if isinstance(ref_code, list): - ref_code = ref_code[0] if ref_code else None - if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0: - ref_code = ref_code.to(torch.long).cpu().contiguous() - if ref_code.ndim == 1: - num_quantizers = int(audio_codes.shape[1]) if audio_codes.ndim == 2 and audio_codes.shape[1] > 0 else 16 - if ref_code.numel() % num_quantizers != 0: - logger.warning( - "Ignoring malformed ref_code with %d elements not divisible by num_quantizers=%d", - ref_code.numel(), - num_quantizers, - ) - ref_code = None - else: - ref_code = ref_code.reshape(-1, num_quantizers) - elif ref_code.ndim != 2: - logger.warning("Ignoring malformed ref_code shape %s", tuple(ref_code.shape)) - ref_code = None - if isinstance(ref_code, torch.Tensor) and ref_code_len > 0 and int(ref_code.shape[0]) > ref_code_len: - logger.warning( - "Trimming ref_code from %d frames to ref_code_len=%d before Code2Wav.", - int(ref_code.shape[0]), - ref_code_len, - ) - ref_code = ref_code[:ref_code_len] - if not isinstance(ref_code, torch.Tensor): - ref_code_len = 0 - else: - ref_code_len = int(ref_code.shape[0]) - audio_codes = torch.cat([ref_code.to(audio_codes.device), audio_codes], dim=0) - else: - ref_code_len = 0 - # Code2Wav expects codebook-major flat: [Q*num_frames] - codec_codes = audio_codes.transpose(0, 1).cpu().reshape(-1).tolist() - additional_information = to_dict( - OmniPayloadStruct( - meta=MetaStruct(left_context_size=ref_code_len) if ref_code_len > 0 else None, - speaker=extract_speaker_from_prompt(prompt, index=i), - language=extract_language_from_prompt(prompt, index=i), - ) - ) - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=codec_codes, - multi_modal_data=None, - mm_processor_kwargs=None, - additional_information=additional_information if additional_information else None, - ) - ) - return code2wav_inputs +def _qwen3_tts_empty_finished_payload(): + """Empty-but-finished stage payload. Returned instead of None on a degenerate + talker take so the Stage-1 wait gate releases (code2wav handles empty codec + input -> 0-sample audio, request finished) rather than the connector silently + dropping the request and Stage-1 hanging to connector_get_max_wait.""" + return { + "codes": {"audio": torch.zeros(0, dtype=torch.long)}, + "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, + } def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None: @@ -340,10 +254,7 @@ def talker2code2wav_async_chunk( def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor: - """Filter zero-padded, out-of-range, and negative-padded codec frames. - - Mirrors the orchestrator-path body in `talker2code2wav` above. - """ + """Filter zero-padded, out-of-range, and negative-padded codec frames.""" if not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: return audio_codes if audio_codes.ndim != 2: @@ -356,10 +267,8 @@ def _filter_audio_codes_qwen3_tts(audio_codes: torch.Tensor) -> torch.Tensor: def _coerce_ref_code_len(raw) -> int: """Coerce mm["meta"]["ref_code_len"] / pooling_output["meta.ref_code_len"] - raw value (Tensor | int | None) into a non-negative int. Mirrors the - extraction inlined in the legacy ``talker2code2wav`` path; clamps any - negative input to 0 since downstream code treats this as a non-negative - frame count.""" + raw value (Tensor | int | None) into a non-negative int; clamps any + negative input to 0 since downstream code treats this as a frame count.""" if isinstance(raw, torch.Tensor): value = int(raw.reshape(-1)[-1].item()) if raw.numel() > 0 else 0 elif raw is None: @@ -370,7 +279,7 @@ def _coerce_ref_code_len(raw) -> int: def _normalize_ref_code(ref_code, num_quantizers: int, ref_code_len: int): - """Coerce ref_code into a [ref_len, Q] tensor or None. Mirrors orchestrator path.""" + """Coerce ref_code into a [ref_len, Q] tensor or None.""" if isinstance(ref_code, list): ref_code = ref_code[0] if ref_code else None if not isinstance(ref_code, torch.Tensor) or ref_code.numel() == 0: @@ -396,8 +305,7 @@ def talker2code2wav_token_only( Sized to the expected codec token count (codebook-major flat: Q * (ref_frames + audio_frames)). Speaker / language metadata are - extracted from `prompt` and threaded via `additional_information` - (orchestrator-style; same as the legacy `talker2code2wav` builder). + extracted from `prompt` and threaded via `additional_information`. Actual codec ids are delivered via the worker connector payload built by `talker2code2wav_full_payload`. """ @@ -462,8 +370,8 @@ def talker2code2wav_full_payload( Reads accumulated codec from `pooling_output["codes.audio"]` (CONCAT across steps via flatten_payload), latest `pooling_output["codes.ref"]` (prefill-emitted), and latest `pooling_output["meta.ref_code_len"]`. - Replicates the orchestrator-path body of `talker2code2wav` (filter, - crop to seq_len, prepend ref, codebook-major flatten). + Filters invalid frames, crops to seq_len, prepends ref, and flattens + codebook-major for code2wav consumption. """ del transfer_manager rid = getattr(request, "request_id", "?") @@ -474,7 +382,7 @@ def talker2code2wav_full_payload( type(pooling_output).__name__, rid, ) - return None + return _qwen3_tts_empty_finished_payload() # codes.audio — try flat dotted first (flatten_payload), then nested fallback. audio = pooling_output.get("codes.audio") @@ -489,7 +397,7 @@ def talker2code2wav_full_payload( list(pooling_output.keys()), rid, ) - return None + return _qwen3_tts_empty_finished_payload() audio = audio.to(torch.long) audio = _filter_audio_codes_qwen3_tts(audio) if audio.numel() == 0: @@ -498,7 +406,7 @@ def talker2code2wav_full_payload( "filter (negative/all-zero/out-of-range rows dropped) for req=%s.", rid, ) - return None + return _qwen3_tts_empty_finished_payload() output_token_ids = list(getattr(request, "output_token_ids", None) or []) seq_len = max(len(output_token_ids) - 1, 0) diff --git a/vllm_omni/model_executor/stage_input_processors/step_audio2.py b/vllm_omni/model_executor/stage_input_processors/step_audio2.py index 5e015aa7c50..c0c0932b27a 100644 --- a/vllm_omni/model_executor/stage_input_processors/step_audio2.py +++ b/vllm_omni/model_executor/stage_input_processors/step_audio2.py @@ -7,6 +7,7 @@ from vllm.inputs import TextPrompt from vllm.logger import init_logger +from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayloadStruct from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import ( @@ -35,14 +36,14 @@ def _ensure_list(x): def thinker2token2wav_async_chunk( transfer_manager: Any, - pooling_output: dict[str, Any], + multimodal_output: dict[str, Any] | None, request: OmniEngineCoreRequest, is_finished: bool = False, -) -> dict[str, Any] | None: +) -> OmniPayloadStruct | None: """ Async chunk processor: stream audio tokens from Thinker to Token2Wav. - Unlike Qwen3-Omni which passes hidden states via pooling_output, + Unlike Qwen3-Omni which passes hidden states via multimodal_output, Step Audio2 generates audio tokens as part of the LLM autoregressive output (token IDs >= audio_start). This processor extracts those audio tokens from the running token stream and sends them in chunks. @@ -62,14 +63,13 @@ def thinker2token2wav_async_chunk( Args: transfer_manager: OmniChunkTransferAdapter instance (has ``code_prompt_token_ids`` defaultdict for state tracking). - pooling_output: Pooler output dict (unused — audio tokens are in + multimodal_output: Multimodal output dict (unused — audio tokens are in the token ID stream, not in hidden states). request: Current engine request with access to all_token_ids. is_finished: Whether the upstream request has finished generating. Returns: - dict with framework-compatible fields, or None if the chunk is - not yet ready. + Structured connector payload, or None if the chunk is not yet ready. """ audio_start = DEFAULT_TOKEN_CONFIG.audio_start audio_eos = DEFAULT_TOKEN_CONFIG.audio_eos @@ -100,18 +100,22 @@ def thinker2token2wav_async_chunk( # Last chunk: send all remaining tokens if available <= 0: # No audio tokens at all (text-only response) — send EOF marker - return { - "code_predictor_codes": [], - "left_context_size": 1, # 1 = last chunk - "finished": torch.tensor(True, dtype=torch.bool), - } + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct( + left_context_size=1, # 1 = last chunk + finished=torch.tensor(True, dtype=torch.bool), + ), + ) remaining_tokens = audio_tokens[consumed:] transfer_manager.code_prompt_token_ids[request_id].extend(remaining_tokens) - return { - "code_predictor_codes": remaining_tokens, - "left_context_size": 1, # 1 = last chunk - "finished": torch.tensor(True, dtype=torch.bool), - } + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor(remaining_tokens, dtype=torch.long)), + meta=MetaStruct( + left_context_size=1, # 1 = last chunk + finished=torch.tensor(True, dtype=torch.bool), + ), + ) else: # Non-last chunk: need chunk_size + pre_lookahead_len tokens required = chunk_size + pre_lookahead_len @@ -120,11 +124,13 @@ def thinker2token2wav_async_chunk( chunk_tokens = audio_tokens[consumed : consumed + required] # Only consume chunk_size tokens; the lookahead portion is re-used transfer_manager.code_prompt_token_ids[request_id].extend(audio_tokens[consumed : consumed + chunk_size]) - return { - "code_predictor_codes": chunk_tokens, - "left_context_size": 0, # 0 = not last chunk - "finished": torch.tensor(False, dtype=torch.bool), - } + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor(chunk_tokens, dtype=torch.long)), + meta=MetaStruct( + left_context_size=0, # 0 = not last chunk + finished=torch.tensor(False, dtype=torch.bool), + ), + ) def thinker2token2wav( diff --git a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py index 2928a7cf31b..af2cb0a257b 100644 --- a/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/voxtral_tts.py @@ -2,7 +2,6 @@ from typing import Any import torch -from vllm.inputs import TextPrompt from vllm.logger import init_logger from vllm_omni.data_entry_keys import ( @@ -11,36 +10,33 @@ OmniPayload, OmniPayloadStruct, ) -from vllm_omni.inputs.data import OmniTokensPrompt logger = init_logger(__name__) -def generator2tokenizer( - source_outputs: list[Any], - _prompt: OmniTokensPrompt | TextPrompt = None, - _requires_multimodal_data: bool = False, -): - tokenizer_inputs = [] - for generator_output in source_outputs: - output = generator_output.outputs[0] - audio_tokens = torch.cat(output.multimodal_output["audio"], dim=-1).flatten().detach().cpu().tolist() - tokenizer_inputs.append( - OmniTokensPrompt( - prompt_token_ids=audio_tokens, - multi_modal_data=None, - mm_processor_kwargs=None, - ) - ) - - return tokenizer_inputs +def _codec_audio_tensors(multimodal_output: OmniPayload | dict[str, Any]) -> list[torch.Tensor] | None: + """Inter-stage codec frames from ``codes.audio``.""" + if not isinstance(multimodal_output, Mapping): + return None + codes = multimodal_output.get("codes") + if not isinstance(codes, Mapping) or "audio" not in codes: + return None + audio = codes["audio"] + if isinstance(audio, torch.Tensor): + return [audio] + if isinstance(audio, list) and audio: + return audio + return None def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torch.Tensor | None: - audio = multimodal_output.get("audio") - if not isinstance(audio, torch.Tensor) or audio.numel() == 0: + audio_tensors = _codec_audio_tensors(multimodal_output) + if not audio_tensors: + return None + frame = audio_tensors[-1] + if not isinstance(frame, torch.Tensor) or frame.numel() == 0: return None - return audio.flatten() + return frame.flatten() def generator2tokenizer_async_chunk( diff --git a/vllm_omni/model_extras/__init__.py b/vllm_omni/model_extras/__init__.py new file mode 100644 index 00000000000..5311f495b03 --- /dev/null +++ b/vllm_omni/model_extras/__init__.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm_omni.model_extras.registry import ( + build_image_to_image_prompt, + build_image_to_video_prompt, + build_text_to_image_prompt, + get_extra_body_params, + get_extra_output_params, + get_model_class_name, + should_init_extra_args_for_non_diffusion_stages, +) + +__all__ = [ + "build_image_to_image_prompt", + "build_image_to_video_prompt", + "build_text_to_image_prompt", + "get_extra_body_params", + "get_extra_output_params", + "get_model_class_name", + "should_init_extra_args_for_non_diffusion_stages", +] diff --git a/vllm_omni/model_extras/audiox.py b/vllm_omni/model_extras/audiox.py new file mode 100644 index 00000000000..220972148d5 --- /dev/null +++ b/vllm_omni/model_extras/audiox.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +# Request params the AudioX pipeline reads from ``sampling_params.extra_args`` +# (see ``pipeline_audiox.py`` ``forward``). Declaring them here routes +# ``extra_body`` (online) / ``--extra-body`` (offline) keys into ``extra_args`` +# instead of silently dropping them. +AUDIOX_EXTRA_BODY_PARAMS = frozenset( + { + "audiox_task", # t2a/t2m/v2a/v2m/tv2a/tv2m + "seconds_start", + "seconds_total", + "sigma_min", + "sigma_max", + "cfg_rescale", + "video_path", # v2*/tv2* conditioning + "audio_path", # optional reference-audio conditioning + } +) +# Echoed back in diffusion output metadata (pipeline_audiox.py). +AUDIOX_EXTRA_OUTPUT_PARAMS = frozenset( + { + "audiox_task", + } +) diff --git a/vllm_omni/model_extras/bagel.py b/vllm_omni/model_extras/bagel.py new file mode 100644 index 00000000000..e8cbcbf11c3 --- /dev/null +++ b/vllm_omni/model_extras/bagel.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +from PIL import Image + + +def build_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + text_prompt: dict[str, Any] = { + "prompt": f"<|im_start|>{prompt}<|im_end|>", + "modalities": ["image"], + "mm_processor_kwargs": { + "target_h": height, + "target_w": width, + "modalities": ["image"], + }, + } + if negative_prompt is not None: + text_prompt["negative_prompt"] = negative_prompt + return text_prompt + + +def build_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + img_prompt: dict[str, Any] = { + "prompt": f"<|fim_middle|><|im_start|>{prompt}<|im_end|>", + "modalities": ["img2img"], + "multi_modal_data": {"img2img": input_image}, + "mm_processor_kwargs": {"modalities": ["img2img"]}, + } + if height is not None: + img_prompt["mm_processor_kwargs"]["target_h"] = height + if width is not None: + img_prompt["mm_processor_kwargs"]["target_w"] = width + if negative_prompt is not None: + img_prompt["negative_prompt"] = negative_prompt + return img_prompt + + +BAGEL_EXTRA_BODY_PARAMS = frozenset( + { + "cfg_text_scale", + "cfg_img_scale", + "cfg_interval", + "cfg_renorm_type", + "cfg_renorm_min", + "negative_prompt", + "think", + "max_think_tokens", + "do_sample", + "text_temperature", + "timestep_shift", + } +) +BAGEL_EXTRA_OUTPUT_PARAMS = frozenset( + { + "text_output", + "think_text", + } +) +BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py new file mode 100644 index 00000000000..a17e8193f00 --- /dev/null +++ b/vllm_omni/model_extras/cosmos3.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + + +def build_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + """Build a Cosmos3 T2I prompt by selecting image output modality. + + Cosmos3 uses the same pipeline class for T2I and video modes. The OpenAI + image endpoint selects T2I by adding ``modalities=["image"]``. ``height`` + and ``width`` are accepted for registry compatibility; sizing travels + through sampling params instead of the prompt payload. + """ + del height, width + text_prompt: dict[str, Any] = { + "prompt": prompt, + "modalities": ["image"], + } + if negative_prompt is not None: + text_prompt["negative_prompt"] = negative_prompt + return text_prompt + + +COSMOS3_EXTRA_BODY_PARAMS = frozenset( + { + "flow_shift", + "max_sequence_length", + "use_resolution_template", + "use_duration_template", + "use_system_prompt", + "system_prompt", + "negative_prompt", + "guardrails", + "condition_frame_indexes_vision", + "condition_video_keep", + "generate_sound", + "sound_gen", + "sound_duration", + "audio_duration", + "action_mode", + "action", + "domain_name", + "domain_id", + "raw_action_dim", + "action_chunk_size", + "action_space", + "action_fps", + "image_height", + "image_width", + "history_length", + "conditioning_fps", + "resolution", + "image_size", + "use_state", + "observation", + "robot_obs", + "deterministic_seed", + "session_id", + } +) +COSMOS3_EXTRA_OUTPUT_PARAMS = frozenset( + { + "action", + "raw_action_dim", + "domain_id", + "action_mode", + } +) diff --git a/vllm_omni/model_extras/helios.py b/vllm_omni/model_extras/helios.py new file mode 100644 index 00000000000..d56c5a65fcf --- /dev/null +++ b/vllm_omni/model_extras/helios.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +# Helios is a text/image/video -> video model. Its model-specific knobs are read +# from sampling_params.extra_args in pipeline_helios.py; declaring them here +# routes request `extra_body` fields into OmniDiffusionSamplingParams.extra_args +# so the model can be driven through the standard task example. +HELIOS_EXTRA_BODY_PARAMS = frozenset( + { + # pyramid / staged sampling + "is_enable_stage2", + "pyramid_num_stages", + "pyramid_num_inference_steps_list", + "is_amplify_first_chunk", + "is_skip_first_chunk", + # CFG-zero* controls + "use_cfg_zero_star", + "use_zero_init", + "zero_steps", + # conditioning inputs + "image", + "video", + "add_noise_to_image_latents", + "image_noise_sigma_min", + "image_noise_sigma_max", + "add_noise_to_video_latents", + "video_noise_sigma_min", + "video_noise_sigma_max", + } +) +HELIOS_EXTRA_OUTPUT_PARAMS = frozenset() diff --git a/vllm_omni/model_extras/magi_human.py b/vllm_omni/model_extras/magi_human.py new file mode 100644 index 00000000000..470a1455b4f --- /dev/null +++ b/vllm_omni/model_extras/magi_human.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +# MagiHuman is a text -> video+audio model. Its model-specific knobs are read +# from sampling_params.extra_args in pipeline_magi_human.py; declaring them here +# routes request `extra_body` fields into OmniDiffusionSamplingParams.extra_args +# so the model can be driven through the standard task example. +MAGI_HUMAN_EXTRA_BODY_PARAMS = frozenset( + { + "seconds", + "audio_path", + "image_path", + "sr_height", + "sr_width", + "sr_num_inference_steps", + } +) +MAGI_HUMAN_EXTRA_OUTPUT_PARAMS = frozenset() diff --git a/vllm_omni/model_extras/mammothmodal2_preview.py b/vllm_omni/model_extras/mammothmodal2_preview.py new file mode 100644 index 00000000000..c17e428a874 --- /dev/null +++ b/vllm_omni/model_extras/mammothmodal2_preview.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +# AR image-grid patch size (pixels per visual token). Mirrors the constant the +# bespoke MammothModa2 example used to convert image dimensions -> AR grid dims. +_PATCH_SIZE = 16 + +_AR_SYSTEM_PROMPT = "You are a helpful image generator." + + +def build_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + """Build the MammothModa2 AR-stage prompt for text-to-image generation. + + Reproduces the prompt string and the structural ``additional_information`` + that the former bespoke example (``run_mammothmoda2_t2i.py``) constructed, + deriving the AR image grid from ``height`` / ``width``. + + Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``, + ``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``. + Config-derived token ids (``eol_token_id``, ``visual_token_start_id``, + ``visual_token_end_id``, ``visual_ids``) are sourced inside the AR stage + rather than passed by the caller -- see the ar2dit stage input processor. + + MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale`` + and has no explicit negative-prompt path, so ``negative_prompt`` is accepted + for signature compatibility but not injected. + """ + h = height or 1024 + w = width or 1024 + ar_height = h // _PATCH_SIZE + ar_width = w // _PATCH_SIZE + + return { + "prompt": ( + f"<|im_start|>system\n{_AR_SYSTEM_PROMPT}<|im_end|>\n" + f"<|im_start|>user\n{prompt}<|im_end|>\n" + "<|im_start|>assistant\n" + f"<|image start|>{ar_width}*{ar_height}<|image token|>" + ), + "additional_information": { + "omni_task": ["t2i"], + "ar_width": [ar_width], + "ar_height": [ar_height], + "image_height": [h], + "image_width": [w], + }, + } + + +MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS = frozenset( + { + "text_guidance_scale", + "cfg_range", + # MammothModa2's DiT stage consumes inputs via the kwargs interface rather + # than OmniDiffusionRequest, so the standard --num-inference-steps flag does + # not reach it; it is routed through extra_body like the CFG knobs. + "num_inference_steps", + } +) +MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset() +MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/ming_flash_omni.py b/vllm_omni/model_extras/ming_flash_omni.py new file mode 100644 index 00000000000..8fece255d7d --- /dev/null +++ b/vllm_omni/model_extras/ming_flash_omni.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from typing import Any + +from PIL import Image + + +def build_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + text_prompt: dict[str, Any] = { + "prompt": prompt, + "modalities": ["image"], + "mm_processor_kwargs": {"modalities": ["image"]}, + } + if height is not None: + text_prompt["mm_processor_kwargs"]["target_h"] = height + if width is not None: + text_prompt["mm_processor_kwargs"]["target_w"] = width + if negative_prompt is not None: + text_prompt["negative_prompt"] = negative_prompt + return text_prompt + + +def build_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + img_prompt: dict[str, Any] = { + "prompt": prompt, + "modalities": ["img2img"], + "multi_modal_data": {"img2img": input_image}, + "mm_processor_kwargs": {"modalities": ["img2img"]}, + } + if height is not None: + img_prompt["mm_processor_kwargs"]["target_h"] = height + if width is not None: + img_prompt["mm_processor_kwargs"]["target_w"] = width + if negative_prompt is not None: + img_prompt["negative_prompt"] = negative_prompt + return img_prompt + + +# For Image-generation (diffusion-stage) +MING_FLASH_OMNI_EXTRA_BODY_PARAMS = frozenset( + { + "height", + "width", + "steps", + "cfg", + "seed", + "byte5_text", + "negative_prompt", + } +) +MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS = frozenset() +MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py new file mode 100644 index 00000000000..79f8b436047 --- /dev/null +++ b/vllm_omni/model_extras/registry.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Callable, Mapping +from typing import Any + +from PIL import Image + +from vllm_omni.model_extras.audiox import ( + AUDIOX_EXTRA_BODY_PARAMS, + AUDIOX_EXTRA_OUTPUT_PARAMS, +) +from vllm_omni.model_extras.bagel import ( + BAGEL_EXTRA_BODY_PARAMS, + BAGEL_EXTRA_OUTPUT_PARAMS, + BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.bagel import ( + build_image_to_image_prompt as build_bagel_image_to_image_prompt, +) +from vllm_omni.model_extras.bagel import ( + build_text_to_image_prompt as build_bagel_text_to_image_prompt, +) +from vllm_omni.model_extras.cosmos3 import ( + COSMOS3_EXTRA_BODY_PARAMS, + COSMOS3_EXTRA_OUTPUT_PARAMS, +) +from vllm_omni.model_extras.cosmos3 import ( + build_text_to_image_prompt as build_cosmos3_text_to_image_prompt, +) +from vllm_omni.model_extras.helios import ( + HELIOS_EXTRA_BODY_PARAMS, + HELIOS_EXTRA_OUTPUT_PARAMS, +) +from vllm_omni.model_extras.magi_human import ( + MAGI_HUMAN_EXTRA_BODY_PARAMS, + MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, +) +from vllm_omni.model_extras.mammothmodal2_preview import ( + MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, + MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS, + MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.mammothmodal2_preview import ( + build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt, +) +from vllm_omni.model_extras.ming_flash_omni import ( + MING_FLASH_OMNI_EXTRA_BODY_PARAMS, + MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS, + MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, +) +from vllm_omni.model_extras.ming_flash_omni import ( + build_image_to_image_prompt as build_ming_flash_omni_image_to_image_prompt, +) +from vllm_omni.model_extras.ming_flash_omni import ( + build_text_to_image_prompt as build_ming_flash_omni_text_to_image_prompt, +) +from vllm_omni.model_extras.sensenova_u1 import ( + SENSENOVA_U1_EXTRA_BODY_PARAMS, + SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, +) +from vllm_omni.model_extras.vace import ( + VACE_EXTRA_BODY_PARAMS, + VACE_EXTRA_OUTPUT_PARAMS, +) +from vllm_omni.model_extras.vace import ( + build_image_to_video_prompt as build_vace_image_to_video_prompt, +) + +TextToImagePromptBuilder = Callable[ + [str, str | None, int | None, int | None], + dict[str, Any], +] +ImageToImagePromptBuilder = Callable[ + [str, str | None, "Image.Image | list[Image.Image]", int | None, int | None], + dict[str, Any], +] +ImageToVideoPromptBuilder = Callable[ + [ + str, + str | None, + "Mapping[str, Any]", + int | None, + int | None, + int | None, + ], + dict[str, Any], +] + + +def default_text_to_image_prompt( + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + result: dict[str, Any] = {"prompt": prompt} + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +def default_image_to_image_prompt( + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + result: dict[str, Any] = { + "prompt": prompt, + "multi_modal_data": {"image": input_image}, + } + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +def default_image_to_video_prompt( + prompt: str, + negative_prompt: str | None, + media_inputs: Mapping[str, Any], + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, +) -> dict[str, Any]: + del height, width, num_frames + if set(media_inputs) != {"image"} or not isinstance(media_inputs["image"], Image.Image): + raise ValueError("This model only supports a single --image input in the shared image-to-video example.") + return default_image_to_image_prompt(prompt, negative_prompt, media_inputs["image"]) + + +_EXTRA_SPECS: dict[str, dict[str, Any]] = { + "AudioXPipeline": { + "extra_body_params": AUDIOX_EXTRA_BODY_PARAMS, + "extra_output_params": AUDIOX_EXTRA_OUTPUT_PARAMS, + }, + "BagelPipeline": { + "extra_body_params": BAGEL_EXTRA_BODY_PARAMS, + "extra_output_params": BAGEL_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": BAGEL_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "text_to_image_prompt_builder": build_bagel_text_to_image_prompt, + "image_to_image_prompt_builder": build_bagel_image_to_image_prompt, + }, + "SenseNovaU1Pipeline": { + "extra_body_params": SENSENOVA_U1_EXTRA_BODY_PARAMS, + "extra_output_params": SENSENOVA_U1_EXTRA_OUTPUT_PARAMS, + }, + "Cosmos3OmniDiffusersPipeline": { + "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS, + "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS, + "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt, + }, + "Cosmos3OmniPipeline": { + "extra_body_params": COSMOS3_EXTRA_BODY_PARAMS, + "extra_output_params": COSMOS3_EXTRA_OUTPUT_PARAMS, + "text_to_image_prompt_builder": build_cosmos3_text_to_image_prompt, + }, + "MagiHumanPipeline": { + "extra_body_params": MAGI_HUMAN_EXTRA_BODY_PARAMS, + "extra_output_params": MAGI_HUMAN_EXTRA_OUTPUT_PARAMS, + }, + "HeliosPipeline": { + "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, + "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS, + }, + "HeliosPyramidPipeline": { + "extra_body_params": HELIOS_EXTRA_BODY_PARAMS, + "extra_output_params": HELIOS_EXTRA_OUTPUT_PARAMS, + }, + "WanVACEPipeline": { + "extra_body_params": VACE_EXTRA_BODY_PARAMS, + "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS, + "image_to_video_prompt_builder": build_vace_image_to_video_prompt, + }, + "MammothModa2DiTPipeline": { + "extra_body_params": MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS, + "extra_output_params": MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "text_to_image_prompt_builder": build_mammothmoda2_text_to_image_prompt, + }, + "MingImagePipeline": { + "extra_body_params": MING_FLASH_OMNI_EXTRA_BODY_PARAMS, + "extra_output_params": MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS, + "init_extra_args_for_non_diffusion_stages": MING_FLASH_OMNI_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES, + "text_to_image_prompt_builder": build_ming_flash_omni_text_to_image_prompt, + "image_to_image_prompt_builder": build_ming_flash_omni_image_to_image_prompt, + }, +} + + +def _get_spec(model_class_name: str | None) -> dict[str, Any] | None: + if not model_class_name: + return None + return _EXTRA_SPECS.get(model_class_name) + + +def get_model_class_name(omni: Any) -> str | None: + """Extract model_class_name from an Omni/AsyncOmni instance. + + This hides the internal ODConfig plumbing from example scripts. + """ + engine = getattr(omni, "engine", None) + if engine is None: + return None + od_config = getattr(engine, "od_config", None) + if od_config is None and hasattr(engine, "get_diffusion_od_config"): + od_config = engine.get_diffusion_od_config() + return getattr(od_config, "model_class_name", None) if od_config else None + + +def get_extra_body_params(model_class_name: str | None) -> frozenset[str]: + spec = _get_spec(model_class_name) + return spec.get("extra_body_params", frozenset()) if spec is not None else frozenset() + + +def get_extra_output_params(model_class_name: str | None) -> frozenset[str]: + spec = _get_spec(model_class_name) + return spec.get("extra_output_params", frozenset()) if spec is not None else frozenset() + + +def should_init_extra_args_for_non_diffusion_stages(model_class_name: str | None) -> bool: + spec = _get_spec(model_class_name) + return bool(spec and spec.get("init_extra_args_for_non_diffusion_stages", False)) + + +def build_text_to_image_prompt( + model_class_name: str | None, + prompt: str, + negative_prompt: str | None, + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + spec = _get_spec(model_class_name) + builder: TextToImagePromptBuilder = ( + spec.get("text_to_image_prompt_builder", default_text_to_image_prompt) + if spec is not None + else default_text_to_image_prompt + ) + return builder(prompt, negative_prompt, height, width) + + +def build_image_to_image_prompt( + model_class_name: str | None, + prompt: str, + negative_prompt: str | None, + input_image: Image.Image | list[Image.Image], + height: int | None = None, + width: int | None = None, +) -> dict[str, Any]: + spec = _get_spec(model_class_name) + builder: ImageToImagePromptBuilder = ( + spec.get("image_to_image_prompt_builder", default_image_to_image_prompt) + if spec is not None + else default_image_to_image_prompt + ) + return builder(prompt, negative_prompt, input_image, height, width) + + +def build_image_to_video_prompt( + model_class_name: str | None, + prompt: str, + negative_prompt: str | None, + media_inputs: Mapping[str, Any], + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, +) -> dict[str, Any]: + spec = _get_spec(model_class_name) + builder: ImageToVideoPromptBuilder = ( + spec.get("image_to_video_prompt_builder", default_image_to_video_prompt) + if spec is not None + else default_image_to_video_prompt + ) + return builder(prompt, negative_prompt, media_inputs, height, width, num_frames) diff --git a/vllm_omni/model_extras/sensenova_u1.py b/vllm_omni/model_extras/sensenova_u1.py new file mode 100644 index 00000000000..29fcc52c4a1 --- /dev/null +++ b/vllm_omni/model_extras/sensenova_u1.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +SENSENOVA_U1_EXTRA_BODY_PARAMS = frozenset( + { + "think", + "cfg_scale", + "cfg_norm", + "timestep_shift", + "t_eps", + "img_cfg_scale", + "max_tokens", + } +) +SENSENOVA_U1_EXTRA_OUTPUT_PARAMS = frozenset( + { + "think_text", + } +) diff --git a/vllm_omni/model_extras/vace.py b/vllm_omni/model_extras/vace.py new file mode 100644 index 00000000000..ed626afbee3 --- /dev/null +++ b/vllm_omni/model_extras/vace.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import numpy as np +from PIL import Image + +# VACE conditioning is expressed through prompt media, and the pipeline does not +# currently read any model-specific values from sampling_params.extra_args. +VACE_EXTRA_BODY_PARAMS: frozenset[str] = frozenset() +VACE_EXTRA_OUTPUT_PARAMS: frozenset[str] = frozenset() + +_MEDIA_KEYS = frozenset({"image", "last_image", "mask", "reference_images"}) + + +def _base_prompt(prompt: str, negative_prompt: str | None) -> dict[str, Any]: + result: dict[str, Any] = {"prompt": prompt} + if negative_prompt is not None: + result["negative_prompt"] = negative_prompt + return result + + +def _require_dimensions( + height: int | None, + width: int | None, + num_frames: int | None, +) -> tuple[int, int, int]: + if height is None or width is None or num_frames is None: + raise ValueError("VACE conditional inputs require height, width, and num_frames.") + if height <= 0 or width <= 0 or num_frames <= 0: + raise ValueError("VACE height, width, and num_frames must be positive.") + return height, width, num_frames + + +def _validate_media_inputs(media_inputs: Mapping[str, Any]) -> None: + unknown = set(media_inputs) - _MEDIA_KEYS + if unknown: + raise ValueError(f"Unsupported VACE media input(s): {', '.join(sorted(unknown))}.") + + has_references = bool(media_inputs.get("reference_images")) + has_frame_conditioning = any(media_inputs.get(key) is not None for key in ("image", "last_image", "mask")) + if has_references and has_frame_conditioning: + raise ValueError("VACE reference images cannot be combined with image, last_image, or mask inputs.") + if media_inputs.get("mask") is not None and media_inputs.get("image") is None: + raise ValueError("VACE mask input requires an image input.") + if media_inputs.get("mask") is not None and media_inputs.get("last_image") is not None: + raise ValueError("VACE mask and last_image inputs cannot be combined.") + if not has_references and not has_frame_conditioning: + raise ValueError("VACE image-to-video generation requires a conditioning media input.") + + +def build_image_to_video_prompt( + prompt: str, + negative_prompt: str | None, + media_inputs: Mapping[str, Any], + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, +) -> dict[str, Any]: + """Build VACE's native video, mask, or reference-image conditioning.""" + _validate_media_inputs(media_inputs) + result = _base_prompt(prompt, negative_prompt) + + reference_images = media_inputs.get("reference_images") + if reference_images: + if not isinstance(reference_images, list) or not all( + isinstance(image, Image.Image) for image in reference_images + ): + raise ValueError("VACE reference_images must be a non-empty list of PIL images.") + result["multi_modal_data"] = {"reference_images": reference_images} + return result + + height, width, num_frames = _require_dimensions(height, width, num_frames) + first_image = media_inputs.get("image") + last_image = media_inputs.get("last_image") + mask_image = media_inputs.get("mask") + for name, image in (("image", first_image), ("last_image", last_image), ("mask", mask_image)): + if image is not None and not isinstance(image, Image.Image): + raise ValueError(f"VACE {name} must be a PIL image.") + + gray = Image.new("RGB", (width, height), (128, 128, 128)) + mask_black = Image.new("L", (width, height), 0) + mask_white = Image.new("L", (width, height), 255) + + if mask_image is not None: + mask = mask_image.convert("L") + source = np.asarray(first_image).copy() + source[np.asarray(mask) > 128] = 128 + frame = Image.fromarray(source) + result["multi_modal_data"] = { + "video": [frame.copy() for _ in range(num_frames)], + "mask": [mask.copy() for _ in range(num_frames)], + } + return result + + if first_image is not None and last_image is not None: + if num_frames < 2: + raise ValueError("VACE first/last-frame conditioning requires at least two frames.") + result["multi_modal_data"] = { + "video": [first_image] + [gray] * (num_frames - 2) + [last_image], + "mask": [mask_black] + [mask_white] * (num_frames - 2) + [mask_black], + } + return result + + if first_image is not None: + result["multi_modal_data"] = { + "video": [first_image] + [gray] * (num_frames - 1), + "mask": [mask_black] + [mask_white] * (num_frames - 1), + } + return result + + result["multi_modal_data"] = { + "video": [gray] * (num_frames - 1) + [last_image], + "mask": [mask_white] * (num_frames - 1) + [mask_black], + } + return result diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs/__init__.py similarity index 94% rename from vllm_omni/outputs.py rename to vllm_omni/outputs/__init__.py index 2dfa1699fc0..34731eecab0 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs/__init__.py @@ -44,17 +44,32 @@ class OmniModelRunnerOutput(ModelRunnerOutput): that may be produced by non-autoregressive stages. Attributes: - multimodal_outputs: Optional per-request list of multimodal output - dicts, indexed by req_index. Each element is a dict mapping - output keys to tensors/values (e.g., [{"audio": tensor}, ...]). + multimodal_outputs: Optional per-request list of client-facing multimodal + output dicts, indexed by req_index. + inter_stage_outputs: Optional per-request list of inter-stage payload dicts + for connector transport (``save_async`` / full_payload). Not forwarded + to the orchestrator output processor. """ multimodal_outputs: list[dict[str, object]] | None = None + inter_stage_outputs: list[dict[str, object]] | None = None # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU. # The Scheduler can safely free the block tables for these requests. kv_extracted_req_ids: list[str] | None = None omni_connector_output: OmniConnectorOutput | None = None + @classmethod + def with_kv_conn_output_only(cls, kv_connector_output: Any) -> "OmniModelRunnerOutput": + return cls( + req_ids=[], + req_id_to_index={}, + sampled_token_ids=[], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + kv_connector_output=kv_connector_output, + ) + @dataclass class OmniRequestOutput: @@ -175,6 +190,7 @@ def from_diffusion( final_output_type: str = "image", stage_durations: dict[str, float] | None = None, peak_memory_mb: float = 0.0, + finished: bool = True, ) -> "OmniRequestOutput": """Create output from diffusion model. @@ -211,7 +227,7 @@ def from_diffusion( _custom_output=custom_output or {}, stage_durations=stage_durations or {}, peak_memory_mb=peak_memory_mb, - finished=True, + finished=finished, ) @property diff --git a/vllm_omni/outputs/mm_outputs.py b/vllm_omni/outputs/mm_outputs.py new file mode 100644 index 00000000000..8c5fe6442b6 --- /dev/null +++ b/vllm_omni/outputs/mm_outputs.py @@ -0,0 +1,261 @@ +"""Multimodal output data structures for vLLM-Omni. + +This module defines structured types for multimodal outputs. + +""" + +from __future__ import annotations + +from collections.abc import Iterator, Mapping +from dataclasses import dataclass, field +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput + +from vllm_omni.outputs.output_modality import TensorAccumulationStrategy +from vllm_omni.outputs.utils import _is_tensor_list, _to_cpu + +logger = init_logger(__name__) + +# Keys whose values are metadata scalars (e.g. audio sample rate) but may +# arrive as 0-d torch.Tensors — from_dict routes all tensors into .tensors, +# so we relocate them to .metadata before consolidation to avoid a bogus +# torch.cat attempt and its warn-and-keep-last fallback. +_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"}) + + +def _cat_tensors( + tensors: list[torch.Tensor], + strategy: TensorAccumulationStrategy, +) -> torch.Tensor: + """Concatenate a list of tensors according to *strategy*.""" + if strategy == TensorAccumulationStrategy.CONCAT_LAST: + return torch.cat(tensors, dim=-1) + if strategy == TensorAccumulationStrategy.REPLACE: + return tensors[-1] + # CONCAT_DIM0 / APPEND_LIST / default + return torch.cat(tensors, dim=0) + + +def _consolidate_tensor_list( + key: str, + tensor_list: list[torch.Tensor], + strategy: TensorAccumulationStrategy, +) -> torch.Tensor: + """Concatenate a deferred tensor list, with fallbacks on shape mismatch.""" + try: + return _cat_tensors(tensor_list, strategy) + except RuntimeError: + # [TODO: this part is for async chunk, not for audio only. can we come up with a design + # that identify this by whether using async chunk instead of modality?] + if key != "audio": + logger.warning("Error concatenating tensor for key %s; keeping last tensor", key) + return tensor_list[-1] + # Audio chunks may have mismatched shapes; retry along the last dim, + # then fall back to flattening each chunk. + try: + return torch.cat(tensor_list, dim=-1) + except RuntimeError: + return torch.cat([chunk.reshape(-1) for chunk in tensor_list], dim=0) + + +def _append_entries(store: dict[str, Any], incoming: dict[str, Any]) -> None: + """Append *incoming* entries into *store*. + + Tensor values are collected into lists (deferred concatenation, avoiding + O(n²) repeated torch.cat); non-tensor values replace the existing entry. + """ + for key, new_value in incoming.items(): + existing = store.get(key) + if isinstance(existing, list) and isinstance(new_value, torch.Tensor): + existing.append(new_value) + elif isinstance(existing, torch.Tensor) and isinstance(new_value, torch.Tensor): + store[key] = [existing, new_value] + else: + store[key] = new_value + + +@dataclass(eq=False) +class MultimodalPayload(Mapping): + """Structured multimodal output payload. + + Implements ``collections.abc.Mapping`` so that ``isinstance(payload, dict)`` + style checks in downstream code can be replaced with duck-typing, and + ``payload.get(key)``, ``payload[key]``, ``key in payload``, ``len(payload)`` + all work seamlessly for both tensors and metadata. + + Attributes: + tensors: Dictionary mapping modality/key names to their tensors. + metadata: Optional dictionary for non-tensor metadata + (e.g., sample rate for audio, image dimensions). + """ + + tensors: dict[str, torch.Tensor] = field(default_factory=dict) + metadata: dict[str, Any] = field(default_factory=dict) + + @property + def primary_tensor(self) -> torch.Tensor | None: + """Return the first tensor in the payload, or None if empty.""" + if self.tensors: + return next(iter(self.tensors.values())) + return None + + @property + def is_empty(self) -> bool: + """Return True if the payload has no tensors and no metadata.""" + return len(self.tensors) == 0 and len(self.metadata) == 0 + + def get(self, key: str, default: Any = None) -> Any: + """Get a value by key, searching tensors first then metadata.""" + if key in self.tensors: + return self.tensors[key] + return self.metadata.get(key, default) + + def __contains__(self, key: str) -> bool: + return key in self.tensors or key in self.metadata + + def __getitem__(self, key: str) -> Any: + """Dict-like indexing: search tensors first, then metadata.""" + if key in self.tensors: + return self.tensors[key] + if key in self.metadata: + return self.metadata[key] + raise KeyError(key) + + def __len__(self) -> int: + return len(self.tensors) + len(self.metadata) + + def __iter__(self) -> Iterator[str]: + yield from self.tensors + yield from self.metadata + + def __bool__(self) -> bool: + return bool(self.tensors) or bool(self.metadata) + + def __eq__(self, other: object) -> bool: + """Support equality with plain dicts and other Mappings.""" + if isinstance(other, MultimodalPayload): + return self.tensors == other.tensors and self.metadata == other.metadata + if isinstance(other, Mapping): + return dict(self) == dict(other) + return NotImplemented + + def merged_with(self, incoming: MultimodalPayload) -> MultimodalPayload: + """Merge *incoming* onto this payload and return the result. + + Tensor values accumulate into lists for deferred concatenation; + non-tensor values are replaced with the latest. When this payload + is empty, *incoming* is returned as-is, so callers should use the + return value: ``accumulated = accumulated.merged_with(incoming)``. + """ + if self.is_empty: + return incoming + _append_entries(self.tensors, incoming.tensors) + _append_entries(self.metadata, incoming.metadata) + return self + + def consolidate_tensors(self, strategy: TensorAccumulationStrategy) -> None: + """Concatenate deferred tensor lists into single tensors. + + Tensors are generated content accumulated as chunks, so lists are + concatenated according to *strategy* (e.g. audio chunks along the + time dimension, latent frames along the batch dimension). + """ + # Relocate scalar metadata tensors (e.g. sample rate) that from_dict + # routed into .tensors, so they take the REPLACE path via .metadata + # instead of failing torch.cat as 0-d tensors. + for key in _METADATA_TENSOR_KEYS: + value = self.tensors.pop(key, None) + if value is None: + continue + self.metadata[key] = value[-1] if isinstance(value, list) else value + + for key, value in list(self.tensors.items()): + if _is_tensor_list(value): + self.tensors[key] = _consolidate_tensor_list(key, value, strategy) + + def consolidate_metadata(self) -> None: + """Resolve deferred tensor lists in metadata by keeping the latest value. + + Metadata values are per-step snapshots (e.g. sample rate), not + content deltas, so the latest value supersedes earlier ones. + Nested dicts (unflattened payloads) are resolved one level down. + """ + for key, value in list(self.metadata.items()): + if _is_tensor_list(value): + self.metadata[key] = value[-1] + elif isinstance(value, dict): + for sub_key, sub_value in list(value.items()): + if _is_tensor_list(sub_value): + value[sub_key] = sub_value[-1] + + def to_dict(self) -> dict[str, Any]: + """Convert back to a plain dict (tensors + metadata merged).""" + result: dict[str, Any] = dict(self.tensors) + result.update(self.metadata) + return result + + @classmethod + def from_dict(cls, data: dict[str, Any] | None) -> MultimodalPayload | None: + """Create a MultimodalPayload from a raw dictionary. + + Separates torch.Tensor values into tensors and everything + else into metadata. + """ + if not data: + return None + tensors: dict[str, torch.Tensor] = {} + metadata: dict[str, Any] = {} + for k, v in data.items(): + if isinstance(v, torch.Tensor): + tensors[k] = v + else: + metadata[k] = v + if not tensors and not metadata: + return None + return cls(tensors=tensors, metadata=metadata) + + @classmethod + def from_raw(cls, payload: Any, modality_key: str) -> MultimodalPayload | None: + """Create a MultimodalPayload from a raw producer payload. + + Accepts a MultimodalPayload (returned as-is), a dict, or a bare + tensor (stored under *modality_key*). Tensors are moved to CPU. + Producer-specific dict keys are remapped to the semantic modality + key (e.g. "audio", "latent"): AR runners produce {"hidden": ...} + and generation runners produce {"model_outputs": ...}. + """ + if isinstance(payload, MultimodalPayload): + return payload + + if not isinstance(payload, dict): + return cls.from_dict({modality_key: _to_cpu(payload)}) + + remapped: dict[str, Any] = {} + for key, value in payload.items(): + is_producer_key = key == "model_outputs" or (key == "hidden" and modality_key != "hidden") + remapped[modality_key if is_producer_key else key] = _to_cpu(value) + return cls.from_dict(remapped) + + +@dataclass +class MultimodalCompletionOutput(CompletionOutput): + """CompletionOutput with multimodal support. + + Inherits all CompletionOutput fields and adds multimodal_output. + As a CompletionOutput subclass, compatible with all existing vLLM consumers. + """ + + def __init__( + self, + multimodal_output: MultimodalPayload | None = None, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.multimodal_output = multimodal_output + + def __repr__(self) -> str: + base = super().__repr__() + return f"{base[:-1]}, multimodal_output={self.multimodal_output!r})" diff --git a/vllm_omni/outputs/output_metadata.py b/vllm_omni/outputs/output_metadata.py new file mode 100644 index 00000000000..f3e364a3462 --- /dev/null +++ b/vllm_omni/outputs/output_metadata.py @@ -0,0 +1,238 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any, Literal, TypeAlias, TypedDict + +import numpy as np +import PIL.Image +import torch + +# Final user-facing output modality returned by pipeline stages. +FinalOutputModalityType: TypeAlias = Literal["text", "image", "audio", "video"] + +# Scalar values that may appear in either payload or metadata. +# Examples: text_output="done", fps=16, action_only_output=True. +DiffusionScalar: TypeAlias = str | int | float | bool | None + +# Shape-like metadata. Examples: [1, 16, 3, 480, 832], (16, 480, 832). +DiffusionShape: TypeAlias = list[int] | tuple[int, ...] + +# Named tensor controls used by transfer-style metadata. +# Example: {"pose": pose_tensor, "depth": depth_tensor}. +DiffusionTensorMap: TypeAlias = Mapping[str, torch.Tensor] + +# Values allowed under metadata groups such as video/audio/actions/transfer. +# Examples: {"fps": 16}, {"shape": [1, 16, 3, 480, 832]}, +# {"controls": {"pose": pose_tensor}}, {"hints": ["pose"]}. +DiffusionMetadataValue: TypeAlias = ( + DiffusionScalar + | DiffusionShape + | DiffusionTensorMap + | torch.Tensor + | list["DiffusionMetadataValue"] + | tuple["DiffusionMetadataValue", ...] + | Mapping[str, "DiffusionMetadataValue"] +) + +# Mutable metadata envelope side. Example: +# {"video": {"fps": 16}, "actions": {"raw_action_dim": 14}}. +DiffusionMetadata: TypeAlias = dict[str, DiffusionMetadataValue] + +# Read-only metadata view for validators and serving readers. +DiffusionMetadataMapping: TypeAlias = Mapping[str, DiffusionMetadataValue] + + +class DiffusionTrajectoryPayload(TypedDict, total=False): + # Optional denoising trajectory payload for RL/debug consumers. + latents: torch.Tensor | dict[str, Any] | None + timesteps: torch.Tensor | dict[str, Any] | None + log_probs: torch.Tensor | dict[str, Any] | None + decoded: list[PIL.Image.Image] | None + + +# Concrete payload values produced by post_process. +# Examples: PIL image, video tensor, audio bytes/tensor, action tensor. +DiffusionPayloadValue: TypeAlias = ( + DiffusionScalar + | bytes + | torch.Tensor + | np.ndarray + | PIL.Image.Image + | DiffusionTrajectoryPayload + | list["DiffusionPayloadValue"] + | tuple["DiffusionPayloadValue", ...] + | Mapping[str, "DiffusionPayloadValue"] +) + +# Mutable payload side. Example: {"video": video_tensor, "action": action_tensor}. +DiffusionPayload: TypeAlias = dict[str, DiffusionPayloadValue] + +# Read-only payload view for code that consumes but does not mutate payload. +DiffusionPayloadMapping: TypeAlias = Mapping[str, DiffusionPayloadValue] + +# Serving-facing multimodal output. Non-metadata keys hold payload values. +# Example: {"video": video_tensor, "metadata": {"video": {"fps": 16}}}. +DiffusionMultimodalOutput: TypeAlias = dict[str, DiffusionPayloadValue | DiffusionMetadata] + + +class DiffusionOutputEnvelope(TypedDict, total=False): + # Canonical post_process return shape: + # {"payload": {"video": video_tensor}, "metadata": {"video": {"fps": 16}}} + payload: DiffusionPayload + metadata: DiffusionMetadata + + +# Raw post_process return before formatter normalization. This keeps backward +# compatibility with direct payload returns while accepting the new envelope. +# Examples: image, {"video": video_tensor}, or +# {"payload": {"video": video_tensor}, "metadata": {"video": {"fps": 16}}}. +DiffusionPostprocessRawOutput: TypeAlias = DiffusionPayloadValue | DiffusionPayload | DiffusionOutputEnvelope + + +def strip_internal_metadata(metadata: DiffusionMetadataMapping) -> DiffusionMetadata: + """Return public metadata, dropping model-private postprocess state.""" + return {key: value for key, value in metadata.items() if key != "internal"} + + +def validate_public_diffusion_metadata(metadata: DiffusionMetadataMapping) -> None: + """Validate metadata that is safe to expose outside model postprocess.""" + if "internal" in metadata: + raise ValueError("Diffusion metadata field 'internal' must not escape public formatting.") + validate_diffusion_metadata(metadata) + + +def validate_diffusion_metadata(metadata: DiffusionMetadataMapping) -> None: + """Validate canonical diffusion output metadata fields. + + Unknown model-specific groups and keys are intentionally tolerated. This + validator only guards fields that are part of the shared output protocol or + consumed by serving code. + """ + _validate_video_metadata(_optional_mapping(metadata, "video")) + _validate_audio_metadata(_optional_mapping(metadata, "audio")) + _validate_actions_metadata(_optional_mapping(metadata, "actions")) + _validate_transfer_metadata(_optional_mapping(metadata, "transfer")) + _validate_text_metadata(_optional_mapping(metadata, "text")) + _validate_common_metadata(_optional_mapping(metadata, "common")) + + +def _validate_video_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_positive_number(metadata, "fps", group="video") + _check_optional_int(metadata, "video_fps_multiplier", group="video", min_value=1) + _check_optional_shape(metadata, "shape", group="video") + + +def _validate_audio_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_int(metadata, "sample_rate", group="audio", min_value=1) + _check_optional_non_empty_str(metadata, "audiox_task", group="audio") + + +def _validate_actions_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_int(metadata, "raw_action_dim", group="actions", min_value=1) + _check_optional_non_empty_str(metadata, "action_mode", group="actions") + _check_optional_int(metadata, "domain_id", group="actions") + + +def _validate_transfer_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + controls = metadata.get("controls") + if controls is not None and not _is_tensor_mapping(controls): + raise TypeError("Diffusion metadata field 'transfer.controls' must be a mapping of str to torch.Tensor.") + hints = metadata.get("hints") + if hints is not None and not _is_str_list(hints): + raise TypeError("Diffusion metadata field 'transfer.hints' must be a list of strings.") + + +def _validate_text_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + _check_optional_str(metadata, "text_output", group="text") + _check_optional_str(metadata, "think_text", group="text") + + +def _validate_common_metadata(metadata: DiffusionMetadataMapping | None) -> None: + if metadata is None: + return + value = metadata.get("action_only_output") + if value is not None and not isinstance(value, bool): + raise TypeError("Diffusion metadata field 'common.action_only_output' must be bool.") + + +def _optional_mapping(metadata: DiffusionMetadataMapping, group: str) -> DiffusionMetadataMapping | None: + value = metadata.get(group) + if value is None: + return None + if not isinstance(value, Mapping): + raise TypeError(f"Diffusion metadata group '{group}' must be a mapping.") + return value + + +def _check_optional_positive_number(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is None: + return + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be a positive number.") + if value <= 0: + raise ValueError(f"Diffusion metadata field '{group}.{key}' must be > 0.") + + +def _check_optional_int( + metadata: DiffusionMetadataMapping, + key: str, + *, + group: str, + min_value: int | None = None, +) -> None: + value = metadata.get(key) + if value is None: + return + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be int.") + if min_value is not None and value < min_value: + raise ValueError(f"Diffusion metadata field '{group}.{key}' must be >= {min_value}.") + + +def _check_optional_str(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is not None and not isinstance(value, str): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be str.") + + +def _check_optional_non_empty_str(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is None: + return + if not isinstance(value, str): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be str.") + if not value: + raise ValueError(f"Diffusion metadata field '{group}.{key}' must be non-empty.") + + +def _check_optional_shape(metadata: DiffusionMetadataMapping, key: str, *, group: str) -> None: + value = metadata.get(key) + if value is None: + return + if not isinstance(value, Sequence) or isinstance(value, (str, bytes, bytearray)): + raise TypeError(f"Diffusion metadata field '{group}.{key}' must be a list or tuple of positive ints.") + if any(isinstance(dim, bool) or not isinstance(dim, int) or dim <= 0 for dim in value): + raise ValueError(f"Diffusion metadata field '{group}.{key}' must contain positive ints.") + + +def _is_tensor_mapping(value: object) -> bool: + return isinstance(value, Mapping) and all( + isinstance(key, str) and isinstance(tensor, torch.Tensor) for key, tensor in value.items() + ) + + +def _is_str_list(value: object) -> bool: + return isinstance(value, list) and all(isinstance(item, str) for item in value) diff --git a/vllm_omni/outputs/output_modality.py b/vllm_omni/outputs/output_modality.py new file mode 100644 index 00000000000..2536a8adedb --- /dev/null +++ b/vllm_omni/outputs/output_modality.py @@ -0,0 +1,120 @@ +"""Output modality types for vLLM-Omni. + +This module defines the OutputModality enum and TensorAccumulationStrategy +for type-safe multimodal output routing and tensor merging. + +""" + +from __future__ import annotations + +import re +from enum import Enum, Flag, StrEnum, auto + +_MODALITY_ALIASES: dict[str, str] = { + "speech": "audio", + "images": "image", + "latents": "latent", + "wav": "audio", + "waveform": "audio", + "pixel_values": "image", + "pixels": "image", + "token_ids": "text", + "tokens": "text", +} + + +class OutputModalityNames(StrEnum): + """Keys for output modalities. + + TODO: (Alex) Integrate this with the big-flag enum below + throughout the code + for better type safety (currently only used for output processor). + """ + + TEXT = "text" + IMAGE = "image" + AUDIO = "audio" + LATENT = "latent" + + +# Specify which output modalities may be drained when handling delta messages. +# For some types, e.g., latents, we need to be careful to ensure the full context +# is passed as the stream yields due to assumptions in the I/O processing and model +# when async chunk isn't enabled. +NON_DRAINABLE_MODALITIES = {OutputModalityNames.TEXT, OutputModalityNames.LATENT} +DRAINABLE_MODALITIES = {mod for mod in OutputModalityNames if mod not in NON_DRAINABLE_MODALITIES} + + +class OutputModality(Flag): + """Bit-flag enum for output modalities. + + Compose freely with ``|`` — no need to enumerate every combination. + + Single: ``OutputModality.TEXT``, ``OutputModality.IMAGE``, ... + Compound: ``OutputModality.TEXT | OutputModality.IMAGE`` (text+image) + + Note: POOLING is intentionally excluded. Pooling/embedding is vLLM's + native path (pooling_output → PoolingRequestOutput), handled entirely + by the base OutputProcessor. vLLM-Omni's layer does not participate. + """ + + TEXT = auto() + IMAGE = auto() + AUDIO = auto() + LATENT = auto() + + @classmethod + def from_string(cls, s: str | None) -> OutputModality: + """Parse a free-text modality string into an OutputModality flag. + + Handles common aliases and compound strings separated by + or ,. + + Examples:: + + OutputModality.from_string("text+image") + # → OutputModality.TEXT | OutputModality.IMAGE + """ + if not s or not s.strip(): + return cls.TEXT + + parts = [p.strip().lower() for p in re.split(r"[+,]", s.strip())] + result = cls(0) + for p in parts: + p = _MODALITY_ALIASES.get(p, p) + try: + result |= cls[p.upper()] + except KeyError: + raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}") + return result + + @property + def has_text(self) -> bool: + return OutputModality.TEXT in self + + @property + def has_multimodal(self) -> bool: + return bool(self & ~OutputModality.TEXT) + + +class TensorAccumulationStrategy(Enum): + """Strategy for merging incremental multimodal tensors.""" + + CONCAT_DIM0 = "concat_dim0" + """Concatenate along dimension 0. Used for image/latent tensors.""" + + CONCAT_LAST = "concat_last" + """Concatenate along the last dimension. Used for audio waveforms.""" + + APPEND_LIST = "append_list" + """Append to a list (no tensor concatenation).""" + + REPLACE = "replace" + """Replace previous tensor entirely with the latest one.""" + + +def get_accumulation_strategy(modality: OutputModality) -> TensorAccumulationStrategy: + """Determine tensor merge strategy from the multimodal flags.""" + if OutputModality.AUDIO in modality: + return TensorAccumulationStrategy.CONCAT_LAST + if OutputModality.IMAGE in modality or OutputModality.LATENT in modality: + return TensorAccumulationStrategy.CONCAT_DIM0 + return TensorAccumulationStrategy.CONCAT_DIM0 # default diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py new file mode 100644 index 00000000000..e39b995b008 --- /dev/null +++ b/vllm_omni/outputs/output_processor.py @@ -0,0 +1,633 @@ +from dataclasses import fields as dataclass_fields +from typing import Any + +import torch +from vllm.logger import init_logger +from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput +from vllm.sampling_params import RequestOutputKind +from vllm.tokenizers import TokenizerLike +from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason +from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor +from vllm.v1.engine.output_processor import ( + OutputProcessorOutput, + RequestOutputCollector, + RequestState, +) +from vllm.v1.engine.parallel_sampling import ParentRequest +from vllm.v1.metrics.stats import IterationStats, RequestStateStats + +from vllm_omni.data_entry_keys import unflatten_payload +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload +from vllm_omni.outputs.output_modality import ( + DRAINABLE_MODALITIES, + OutputModality, + get_accumulation_strategy, +) + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Module-level helpers +# --------------------------------------------------------------------------- + + +def _modality_to_type_string(modality: OutputModality) -> str: + """Convert an OutputModality flag to a lowercase type string.""" + try: + if OutputModality.AUDIO in modality: + return "audio" + if OutputModality.IMAGE in modality: + return "image" + if OutputModality.LATENT in modality: + return "latent" + except TypeError: + # Flag identity mismatch (e.g. after module reload in tests). + name = getattr(modality, "name", "") or "" + lowered = name.lower() + if "audio" in lowered: + return "audio" + if "image" in lowered: + return "image" + if "latent" in lowered: + return "latent" + return "text" + + +class OmniRequestState(RequestState): + """Request state for omni models, tracking multimodal outputs. + + Extends the base RequestState with support for accumulating + multimodal tensor outputs (e.g., images, audio, latents) that + are produced incrementally during generation. + """ + + def __init__( + self, + *args, + **kwargs, + ): + # arrival_time is always passed as a keyword argument by the sole caller + # (from_new_request). The *args fallback via positional index was removed + # because upstream RequestState.__init__ reorders parameters between + # releases, making args[12] fragile. + arrival_time = kwargs.get("arrival_time") + super().__init__(*args, **kwargs) + self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) + # Omni-specific: multimodal output accumulation + # TODO: mm_type is per-request, not per-key. If a model ever produces + # both audio and latent outputs, the modality type would flip on each + # add_multimodal_tensor() call. Consider tracking per-key modality + # types (e.g. dict[str, str]) for future multi-output models. + self.mm_type: str | None = None + self.mm_accumulated: MultimodalPayload = MultimodalPayload() + + def apply_streaming_update(self, update) -> None: + super().apply_streaming_update(update) + self.native_text_stats.arrival_time = update.arrival_time + + def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: + """Accumulate a multimodal tensor payload into the request state. + + Normalizes incoming payloads (dict or raw tensor) into a + MultimodalPayload and merges with any previously accumulated data. + Uses list-based deferred concatenation to avoid O(n²) repeated + torch.cat calls. + """ + if payload is None: + return + try: + if mm_type: + self.mm_type = mm_type.lower() + modality_key = self.mm_type or "hidden" + + incoming = MultimodalPayload.from_raw(payload, modality_key) + if incoming is not None: + self.mm_accumulated = self.mm_accumulated.merged_with(incoming) + except (ValueError, TypeError, RuntimeError): + logger.exception("Error accumulating multimodal tensor") + + def _consolidate_multimodal_tensors(self) -> None: + """Consolidate accumulated tensor lists into single tensors. + + Uses TensorAccumulationStrategy derived from the output modality + to determine concatenation behavior. Metadata values always use + REPLACE (keep latest). + """ + if self.mm_accumulated.is_empty: + return + + try: + modality = OutputModality.from_string(self.mm_type) + except (ValueError, KeyError): + modality = OutputModality.TEXT + strategy = get_accumulation_strategy(modality) + + try: + self.mm_accumulated.consolidate_tensors(strategy) + self.mm_accumulated.consolidate_metadata() + except (RuntimeError, TypeError, KeyError): + logger.exception("Error consolidating multimodal tensors") + + # Override: do not route to pooling-only path; always create completion + # outputs, and attach pooling_result into the CompletionOutput. + def make_request_output( + self, + new_token_ids: list[int], + pooling_output: torch.Tensor | None, + finish_reason: FinishReason | None, + stop_reason: int | str | None, + kv_transfer_params: dict[str, Any] | None = None, + routed_experts: Any = None, + ) -> OmniRequestOutput | PoolingRequestOutput | None: + """Create a request output from generation results. + + Creates a RequestOutput or PoolingRequestOutput from the generated + tokens and accumulated multimodal outputs. Attaches multimodal + tensors to the completion output if available. + + Args: + new_token_ids: List of newly generated token IDs + pooling_output: Optional pooling output tensor + finish_reason: Optional finish reason indicating why generation stopped + stop_reason: Optional stop reason (token ID or stop string) + kv_transfer_params: Optional KV cache transfer parameters + routed_experts: Optional MoE routed-expert ids for this step, + attached to the completion output for generation stages + + Returns: + OmniRequestOutput or PoolingRequestOutput if output should be + emitted (based on finish status and output kind), None otherwise + """ + # Pooling-only requests should follow base behavior. + if self.detokenizer is None and pooling_output is not None: + return super().make_request_output( + new_token_ids, + pooling_output, + finish_reason, + stop_reason, + kv_transfer_params, + ) + + finished = finish_reason is not None + final_only = self.output_kind == RequestOutputKind.FINAL_ONLY + is_delta = self.output_kind == RequestOutputKind.DELTA + + if not finished and final_only: + return None + + # Consolidate accumulated tensors when finishing, or for every + # CUMULATIVE step so the consumer always sees a single tensor. + if finished or not is_delta: + self._consolidate_multimodal_tensors() + + if self.stream_interval > 1 and self.detokenizer is not None: + # Send output request only when + # 1. It has finished, or + # 2. It is the first token, or + # 3. It has reached the stream interval number of tokens + if not ( + finished + or self.sent_tokens_offset == 0 + or self.detokenizer.num_output_tokens() - self.sent_tokens_offset >= self.stream_interval + ): + return None + + if self.output_kind == RequestOutputKind.DELTA: + # Send tokens from the offset in DELTA mode, otherwise all + # tokens are sent. + new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] + self.sent_tokens_offset = self.detokenizer.num_output_tokens() + + external_req_id = self.external_req_id + + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, routed_experts) + + if self.parent_req is None: + outputs = [output] + else: + outputs, finished = self.parent_req.get_outputs(self.request_id, output) + if not outputs: + return None + external_req_id = self.parent_req.external_req_id + + return self._new_request_output( + external_req_id, + outputs, + finished, + kv_transfer_params, + ) + + def _new_completion_output( + self, + token_ids: list[int], + finish_reason: FinishReason | None, + stop_reason: int | str | None, + routed_experts: Any = None, + ) -> MultimodalCompletionOutput | CompletionOutput: + """Create a completion output with multimodal data attached. + + Returns a MultimodalCompletionOutput when multimodal data has been + accumulated, otherwise returns the base CompletionOutput. Snapshots + the accumulated payload before draining DELTA-mode modality keys so + that callers receive an immutable view. + """ + # When there is no detokenizer (generation stages), build a minimal + # CompletionOutput directly since upstream asserts detokenizer != None. + if self.detokenizer is None: + finished = finish_reason is not None + base_output = CompletionOutput( + index=self.request_index, + text="", + token_ids=token_ids, + logprobs=None, + cumulative_logprob=None, + finish_reason=str(finish_reason) if finished else None, + stop_reason=stop_reason if finished else None, + routed_experts=routed_experts, + ) + else: + base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) + + # Always provide cumulative token IDs for inter-stage processors. + if self.detokenizer is not None: + base_output.cumulative_token_ids = list(self.detokenizer.output_token_ids) + else: + base_output.cumulative_token_ids = list(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 + + try: + if self.mm_accumulated and not self.mm_accumulated.is_empty: + # Snapshot: copy current tensors/metadata so drain doesn't + # mutate the payload already handed to the caller. + # Unflatten dotted keys (e.g. "hidden_states.layer_0") back + # to nested dicts so downstream consumers (thinker2talker etc.) + # can access with .get("hidden_states", {}).get("layers", {}). + merged = {**self.mm_accumulated.tensors, **self.mm_accumulated.metadata} + nested = unflatten_payload(merged) + snapshot = MultimodalPayload.from_dict(nested) or MultimodalPayload() + kwargs = {f.name: getattr(base_output, f.name) for f in dataclass_fields(CompletionOutput)} + output = MultimodalCompletionOutput( + multimodal_output=snapshot, + **kwargs, + ) + output.cumulative_token_ids = base_output.cumulative_token_ids + if hasattr(base_output, "cumulative_text"): + output.cumulative_text = base_output.cumulative_text + + # DELTA mode: drain modality keys (e.g. audio) so the next + # step only sees freshly accumulated data for those keys. + if self.output_kind == RequestOutputKind.DELTA: + for modality_key in DRAINABLE_MODALITIES: + self.mm_accumulated.tensors.pop(modality_key, None) + + return output + except (RuntimeError, TypeError, AttributeError): + logger.exception("Error creating MultimodalCompletionOutput") + return base_output + + def _new_request_output( + self, + external_req_id: str, + outputs: list, + finished: bool, + kv_transfer_params: dict[str, Any] | None = None, + ) -> RequestOutput | PoolingRequestOutput: + """Create request output, handling no-detokenizer generation stages. + + Upstream asserts ``self.logprobs_processor is not None`` which fails + for generation stages (talker, code2wav) that have no tokenizer. + When the logprobs processor is absent we build the RequestOutput + directly with ``prompt_logprobs=None``. + """ + if self.logprobs_processor is not None: + return super()._new_request_output( + external_req_id, + outputs, + finished, + kv_transfer_params, + ) + + # No-detokenizer path: build RequestOutput directly. + prompt_token_ids = self.prompt_token_ids + if prompt_token_ids is None and self.prompt_embeds is not None: + prompt_token_ids = [0] * len(self.prompt_embeds) + if prompt_token_ids is None: + prompt_token_ids = [] + + return RequestOutput( + request_id=external_req_id, + lora_request=self.lora_request, + prompt=self.prompt, + prompt_token_ids=prompt_token_ids, + prompt_logprobs=None, + outputs=outputs, + finished=finished, + kv_transfer_params=kv_transfer_params, + num_cached_tokens=self.num_cached_tokens, + metrics=self.stats, + ) + + +class MultimodalOutputProcessor(VLLMOutputProcessor): + """Handles multimodal output processing. + + Captures multimodal outputs from OmniEngineCoreOutput and accumulates + them as MultimodalPayload in OmniRequestState, before delegating to + the base vLLM OutputProcessor for text handling. + + The data flow is: + 1. For each EngineCoreOutput with multimodal_output: + - Capture into OmniRequestState.add_multimodal_tensor() + 2. Base vLLM OutputProcessor handles text detokenization + 3. On finish, _consolidate_multimodal_tensors() concatenates accumulated + tensors using strategy-based dispatch + 4. _new_completion_output() returns MultimodalCompletionOutput + """ + + def __init__( + self, + tokenizer: TokenizerLike | None, + *, + log_stats: bool, + stream_interval: int = 1, + tracing_enabled: bool = False, + engine_core_output_type: str | None = None, + output_modality: OutputModality = OutputModality.TEXT, + ): + """Initialize the multimodal output processor. + + Args: + tokenizer: Tokenizer for detokenizing text outputs + log_stats: Whether to log statistics + stream_interval: Stream interval for output generation + engine_core_output_type: Optional output type string (e.g., + "image", "audio", "latent"). Converted to OutputModality + internally. Kept for backward compatibility with + stage_init_utils. + output_modality: Type-safe output modality flag. Used to tag + multimodal outputs with the correct modality key when + per-output type info is unavailable. + """ + super().__init__( + tokenizer=tokenizer, + log_stats=log_stats, + stream_interval=stream_interval, + tracing_enabled=tracing_enabled, + ) + # Convert string-based engine_core_output_type to OutputModality + if engine_core_output_type is not None: + self.output_modality = OutputModality.from_string(engine_core_output_type) + else: + self.output_modality = output_modality + self.engine_core_output_type = engine_core_output_type + self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {} + + def _native_text_metric_record(self, request_id: str) -> dict[str, Any]: + return self._native_text_metrics_by_request.setdefault( + request_id, + { + "vllm_ttft_ms": 0.0, + "vllm_tpot_ms": 0.0, + "vllm_itl_ms": 0.0, + "vllm_itls_ms": [], + }, + ) + + def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: + return self._native_text_metrics_by_request.pop(request_id, {}) + + def abort_requests(self, request_ids, internal: bool) -> list[str]: + request_ids = list(request_ids) + for request_id in request_ids: + if internal: + req_state = self.request_states.get(request_id) + if req_state is not None: + self._native_text_metrics_by_request.pop(req_state.external_req_id, None) + else: + self._native_text_metrics_by_request.pop(request_id, None) + return super().abort_requests(request_ids, internal) + + def add_request( + self, + request: EngineCoreRequest, + prompt: str | None, + parent_req: ParentRequest | None = None, + request_index: int = 0, + queue: RequestOutputCollector | None = None, + ) -> None: + """Add a new request to be processed. + + Creates an OmniRequestState for the request and registers it + for output processing. + + Args: + request: Engine core request to add + prompt: Optional prompt string for the request + parent_req: Optional parent request for parallel sampling + request_index: Index of the request in the batch + queue: Optional queue for collecting outputs + + Raises: + ValueError: If the request ID is already registered + """ + request_id = request.request_id + req_state = self.request_states.get(request_id) + if req_state is not None: + self._update_streaming_request_state(req_state, request, prompt) + return + + req_state = OmniRequestState.from_new_request( + tokenizer=self.tokenizer, + request=request, + prompt=prompt, + parent_req=parent_req, + request_index=request_index, + queue=queue, + log_stats=self.log_stats, + stream_interval=self.stream_interval, + ) + self.request_states[request_id] = req_state + if parent_req: + self.parent_requests[parent_req.request_id] = parent_req + self.external_req_ids[req_state.external_req_id].append(request_id) + + def remove_request(self, request_id: str) -> None: + """Rollback one previously registered request if it was never submitted.""" + req_state = self.request_states.pop(request_id, None) + if req_state is None: + return + + external_req_id = getattr(req_state, "external_req_id", None) + if external_req_id is not None: + self._native_text_metrics_by_request.pop(external_req_id, None) + request_ids = self.external_req_ids.get(external_req_id) + if request_ids is not None: + self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id] + if not self.external_req_ids[external_req_id]: + self.external_req_ids.pop(external_req_id, None) + + parent_req = getattr(req_state, "parent_req", None) + if parent_req is not None: + self.parent_requests.pop(parent_req.request_id, None) + + def process_outputs( + self, + engine_core_outputs: list[EngineCoreOutput], + engine_core_timestamp: float | None = None, + iteration_stats: IterationStats | None = None, + ) -> OutputProcessorOutput: + default_mm_type = _modality_to_type_string(self.output_modality) + + # Separate outputs that upstream can handle (has detokenizer or + # pooling) from multimodal-only outputs (no detokenizer, no pooling) + # that would trigger upstream's `assert detokenizer is not None`. + upstream_outputs: list[EngineCoreOutput] = [] + mm_only_outputs: list[EngineCoreOutput] = [] + + for eco in engine_core_outputs: + req_state = self.request_states.get(eco.request_id) + if req_state is None: + continue + + # Accumulate multimodal tensors regardless of path. + if isinstance(req_state, OmniRequestState): + mm_output = getattr(eco, "multimodal_output", None) + if mm_output is not None: + mm_type = getattr(eco, "output_type", None) or default_mm_type + req_state.add_multimodal_tensor(mm_output, mm_type) + + # Route: if no detokenizer and no pooling output, handle locally + # to avoid upstream's assert on detokenizer. + if req_state.detokenizer is None and eco.pooling_output is None: + mm_only_outputs.append(eco) + else: + upstream_outputs.append(eco) + + # Handle multimodal-only outputs (generation stages) locally. + self._process_mm_only_outputs(mm_only_outputs) + + # Delegate text/pooling outputs to upstream. + return super().process_outputs( + upstream_outputs, + engine_core_timestamp=engine_core_timestamp, + iteration_stats=iteration_stats, + ) + + def _process_mm_only_outputs( + self, + engine_core_outputs: list[EngineCoreOutput], + ) -> None: + """Handle outputs from generation stages that have no detokenizer. + + These cannot go through upstream process_outputs because it asserts + detokenizer is not None when pooling_output is None. + """ + for eco in engine_core_outputs: + req_state = self.request_states.get(eco.request_id) + if req_state is None or not isinstance(req_state, OmniRequestState): + continue + + new_token_ids = eco.new_token_ids + finish_reason = eco.finish_reason + stop_reason = eco.stop_reason + kv_transfer_params = eco.kv_transfer_params + routed_experts = eco.routed_experts + req_state.num_cached_tokens = eco.num_cached_tokens + req_state.is_prefilling = False + + if request_output := req_state.make_request_output( + new_token_ids, + None, # pooling_output + finish_reason, + stop_reason, + kv_transfer_params, + routed_experts, + ): + if req_state.queue is not None: + req_state.queue.put(request_output) + + if finish_reason is not None: + self._finish_request(req_state) + + def _update_stats_from_output( + self, + req_state: RequestState, + engine_core_output: EngineCoreOutput, + engine_core_timestamp: float | None, + iteration_stats: IterationStats | None, + ): + was_prefilling = req_state.is_prefilling + native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None + previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 + + # NOTE: We pass ``None`` for *iteration_stats* to the parent so that + # the upstream's ``_update_stats_from_output`` logs stats via its own + # ``RequestStateStats`` (req_state.stats) without touching the omni + # ``IterationStats`` object. Below we re-apply the same update using + # *native_stats* (the omni-specific ``RequestStateStats`` attached to + # ``OmniRequestState``), which correctly feeds the omni metrics path. + # If the upstream adds new logic in ``_update_stats_from_output`` (e.g. + # additional stat fields), this override must be kept in sync. + super()._update_stats_from_output( + req_state, + engine_core_output, + engine_core_timestamp, + None, + ) + + if iteration_stats is None or engine_core_timestamp is None or native_stats is None: + return + + iteration_stats.update_from_output( + engine_core_output, + engine_core_timestamp, + was_prefilling, + native_stats, + self.lora_states, + req_state.lora_name, + ) + record = self._native_text_metric_record(req_state.external_req_id) + if was_prefilling: + record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0) + return + + if previous_last_token_ts > 0: + itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0) + itls_ms = record.setdefault("vllm_itls_ms", []) + itls_ms.append(itl_ms) + record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms)) + + def _update_stats_from_finished( + self, + req_state: RequestState, + finish_reason: FinishReason | None, + iteration_stats: IterationStats | None, + ): + super()._update_stats_from_finished(req_state, finish_reason, None) + + native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None + if iteration_stats is None or native_stats is None or finish_reason is None: + return + + iteration_stats.update_from_finished_request( + finish_reason=finish_reason, + request_id=req_state.external_req_id, + num_prompt_tokens=req_state.prompt_len, + max_tokens_param=req_state.max_tokens_param, + req_stats=native_stats, + num_cached_tokens=req_state.num_cached_tokens, + ) + if not iteration_stats.finished_requests: + return + + finished_request = iteration_stats.finished_requests[-1] + if finished_request.request_id != req_state.external_req_id: + return + self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = ( + float(finished_request.mean_time_per_output_token) * 1000.0 + ) diff --git a/vllm_omni/outputs/utils.py b/vllm_omni/outputs/utils.py new file mode 100644 index 00000000000..88fc11fb5c2 --- /dev/null +++ b/vllm_omni/outputs/utils.py @@ -0,0 +1,24 @@ +"""Shared helpers for multimodal output handling.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def _to_cpu(x: Any) -> Any: + """Move a tensor to CPU; pass through non-tensors unchanged.""" + if isinstance(x, torch.Tensor): + try: + return x.detach().to("cpu", non_blocking=True).contiguous() + except (RuntimeError, AttributeError): + return x + if isinstance(x, dict): + return {str(sub_key): _to_cpu(sub_value) for sub_key, sub_value in x.items()} + return x + + +def _is_tensor_list(value: Any) -> bool: + """Return True if *value* is a non-empty list of tensors (deferred concat).""" + return isinstance(value, list) and bool(value) and isinstance(value[0], torch.Tensor) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 373246d71c2..4f244129461 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,7 +1,9 @@ import logging +import os import sys from functools import cached_property +import torch from aenum import extend_enum from vllm.config import ModelConfig as _OriginalModelConfig from vllm.inputs import TokensPrompt as _OriginalTokensPrompt @@ -21,6 +23,8 @@ from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding from vllm_omni.request import OmniRequest, OmniStreamingUpdate +_PATCH_LOGGER = logging.getLogger("vllm_omni.patch") + # ============================================================================= # Patch ModelConfig.is_mm_prefix_lm to support omni-specific models # ============================================================================= @@ -71,6 +75,185 @@ def _patched_is_mm_prefix_lm(self): "for HunyuanImage3 will not work. Check vLLM ModelConfig changes." ) +# ============================================================================= +# Patch ModelOptNvFp4LinearMethod.process_weights_after_loading to clamp NaN +# bytes in per-block FP8 weight_scale tensors at load time. +# ============================================================================= +# WHY: ModelOpt 0.44's float32 -> torch.float8_e4m3fn cast of per-block weight +# scales occasionally emits NaN bytes (E4M3 encoding 0x7F / 0xFF) when the +# pre-cast scale rounds above the FP8 E4M3 max of 448 after the global-scale +# division. Any single NaN byte in weight_scale propagates through the +# FlashInfer FP4 GEMM and the served model collapses to `!!!!`. NVFP4 W4A4 +# Qwen3-Omni checkpoints published by the community — exported with stock +# ModelOpt 0.44 — currently fail this way the moment they're served. +# +# WHY here (vllm-omni, not upstream): the root cause sits in ModelOpt; the +# analogous fix in vLLM's own PWAL is also pending. Both are planned as +# follow-up PRs. Until those land in vllm-omni's pinned vLLM, this load-time +# clamp keeps the user-visible failure from looking like a vllm-omni +# inference bug. Newly calibrated, clean checkpoints pay no runtime cost +# (the clamp is a no-op when no NaN bytes are present). +# +# SCOPE: ModelOptNvFp4LinearMethod (W4A4 NVFP4 Linear) only. NvFp4FusedMoE / +# NvFp4W4A16 / CompressedTensors / Quark NVFP4 paths are not covered. +# +# SELF-EXTINGUISH: `_already_patched_upstream` heuristically detects when +# vLLM's own PWAL contains an in-place `masked_fill_` against `weight_scale` +# / `isnan` — the structure the upstream fix is expected to take when it is +# filed (planned as a follow-up PR after this one merges). Once vllm-omni's +# vllm pin moves to a release with that upstream fix, the override is +# skipped at import and this block can be deleted. NOTE: the heuristic only +# matches "vLLM PWAL clamps NaN in-place with `masked_fill_`"; if the +# upstream fix lands as `nan_to_num_` or as a clamp before the FP32→FP8 +# cast, the check won't fire and this override stays active. The override +# is idempotent so the overlap is a warning log, not a correctness issue — +# but the heuristic should be revisited when the upstream PR is filed. +# +# ORDERING: the clamp must run BEFORE the original PWAL. The non-Blackwell +# Marlin fallback (sm_<100) casts weight_scale FP8 -> bf16/fp16 and permutes +# inside its kernel PWAL; an after-PWAL clamp would either trip the +# byte-view shape assertion (FP8 1B/elem -> bf16 2B/elem) or operate on +# already-transformed bytes where NaN has propagated through permute. On +# Blackwell (Cutlass/TRTLLM/cuDNN) the kernel PWAL only swizzles/shuffles +# FP8 bytes, so clamp-before is equivalent to clamp-after there. + + +def _clamp_nvfp4_weight_scale_nans(layer) -> int: + """Scan ``layer.weight_scale`` for NaN bytes and clamp them to the FP8 + E4M3 max byte (0x7E) in place. Returns the number of bytes clamped. + + Scoped to NVFP4 ``float8_e4m3fn`` weight_scale only — a future MXFP4 + weight-scale clamp would need its own helper (MXFP4 uses + ``float8_e8m0fnu`` for scales, and the byte encoding of NaN differs). + + Exposed at module scope so tests can exercise the clamp directly + without monkey-patching + reloading the parent ``vllm_omni.patch`` + module. + """ + # Defensive: a stray subclass without weight_scale would crash inside + # the override; just bail out. + weight_scale = getattr(layer, "weight_scale", None) + if weight_scale is None: + return 0 + # Re-entry safety: this clamp expects raw on-disk FP8 weight_scale (1 + # byte/elem). If we are ever invoked after the original PWAL ran (e.g. + # a caller re-issues process_weights_after_loading on the same layer), + # weight_scale may have been cast to bf16/fp16 by the Marlin path, in + # which case the uint8 byte view below would have a different shape + # than the NaN mask. Bail out cleanly rather than trip the assert. + if weight_scale.dtype != torch.float8_e4m3fn: + return 0 + # `.view(torch.uint8)` below requires a contiguous tensor and raises on a + # non-contiguous one. A checkpoint whose per-block scale tensor isn't + # contiguous would otherwise crash here. Materialize a contiguous copy and + # write it back to the Parameter so the in-place byte clamp lands on the + # tensor the kernel will actually read (a bare `.contiguous()` copy would + # be discarded and the original left un-clamped). + weight_scale_data = weight_scale.data + if not weight_scale_data.is_contiguous(): + weight_scale_data = weight_scale_data.contiguous() + weight_scale.data = weight_scale_data + nan_mask = torch.isnan(weight_scale_data) + if not nan_mask.any(): + return 0 + # PyTorch does not implement masked_fill_ for float8_e4m3fn, so view + # the storage as uint8 and write the FP8 max byte (0x7E) directly. + # float8_e4m3fn is 1 byte per element so the byte view is shape- + # for-shape with the original tensor; pin that here so a future + # packed scale format would fail loudly rather than silently corrupt + # unrelated bytes. + byte_view = weight_scale_data.view(torch.uint8) + if byte_view.shape != nan_mask.shape: + raise RuntimeError( + f"NVFP4 weight_scale uint8 view shape {byte_view.shape} != NaN mask shape " + f"{nan_mask.shape}; per-block scale layout changed — recheck the clamp." + ) + fp8_max_byte = ( + torch.tensor(torch.finfo(torch.float8_e4m3fn).max, dtype=torch.float8_e4m3fn).view(torch.uint8).item() + ) + byte_view.masked_fill_(nan_mask, fp8_max_byte) + n = int(nan_mask.sum().item()) + _PATCH_LOGGER.warning("Clamped %d NaN entries in NVFP4 per-block weight_scale.", n) + return n + + +# Module-level defaults so downstream code (and tests) can import these names +# without guarding for the import-failure / escape-hatch branches below. +# `_already_patched_upstream` = upstream PWAL contains its own NaN clamp. +# `_clamp_installed` = our wrapper was installed on the upstream class. +# These are independent: the env-var escape hatch and the import-failure path +# both leave the wrapper uninstalled WITHOUT upstream being patched, so the +# right check for "we own NaN-clamp behavior" is `_clamp_installed`. +_already_patched_upstream = False +_clamp_installed = False + +try: + # Escape hatch — set VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP=1 to skip installing + # the patch (e.g. to confirm a `!!!!` failure is the NaN-byte case). + # The escape-hatch deliberately raises ImportError so the not-installed + # warning below logs through the same path a real ImportError would. + # Use the repo-wide bool-env idiom so values like `0`, `false`, `no`, + # `off` correctly mean "do not skip" rather than tripping naive + # truthiness on the non-empty string. + if os.environ.get("VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP", "").lower() in ("1", "true", "yes", "on"): + raise ImportError("VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP is set; skipping NaN-clamp install") + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod as _OriginalModelOptNvFp4LinearMethod, + ) +except ImportError as _nan_clamp_import_err: + _PATCH_LOGGER.warning( + "NVFP4 weight_scale NaN-clamp patch could NOT install: %s. NVFP4 W4A4 " + "checkpoints with NaN bytes in per-block weight_scale will serve `!!!!`.", + _nan_clamp_import_err, + ) +else: + _current_nvfp4_pwal = _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading + # Reload idempotency: on a module reload (importlib.reload in a test, or a + # second import path) the class attribute already holds OUR wrapper, so + # capturing it as the "original" and re-wrapping would nest the clamp and + # run it twice per load. Recover the genuine upstream method from the + # sentinel we stamp below; the heuristic + a fresh single-level wrapper are + # then computed against the real upstream, not against our own wrapper. + _original_nvfp4_pwal = getattr(_current_nvfp4_pwal, "_vllm_omni_wrapped_pwal", _current_nvfp4_pwal) + _upstream_pwal_names = set(_original_nvfp4_pwal.__code__.co_names or ()) + # Require ALL three names — `weight_scale` alone is too loose because the + # current upstream PWAL already references `weight_scale_2` (close prefix + # collisions aside, future PWALs may legitimately reference `weight_scale` + # without a NaN clamp). `masked_fill_` + `isnan` + `weight_scale` together + # are the structural signature of an in-place NaN clamp on weight_scale. + # KNOWN BLIND SPOTS (false-negatives, safe direction — we install a + # redundant but idempotent clamp): upstream using + # `getattr(layer, "weight_scale")` (string lands in co_consts, not + # co_names), or factoring the clamp into a helper called from PWAL + # (none of the three names appear in the top-level co_names). Revisit + # this heuristic when the upstream PR is actually filed. + _already_patched_upstream = all(n in _upstream_pwal_names for n in ("masked_fill_", "isnan", "weight_scale")) + + def _patched_nvfp4_pwal(self, layer, *args, **kwargs): + # Clamp BEFORE the original PWAL — see ORDERING note above. + _clamp_nvfp4_weight_scale_nans(layer) + _original_nvfp4_pwal(self, layer, *args, **kwargs) + + # Sentinel so a later reload of this module recovers the genuine upstream + # method above instead of wrapping this wrapper (see reload-idempotency note). + _patched_nvfp4_pwal._vllm_omni_wrapped_pwal = _original_nvfp4_pwal + + if not _already_patched_upstream: + _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading = _patched_nvfp4_pwal + # Fail loudly if install dropped silently (e.g. another plugin + # patched the same class) rather than degrade to `!!!!` at decode. + # Use raise, not assert: asserts are compiled out under `python -O` / + # PYTHONOPTIMIZE, which would silently disable this guard in exactly + # the optimized runs where a conflicting plugin must still be caught. + if _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading is not _patched_nvfp4_pwal: + raise RuntimeError("NVFP4 weight_scale NaN-clamp install failed — check for conflicting plugins.") + _clamp_installed = True + + _PATCH_LOGGER.info( + "NVFP4 W4A4 weight_scale NaN-clamp: %s.", + "skipped (upstream already patched)" if _already_patched_upstream else "installed", + ) + # ============================================================================= # Patch GlmImageTextConfig to expose mrope_section in rope_parameters # ============================================================================= @@ -130,39 +313,6 @@ def _patched_glm_image_text_config_init(self, *args, **kwargs): if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest: module.EngineCoreRequest = OmniEngineCoreRequest -# ============================================================================= -# Patch unregister_vllm_metrics to skip vllm:omni_* collectors -# ============================================================================= -# WHY: Upstream unregister_vllm_metrics() uses `"vllm" in collector._name` to -# strip collectors before each new PrometheusStatLogger registers, which is -# how it avoids "Duplicated timeseries" when the same process spawns multiple -# engines / orchestrators. But its substring match also wipes our vllm:omni_* -# families, so we must replace it with a scoped version that keeps ours. -# -# REMOVAL: Remove this patch once upstream vLLM adds -# _STAT_LOGGER_METRIC_NAMES to vllm.v1.metrics.prometheus and scopes -# unregister_vllm_metrics() to that set. Track: -# https://github.com/vllm-project/vllm/pull/42331 -import vllm.v1.metrics.prometheus as _vllm_prometheus # noqa: E402 - -_logger = logging.getLogger(__name__) - - -def _scoped_unregister_vllm_metrics(): - """Drop upstream vllm:* collectors but preserve vllm:omni_* (ours).""" - from prometheus_client import REGISTRY - - for collector in list(REGISTRY._collector_to_names): - name = getattr(collector, "_name", "") - if "vllm" not in name: - continue - if name.startswith("vllm:omni_") or name.startswith("vllm_omni"): - continue - REGISTRY.unregister(collector) - - -_vllm_prometheus.unregister_vllm_metrics = _scoped_unregister_vllm_metrics - # Patch: add qwen3_omni_moe to vllm's chat template fallback registry. # Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone @@ -174,11 +324,13 @@ def _patch_chat_template_registry(): try: from vllm.transformers_utils.chat_templates.registry import ( _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK, - _get_qwen_chat_template_fallback, + CHAT_TEMPLATES_DIR, ) if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK: - _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = _get_qwen_chat_template_fallback + _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = ( + lambda _: CHAT_TEMPLATES_DIR / "template_chatml.jinja" + ) except ImportError: pass @@ -265,3 +417,70 @@ def _patch_fp8_use_quack_fused_bias(): _patch_fp8_use_quack_fused_bias() + + +# ============================================================================= +# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown +# ============================================================================= +# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free +# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm +# gets the safe empty-handle return. On CUDA, the callback falls through and +# returns the original handle, causing cuMemRelease on already-freed memory +# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup. +# +# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL +# platforms, then sets ``is_asleep = True``. When the atexit handler +# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free +# callback, the asleep entries must return an empty chunk list so the C +# extension skips ``cuMemRelease`` — exactly the same logic already used by +# the ROCm guard. +# +# The fix removes the ``current_platform.is_rocm()`` condition so the guard +# applies to CUDA (and any future platform that implements cumem). +# +# FRAGILITY: Relies on ``_python_free_callback`` being a regular method +# (not a slot or C extension). If CuMemAllocator is rewritten in C/Cython, +# this monkey-patch will silently become a no-op. +def _patch_cumem_free_callback_cuda() -> None: + try: + from vllm.device_allocator.cumem import CuMemAllocator + except ImportError: + _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch") + return + + _original_free_callback = CuMemAllocator._python_free_callback + + if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False): + return + + # The upstream bug: `_python_free_callback` only skips the double-free + # for ROCm (line ~206). We wrap the method to extend the guard to all + # platforms. + def _patched_free_callback(self, ptr: int) -> tuple: + data = self.pointer_to_data.pop(ptr) + if data.cpu_backup_tensor is not None: + data.cpu_backup_tensor = None + if data.is_asleep: + # sleep() already called unmap_and_release() on this allocation. + # Return an empty chunk list so the C extension skips + # cuMemRelease, avoiding a double-free. Same logic as the + # existing ROCm guard, but applied to all platforms. + device, size, d_mem, _ = data.handle + result = (device, size, d_mem, []) + _PATCH_LOGGER.debug( + "[cumem-cuda] Free callback: asleep entry %s -> empty handle", + ptr, + ) + return result + # Drain pending kernels before the C extension's cuMemUnmap. + torch.accelerator.synchronize(data.handle[0]) + return data.handle + + _patched_free_callback._omni_cumem_cuda_patched = True + CuMemAllocator._python_free_callback = _patched_free_callback + _PATCH_LOGGER.info( + "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms." + ) + + +_patch_cumem_free_callback_cuda() diff --git a/vllm_omni/platforms/__init__.py b/vllm_omni/platforms/__init__.py index 64a7cdb16f0..942eb7c24bd 100644 --- a/vllm_omni/platforms/__init__.py +++ b/vllm_omni/platforms/__init__.py @@ -7,7 +7,6 @@ from typing import TYPE_CHECKING from vllm.utils.import_utils import resolve_obj_by_qualname -from vllm.utils.torch_utils import supports_xccl from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum from vllm_omni.plugins import ( @@ -86,7 +85,7 @@ def xpu_omni_platform_plugin() -> str | None: try: import torch - if supports_xccl(): + if torch.distributed.is_xccl_available(): dist_backend = "xccl" else: dist_backend = "ccl" diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index c96da69fca5..7be13535229 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -36,7 +36,7 @@ def get_omni_generation_worker_cls(cls) -> str: @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/model_executor/stage_configs" + return "vllm_omni/deploy" @classmethod def has_flash_attn_package(cls) -> bool: @@ -121,6 +121,31 @@ def get_diffusion_attn_backend_cls( if selected_backend is not None: backend_upper = selected_backend.upper() + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + try: + importlib.import_module("kernels") + logger.info("Using HuggingFace kernels-backed attention backend '%s'", backend_upper) + except ImportError: + if backend_upper == "FLASH_ATTN_HUB": + logger.warning( + "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN." + ) + backend_upper = "FLASH_ATTN" + elif backend_upper == "FLASH_ATTN_3_HUB": + logger.warning( + "HuggingFace `kernels` library is not available. Falling back to local FLASH_ATTN." + ) + backend_upper = "FLASH_ATTN" + + if backend_upper == "FLASH_ATTN_3_HUB": + fa3_hub_supported = compute_capability is not None and compute_capability.major >= 9 + if not fa3_hub_supported: + logger.warning( + "FLASH_ATTN_3_HUB requires a Hopper-class GPU with compute capability >= 9.0. " + "Falling back to FLASH_ATTN_HUB." + ) + backend_upper = "FLASH_ATTN_HUB" + if backend_upper == "FLASH_ATTN" and not flash_attn_supported: if not compute_supported: logger.warning( @@ -213,14 +238,19 @@ def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.cuda.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.cuda.mem_get_info(device) + return free, total + @classmethod def get_device_name(cls, device_id: int = 0) -> str: return torch.cuda.get_device_name(device_id) @classmethod def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig: - """Copied from vllm/platforms/cuda/platform.py v0.20.0 with force using vllm_c kernels""" - default = ["vllm_c", "native"] # Originally using "native" here when compiling + """Prefer ``vllm_c`` CUDA kernels over ``native`` for diffusion IR ops.""" + default = ["vllm_c", "native"] # Use oink if enabled for rms_norm # TODO(Laurawly/luka): remove this env var, diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 14aa93b20fb..54725b553dd 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -136,6 +136,11 @@ def get_diffusion_model_runner_cls(cls) -> str: """ return "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner" + @classmethod + def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None: + """Initialize platform-specific state for diffusion worker VllmConfig.""" + return None + @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: raise NotImplementedError @@ -156,6 +161,10 @@ def synchronize(cls) -> None: def get_free_memory(cls, device: torch.device | None = None) -> int: raise NotImplementedError + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + raise NotImplementedError + @classmethod def create_autocast_context( cls, @@ -243,6 +252,10 @@ class UnspecifiedOmniPlatform(OmniPlatform): _enum = PlatformEnum.UNSPECIFIED device_type = "cpu" + @classmethod + def get_torch_device(cls, local_rank: int | None = None) -> torch.device: + return torch.device("cpu") + @classmethod def get_device_count(cls) -> int: return 0 diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 75f20a5a725..818338da496 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -32,7 +32,7 @@ def get_omni_generation_worker_cls(cls) -> str: @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/model_executor/stage_configs" + return "vllm_omni/deploy" @classmethod def has_flash_attn_package(cls) -> bool: @@ -78,6 +78,14 @@ def get_diffusion_attn_backend_cls( if selected_backend is not None: backend_upper = selected_backend.upper() + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + logger.warning( + "HuggingFace kernels-backed FlashAttention is " + "not supported on MUSA. Falling back to local " + "FLASH_ATTN." + ) + backend_upper = "FLASH_ATTN" + if backend_upper == "FLASH_ATTN" and not flash_attn_supported: if not compute_supported: logger.warning( @@ -146,17 +154,14 @@ def synchronize(cls) -> None: @classmethod def get_free_memory(cls, device: torch.device | None = None) -> int: - """Get the free memory on the MUSA device. - - Args: - device: Optional device to query. If None, uses current device. - - Returns: - Free memory in bytes. - """ free, _ = torch.musa.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.musa.mem_get_info(device) + return free, total + @classmethod def get_device_name(cls, device_id: int = 0) -> str: return torch.musa.get_device_name(device_id) diff --git a/vllm_omni/platforms/npu/_310p/patch/__init__.py b/vllm_omni/platforms/npu/_310p/patch/__init__.py index b3ab637ff2e..5bfd424294f 100644 --- a/vllm_omni/platforms/npu/_310p/patch/__init__.py +++ b/vllm_omni/platforms/npu/_310p/patch/__init__.py @@ -7,6 +7,7 @@ _WORKER_PATCHED = False _QWEN3_TTS_TALKER_ARCH = "Qwen3TTSTalkerForConditionalGeneration" +_QWEN3_TTS_CODE2WAV_ARCH = "Qwen3TTSCode2Wav" def apply_patches() -> None: @@ -22,9 +23,13 @@ def apply_patches() -> None: def apply_model_patches(model_config) -> None: - if getattr(model_config, "model_arch", None) != _QWEN3_TTS_TALKER_ARCH: - return + model_arch = getattr(model_config, "model_arch", None) + if model_arch == _QWEN3_TTS_TALKER_ARCH: + from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_talker_patches - from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_talker_patches + apply_talker_patches() + return + elif model_arch == _QWEN3_TTS_CODE2WAV_ARCH: + from vllm_omni.platforms.npu._310p.patch.qwen3_tts import apply_code2wav_patches - apply_talker_patches() + apply_code2wav_patches() diff --git a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py index c43cdd3cf42..713b67f8136 100644 --- a/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py +++ b/vllm_omni/platforms/npu/_310p/patch/qwen3_tts.py @@ -5,26 +5,30 @@ from __future__ import annotations +from collections.abc import Sequence + import numpy as np import torch -from transformers.models.mimi import modeling_mimi +import torch.nn as nn +import torch.nn.functional as F +import torch_npu from vllm.multimodal.audio import AudioResampler +from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310 +from vllm_ascend.sample.sampler import apply_top_k_top_p, random_sample +from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, aligned_16, maybe_trans_nz, nd_to_nz_2d from vllm_omni.model_executor.models.common import qwen3_code_predictor -from vllm_omni.model_executor.models.qwen3_tts import prompt_embeds_builder, qwen3_tts_talker +from vllm_omni.model_executor.models.qwen3_tts import ( + prompt_embeds_builder, + qwen3_tts_code2wav, + qwen3_tts_code_predictor_vllm, + qwen3_tts_talker, +) _RUNTIME_DTYPE = torch.float16 _CPU_DEVICE = torch.device("cpu") - - -class _MimiEuclideanCodebook310P(modeling_mimi.MimiEuclideanCodebook): - def quantize(self, hidden_states): - # 310P does not support torch.cdist on NPU. - device = hidden_states.device - dists = torch.cdist( - hidden_states[None].to(_CPU_DEVICE, torch.float32), self.embed[None].to(_CPU_DEVICE, torch.float32), p=2 - )[0] - return dists.argmin(dim=-1).to(device=device) +_PATCHED = False +_CODE2WAV_PATCHED = False class _Qwen3TTSTalker310P(qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration): @@ -32,10 +36,18 @@ def __init__(self, *, vllm_config, prefix: str = "") -> None: super().__init__(vllm_config=vllm_config, prefix=prefix) self._embedding_dtype = _RUNTIME_DTYPE self._prompt_builder._embedding_dtype = _RUNTIME_DTYPE + # 310P random sampling operators cannot be captured by ACL graphs. + # Keep Talker-MTP eager so CodePredictor can replay its own NPU graphs + # without changing sampling semantics. + self.talker_mtp_graph_safe = False + self.talker_mtp_accepts_per_row_generators = True def load_weights(self, weights): loaded = super().load_weights(weights) - self.encoder.to(dtype=_RUNTIME_DTYPE) + # The Mimi tokenizer encoder is used only while building ref_audio + # prompts. Keep this preprocessing-only module on CPU because the 310P + # NPU path is sensitive to changing reference-audio shapes. + self.encoder.to(device=_CPU_DEVICE, dtype=torch.float32) return loaded def _encode_ref_audio_batch( @@ -51,7 +63,7 @@ def _encode_ref_audio_batch( resampler = AudioResampler(target_sr=target_sr) wavs = [resampler.resample(w.astype(np.float32), orig_sr=int(sr)) for w in wavs] - inputs = fe(raw_audio=wavs, sampling_rate=target_sr, return_tensors="pt").to(device).to(_RUNTIME_DTYPE) + inputs = fe(raw_audio=wavs, sampling_rate=target_sr, return_tensors="pt").to(torch.float32) with torch.inference_mode(): encoded = self.encoder.encode( @@ -100,10 +112,56 @@ def extract_speaker_embedding(self, wav: np.ndarray, sr: int) -> torch.Tensor: return spk.to(dtype=dtype) +# =================================================================== +# Code2Wav layer patches +# =================================================================== +# +# Code2Wav runs under the 310P graph path after the stage-0 Talker has +# produced codec tokens. Common NPU weight packing and fused tokenizer ops +# live in the shared model code; the 310P patch only selects the runtime dtype. + + +class _Qwen3TTSCode2Wav310P(qwen3_tts_code2wav.Qwen3TTSCode2Wav): + """Qwen3-TTS Code2Wav specialized for the 310P NPU path.""" + + def _npu_decoder_runtime_dtype(self, device: torch.device) -> torch.dtype: + return _RUNTIME_DTYPE + + +# =================================================================== +# CodePredictor layer patches +# =================================================================== +# +# Keep the portable implementation in common/qwen3_code_predictor.py. +# The overrides below are installed only by the 310P platform patch because +# the short CodePredictor loop is graph-captured on 310P and needs the 310P +# flash-attention mask layout plus loop-local projection and sampling changes. + + class _Qwen3CodePredictorAttention310P(qwen3_code_predictor.CodePredictorAttention): + """Attention override using 310P RoPE and flash-attention kernels. + + The shared attention path is written in portable PyTorch. This override + keeps the 310P-specific RoPE op, token alignment, and FRACTAL_NZ mask path + scoped to the platform patch. + """ + def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self._buffers.pop("_fusion_causal_mask", None) + self._q_size = self.num_heads * self.head_dim + self._kv_size = self.num_kv_heads * self.head_dim + self._fused_qkv_weight = None + self._fused_qkv_bias = None + + def prepare_qkv_weights(self) -> None: + # Pack QKV once so each graph replay uses one matmul and consumes the + # weight directly in the 310P matmul layout. + self._fused_qkv_weight = maybe_trans_nz( + torch.cat((self.q_proj.weight, self.k_proj.weight, self.v_proj.weight), dim=0).contiguous() + ) + if self.q_proj.bias is not None: + self._fused_qkv_bias = torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias), dim=0) def forward( self, @@ -111,32 +169,29 @@ def forward( position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: torch.Tensor | None = None, ) -> torch.Tensor: - if hidden_states.device.type != "npu" or attention_mask is None: - return super().forward(hidden_states, position_embeddings) - bsz, seq_len, _ = hidden_states.shape - hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) - hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) - - q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2) - k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2) - v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2) + qkv = F.linear(hidden_states, self._fused_qkv_weight, self._fused_qkv_bias) + q, k, v = qkv.split((self._q_size, self._kv_size, self._kv_size), dim=-1) + q = self.q_norm(q.view(bsz, seq_len, self.num_heads, self.head_dim)).transpose(1, 2) + k = self.k_norm(k.view(bsz, seq_len, self.num_kv_heads, self.head_dim)).transpose(1, 2) + v = v.view(bsz, seq_len, self.num_kv_heads, self.head_dim).transpose(1, 2) cos, sin = position_embeddings cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) - q = (q * cos) + (qwen3_code_predictor._rotate_half(q) * sin) - k = (k * cos) + (qwen3_code_predictor._rotate_half(k) * sin) + # Use the fused Ascend RoPE op instead of expanding RoPE into + # elementwise mul/add/rotate-half kernels. + q = torch_npu.npu_rotary_mul(q, cos, sin) + k = torch_npu.npu_rotary_mul(k, cos, sin) real_tokens = int(bsz) * int(seq_len) output_dtype = q.dtype - import torch_npu - from vllm_ascend.utils import aligned_16 - - q_f = aligned_16(q.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_heads, self.head_dim)) - k_f = aligned_16(k.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim)) - v_f = aligned_16(v.to(torch.float16).transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim)) + # 310P flash attention consumes token-major fp16 inputs with 16-token + # alignment; seq_lens carries the padding information. + q_f = aligned_16(q.transpose(1, 2).reshape(real_tokens, self.num_heads, self.head_dim)) + k_f = aligned_16(k.transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim)) + v_f = aligned_16(v.transpose(1, 2).reshape(real_tokens, self.num_kv_heads, self.head_dim)) aligned_tokens = int(q_f.shape[0]) seq_lens = torch.full((int(bsz),), int(seq_len), dtype=torch.int32, device="cpu") @@ -160,23 +215,34 @@ def forward( class _Qwen3CodePredictorDecoderLayer310P(qwen3_code_predictor.CodePredictorDecoderLayer): + """Decoder layer override that passes the 310P attention mask.""" + def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None = None, + attention_mask: torch.Tensor, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states = self.self_attn(hidden_states, position_embeddings, attention_mask=attention_mask) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states, _, residual = torch_npu.npu_add_rms_norm( + hidden_states, + residual, + self.post_attention_layernorm.weight, + self.post_attention_layernorm.variance_epsilon, + ) return residual + self.mlp(hidden_states) class _Qwen3CodePredictorBaseModel310P(qwen3_code_predictor.CodePredictorBaseModel): + """Base model override with a cached 310P causal mask. + + The 310P flash-attention path consumes the additive causal mask in + FRACTAL_NZ format. The CodePredictor sequence length is fixed and short, + so the mask is built once and reused across graph replays. + """ + def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) self._attention_mask_310p = None @@ -188,14 +254,9 @@ def forward( inputs_embeds: torch.Tensor, position_ids: torch.Tensor, ) -> torch.Tensor: - if inputs_embeds.device.type != "npu": - return super().forward(inputs_embeds, position_ids) - if self._attention_mask_310p is None or self._attention_mask_310p_device != inputs_embeds.device: - import torch_npu - from vllm_ascend._310p.attention.attention_mask import AttentionMaskBuilder310 - from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, nd_to_nz_2d - + # Store the additive causal mask in the format consumed by the + # 310P flash-attention kernel. mask = AttentionMaskBuilder310.gen_causal_additive_mask( self._attention_mask_310p_max_seq, inputs_embeds.device, @@ -205,23 +266,239 @@ def forward( input_dtype = inputs_embeds.dtype hidden_states = inputs_embeds - with torch.amp.autocast(inputs_embeds.device.type, enabled=False, dtype=torch.float32): - position_embeddings = self.rotary_emb(hidden_states, position_ids) - for layer in self.layers: - hidden_states = layer( - hidden_states, - position_embeddings, - attention_mask=self._attention_mask_310p, - ) - hidden_states = self.norm(hidden_states) + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer( + hidden_states, + position_embeddings, + attention_mask=self._attention_mask_310p, + ) + hidden_states = self.norm(hidden_states) return hidden_states.to(input_dtype) +class _Qwen3TTSTalkerCodePredictor310P( + qwen3_tts_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM +): + """Qwen3-TTS code predictor specialized for the 310P NPU path.""" + + def __init__( + self, + *, + vllm_config, + config, + talker_config, + prefix: str = "code_predictor", + ) -> None: + super().__init__( + vllm_config=vllm_config, + config=config, + talker_config=talker_config, + prefix=prefix, + ) + self._projected_codec_embed_weight = None + + def _prepare_npu_weights(self) -> None: + qkv_projections = set() + with torch.no_grad(): + for layer in self.model.layers: + attention = layer.self_attn + attention.prepare_qkv_weights() + qkv_projections.update((attention.q_proj, attention.k_proj, attention.v_proj)) + + for module in self.modules(): + if isinstance(module, nn.Linear) and module not in qkv_projections: + module.weight.data = maybe_trans_nz(module.weight.data) + + def load_weights(self, weights): + loaded = super().load_weights(weights) + with torch.no_grad(): + self._projected_codec_embed_weight = torch.stack( + [self.small_to_mtp_projection(embed.weight).detach() for embed in self.model.codec_embedding], + dim=0, + ).contiguous() + return loaded + + @torch.inference_mode() + def forward( + self, + layer0_code: torch.Tensor, + layer0_embed: torch.Tensor, + last_talker_hidden: torch.Tensor, + do_sample: bool = True, + temperature: float = 0.9, + top_k: int = 50, + top_p: float = 1.0, + generator: torch.Generator | None = None, + generators: Sequence[torch.Generator | None] | None = None, + ) -> torch.Tensor: + bsz = int(layer0_code.shape[0]) + if generators is not None and len(generators) != bsz: + raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}") + num_groups = self._num_groups + device = layer0_code.device + + self._setup_compile() + dtype = self._model_dtype + + padded_bsz = self._padded_bsz(bsz) + self._ensure_buffers(device, dtype, padded_bsz) + + proj_buf = self._proj_buf + max_seq = num_groups + 1 + projection = self.small_to_mtp_projection + model_fwd = self._compiled_model_fwd + lm_heads = self._lm_heads_list + if generators is not None: + npu_generators = {i: row_generator for i, row_generator in enumerate(generators) if row_generator} + elif generator is not None: + npu_generators = {i: generator for i in range(bsz)} + else: + npu_generators = {} + + proj_buf[:padded_bsz].zero_() + initial_embeds = torch.cat( + ( + last_talker_hidden.reshape(bsz, 1, -1), + layer0_embed.reshape(bsz, 1, -1), + ), + dim=1, + ) + if initial_embeds.dtype != dtype: + initial_embeds = initial_embeds.to(dtype) + proj_buf[:bsz, :2, :].copy_(projection(initial_embeds)) + + stored_mode = self._wrapper_config.sampling_mode == "stored" + if stored_mode: + s_top_k = self._top_k + s_top_p = self._top_p + else: + use_sampling = do_sample and temperature > 0 + inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0 + if use_sampling and top_p != 1.0: + raise NotImplementedError( + "top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0." + ) + + top_k_tensor = None + top_p_tensor = None + if stored_mode: + top_k_hint = s_top_k if s_top_k > 0 else None + if s_top_k > 0: + top_k_tensor = torch.full((bsz,), s_top_k, dtype=torch.int32, device=device) + if s_top_p < 1.0: + top_p_tensor = torch.full((bsz,), s_top_p, dtype=dtype, device=device) + elif use_sampling: + top_k_hint = top_k if top_k > 0 else None + if top_k > 0: + top_k_tensor = torch.full((bsz,), top_k, dtype=torch.int32, device=device) + else: + top_k_hint = None + + all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) + all_codes[:, 0] = layer0_code.reshape(bsz) + + for step in range(1, num_groups): + graph_key: int | tuple[int, int] = padded_bsz + seq_len = max_seq + if self._prefix_graphs_enabled: + prefix_key = (padded_bsz, step + 1) + if prefix_key in self._device_graphs: + graph_key = prefix_key + seq_len = step + 1 + pos_ids = self._bucket_pos_ids.get(graph_key) + if pos_ids is None: + pos_ids = ( + torch.arange(seq_len, device=device, dtype=torch.long) + .unsqueeze(0) + .expand(padded_bsz, -1) + .contiguous() + ) + + device_graph_entry = self._device_graphs.get(graph_key) + if device_graph_entry is not None: + device_graph_entry[0].replay() + hidden_out = device_graph_entry[1] + else: + hidden_out = model_fwd(proj_buf[:padded_bsz, :seq_len, :], pos_ids) + + logits = lm_heads[step - 1](hidden_out[:bsz, step, :]) + + if stored_mode: + if top_k_tensor is not None or top_p_tensor is not None: + logits = apply_top_k_top_p(logits, p=top_p_tensor, k=top_k_tensor, top_k=top_k_hint) + candidate_indices = None + if isinstance(logits, tuple): + logits, candidate_indices = logits + probs = F.softmax(logits, dim=-1, dtype=torch.float32) + code = random_sample(probs, npu_generators) + if candidate_indices is not None: + code = candidate_indices.gather(1, code.unsqueeze(1)).squeeze(1) + else: + if use_sampling: + scaled = logits * inv_temperature + if top_k_tensor is not None: + scaled = apply_top_k_top_p(scaled, p=None, k=top_k_tensor, top_k=top_k_hint) + candidate_indices = None + if isinstance(scaled, tuple): + scaled, candidate_indices = scaled + probs = F.softmax(scaled, dim=-1, dtype=torch.float32) + code = random_sample(probs, npu_generators) + if candidate_indices is not None: + code = candidate_indices.gather(1, code.unsqueeze(1)).squeeze(1) + else: + code = logits.argmax(dim=-1, keepdim=True) + + all_codes[:, step] = code.reshape(bsz) + if step < num_groups - 1: + proj_buf[:bsz, step + 1, :].copy_( + F.embedding(code.reshape(-1), self._projected_codec_embed_weight[step - 1]) + ) + + return all_codes + + +# =================================================================== +# Patch registration +# =================================================================== + + def apply_talker_patches() -> None: - modeling_mimi.MimiEuclideanCodebook = _MimiEuclideanCodebook310P + """Install Qwen3-TTS Talker and CodePredictor 310P patches. + + The generic model modules stay unchanged. Patch registration swaps in the + 310P-specialized CodePredictor classes and wrapper methods only when the + 310P platform applies the Talker patch. + """ + + global _PATCHED + + if _PATCHED: + return + qwen3_tts_talker.Qwen3TTSTalkerForConditionalGeneration = _Qwen3TTSTalker310P qwen3_tts_talker.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilder310P + qwen3_tts_talker.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM = _Qwen3TTSTalkerCodePredictor310P + qwen3_tts_code_predictor_vllm.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM = ( + _Qwen3TTSTalkerCodePredictor310P + ) + qwen3_tts_code_predictor_vllm.CodePredictorBaseModel = _Qwen3CodePredictorBaseModel310P + qwen3_tts_code_predictor_vllm.Qwen3TTSTalkerCodePredictorModelVLLM = _Qwen3CodePredictorBaseModel310P prompt_embeds_builder.Qwen3TTSPromptEmbedsBuilder = _Qwen3TTSPromptEmbedsBuilder310P qwen3_code_predictor.CodePredictorAttention = _Qwen3CodePredictorAttention310P qwen3_code_predictor.CodePredictorDecoderLayer = _Qwen3CodePredictorDecoderLayer310P qwen3_code_predictor.CodePredictorBaseModel = _Qwen3CodePredictorBaseModel310P + + _PATCHED = True + + +def apply_code2wav_patches() -> None: + """Install the 310P Code2Wav runtime patch.""" + global _CODE2WAV_PATCHED + + if _CODE2WAV_PATCHED: + return + + qwen3_tts_code2wav.Qwen3TTSCode2Wav = _Qwen3TTSCode2Wav310P + + _CODE2WAV_PATCHED = True diff --git a/vllm_omni/platforms/npu/_310p/patch/worker.py b/vllm_omni/platforms/npu/_310p/patch/worker.py index 3116b54da43..49ee6052cb6 100644 --- a/vllm_omni/platforms/npu/_310p/patch/worker.py +++ b/vllm_omni/platforms/npu/_310p/patch/worker.py @@ -12,6 +12,9 @@ from __future__ import annotations +from vllm.v1.sample.sampler import Sampler +from vllm_ascend._310p.sample.sampler import AscendSampler310 + from vllm_omni.platforms.npu._310p import disable_jit_compile from vllm_omni.platforms.npu.worker import base as worker_base @@ -24,4 +27,6 @@ def _init_device(self): def apply_patch() -> None: + # Triton-Ascend does not target 310P; use vLLM's native penalty path. + AscendSampler310.apply_penalties = staticmethod(Sampler.apply_penalties) worker_base.OmniNPUWorkerBase = _OmniNPUWorkerBase310P diff --git a/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py b/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py new file mode 100644 index 00000000000..7e1d5d54019 --- /dev/null +++ b/vllm_omni/platforms/npu/models/cosyvoice2_dit_attn.py @@ -0,0 +1,169 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NPU patches for CosyVoice2 / Token2Wav DiT attention. + +CosyVoice2's DiT ``Attention.forward`` builds a key-padding style mask +``(B, 1, 1, S)`` via ``mask.unsqueeze(1)`` and feeds it to +``F.scaled_dot_product_attention``. On Ascend that call is routed to +``npu_fusion_attention`` / ``aclnnFlashAttentionScore``, which only accepts +mask shapes ``[B,N,S,S] / [B,1,S,S] / [1,1,S,S] / [S,S]`` and fails with +error **161001** (tiling / parameter invalid) for ``[B,1,1,S]``. + +This module: +1. Expands DiT attention masks to ``[B, 1, S, S]`` before SDPA. +2. Provides a MATH-backend SDPA context so inference can avoid the fused FA + kernel entirely when the platform still incorrectly routes SDPA. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext + +import torch +import torch.nn.functional as F +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False + + +def _expand_attn_mask_for_npu( + attn_mask: torch.Tensor | None, + q_len: int, + kv_len: int | None = None, +) -> torch.Tensor | None: + """Expand CosyVoice key-padding masks to Ascend FA-compatible shapes.""" + if attn_mask is None: + return None + kv_len = kv_len if kv_len is not None else q_len + + # (B, 1, S) key-padding -> (B, S_q, S_kv) then unsqueeze heads below. + if attn_mask.dim() == 3 and attn_mask.shape[-2] == 1: + attn_mask = attn_mask.expand(-1, q_len, -1) + if attn_mask.dim() == 3: + # (B, S_q, S_kv) -> (B, 1, S_q, S_kv) + attn_mask = attn_mask.unsqueeze(1) + if attn_mask.dim() == 4 and attn_mask.shape[-2] == 1 and kv_len > 1: + # (B, 1, 1, S) -> (B, 1, S_q, S) + attn_mask = attn_mask.expand(-1, -1, q_len, -1) + return attn_mask + + +def _patched_attention_forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: + b, t, c = x.shape + + q = self.to_heads(self.to_q(x)) + k = self.to_heads(self.to_k(x)) + v = self.to_heads(self.to_v(x)) + + q = self.q_norm(q) + k = self.k_norm(k) + + attn_mask = _expand_attn_mask_for_npu(attn_mask, q_len=t, kv_len=t) + x = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + dropout_p=self.attn_drop.p if self.training else 0.0, + ) + x = x.transpose(1, 2).reshape(b, t, -1) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +def _patched_attention_forward_chunk( + self, + x: torch.Tensor, + att_cache: torch.Tensor | None = None, + attn_mask: torch.Tensor | None = None, +): + b, t, c = x.shape + + q = self.to_heads(self.to_q(x)) + k = self.to_heads(self.to_k(x)) + v = self.to_heads(self.to_v(x)) + + q = self.q_norm(q) + k = self.k_norm(k) + + if att_cache is not None: + k_cache, v_cache = att_cache.chunk(2, dim=3) + k = torch.cat([k, k_cache], dim=2) + v = torch.cat([v, v_cache], dim=2) + + new_att_cache = torch.cat([k, v], dim=3) + kv_len = k.shape[2] + if attn_mask is not None: + attn_mask = _expand_attn_mask_for_npu(attn_mask, q_len=t, kv_len=kv_len) + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) + x = x.transpose(1, 2).reshape(b, t, -1) + x = self.proj(x) + x = self.proj_drop(x) + return x, new_att_cache + + +@contextmanager +def npu_math_sdpa_context() -> Iterator[None]: + """Force SDPA MATH backend so Ascend does not call fused FA.""" + try: + from torch.nn.attention import SDPBackend, sdpa_kernel + + with sdpa_kernel(SDPBackend.MATH): + yield + except Exception: + # Older torch / missing backend enum — just run as-is. + with nullcontext(): + yield + + +def _disable_upsample_encoder_compile() -> None: + """Run CosyVoice2 UpsampleConformerEncoderV2 eagerly on Ascend. + + ``forward_chunk`` is decorated with ``@torch.compile(dynamic=True, + backend="eager")`` upstream. The backend gives no speedup (still eager), + but Dynamo tracing of the relative-position attention with a KV cache + creates a symbolic add ``matrix_ac(..., chunk+cache) + matrix_bd(..., + (pos//2)+1)`` whose sizes Ascend torch's fake-tensor ``infer_size`` + cannot reconcile, crashing on the 2nd+ streaming chunk. Unwrapping the + compile and disabling Dynamo keeps concrete shapes (always equal). + """ + try: + from cosyvoice2.transformer import upsample_encoder_v2 + except ImportError: + return + + enc_cls = getattr(upsample_encoder_v2, "UpsampleConformerEncoderV2", None) + if enc_cls is None: + return + + fn = enc_cls.forward_chunk + orig = getattr(fn, "_torchdynamo_orig_callable", None) or getattr(fn, "__wrapped__", fn) + enc_cls.forward_chunk = torch._dynamo.disable(orig) # type: ignore[method-assign] + logger.info("Disabled torch.compile on CosyVoice2 UpsampleConformerEncoderV2.forward_chunk (Ascend eager)") + + +def apply_cosyvoice2_dit_attn_npu_patch() -> None: + """Monkey-patch CosyVoice2 DiT Attention for Ascend FA mask constraints.""" + global _PATCHED + if _PATCHED: + return + + try: + from cosyvoice2.flow import decoder_dit + except ImportError: + logger.debug("cosyvoice2 not installed; skip DiT attn NPU patch") + return + + attn_cls = getattr(decoder_dit, "Attention", None) + if attn_cls is None: + return + + attn_cls.forward = _patched_attention_forward # type: ignore[method-assign] + attn_cls.forward_chunk = _patched_attention_forward_chunk # type: ignore[method-assign] + _disable_upsample_encoder_compile() + _PATCHED = True + logger.info("Applied CosyVoice2 DiT Attention NPU patch (expand attn_mask to Bx1xSxS for aclnnFlashAttentionScore)") diff --git a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py b/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py index c53fad2e2d5..251eac2cff6 100644 --- a/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py +++ b/vllm_omni/platforms/npu/models/hunyuan_fused_moe.py @@ -7,7 +7,6 @@ import vllm.forward_context as _vllm_fc from vllm.config import VllmConfig from vllm.distributed import get_ep_group -from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size from vllm.distributed.parallel_state import ( init_model_parallel_group as vllm_init_model_parallel_group, ) @@ -17,7 +16,7 @@ from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type from vllm_omni.diffusion.distributed.parallel_state import ( - get_data_parallel_world_size, + get_expert_parallel_group_ranks, get_world_group, ) from vllm_omni.diffusion.forward_context import get_forward_context as omni_get_ctx @@ -42,20 +41,14 @@ def _set_hunyuan_fused_moe_forward_context(num_tokens: int) -> None: def _init_mc2_group_for_diffusion( - world_size: int, - data_parallel_size: int, - tensor_parallel_size: int, backend: str, local_rank: int, + group_ranks: list[list[int]], ) -> None: import vllm_ascend.distributed.parallel_state as vllm_ascend_parallel_state if getattr(vllm_ascend_parallel_state, "_MC2", None) is not None: return - all_ranks = torch.arange(world_size).reshape(-1, data_parallel_size * tensor_parallel_size) - group_ranks = all_ranks.unbind(0) - group_ranks = [x.tolist() for x in group_ranks] - vllm_ascend_parallel_state._MC2 = vllm_init_model_parallel_group( group_ranks, local_rank, @@ -82,20 +75,17 @@ def _select_moe_comm_method(vllm_config: VllmConfig) -> MoECommType | None: def prepare_hunyuan_fused_moe_runtime() -> None: - world_size = torch.distributed.get_world_size() - data_parallel_size = get_data_parallel_world_size() - tensor_parallel_size = get_tensor_model_parallel_world_size() - backend = torch.distributed.get_backend(get_world_group().device_group) - local_rank = get_world_group().local_rank - _init_mc2_group_for_diffusion( - world_size=world_size, - data_parallel_size=data_parallel_size, - tensor_parallel_size=tensor_parallel_size, - backend=backend, - local_rank=local_rank, - ) - - moe_comm_type = _select_moe_comm_method(vllm_config=omni_get_ctx().vllm_config) + vllm_config = omni_get_ctx().vllm_config + if vllm_config.parallel_config.enable_expert_parallel: + backend = torch.distributed.get_backend(get_world_group().device_group) + local_rank = get_world_group().local_rank + _init_mc2_group_for_diffusion( + backend=backend, + local_rank=local_rank, + group_ranks=get_expert_parallel_group_ranks(), + ) + + moe_comm_type = _select_moe_comm_method(vllm_config=vllm_config) _ensure_forward_context_attr("num_tokens", int | None, None) _ensure_forward_context_attr("in_profile_run", bool, False) _ensure_forward_context_attr("moe_comm_type", MoECommType | None, moe_comm_type) diff --git a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py b/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py index 282689bca47..bd74e9dc6ad 100644 --- a/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py +++ b/vllm_omni/platforms/npu/models/qwen3_tts_code2wav.py @@ -1,15 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Monkey-patch ``Qwen3TTSCode2Wav.__init__`` for NPU Code2Wav runtime knobs.""" +"""Patch Qwen3-TTS Code2Wav NPU runtime setup and weight preparation.""" from __future__ import annotations from typing import TYPE_CHECKING import torch +import torch.nn as nn +import torch_npu from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm_ascend.utils import maybe_trans_nz if TYPE_CHECKING: pass @@ -18,6 +21,8 @@ _PATCHED = False _original_init = None +_original_load_weights = None +_ACL_FORMAT_FRACTAL_Z = 4 def _prepare_npu_code2wav_runtime() -> None: @@ -35,14 +40,43 @@ def _patched_init(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: _original_init(self, vllm_config=vllm_config, prefix=prefix) +def _prepare_npu_decoder_weights(decoder: nn.Module) -> None: + linear_count = 0 + conv_count = 0 + with torch.no_grad(): + for module in decoder.modules(): + if isinstance(module, nn.Linear): + module.weight.data = maybe_trans_nz(module.weight.data) + linear_count += 1 + elif isinstance(module, (nn.Conv1d, nn.ConvTranspose1d)) and module.groups == 1: + module.weight.data = torch_npu.npu_format_cast(module.weight.data.contiguous(), _ACL_FORMAT_FRACTAL_Z) + conv_count += 1 + + logger.info("Prepared NPU Code2Wav weights: linear=%d conv=%d", linear_count, conv_count) + + +def _patched_load_weights(self, weights): + assert _original_load_weights is not None + loaded = _original_load_weights(self, weights) + device = self.vllm_config.device_config.device + runtime_dtype = getattr(self, "_npu_decoder_runtime_dtype", lambda _: torch.float32)(device) + self.decoder.to(device=device, dtype=runtime_dtype) + _prepare_npu_decoder_weights(self.decoder) + if runtime_dtype != torch.float32 and hasattr(self.decoder, "precompute_snake_caches"): + self.decoder.precompute_snake_caches() + return loaded + + def apply_qwen3_tts_code2wav_patch() -> None: - global _PATCHED, _original_init + global _PATCHED, _original_init, _original_load_weights if _PATCHED: return from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav _original_init = Qwen3TTSCode2Wav.__init__ + _original_load_weights = Qwen3TTSCode2Wav.load_weights Qwen3TTSCode2Wav.__init__ = _patched_init # type: ignore[method-assign] + Qwen3TTSCode2Wav.load_weights = _patched_load_weights # type: ignore[method-assign] _PATCHED = True - logger.debug("Applied NPU patch for Qwen3TTSCode2Wav.__init__") + logger.debug("Applied NPU patch for Qwen3TTSCode2Wav") diff --git a/vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py b/vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py new file mode 100644 index 00000000000..f89ecded9e7 --- /dev/null +++ b/vllm_omni/platforms/npu/models/qwen3_tts_tokenizer_v2.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""NPU patches for the Qwen3-TTS 12Hz tokenizer decoder.""" + +from __future__ import annotations + +import torch +import torch_npu +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False + + +def _apply_rotary_pos_emb_npu( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + position_ids=None, + unsqueeze_dim: int = 1, +) -> tuple[torch.Tensor, torch.Tensor]: + del position_ids + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + return torch_npu.npu_rotary_mul(q, cos, sin), torch_npu.npu_rotary_mul(k, cos, sin) + + +def _rms_norm_forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=self.variance_epsilon)[0] + + +def apply_qwen3_tts_tokenizer_v2_patch() -> None: + global _PATCHED + if _PATCHED: + return + + from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import ( + modeling_qwen3_tts_tokenizer_v2, + ) + + modeling_qwen3_tts_tokenizer_v2.apply_rotary_pos_emb = _apply_rotary_pos_emb_npu + modeling_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderRMSNorm.forward = _rms_norm_forward_npu + _PATCHED = True + logger.debug("Applied NPU patch for Qwen3-TTS 12Hz tokenizer decoder") diff --git a/vllm_omni/platforms/npu/models/step_audio2_token2wav.py b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py new file mode 100644 index 00000000000..04768d04dfd --- /dev/null +++ b/vllm_omni/platforms/npu/models/step_audio2_token2wav.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""NPU patches for Step-Audio2 / MiniCPM Token2Wav. + +Ascend-specific workarounds that must not live in the shared GPU model file: + +1. HiFT vocoder CPU offload — STFT/ISTFT and sine-source generation are + unstable on Ascend (aicore 507015). HiFT is tiny, so running it on CPU + is acceptable; inputs/outputs are moved transparently. +2. CosyVoice2 DiT SDPA — force MATH backend (+ DiT attn mask expand) to + avoid fused FA rejecting CosyVoice ``(B,1,1,S)`` masks (error 161001). +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager, nullcontext +from types import MethodType + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False +_original_ensure_models_loaded = None +_original_forward = None +_original_stream_chunk_for = None + + +def patch_hift_module_for_npu(hift: torch.nn.Module) -> None: + """Run the entire HiFT vocoder on CPU when on Ascend NPU. + + HiFT is tiny, so keeping it on NPU yields negligible speedup while + several of its operators are unstable on Ascend and trigger aicore + exceptions (runtime error 507015): the STFT/ISTFT as well as the + sine-source generation (``f0_upsamp`` + ``m_source`` cumsum/sin over + the full waveform length). Running the whole module on CPU sidesteps + all of them. Inputs are copied to CPU and outputs are moved back to + the original device transparently, so callers stay unchanged. + """ + if getattr(hift, "_npu_cpu_offload_patched", False): + return + + hift.to("cpu") + original_forward = hift.forward + + def _forward_on_cpu(_module, *args, **kwargs): + output_device: torch.device | None = None + + def to_cpu(value): + nonlocal output_device + if isinstance(value, torch.Tensor): + if output_device is None and value.device.type != "cpu": + output_device = value.device + return value.cpu() + return value + + cpu_args = tuple(to_cpu(a) for a in args) + cpu_kwargs = {k: to_cpu(v) for k, v in kwargs.items()} + output = original_forward(*cpu_args, **cpu_kwargs) + + if output_device is None: + return output + + def to_device(value): + if isinstance(value, torch.Tensor): + return value.to(output_device) + return value + + if isinstance(output, tuple): + return tuple(to_device(v) for v in output) + return to_device(output) + + hift.forward = MethodType(_forward_on_cpu, hift) + hift._npu_cpu_offload_patched = True + + +@contextmanager +def npu_token2wav_sdpa_context() -> Iterator[None]: + """Expand CosyVoice masks + force MATH SDPA to avoid FA 161001.""" + try: + from vllm_omni.platforms.npu.models.cosyvoice2_dit_attn import ( + apply_cosyvoice2_dit_attn_npu_patch, + npu_math_sdpa_context, + ) + + apply_cosyvoice2_dit_attn_npu_patch() + with npu_math_sdpa_context(): + yield + except Exception: + with nullcontext(): + yield + + +def _patched_ensure_models_loaded(self) -> None: + assert _original_ensure_models_loaded is not None + was_loaded = self._models_loaded + _original_ensure_models_loaded(self) + if was_loaded or self.device.type != "npu" or self._hift is None: + return + patch_hift_module_for_npu(self._hift) + + +def _patched_forward(self, generated_speech_tokens, prompt_wav, return_bytes=True): + assert _original_forward is not None + if self.device.type != "npu": + return _original_forward(self, generated_speech_tokens, prompt_wav, return_bytes) + with npu_token2wav_sdpa_context(): + return _original_forward(self, generated_speech_tokens, prompt_wav, return_bytes) + + +def _patched_stream_chunk_for(self, audio_tokens, prompt_wav, last_chunk, state): + assert _original_stream_chunk_for is not None + if self.device.type != "npu": + return _original_stream_chunk_for(self, audio_tokens, prompt_wav, last_chunk, state) + with npu_token2wav_sdpa_context(): + return _original_stream_chunk_for(self, audio_tokens, prompt_wav, last_chunk, state) + + +def apply_step_audio2_token2wav_npu_patch() -> None: + """Monkey-patch StepAudio2Token2WavCore for Ascend NPU. + + Import is deferred and optional: platform bootstrap (e.g. resolving + ``current_omni_platform`` from rotary embedding) must not require + Token2Wav optional deps such as ``librosa``. + """ + global _PATCHED, _original_ensure_models_loaded, _original_forward, _original_stream_chunk_for + if _PATCHED: + return + + try: + from vllm_omni.model_executor.models.step_audio2.step_audio2_token2wav import ( + StepAudio2Token2WavCore, + ) + except ImportError as e: + logger.debug("step_audio2 token2wav deps unavailable; skip NPU patch: %s", e) + return + + _original_ensure_models_loaded = StepAudio2Token2WavCore._ensure_models_loaded + _original_forward = StepAudio2Token2WavCore.forward + _original_stream_chunk_for = StepAudio2Token2WavCore.stream_chunk_for + + StepAudio2Token2WavCore._ensure_models_loaded = _patched_ensure_models_loaded # type: ignore[method-assign] + StepAudio2Token2WavCore.forward = _patched_forward # type: ignore[method-assign] + StepAudio2Token2WavCore.stream_chunk_for = _patched_stream_chunk_for # type: ignore[method-assign] + + _PATCHED = True + logger.debug("Applied NPU patch for StepAudio2Token2WavCore") diff --git a/vllm_omni/platforms/npu/platform.py b/vllm_omni/platforms/npu/platform.py index 6f2eb99aad3..5133fa98eaa 100644 --- a/vllm_omni/platforms/npu/platform.py +++ b/vllm_omni/platforms/npu/platform.py @@ -37,13 +37,23 @@ def __init__(self) -> None: from vllm_omni.platforms.npu.models.qwen3_tts_code2wav import ( apply_qwen3_tts_code2wav_patch, ) + from vllm_omni.platforms.npu.models.qwen3_tts_tokenizer_v2 import ( + apply_qwen3_tts_tokenizer_v2_patch, + ) apply_qwen3_tts_code2wav_patch() + apply_qwen3_tts_tokenizer_v2_patch() apply_310p_patches() @classmethod def set_device(cls, device: torch.device) -> None: super().set_device(device) + + # Register vllm_ascend custom ops (torch.ops._C_ascend.*). + from vllm_ascend.utils import enable_custom_op + + enable_custom_op() + # Ascend quantized weights are converted from ND to FRACTAL_NZ # after loading. Enable internal format so the NZ storage layout # is preserved for fused NPU kernels. @@ -57,9 +67,15 @@ def get_omni_ar_worker_cls(cls) -> str: def get_omni_generation_worker_cls(cls) -> str: return "vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker" + @classmethod + def init_diffusion_worker_vllm_config(cls, vllm_config: Any) -> None: + from vllm_ascend.ascend_config import init_ascend_config + + init_ascend_config(vllm_config) + @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/platforms/npu/stage_configs" + return "vllm_omni/deploy" @classmethod def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: @@ -95,6 +111,14 @@ def get_diffusion_attn_backend_cls( if selected_backend is not None: backend_upper = selected_backend.upper() + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + logger.warning( + "HuggingFace kernels-backed FlashAttention is " + "not supported on NPU. Falling back to local " + "FLASH_ATTN." + ) + backend_upper = "FLASH_ATTN" + backend = DiffusionAttentionBackendEnum[backend_upper] logger.debug("Using diffusion attention backend '%s'", backend_upper) return backend.get_path() @@ -137,6 +161,11 @@ def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.npu.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.npu.mem_get_info(device) + return free, total + @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: device_props = torch.npu.get_device_properties(device_id) diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 38808648e93..9a8a5570978 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -45,7 +45,41 @@ from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner -from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element +from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids + + +def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]: + """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``. + + Non-tensor scalars (int, float, bool) are wrapped with ``torch.tensor()``. + Values that cannot be safely converted are dropped with a warning. + This enforces the tensor-only invariant required by the + ``OmniEngineCoreOutput.multimodal_output`` wire field and msgspec + serialization. Mirrors ``gpu_ar_model_runner._ensure_tensor_values``. + """ + result: dict[str, torch.Tensor] = {} + for key, val in payload.items(): + if isinstance(val, torch.Tensor): + result[key] = val + elif isinstance(val, (int, float, bool)): + result[key] = torch.tensor(val) + elif isinstance(val, (list, tuple)): + try: + result[key] = torch.tensor(val) + except (ValueError, TypeError, RuntimeError): + logger.warning( + "Dropping non-tensorizable multimodal output key '%s' (type=%s) from wire payload.", + key, + type(val).__name__, + ) + else: + logger.warning( + "Dropping non-tensor multimodal output key '%s' (type=%s) from wire payload.", + key, + type(val).__name__, + ) + return result class ExecuteModelState(NamedTuple): @@ -77,6 +111,7 @@ def __init__(self, *args, **kwargs): self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) + self._async_chunk = getattr(self.model_config, "async_chunk", False) self._downstream_payload_cache: dict[str, bool] = {} def _make_buffer(self, *size, dtype, numpy=True): @@ -216,6 +251,26 @@ def _resolve_req_hidden_states( return combined_hidden_states[rid] return hidden_states_cpu[start:end] + + def _build_multimodal_outputs( + self, + per_req_payloads: list[dict[str, object] | None] | None, + ) -> list[dict[str, torch.Tensor] | None] | None: + if self.vllm_config.model_config.engine_output_type == "text": + return None + if per_req_payloads is None: + return None + wire_payloads: list[dict[str, torch.Tensor] | None] = [] + for payload in per_req_payloads: + if not payload: + wire_payloads.append(None) + else: + wire_payloads.append(_ensure_tensor_values(payload)) + if all(item is None for item in wire_payloads): + return None + return wire_payloads + + def _request_final_stage_id(self, req_id: str) -> int | None: info = self.model_intermediate_buffer.get(req_id) if not isinstance(info, dict): @@ -425,7 +480,6 @@ def execute_model( logits_indices, spec_decode_metadata, total_num_scheduled_tokens, - num_scheduled_tokens_compressed_list, ) = self._prepare_inputs( scheduler_output, num_scheduled_tokens_np, @@ -522,7 +576,12 @@ def execute_model( # Another possible condition is num_tokens_padded != num_tokens_unpadded # but this scope is way too big and the consequences are unpredictable num_reqs_padded = self._pad_query_start_loc_for_fia( - num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_mode, + batch_desc.num_reqs, ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( @@ -563,8 +622,8 @@ def execute_model( positions=positions, inputs_embeds=inputs_embeds, req_ids=req_ids[:num_reqs], - num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)], - num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)], + num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs], + num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], input_ids_buffer=self.input_ids.gpu[:num_tokens_padded], ) # -------------------------------------- Omni-new ------------------------------------------------- @@ -632,8 +691,9 @@ def execute_model( with record_function_or_nullcontext("post process"): # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Map pending ropes metadata to req_ids. - if hasattr(self.model, "flush_pending_metadata"): - self.model.flush_pending_metadata(list(req_ids)) + flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) + if callable(flush_pending_metadata): + flush_pending_metadata(req_ids[:num_reqs]) hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) @@ -843,7 +903,8 @@ def sample_tokens( ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None - seq_len = hidden_states.shape[0] + hidden_seq_len = int(hidden_states.shape[0]) + scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens) # Apply structured output bitmasks if present. if grammar_output is not None: @@ -862,6 +923,15 @@ def sample_tokens( logits_vocab = logits.shape[-1] if self.input_batch.vocab_size > logits_vocab: smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + + # Drop min-tokens stop ids the head cannot emit (e.g. the text + # tokenizer EOS folded into all_stop_token_ids on a narrow codec + # talker head); they would index_put_ out of bounds (#4962). + if logits is not None: + sanitize_min_tokens_stop_ids( + self.input_batch.sampling_metadata.logitsprocs, + logits.shape[-1], + ) # -------------------------------------- Omni-new ------------------------------------------------- @@ -1082,18 +1152,31 @@ def _unwrap_lists(v): start=start, end=end, pass_lists_through=False, - seq_len=seq_len, + seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, ) payload.update(mm_payload) pooler_output.append(flatten_payload(payload)) + pooler_output = pooler_output or [] + if self._async_chunk: + pooler_inter, pooler_client = partition_payload_list(pooler_output) + else: + # Non-async-chunk ships the full payload to the next stage via + # inter_stage_outputs (the NPU runner has no separate full-payload + # accumulate). #4527's (None, pooler_output) starved it. (PR #4792) + pooler_inter, pooler_client = pooler_output, pooler_output + inter_stage_outputs = self._build_multimodal_outputs(pooler_inter) + multimodal_outputs = self._build_multimodal_outputs(pooler_client) model_runner_output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, sampled_token_ids=valid_sampled_token_ids, logprobs=logprobs_lists, prompt_logprobs_dict=prompt_logprobs_dict, - pooler_output=(pooler_output if engine_output_type != "text" and needs_pooler_payload else None), + pooler_output=None, + multimodal_outputs=multimodal_outputs, + inter_stage_outputs=inter_stage_outputs, kv_connector_output=kv_connector_output, ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, cudagraph_stats=cudagraph_stats, diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index 341b79d36e3..ea03842b40b 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -33,13 +33,18 @@ from vllm_ascend.worker.model_runner_v1 import SEQ_LEN_WITH_MAX_PA_WORKSPACE from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState +from vllm_omni.platforms.npu.worker.npu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner +from vllm_omni.utils.mm_outputs import partition_payload_list class NPUGenerationModelRunner(OmniNPUModelRunner): """Generation model runner for vLLM-omni on NPU (non-autoregressive).""" + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._async_chunk = getattr(self.model_config, "async_chunk", False) + def _update_request_states(self, scheduler_output: SchedulerOutput): # remove requests for req_id in scheduler_output.finished_req_ids: @@ -159,7 +164,6 @@ def execute_model( logits_indices, spec_decode_metadata, total_num_scheduled_tokens, - num_scheduled_tokens_compressed_list, ) = self._prepare_inputs( scheduler_output, num_scheduled_tokens_np, @@ -256,7 +260,12 @@ def execute_model( # Another possible condition is num_tokens_padded != num_tokens_unpadded # but this scope is way too big and the consequences are unpredictable num_reqs_padded = self._pad_query_start_loc_for_fia( - num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_mode, batch_desc.num_reqs + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_mode, + batch_desc.num_reqs, ) (attn_metadata, spec_decode_common_attn_metadata) = self._build_attention_metadata( @@ -427,33 +436,35 @@ def sample_tokens( ec_connector_output, cudagraph_stats, _batch_desc, - multimodal_outputs, # Omni-Specific + multimodal_outputs_raw, # Omni-Specific ) = self.execute_model_state # Clear ephemeral state. self.execute_model_state = None # -------------------------------------- Omni-new ------------------------------------------------- - pooler_output: list[object] = [] - if isinstance(multimodal_outputs, torch.Tensor): - assert multimodal_outputs.shape[0] == 1, ( + # Build per-request multimodal_outputs list (dedicated channel). + # pooler_output is no longer used for multimodal data. + per_req_payloads: list[dict[str, object]] = [] + if isinstance(multimodal_outputs_raw, torch.Tensor): + assert multimodal_outputs_raw.shape[0] == 1, ( "model should return a single tensor, to return multiple tensors, use a dict" ) - assert multimodal_outputs.shape[0] == self.input_batch.num_reqs + assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): - pooler_output.append({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()}) - elif isinstance(multimodal_outputs, list): - assert len(multimodal_outputs) == 1, ( + per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) + elif isinstance(multimodal_outputs_raw, list): + assert len(multimodal_outputs_raw) == 1, ( "model should return a single list, to return multiple lists, use a dict" ) - for out in multimodal_outputs: - pooler_output.append( + for out in multimodal_outputs_raw: + per_req_payloads.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) - elif isinstance(multimodal_outputs, Mapping): + elif isinstance(multimodal_outputs_raw, Mapping): num_reqs = self.input_batch.num_reqs for i in range(num_reqs): mm_payload = {} - for key, out in multimodal_outputs.items(): + for key, out in multimodal_outputs_raw.items(): if isinstance(out, list): if len(out) != num_reqs: raise ValueError( @@ -465,9 +476,16 @@ def sample_tokens( mm_payload[key] = out.detach().to("cpu").contiguous() else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") - pooler_output.append(mm_payload) + per_req_payloads.append(_ensure_tensor_values(mm_payload)) else: raise RuntimeError("Unsupported diffusion output type") + + if self._async_chunk: + inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads) + else: + # See npu_ar_model_runner: non-async-chunk ships the full payload to the next + # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792) + inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() @@ -482,10 +500,13 @@ def sample_tokens( sampled_token_ids=[], logprobs=None, prompt_logprobs_dict={}, - pooler_output=pooler_output, + pooler_output=None, + multimodal_outputs=multimodal_outputs, + inter_stage_outputs=inter_stage_outputs, kv_connector_output=kv_connector_output, - ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, + num_nans_in_logits={}, cudagraph_stats=cudagraph_stats, + ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) output.routed_experts = routed_experts_lists # -------------------------------------- Omni-new ------------------------------------------------- @@ -712,7 +733,12 @@ def _dummy_run( if not profile_cpp: num_reqs_padded = self._pad_query_start_loc_for_fia( - num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_runtime_mode, + batch_desc.num_reqs, ) pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL @@ -738,7 +764,15 @@ def _dummy_run( ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - if self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: + model_kwargs = self._init_model_kwargs() + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: @@ -746,7 +780,6 @@ def _dummy_run( inputs_embeds = None # -------------------------------------- Omni-new ------------------------------------------------- - model_kwargs = self._init_model_kwargs() # Some generation-stage models (e.g. MammothModa2DiTPipeline) require # model-specific runtime information (such as image size and conditioning # embeddings) even during the dummy profiling run that vLLM uses to diff --git a/vllm_omni/platforms/npu/worker/npu_model_runner.py b/vllm_omni/platforms/npu/worker/npu_model_runner.py index 5a50f523d02..671c7b7e53a 100644 --- a/vllm_omni/platforms/npu/worker/npu_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_model_runner.py @@ -193,7 +193,12 @@ def _dummy_run( if not profile_cpp: num_reqs_padded = self._pad_query_start_loc_for_fia( - num_tokens_padded, num_reqs_padded, num_reqs, cudagraph_runtime_mode, batch_desc.num_reqs + self.query_start_loc, + num_tokens_padded, + num_reqs_padded, + num_reqs, + cudagraph_runtime_mode, + batch_desc.num_reqs, ) pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL @@ -219,7 +224,10 @@ def _dummy_run( ): # Make sure padding doesn't exceed max_num_tokens assert num_tokens_padded <= self.max_num_tokens - if self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: + if getattr(getattr(self, "model", None), "has_preprocess", False): + input_ids = self.input_ids.gpu[:num_tokens_padded] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + elif self.supports_mm_inputs and not self.model_config.is_encoder_decoder or self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 1b50c328b80..38a14c3258e 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -88,6 +88,14 @@ def get_diffusion_attn_backend_cls( if selected_backend is not None: backend_upper = selected_backend.upper() + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + logger.warning( + "HuggingFace kernels-backed FlashAttention is " + "not supported on ROCm. Falling back to local " + "FLASH_ATTN." + ) + backend_upper = "FLASH_ATTN" + if backend_upper == "FLASH_ATTN" and not aiter_supported: logger.warning( "Flash Attention requires `aiter` library which is only supported " @@ -114,7 +122,7 @@ def supports_torch_inductor(cls) -> bool: @classmethod def get_default_stage_config_path(cls) -> str: - return "vllm_omni/platforms/rocm/stage_configs" + return "vllm_omni/deploy" @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: @@ -142,6 +150,11 @@ def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.cuda.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.cuda.mem_get_info(device) + return free, total + @classmethod def set_device_control_env_var(cls, devices: str | int | None) -> None: import os @@ -158,10 +171,18 @@ def unset_device_control_env_var(cls) -> None: @classmethod def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig: - """Copied from vllm/platforms/rocm/platform.py v0.20.0 with force using vllm_c kernels""" + """Copied from upstream RocmPlatform with inductor-aware logic. + + When inductor is active (compiling) use native as the default; + otherwise prefer vllm_c kernels where available. + Preserves omni-specific is_custom_op_enabled('rms_norm') check. + """ + from vllm.config.compilation import CompilationMode + # TODO(luka/TJ) use aiter, vllm_c, native by default on ROCm cc = vllm_config.compilation_config - default = ["vllm_c", "native"] # Originally using "native" here when compiling + using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE + default = ["native"] if using_inductor else ["vllm_c", "native"] # This (mostly) preserves previous CustomOp behavior # Necessary on ROCm because it's common that users diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 5500e013f3c..0bc782e0374 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/vllm_omni/platforms/xpu/platform.py @@ -42,6 +42,14 @@ def get_diffusion_attn_backend_cls( if selected_backend is not None: backend_upper = selected_backend.upper() + if backend_upper in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB"): + logger.warning( + "HuggingFace kernels-backed FlashAttention is " + "not supported on XPU. Falling back to local " + "FLASH_ATTN." + ) + backend_upper = "FLASH_ATTN" + backend = DiffusionAttentionBackendEnum[backend_upper] logger.debug("Using diffusion attention backend '%s'", backend_upper) return backend.get_path() @@ -85,6 +93,11 @@ def get_free_memory(cls, device: torch.device | None = None) -> int: free, _ = torch.xpu.mem_get_info(device) return free + @classmethod + def get_device_memory(cls, device: torch.device | None = None) -> tuple[int, int]: + free, total = torch.xpu.mem_get_info(device) + return free, total + @classmethod def get_profiler_cls(cls) -> str: """Return XPU-specific profiler that handles XPU events.""" @@ -92,7 +105,15 @@ def get_profiler_cls(cls) -> str: @classmethod def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig: - """Copied from vllm/platforms/xpu/platform.py v0.20.0 with force using xpu_kernels kernels""" - default = ["xpu_kernels", "native"] # Originally using "native" here when compiling + """Copied from upstream XPUPlatform with inductor-aware logic. + + When inductor is active (compiling) use native as the default; + otherwise prefer xpu_kernels where available. + """ + from vllm.config.compilation import CompilationMode + + cc = vllm_config.compilation_config + using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE + default = ["native"] if using_inductor else ["xpu_kernels", "native"] return IrOpPriorityConfig.with_default(default) diff --git a/vllm_omni/plugins/__init__.py b/vllm_omni/plugins/__init__.py index 41bf398cabf..240394b2c72 100644 --- a/vllm_omni/plugins/__init__.py +++ b/vllm_omni/plugins/__init__.py @@ -68,6 +68,10 @@ def load_omni_general_plugins() -> None: return omni_plugins_loaded = True + from vllm.plugins import load_general_plugins + + load_general_plugins() + plugins = load_omni_plugins_by_group(group=OMNI_DEFAULT_PLUGINS_GROUP) # General plugins: we only need to execute the loaded functions. for func in plugins.values(): diff --git a/vllm_omni/quantization/__init__.py b/vllm_omni/quantization/__init__.py index d8cad3a98e3..8896903e886 100644 --- a/vllm_omni/quantization/__init__.py +++ b/vllm_omni/quantization/__init__.py @@ -12,10 +12,10 @@ """ from .component_config import ComponentQuantizationConfig -from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config +from .factory import SUPPORTED_QUANTIZATION_METHODS, build_quant_config, register_quantization_override from .inc_config import OmniINCConfig -# Heavy configs (GGUF, MXFP8) are NOT imported here to avoid pulling in +# Heavy configs are NOT imported here to avoid pulling in # optional dependencies (pynvml, torch_npu) at module load time. # Import them directly when needed: # from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config @@ -25,4 +25,5 @@ "ComponentQuantizationConfig", "OmniINCConfig", "SUPPORTED_QUANTIZATION_METHODS", + "register_quantization_override", ] diff --git a/vllm_omni/quantization/bitsandbytes_config.py b/vllm_omni/quantization/bitsandbytes_config.py new file mode 100644 index 00000000000..7d8e7a68425 --- /dev/null +++ b/vllm_omni/quantization/bitsandbytes_config.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""BitsAndBytes 4-bit quantization config for diffusion transformers. + +Supports online (dynamic) NF4/FP4 weight-only quantization from BF16/FP16 +checkpoints on CUDA GPUs. +""" + +from typing import TYPE_CHECKING, Any, Optional + +import torch +from torch.nn import Module +from vllm.model_executor.layers.linear import ( + LinearBase, + LinearMethodBase, + UnquantizedLinearMethod, +) +from vllm.model_executor.layers.quantization import QuantizationMethods +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from vllm.model_executor.layers.quantization.fp8 import _copy_missing_attrs +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.model_loader.weight_utils import initialize_single_dummy_weight +from vllm.model_executor.parameter import ModelWeightParameter +from vllm.model_executor.utils import replace_parameter + +from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.int8_config import LazyWeightMixin + +if TYPE_CHECKING: + from vllm.model_executor.models.utils import WeightsMapper + + +class DiffusionBitsAndBytesConfig(QuantizationConfig): + """BitsAndBytes 4-bit weight-only config for diffusion transformers. + + Supports online (dynamic) quantization from BF16/FP16 checkpoints. + Works on CUDA GPUs with the optional ``bitsandbytes`` package installed. + """ + + def __init__( + self, + quant_type: str = "nf4", + compress_statistics: bool = True, + ignored_layers: list[str] | None = None, + ) -> None: + super().__init__() + + if quant_type not in ("nf4", "fp4"): + raise ValueError(f"Unsupported quant_type {quant_type!r}; expected 'nf4' or 'fp4'") + self.quant_type = quant_type + self.compress_statistics = compress_statistics + self.ignored_layers = ignored_layers or [] + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "bitsandbytes" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + return 75 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"): + if self.ignored_layers is not None: + self.ignored_layers = hf_to_vllm_mapper.apply_list(self.ignored_layers) + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "DiffusionBitsAndBytesConfig": + quant_type = cls.get_from_keys_or(config, ["quant_type"], "nf4") + compress_statistics = cls.get_from_keys_or(config, ["compress_statistics"], True) + ignored_layers = cls.get_from_keys_or(config, ["ignored_layers"], None) + + if not ignored_layers: + ignored_layers = cls.get_from_keys_or(config, ["modules_to_not_convert"], None) + return cls( + quant_type=quant_type, + compress_statistics=compress_statistics, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional["QuantizeMethodBase"]: + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if current_omni_platform.is_cuda(): + return BnBOnlineLinearMethod(self) + raise NotImplementedError("BitsAndBytes online quantization is only supported on CUDA.") + return None + + +class BnBOnlineLinearMethod(LazyWeightMixin, LinearMethodBase): + """Online BitsAndBytes 4-bit linear method. + + Loads BF16/FP16 checkpoint weights and quantizes them during loading. + """ + + def __init__(self, quant_config: DiffusionBitsAndBytesConfig): + self.quant_config = quant_config + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + if layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + import bitsandbytes.functional as bnb_F + + weight = layer.weight.data.contiguous() + if not weight.is_cuda: + weight = weight.cuda() + + original_shape = tuple(weight.shape) + qweight, quant_state = bnb_F.quantize_4bit( + weight, + quant_type=self.quant_config.quant_type, + compress_statistics=self.quant_config.compress_statistics, + ) + + replace_parameter( + layer, + "weight", + torch.nn.Parameter(qweight, requires_grad=False), + ) + layer.quant_state = quant_state + layer.bnb_shape = original_shape + + layer._already_called_process_weights_after_loading = True + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + import bitsandbytes as bnb + + ori_shape = x.shape + ori_dtype = x.dtype + x_2d = x.reshape(-1, ori_shape[-1]) + + out = bnb.matmul_4bit( + x_2d, + layer.weight.t(), + quant_state=layer.quant_state, + bias=bias, + ) + return out.reshape(*ori_shape[:-1], -1).to(ori_dtype) diff --git a/vllm_omni/quantization/component_config.py b/vllm_omni/quantization/component_config.py index 25d1abf1e72..28c33d63771 100644 --- a/vllm_omni/quantization/component_config.py +++ b/vllm_omni/quantization/component_config.py @@ -48,6 +48,30 @@ def resolve_encoder_quant_config( return quant_config +def safe_quant_config( + quant_config: QuantizationConfig | None, +) -> QuantizationConfig | None: + """Return *quant_config* only if it is safe for norm/modulation layers. + + Norm and modulation layers (LayerNorm, RMSNorm, AdaLayerNorm, img_mod, + txt_mod, etc.) produce precision-sensitive shift/scale/gate values and + should not receive FP8 quant configs (see #2728). Pre-quantized methods + like INC/AutoRound W4A16 need the config propagated so packed weights + load correctly. + + This is the inverse of :func:`resolve_encoder_quant_config`: that function + strips pre-quantized configs from encoders, while this one strips + *every config except* pre-quantized configs from norm/mod layers. + """ + if quant_config is None: + return None + from vllm.model_executor.layers.quantization.inc import INCConfig + + if isinstance(quant_config, INCConfig): + return quant_config + return None + + class ComponentQuantizationConfig(QuantizationConfig): """Routes quantization to different configs by layer prefix.""" @@ -56,6 +80,7 @@ def __init__( component_configs: dict[str, QuantizationConfig | None], default_config: QuantizationConfig | None = None, ) -> None: + super().__init__() self._components = component_configs self._default = default_config self._sorted_prefixes = sorted(self._components.keys(), key=len, reverse=True) diff --git a/vllm_omni/quantization/factory.py b/vllm_omni/quantization/factory.py index d9a7a31bfb8..b5d0a8ad794 100644 --- a/vllm_omni/quantization/factory.py +++ b/vllm_omni/quantization/factory.py @@ -3,8 +3,8 @@ """Factory for building quantization configs. build_quant_config() delegates to vLLM's quantization registry. -The only extension point is _OVERRIDES for methods that need a -different QuantizationConfig subclass in the OMNI context (e.g. GGUF). +Out-of-tree integrations can register Omni-specific builders with +register_quantization_override(). """ from __future__ import annotations @@ -81,13 +81,6 @@ def _register_humming_stubs() -> None: logger = init_logger(__name__) -def _build_gguf(**kw: Any) -> QuantizationConfig: - """Lazy import to avoid pulling in CUDA/pynvml at module load time.""" - from .gguf_config import DiffusionGGUFConfig - - return DiffusionGGUFConfig(**kw) - - def _build_int8(**kw: Any) -> QuantizationConfig: """Lazy import for Int8 diffusion config (supports CUDA + NPU).""" from .int8_config import DiffusionInt8Config @@ -95,6 +88,13 @@ def _build_int8(**kw: Any) -> QuantizationConfig: return DiffusionInt8Config(**kw) +def _build_bitsandbytes(**kw: Any) -> QuantizationConfig: + """Lazy import for BitsAndBytes 4-bit diffusion config (CUDA only).""" + from .bitsandbytes_config import DiffusionBitsAndBytesConfig + + return DiffusionBitsAndBytesConfig(**kw) + + def _build_mxfp8(**kw: Any) -> QuantizationConfig: """Lazy import for W8A8 MXFP8 diffusion config (NPU only).""" from .mxfp8_config import DiffusionMXFP8Config @@ -140,8 +140,8 @@ def _build_inc(**kw: Any) -> QuantizationConfig: _OVERRIDES: dict[str, Callable[..., QuantizationConfig]] = { - "gguf": _build_gguf, "int8": _build_int8, + "bitsandbytes": _build_bitsandbytes, "mxfp8": _build_mxfp8, "mxfp4": _build_mxfp4, "mxfp4_dualscale": _build_mxfp4_dualscale, @@ -150,7 +150,24 @@ def _build_inc(**kw: Any) -> QuantizationConfig: "auto_round": _build_inc, } -SUPPORTED_QUANTIZATION_METHODS: list[str] = list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys()))) + +def _compute_supported_quantization_methods() -> list[str]: + return list(dict.fromkeys(QUANTIZATION_METHODS + list(_OVERRIDES.keys()))) + + +SUPPORTED_QUANTIZATION_METHODS: list[str] = _compute_supported_quantization_methods() + + +def _refresh_registered_methods() -> None: + SUPPORTED_QUANTIZATION_METHODS[:] = _compute_supported_quantization_methods() + global _CACHED_ALIAS_MAP + _CACHED_ALIAS_MAP = None + + +def register_quantization_override(method: str, builder: Callable[..., QuantizationConfig]) -> None: + """Register an Omni-specific quantization config builder.""" + _OVERRIDES[_normalize_method_name(method)] = builder + _refresh_registered_methods() def _build_reverse_alias_map() -> dict[str, str]: diff --git a/vllm_omni/quantization/gguf_config.py b/vllm_omni/quantization/gguf_config.py deleted file mode 100644 index 42b2b2fe214..00000000000 --- a/vllm_omni/quantization/gguf_config.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""GGUF quantization config for diffusion transformers. - -Uses dequant+GEMM instead of the fused kernel path (which expects 2D inputs). -""" - -from __future__ import annotations - -import gguf -import torch -from vllm import _custom_ops as ops -from vllm.model_executor.layers.quantization.gguf import ( - UNQUANTIZED_TYPES, - GGUFConfig, - GGUFLinearMethod, - LinearBase, - QuantizeMethodBase, - UnquantizedLinearMethod, - is_layer_skipped_gguf, -) - - -def dequant_gemm_gguf(x: torch.Tensor, qweight: torch.Tensor, qweight_type: int) -> torch.Tensor: - if qweight_type in UNQUANTIZED_TYPES: - return x @ qweight.T - block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type] - shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size) - weight = ops.ggml_dequantize(qweight, qweight_type, *shape, x.dtype) - return x @ weight.T - - -class DiffusionGGUFLinearMethod(GGUFLinearMethod): - """GGUF linear method using dequant+GEMM for N-D diffusion tensors.""" - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - shard_id = getattr(layer.qweight, "shard_id", None) - - if shard_id: - shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id - qweight = layer.qweight - result = [] - for idx in shard_id: - start, end, offset = layer.qweight.shard_offset_map[idx] - qweight_type = layer.qweight_type.shard_weight_type[idx] - result.append(dequant_gemm_gguf(x, qweight[start:end, :offset].contiguous(), qweight_type)) - out = torch.cat(result, axis=-1) - else: - qweight = layer.qweight - qweight_type = layer.qweight_type.weight_type - out = dequant_gemm_gguf(x, qweight, qweight_type) - - if bias is not None: - out.add_(bias) - return out - - -class DiffusionGGUFConfig(GGUFConfig): - """GGUF config that carries gguf_model path and uses dequant+GEMM.""" - - def __init__( - self, - gguf_model: str | None = None, - unquantized_modules: list[str] | None = None, - ) -> None: - super().__init__(unquantized_modules=unquantized_modules or []) - self.gguf_model = gguf_model - - def get_quant_method(self, layer: torch.nn.Module, prefix: str) -> QuantizeMethodBase | None: - if isinstance(layer, LinearBase): - if is_layer_skipped_gguf(prefix, self.unquantized_modules, self.packed_modules_mapping): - return UnquantizedLinearMethod() - return DiffusionGGUFLinearMethod(self) - return None diff --git a/vllm_omni/quantization/mxfp4_config.py b/vllm_omni/quantization/mxfp4_config.py index 6a02d5edb87..a6b29ea7937 100644 --- a/vllm_omni/quantization/mxfp4_config.py +++ b/vllm_omni/quantization/mxfp4_config.py @@ -7,6 +7,8 @@ MXFPLinearMethodBase – platform-agnostic skeleton (imported from mxfp8_config) NPUMxfp4LinearMethod – NPU single-scale offline (W4A4 MXFP4) NPUMxfp4OnlineLinearMethod – NPU single-scale online (BF16 → FP4) + ROCmMxfp4LinearMethod – ROCm base class (AITER quant + shuffle; online-only) + ROCmMxfp4OnlineLinearMethod – ROCm online (BF16 → FP4 via AITER) NPUMxfp4DualScaleLinearMethod – NPU dual-scale offline (W4A4 MXFP4 DualScale) NPUMxfp4DualScaleOnlineLinearMethod – NPU dual-scale online (BF16 → FP4) @@ -153,9 +155,16 @@ def get_quant_method( if self.is_checkpoint_mxfp4_serialized: return NPUMxfp4LinearMethod(self) return NPUMxfp4OnlineLinearMethod(self) + if current_omni_platform.is_rocm(): + gcn_arch = torch.cuda.get_device_properties(torch.accelerator.current_device_index()).gcnArchName + if "gfx950" not in gcn_arch: + raise NotImplementedError(f"MXFP4 on ROCm requires gfx950 (MI355X). Detected: {gcn_arch}") + if self.is_checkpoint_mxfp4_serialized: + raise NotImplementedError("Pre-quantized MXFP4 checkpoints are not yet supported on ROCm.") + return ROCmMxfp4OnlineLinearMethod(self) raise NotImplementedError( "DiffusionMXFP4Config (W4A4 MXFP4) is currently only supported " - "on NPU (Ascend) platforms. CUDA support is not yet implemented." + "on NPU (Ascend) and ROCm (AMD, gfx950) platforms." ) return None @@ -332,6 +341,151 @@ def process_weights_after_loading(self, layer: Module) -> None: layer._already_called_process_weights_after_loading = True +# --------------------------------------------------------------------------- +# ROCm MXFP4 base method (AITER) — used as base class for the online subclass. +# Online-only for now; offline (pre-quantized checkpoint) not yet supported. +# --------------------------------------------------------------------------- + + +def _register_rocm_mxfp4_op() -> None: + """Register the vllm_omni::rocm_mxfp4_gemm custom op for torch.compile. + + Wraps activation quantization + GEMM in a single opaque op so that + torch.compile/inductor doesn't try to trace through AITER internals. + """ + import aiter + + @torch.library.custom_op("vllm_omni::rocm_mxfp4_gemm", mutates_args=()) + def _rocm_mxfp4_gemm( + a: torch.Tensor, + w_quant: torch.Tensor, + w_scale: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32) + a_quant, a_scale = quant_func(a, shuffle=True) + return aiter.gemm_a4w4( + a_quant, + w_quant, + a_scale, + w_scale, + bpreshuffle=True, + bias=bias, + ) + + @_rocm_mxfp4_gemm.register_fake + def _rocm_mxfp4_gemm_fake( + a: torch.Tensor, + w_quant: torch.Tensor, + w_scale: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + M, _ = a.shape + N, _ = w_quant.shape + return torch.empty(M, N, dtype=a.dtype, device=a.device) + + +class ROCmMxfp4LinearMethod(MXFPLinearMethodBase): + """ROCm W4A4 MXFP4 linear method using AITER (gemm_a4w4). + + Weight buffers after process_weights_after_loading: + weight_shuffle : FP4 quantized + shuffled via shuffle_weight(layout=(16,16)) + weight_scale : per-group-of-32 scales from AITER per_1x32 + + Forward path: + _quantize_activation: pass-through (activation quant is inside the custom op) + _quant_matmul: torch.ops.vllm_omni.rocm_mxfp4_gemm wraps activation quant + then gemm_a4w4 in one custom op + """ + + def __init__(self, quant_config: DiffusionMXFP4Config) -> None: + self.quant_config = quant_config + self.out_dtype = torch.get_default_dtype() + if not hasattr(torch.ops.vllm_omni, "rocm_mxfp4_gemm"): + _register_rocm_mxfp4_op() + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + import aiter + from aiter.ops.shuffle import shuffle_weight + + quant_func = aiter.get_hip_quant(aiter.QuantType.per_1x32) + weight_quant, weight_scale = quant_func(layer.weight.data, shuffle=True) + weight_shuffled = shuffle_weight(weight_quant, layout=(16, 16)) + + # Store quantized tensors as non-parameter buffers; delete original weight. + layer.register_buffer("weight_shuffle", weight_shuffled, persistent=True) + layer.register_buffer("weight_scale", weight_scale, persistent=True) + + # Remove the original weight parameter to free memory. + if hasattr(layer, "weight") and isinstance(layer.weight, torch.nn.Parameter): + delattr(layer, "weight") + layer.register_parameter("weight", None) + + layer._already_called_process_weights_after_loading = True + + def _quantize_activation(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # Activation quantization is inside the custom op called by _quant_matmul, + # so pass through the raw activation here. + return x, None + + def _quant_matmul( + self, + x_q: torch.Tensor, + x_scale: torch.Tensor, + layer: torch.nn.Module, + bias: torch.Tensor | None, + ori_dtype: torch.dtype, + ) -> torch.Tensor: + output = torch.ops.vllm_omni.rocm_mxfp4_gemm( + x_q, + layer.weight_shuffle, + layer.weight_scale, + bias, + ) + if output.dtype != ori_dtype: + output = output.to(ori_dtype) + return output + + +# --------------------------------------------------------------------------- +# ROCm MXFP4 online method (BF16 checkpoint → quantize at load time) +# --------------------------------------------------------------------------- + + +class ROCmMxfp4OnlineLinearMethod(_LazyWeightMixin, ROCmMxfp4LinearMethod): + """ROCm W4A4 MXFP4 online linear method using AITER. + + MRO: ROCmMxfp4OnlineLinearMethod → _LazyWeightMixin → ROCmMxfp4LinearMethod + → MXFPLinearMethodBase → LinearMethodBase + + create_weights : _LazyWeightMixin (meta device + patched loader) + process_weights : ROCmMxfp4OnlineLinearMethod (meta → materialize, then AITER quant + shuffle) + apply / ops : ROCmMxfp4LinearMethod / MXFPLinearMethodBase (shared) + """ + + def process_weights_after_loading(self, layer: Module) -> None: + if getattr(layer, "_already_called_process_weights_after_loading", False): + return + + # Materialise from meta device if needed (same pattern as NPU online). + if layer.weight is not None and layer.weight.device == torch.device("meta"): + weight = ModelWeightParameter( + data=torch.empty_like(layer.weight, device=layer._load_device), + input_dim=1, + output_dim=0, + weight_loader=layer.weight.weight_loader, + ) + _copy_missing_attrs(layer.weight, weight) + layer.register_parameter("weight", weight) + initialize_single_dummy_weight(layer.weight) + + # Delegate to the base class which does AITER quant + shuffle. + ROCmMxfp4LinearMethod.process_weights_after_loading(self, layer) + + # --------------------------------------------------------------------------- # NPU MXFP4 dual-scale offline method (W4A4_MXFP4_DUALSCALE checkpoint) # --------------------------------------------------------------------------- diff --git a/vllm_omni/quantization/quack_fp8.py b/vllm_omni/quantization/quack_fp8.py index ba68d078dd1..86595838696 100644 --- a/vllm_omni/quantization/quack_fp8.py +++ b/vllm_omni/quantization/quack_fp8.py @@ -15,11 +15,21 @@ _gemm_interface = None -def _is_blackwell() -> bool: +def _is_quack_capable() -> bool: + """quack's CuteDSL FP8 / block-scaled MMA is built on the 5th-gen tensor-core + ``tcgen05`` instruction family, which is datacenter-Blackwell only + (``sm_100a`` / ``sm_101a`` / ``sm_103a``, compute capability ``10.x``). + + Workstation/consumer Blackwell (``sm_120`` / ``sm_121``, compute capability + ``12.x``, e.g. RTX PRO 6000 / RTX 50-series) lacks ``tcgen05``, so quack can + never run there — CuteDSL rejects the arch and every GEMM falls back to + FlashInfer one call at a time, which is catastrophically slow. Those GPUs + have working native FlashInfer FP8 kernels, so default quack off for them. + """ try: if not torch.cuda.is_available(): return False - return torch.cuda.get_device_capability()[0] >= 10 + return torch.cuda.get_device_capability()[0] == 10 except Exception: # noqa: BLE001 return False @@ -28,7 +38,7 @@ def quack_enabled() -> bool: override = os.environ.get("VLLM_OMNI_USE_QUACK_FP8") if override is not None: return override.lower() in _TRUTHY - return _is_blackwell() + return _is_quack_capable() def _set_persistent_cache_dir() -> None: @@ -84,7 +94,7 @@ def quack_scaled_fp8_mm( return None out = torch.empty(a.shape[0], b.shape[1], device=a.device, dtype=out_dtype) alpha = scale_a.reshape(1).float() * scale_b.reshape(1).float() - gemm.gemm_tuned(a, b, out, bias=bias, alpha=alpha) + gemm.gemm(a, b, out=out, bias=bias, alpha=alpha, tuned=True) return out diff --git a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py index 3d4f68d23f0..77b1917383f 100644 --- a/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py +++ b/vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py @@ -280,7 +280,6 @@ def _build_quantization_config( quantization: str | None, quantization_config_json: str | None, ignored_layers: str | None, - gguf_model: str | None, ) -> dict[str, Any] | None: explicit = _load_json_arg(quantization_config_json) if explicit is not None: @@ -292,8 +291,6 @@ def _build_quantization_config( ignored = _split_csv(ignored_layers) if ignored: config["ignored_layers"] = ignored - if gguf_model is not None: - config["gguf_model"] = gguf_model return config @@ -305,7 +302,6 @@ def _build_variant_config(args: argparse.Namespace, variant: str) -> VariantConf quantization=quantization, quantization_config_json=getattr(args, f"{variant}_quantization_config_json"), ignored_layers=ignored_layers, - gguf_model=getattr(args, f"{variant}_gguf_model"), ) return VariantConfig( label=variant, @@ -554,8 +550,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--reference-ignored-layers", help="Comma-separated reference ignored layer patterns.") parser.add_argument("--candidate-ignored-layers", help="Comma-separated candidate ignored layer patterns.") parser.add_argument("--ignored-layers", help="Alias for --candidate-ignored-layers.") - parser.add_argument("--reference-gguf-model", help="Reference GGUF file or HF reference when using gguf.") - parser.add_argument("--candidate-gguf-model", help="Candidate GGUF file or HF reference when using gguf.") parser.add_argument("--reference-load-format", default="default", help="Reference diffusion load format.") parser.add_argument("--candidate-load-format", default="default", help="Candidate diffusion load format.") parser.add_argument("--prompt", default="a cup of coffee on the table") diff --git a/vllm_omni/request.py b/vllm_omni/request.py index f936069be2d..b094f2d39fa 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -30,15 +30,16 @@ class OmniRequest(Request): def __init__( self, + *args, prompt_embeds: PromptEmbedsPayload | torch.Tensor | None = None, # Optional external request ID for tracking external_req_id: str | None = None, additional_information: AdditionalInformationPayload | None = None, - *args, **kwargs, ): - prompt_embeds_tensor = self._maybe_decode_prompt_embeds(prompt_embeds) - super().__init__(prompt_embeds=prompt_embeds_tensor, *args, **kwargs) + if prompt_embeds is not None: + kwargs["prompt_embeds"] = self._maybe_decode_prompt_embeds(prompt_embeds) + super().__init__(*args, **kwargs) # Preserve serialized prompt embeddings payload (optional) self.prompt_embeds_payload: PromptEmbedsPayload | None = ( prompt_embeds if isinstance(prompt_embeds, PromptEmbedsPayload) else None diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index ba6d329eab8..88a23dc0c52 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -10,6 +10,7 @@ import importlib _CLASS_TO_MODULE: dict[str, str] = { + "HiggsAudioV3Config": "vllm_omni.transformers_utils.configs.higgs_audio_v3", "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2", @@ -31,6 +32,7 @@ } __all__ = [ + "HiggsAudioV3Config", "Mammothmoda2Config", "Mammothmoda2Qwen2_5_VLConfig", "Mammothmoda2Qwen2_5_VLTextConfig", @@ -70,6 +72,7 @@ def __dir__(): from vllm_omni.transformers_utils.configs import cosyvoice3 as _cosyvoice3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni # noqa: F401, E402 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice # noqa: F401, E402 diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py b/vllm_omni/transformers_utils/configs/higgs_audio_v3.py similarity index 100% rename from vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py rename to vllm_omni/transformers_utils/configs/higgs_audio_v3.py diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py index 4acd098e211..1c6352264ca 100644 --- a/vllm_omni/transformers_utils/configs/mammoth_moda2.py +++ b/vllm_omni/transformers_utils/configs/mammoth_moda2.py @@ -2,13 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from typing import ClassVar, Literal -from transformers import AutoConfig, PretrainedConfig +from transformers import AutoConfig, AutoTokenizer, PretrainedConfig from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import ( Qwen2_5_VLConfig, Qwen2_5_VLTextConfig, Qwen2_5_VLVisionConfig, ) +from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer + __all__ = [ "Mammothmoda2Config", "Mammothmoda2Qwen2_5_VLConfig", @@ -229,7 +231,8 @@ def __init__( architectures: list[str] | None = None, **kwargs, ) -> None: - super().__init__(**kwargs) + # Must set llm_config before super().__init__() because the parent's + # validation calls get_text_config() which accesses self.llm_config self.llm_config = AutoConfig.for_model(**llm_config) if llm_config is not None else None self.gen_vae_config = gen_vae_config self.gen_dit_config = gen_dit_config @@ -243,8 +246,13 @@ def __init__( self.tokenizer_class = "MammothUTokenizer" self.architectures = ["Mammothmoda2Model"] + super().__init__(**kwargs) + def get_text_config(self, decoder: bool = False) -> PretrainedConfig: # noqa: ARG002 - return self.llm_config + if self.llm_config is None: + return None + # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config + return self.llm_config.text_config def _require_llm_config(self) -> PretrainedConfig: if self.llm_config is None: @@ -281,3 +289,6 @@ def vision_end_token_id(self) -> int: AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig) AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig) AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig) +# Register tokenizer_type -> the configs & AutoTokenizer +AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer) +AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer) diff --git a/vllm_omni/transformers_utils/processors/ming.py b/vllm_omni/transformers_utils/processors/ming.py index 685909029ad..b32f10d1204 100644 --- a/vllm_omni/transformers_utils/processors/ming.py +++ b/vllm_omni/transformers_utils/processors/ming.py @@ -143,7 +143,6 @@ def __call__( # (conv1 has stride=1 and does not change T) # 2. AudioProjector Conv1d: kernel=3, stride=2, padding=1 # Combined: T → ((T-1)//2 + 1 - 1)//2 + 1 - # See also: AudioProjector.compute_output_length() encoder_feats_lengths = ((audio_feats_lengths - 3 + 2 * 1) // 2 + 1 - 3 + 2 * 1) // 2 + 1 audio_feats = torch.cat(audio_feat_list, dim=0).unsqueeze(0) # [1, T_total, n_mels] diff --git a/vllm_omni/utils/forced_aligner.py b/vllm_omni/utils/forced_aligner.py new file mode 100644 index 00000000000..2a97df3d829 --- /dev/null +++ b/vllm_omni/utils/forced_aligner.py @@ -0,0 +1,399 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared forced aligner for streaming TTS word timestamps. + +Hosts one lazy-loaded, in-process pooling ``vllm.LLM`` shared by the TTS +frontend. ``llm.encode`` is sync/blocking, so :func:`align` runs it in +``asyncio.to_thread``, serialized on a lock since the offline LLM is not +thread-safe. See :func:`align` for the return-value contract. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import threading +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np +import yaml + +from vllm_omni.utils import qwen3_force_align_processor as _processor + +logger = logging.getLogger(__name__) + + +_DEFAULT_CONFIG_PATH = Path(__file__).resolve().parents[1] / "deploy" / "qwen3_tts_forced_aligner.yaml" + + +@dataclass(frozen=True, slots=True) +class WordTimestamp: + """Internal alignment record. Serialized to a plain JSON object + (``{"word", "start_ms", "end_ms"}``) at the WebSocket boundary. + """ + + word: str + start_ms: int + end_ms: int + + +@dataclass(frozen=True, slots=True) +class ForcedAlignerConfig: + """Plain-data config for constructing the aligner ``vllm.LLM`` (lazy-loaded).""" + + model: str + runner: str | None = None + architecture: str | None = None + pooling_task: str | None = None + gpu_memory_utilization: float | None = None + dtype: str | None = None + max_model_len: int | None = None + trust_remote_code: bool | None = None + extra_llm_kwargs: dict[str, Any] = field(default_factory=dict) + + +def _load_forced_aligner_yaml(path: str | Path) -> dict[str, Any]: + with Path(path).open(encoding="utf-8") as f: + raw = yaml.safe_load(f) or {} + return dict(raw.get("forced_aligner", raw)) + + +def build_forced_aligner_config(args: Any) -> ForcedAlignerConfig | None: + """Build a config from CLI args, or ``None`` when the feature is off. + + Precedence (lowest to highest): the packaged default YAML + (:data:`_DEFAULT_CONFIG_PATH`, Qwen deploy defaults) -> a user YAML passed + via ``--forced-aligner-config`` -> the ``--forced-aligner`` model path. The + feature is off (returns ``None``) unless a model resolves from this chain. + Per-field overrides such as ``gpu_memory_utilization`` live in the YAML. + """ + config_path = getattr(args, "forced_aligner_config", None) + config_data: dict[str, Any] = {} + model = getattr(args, "forced_aligner", None) + if config_path or model: + config_data.update(_load_forced_aligner_yaml(_DEFAULT_CONFIG_PATH)) + if config_path: + config_data.update(_load_forced_aligner_yaml(config_path)) + if model: + config_data["model"] = str(model) + model = config_data.get("model") + if not model: + return None + allowed = set(ForcedAlignerConfig.__dataclass_fields__) + return ForcedAlignerConfig(**{k: v for k, v in config_data.items() if k in allowed}) + + +class ForcedAlignerLoadError(RuntimeError): + """The aligner model/config could not be loaded (permanent until restart). + + Distinguished from a per-request alignment failure so the streaming layer + can report a clear reason once rather than degrading every request to + ``timestamps: null``. + """ + + +# --- Singleton state --- +# A single LLM serves the whole API server. ``_lock`` guards the lazy +# constructor; ``_encode_lock`` serializes ``encode`` because the offline +# ``vllm.LLM`` is a sync batch interface that is not safe to call from multiple +# threads at once (concurrent WebSocket sessions otherwise race the engine). +_lock = threading.Lock() +_encode_lock = asyncio.Lock() +_llm: Any = None +_classify_num: int | None = None +_timestamp_token_id: int | None = None +_timestamp_segment_time_ms: float | None = None +_loaded_config: ForcedAlignerConfig | None = None + + +async def align( + *, + audio: bytes, + text: str, + sample_rate: int, + config: ForcedAlignerConfig, + language: str | None = None, +) -> list[WordTimestamp] | None: + """Run one forced-alignment pass. + + Args: + audio: Signed-int16 little-endian mono PCM bytes. + text: Ground-truth text whose words to align. + sample_rate: Sample rate of ``audio`` in Hz. + config: Aligner config (same instance for every call across the + server's lifetime; reload requires a server restart). + language: Optional language hint for word segmentation. ``None`` / + ``"auto"`` use the space + Chinese-mixed path; ``"japanese"`` / + ``"korean"`` (or their codes) need ``qwen_asr`` installed for + faithful tokenisation, else they degrade to whitespace splitting. + + Returns: + List of :class:`WordTimestamp` on success (possibly empty for + silence / no aligned tokens), ``None`` on a per-request alignment + failure (decode error). + + Raises: + ForcedAlignerLoadError: the model/config could not be loaded. This is + permanent until restart, so callers surface it once instead of + degrading every request to ``None``. + """ + async with _encode_lock: + try: + await asyncio.to_thread(_ensure_loaded, config) + except Exception as exc: # noqa: BLE001 + logger.exception("Forced aligner failed to load") + raise ForcedAlignerLoadError(str(exc)) from exc + try: + return await asyncio.to_thread(_align_sync, audio, text, sample_rate, config, language) + except Exception: # noqa: BLE001 + logger.exception("Forced aligner failed for text=%r", text) + return None + + +def _align_sync( + audio: bytes, + text: str, + sample_rate: int, + config: ForcedAlignerConfig, + language: str | None = None, +) -> list[WordTimestamp]: + # Precondition: align() has already run _ensure_loaded(config) under + # _encode_lock and owns the load-failure -> ForcedAlignerLoadError mapping. + # Don't reload here, so a load error can't be swallowed as a per-request None. + audio_arr = _pcm_bytes_to_float32(audio) + if audio_arr.size == 0: + return [] + audio_duration_ms = (audio_arr.size / sample_rate) * 1000.0 + + # Segment once and reuse for both the prompt and the decode: the word + # units MUST match between the two or the markers drift out of sync. + words = _processor.segment_words(text, language) + prompt = _processor.build_prompt(words) + request = { + "prompt": prompt, + "multi_modal_data": {"audio": (audio_arr, sample_rate)}, + } + + from vllm.pooling_params import PoolingParams + + outputs = _llm.encode( # type: ignore[union-attr] + [request], + pooling_params=PoolingParams(), + pooling_task=config.pooling_task or "token_classify", + use_tqdm=False, + ) + if not outputs: + return [] + + result = outputs[0] + logits = result.outputs.data # [n_token, classify_num] + prompt_token_ids = list(result.prompt_token_ids) + timestamp_positions = [i for i, tid in enumerate(prompt_token_ids) if tid == _timestamp_token_id] + if not timestamp_positions: + logger.warning( + "No <|timestamp|> tokens found in prompt for text=%r; aligner returned %d rows.", + text, + logits.shape[0] if hasattr(logits, "shape") else len(logits), + ) + return [] + + return _decode_timestamps( + logits=logits, + words=words, + timestamp_positions=timestamp_positions, + classify_num=_classify_num, + timestamp_segment_time_ms=_timestamp_segment_time_ms, + audio_duration_ms=audio_duration_ms, + ) + + +def _ensure_loaded(config: ForcedAlignerConfig) -> None: + """Lazy-load the singleton ``vllm.LLM`` under lock; idempotent.""" + global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config + + if _llm is not None: + if _loaded_config is not None and _loaded_config.model != config.model: + # Multiple configs from different requests — refuse rather + # than swap models silently. A server restart is required. + raise RuntimeError( + f"Forced aligner already loaded with config={_loaded_config!r}; " + f"cannot serve a request that asks for {config!r}. " + "Restart the server to change the aligner config." + ) + return + + with _lock: + if _llm is not None: + return # raced; another caller did the load + + # Lazy import: vllm pulls torch + CUDA, which we want to avoid at + # module import time. + from vllm import LLM + + logger.info( + "Loading forced aligner %s (gpu_memory_utilization=%s)", + config.model, + config.gpu_memory_utilization if config.gpu_memory_utilization is not None else "default", + ) + llm_kwargs: dict[str, Any] = {"model": config.model} + if config.runner is not None: + llm_kwargs["runner"] = config.runner + if config.architecture is not None: + llm_kwargs["hf_overrides"] = {"architectures": [config.architecture]} + if config.gpu_memory_utilization is not None: + llm_kwargs["gpu_memory_utilization"] = config.gpu_memory_utilization + if config.trust_remote_code is not None: + llm_kwargs["trust_remote_code"] = config.trust_remote_code + if config.dtype is not None: + llm_kwargs["dtype"] = config.dtype + if config.max_model_len is not None: + llm_kwargs["max_model_len"] = config.max_model_len + llm_kwargs.update(config.extra_llm_kwargs) + + # Force the aligner engine in-process: in multiprocessing mode the + # token_classify ``pooling_output`` fails to round-trip over the + # engine-core IPC on current vLLM. Hard-set (not setdefault) so a + # globally exported ``=1`` can't push it into a subprocess; save/restore + # to avoid leaking the override to anything built later in this process. + prev_mp = os.environ.get("VLLM_ENABLE_V1_MULTIPROCESSING") + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" + try: + llm = LLM(**llm_kwargs) + finally: + if prev_mp is None: + os.environ.pop("VLLM_ENABLE_V1_MULTIPROCESSING", None) + else: + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = prev_mp + + thinker_config = getattr(llm.llm_engine.model_config.hf_config, "thinker_config", None) + if thinker_config is None or not hasattr(thinker_config, "classify_num"): + raise RuntimeError( + "Loaded aligner has no thinker_config.classify_num; " + "expected a Qwen3ASRForcedAlignerForTokenClassification checkpoint." + ) + timestamp_segment_time_ms = getattr(llm.llm_engine.model_config.hf_config, "timestamp_segment_time", None) + if timestamp_segment_time_ms is None: + raise RuntimeError( + "Loaded aligner has no timestamp_segment_time; expected a Qwen3ASR forced aligner checkpoint." + ) + + tokenizer = llm.get_tokenizer() + timestamp_token_id = _processor.resolve_timestamp_token_id(tokenizer) + + # Publish in this order so a concurrent reader either sees + # _llm == None (will block on the lock) or sees a fully + # initialized aligner. + _classify_num = int(thinker_config.classify_num) + _timestamp_token_id = timestamp_token_id + _timestamp_segment_time_ms = float(timestamp_segment_time_ms) + _loaded_config = config + _llm = llm + + logger.info( + "Forced aligner ready: timestamp_token_id=%d, classify_num=%d, timestamp_segment_time_ms=%.1f", + timestamp_token_id, + _classify_num, + _timestamp_segment_time_ms, + ) + + +# --- pure helpers (testable without a GPU / vllm) --- +# +# Word segmentation, prompt building, timestamp repair and marker-token +# lookup are Qwen-specific and live in +# :mod:`vllm_omni.utils.qwen3_force_align_processor`. This module owns only +# the generic vLLM orchestration. + + +def _pcm_bytes_to_float32(audio: bytes) -> np.ndarray: + """Decode signed-int16 mono PCM bytes into a [-1, 1] float32 array.""" + if not audio: + return np.zeros(0, dtype=np.float32) + if len(audio) % 2 != 0: + # Drop a trailing odd byte rather than raise; keeps streaming + # robust against off-by-one chunk boundaries. + audio = audio[:-1] + pcm = np.frombuffer(audio, dtype=np.int16) + return (pcm.astype(np.float32) / 32768.0).copy() + + +def _decode_timestamps( + *, + logits: Any, + words: list[str], + timestamp_positions: list[int], + classify_num: int, + audio_duration_ms: float, + timestamp_segment_time_ms: float | None = None, +) -> list[WordTimestamp]: + """Translate ``[n_token, classify_num]`` logits into word timestamps. + + ``words`` must be the exact segmentation used to build the prompt (see + :func:`qwen3_force_align_processor.segment_words`); each word owns two + consecutive markers. + """ + arr = logits.detach().cpu().numpy() if hasattr(logits, "detach") else np.asarray(logits) + if arr.ndim != 2: + raise ValueError(f"Expected 2D logits [n_token, classify_num]; got shape {arr.shape}") + if arr.shape[1] != classify_num: + raise ValueError( + f"Logits last dim {arr.shape[1]} != classify_num {classify_num}; " + "model config and prompt template may be out of sync." + ) + + expected = len(words) * 2 + if len(timestamp_positions) != expected: + logger.warning( + "Got %d timestamp positions but text has %d words (expected %d start/end markers); " + "returning empty alignment.", + len(timestamp_positions), + len(words), + expected, + ) + return [] + + marker_logits = arr[timestamp_positions, :] + bin_idx = marker_logits.argmax(axis=-1) + bin_size_ms = ( + float(timestamp_segment_time_ms) + if timestamp_segment_time_ms is not None + else (audio_duration_ms / classify_num if classify_num > 0 else 0.0) + ) + + # Repair non-monotonic bins across the whole start/end sequence before + # pairing, so each (start, end) pair stays ordered and one bad bin can't + # corrupt its neighbours. + marker_ms = _processor.fix_timestamp([int(round(int(b) * bin_size_ms)) for b in bin_idx]) + + # Pair markers into per-word intervals, enforcing the wire contract here so + # callers don't have to: monotonic and non-overlapping across words, with + # start <= end and nothing past the audio length. + duration_ms = int(round(audio_duration_ms)) + out: list[WordTimestamp] = [] + prev_end_ms = 0 + for i, word in enumerate(words): + start_ms = min(max(marker_ms[i * 2], prev_end_ms), duration_ms) + end_ms = min(max(marker_ms[i * 2 + 1], start_ms), duration_ms) + out.append(WordTimestamp(word=word, start_ms=start_ms, end_ms=end_ms)) + prev_end_ms = end_ms + return out + + +# Test hooks --------------------------------------------------------------- +# Tests need a way to reset module state without restarting Python. Not +# part of the public API; do not call in production code. + + +def _reset_for_tests() -> None: + """Drop the cached aligner state so the next call reloads.""" + global _llm, _classify_num, _timestamp_token_id, _timestamp_segment_time_ms, _loaded_config + with _lock: + _llm = None + _classify_num = None + _timestamp_token_id = None + _timestamp_segment_time_ms = None + _loaded_config = None + _processor._reset_for_tests() diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py index 3009e21d4d8..189cc00aa19 100644 --- a/vllm_omni/utils/mm_outputs.py +++ b/vllm_omni/utils/mm_outputs.py @@ -9,6 +9,55 @@ logger = init_logger(__name__) +# Flat payload keys partitioned at worker output into inter-stage connector +# payloads vs client-facing multimodal outputs. Only final output roots are +# listed here; everything else remains available for stage-to-stage transport. +_CLIENT_MM_ROOT_KEYS: frozenset[str] = frozenset( + { + "model_outputs", + "sr", + "audio", + "image", + "images", + "video", + "videos", + "trajectory_latents", + "latents", + } +) + + +def partition_flat_payload( + payload: Mapping[str, object], +) -> tuple[dict[str, object], dict[str, object]]: + """Split a flattened per-request payload into inter-stage vs client mm dicts.""" + if not payload: + return {}, {} + inter_stage: dict[str, object] = {} + client_mm: dict[str, object] = {} + for key, value in payload.items(): + root = key.split(".", 1)[0] + if root in _CLIENT_MM_ROOT_KEYS: + client_mm[key] = value + else: + inter_stage[key] = value + return inter_stage, client_mm + + +def partition_payload_list( + payloads: list[dict[str, object]], +) -> tuple[list[dict[str, object] | None] | None, list[dict[str, object] | None] | None]: + inter_stage_list: list[dict[str, object] | None] = [] + client_mm_list: list[dict[str, object] | None] = [] + for payload in payloads: + inter_stage, client_mm = partition_flat_payload(payload) + inter_stage_list.append(inter_stage or None) + client_mm_list.append(client_mm or None) + return ( + None if all(item is None for item in inter_stage_list) else inter_stage_list, + None if all(item is None for item in client_mm_list) else client_mm_list, + ) + def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]: """Pre-copies multimodal tensor to CPU once (not per-request) to avoid @@ -20,6 +69,9 @@ def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]: Args: multimodal_outputs: Multimodal dict mapping strings to objects. """ + if not multimodal_outputs: + return {} + # Pre-copy multimodal tensors to CPU once (not per-request) to avoid # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU. mm_cpu: dict[str, object] = {} @@ -28,11 +80,10 @@ def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]: if not isinstance(multimodal_outputs, Mapping): logger.warning("Multimodal outputs are not a dict and will not be passed") - if multimodal_outputs: - for k, v in multimodal_outputs.items(): - cpu_v = _to_cpu(v) - if cpu_v is not None: - mm_cpu[k] = cpu_v + for k, v in multimodal_outputs.items(): + cpu_v = _to_cpu(v) + if cpu_v is not None: + mm_cpu[k] = cpu_v return mm_cpu @@ -55,7 +106,13 @@ def _to_cpu(value): def to_payload_element( - element: object, idx: int, start: int, end: int, pass_lists_through: bool = False, seq_len: int | None = None + element: object, + idx: int, + start: int, + end: int, + pass_lists_through: bool = False, + seq_len: int | None = None, + scheduled_seq_len: int | None = None, ): """Build an mm payload element corresponding to one request index from an element containing 0 or more CPU tensors. @@ -69,21 +126,39 @@ def to_payload_element( passthrough data; this should be False in normal cases, but True if we need to avoid splitting nonempty lists prior to calling postprocess, which is the case for prefix cache. - seq_len: Optional sequence length (i.e., dim 0 of hidden states). - When set, a tensor whose first dimension equals seq_len is - sliced per request. The prefix cache passthrough also passes - the total scheduled token count here so 1D (seq_len,) metadata - that is intentionally not cached is still split per request. + seq_len: Optional hidden-aligned batch length (``hidden_states.shape[0]``). + When a tensor's first dimension equals this value, it is sliced + per request using ``start:end``. + scheduled_seq_len: Optional scheduler-aligned batch length + (``scheduler_output.total_num_scheduled_tokens``). Some full-payload + mm tensors (e.g. batched ``codes.audio`` with tail-only hidden states) + are laid out by scheduled tokens instead of the hidden tail shape. + When omitted, ``seq_len`` is reused for backward compatibility. """ - # Cached per-token tensors are merged elsewhere; here a first dim - # equal to seq_len means a per-request slice is required. - if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len: + if scheduled_seq_len is None: + scheduled_seq_len = seq_len + + # Cached per-token tensors are merged elsewhere; here a first dim equal to + # either the hidden-aligned or scheduler-aligned batch length means a + # per-request slice is required. + if isinstance(element, torch.Tensor) and ( + (seq_len is not None and element.shape[0] == seq_len) + or (scheduled_seq_len is not None and element.shape[0] == scheduled_seq_len) + ): return element[start:end].contiguous() # Every other case is shared between prefix cache (passthrough data) # and running a model without prefix caching. elif isinstance(element, dict): return { - sk: to_payload_element(sv, idx, start, end, pass_lists_through=pass_lists_through, seq_len=seq_len) + sk: to_payload_element( + sv, + idx, + start, + end, + pass_lists_through=pass_lists_through, + seq_len=seq_len, + scheduled_seq_len=scheduled_seq_len, + ) for sk, sv in element.items() } elif isinstance(element, list): diff --git a/vllm_omni/utils/qwen3_force_align_processor.py b/vllm_omni/utils/qwen3_force_align_processor.py new file mode 100644 index 00000000000..66a084b4488 --- /dev/null +++ b/vllm_omni/utils/qwen3_force_align_processor.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from: +# https://github.com/QwenLM/Qwen3-ASR/blob/main/qwen_asr/inference/qwen3_forced_aligner.py +# +# Copyright 2026 The Alibaba Qwen team. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Qwen3 forced-aligner text/timestamp processor. + +This is the model-specific half of upstream's ``Qwen3ForceAlignProcessor``: +it turns text into the aligner's word units and prompt, repairs the predicted +timestamp bins, and resolves the marker token id. The generic "run a vLLM +pooling model" half lives in :mod:`vllm_omni.utils.forced_aligner`, which +calls into this module. + +Keeping the Qwen-specific pieces here marks the seam for the model-agnostic +aligner the issue asks for: a different aligner family would supply its own +processor exposing the same small surface — :func:`segment_words`, +:func:`build_prompt`, :func:`fix_timestamp`, :func:`resolve_timestamp_token_id`. + +Word segmentation prefers qwen_asr's official ``Qwen3ForceAlignProcessor`` +when installed (full multilingual fidelity, incl. Japanese/Korean) and +otherwise uses the faithful port below: exact for whitespace-delimited and +Chinese-mixed text; Japanese/Korean degrade to whitespace splitting. +""" + +from __future__ import annotations + +import logging +import unicodedata +from typing import Any + +logger = logging.getLogger(__name__) + + +# Prompt tokens for the Qwen3 forced aligner. Baked in (not config) because the +# prompt template and the word segmentation are inherently Qwen-specific: a +# different aligner family needs a different processor module regardless. +AUDIO_PLACEHOLDER = "<|audio_start|><|audio_pad|><|audio_end|>" +TIMESTAMP_TOKEN = "<timestamp>" + +_CJK_RANGES = ( + (0x4E00, 0x9FFF), + (0x3400, 0x4DBF), + (0x20000, 0x2A6DF), + (0x2A700, 0x2B73F), + (0x2B740, 0x2B81F), + (0x2B820, 0x2CEAF), + (0xF900, 0xFAFF), +) + +# Map common language codes to the names qwen_asr's processor expects. +_LANG_ALIASES = { + "ja": "japanese", + "jp": "japanese", + "jpn": "japanese", + "ko": "korean", + "kr": "korean", + "kor": "korean", +} + + +def build_prompt(words: list[str]) -> str: + """Wrap segmented words in the Qwen3 aligner prompt template. + + Each word gets two trailing ``<timestamp>`` markers (start + end); the + model classifies each marker into an audio time bin. + """ + if not words: + # Pad with one timestamp so the decoder always has something to + # read; an empty result still surfaces as "[]" upstream. + body = TIMESTAMP_TOKEN + else: + body = f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}".join(words) + f"{TIMESTAMP_TOKEN}{TIMESTAMP_TOKEN}" + return f"<|im_start|>user\n{AUDIO_PLACEHOLDER}{body}<|im_end|>\n<|im_start|>assistant\n" + + +# --- word segmentation (port of Qwen3ForceAlignProcessor) ---------------- +# +# The marker count fed to the model must match the segmentation the aligner +# was trained with, or the decode drifts and returns empty. + + +def _is_cjk_char(ch: str) -> bool: + code = ord(ch) + return any(lo <= code <= hi for lo, hi in _CJK_RANGES) + + +def _is_kept_char(ch: str) -> bool: + """Keep letters/numbers (any script) and the apostrophe; drop punctuation.""" + if ch == "'": + return True + cat = unicodedata.category(ch) + return cat.startswith("L") or cat.startswith("N") + + +def _clean_token(token: str) -> str: + return "".join(ch for ch in token if _is_kept_char(ch)) + + +def _split_segment_with_chinese(seg: str) -> list[str]: + """Split a cleaned segment so each CJK char is its own token.""" + tokens: list[str] = [] + buf: list[str] = [] + + def flush() -> None: + if buf: + tokens.append("".join(buf)) + buf.clear() + + for ch in seg: + if _is_cjk_char(ch): + flush() + tokens.append(ch) + else: + buf.append(ch) + flush() + return tokens + + +def _tokenize_space_lang(text: str) -> list[str]: + """Whitespace + Chinese-mixed segmentation (Qwen3ForceAlignProcessor port). + + Splits on whitespace, strips punctuation from each segment (keeping + cross-script letters/digits and the apostrophe), then peels off CJK + characters as individual tokens. + """ + tokens: list[str] = [] + for seg in text.split(): + cleaned = _clean_token(seg) + if cleaned: + tokens.extend(_split_segment_with_chinese(cleaned)) + return tokens + + +# Cached qwen_asr processor (or a sentinel that it is unavailable). Importing +# it eagerly pulls nagisa + the Qwen3ASR transformers backend, so we only try +# once, lazily, and remember the outcome. +_official_processor: Any = None +_official_processor_unavailable = False + + +def _get_official_processor() -> Any | None: + global _official_processor, _official_processor_unavailable + if _official_processor is not None: + return _official_processor + if _official_processor_unavailable: + return None + try: + from qwen_asr.inference.qwen3_forced_aligner import ( # type: ignore[import-not-found] + Qwen3ForceAlignProcessor, + ) + + _official_processor = Qwen3ForceAlignProcessor() + except Exception: # noqa: BLE001 - any import/init error means "use the port" + _official_processor_unavailable = True + logger.info( + "qwen_asr not available; using the built-in word segmentation " + "(exact for whitespace/Chinese-mixed text; Japanese/Korean fall " + "back to whitespace splitting). Install qwen-asr for full fidelity." + ) + return None + return _official_processor + + +def segment_words(text: str, language: str | None = None) -> list[str]: + """Split ``text`` into the aligner's word units. + + Prefers qwen_asr's official processor (full multilingual fidelity); falls + back to the built-in port otherwise. ``language`` follows the official + naming — ``"japanese"`` / ``"korean"`` (or codes like ``ja`` / ``ko``) + trigger language-specific tokenisers; anything else (incl. ``None`` / + ``"auto"``) uses the whitespace + Chinese-mixed path. + """ + lang = (language or "auto").strip().lower() + lang = _LANG_ALIASES.get(lang, lang) + processor = _get_official_processor() + if processor is not None: + try: + word_list, _ = processor.encode_timestamp(text, lang) + return list(word_list) + except Exception: # noqa: BLE001 + logger.warning( + "Official Qwen3ForceAlignProcessor failed for language=%r; falling back to built-in segmentation.", + lang, + exc_info=True, + ) + if lang in ("japanese", "korean"): + logger.warning( + "language=%r needs qwen_asr (nagisa/soynlp) for faithful word " + "segmentation; falling back to whitespace splitting.", + lang, + ) + return _tokenize_space_lang(text) + + +def fix_timestamp(values: list[int]) -> list[int]: + """Repair non-monotonic timestamp bins via LIS + interpolation. + + The model occasionally emits out-of-order bins; this snaps anomalies back + onto the longest non-decreasing subsequence (interpolating longer runs) + so paired start/end markers stay ordered. Port of + ``Qwen3ForceAlignProcessor.fix_timestamp`` (Apache-2.0). + """ + n = len(values) + if n == 0: + return [] + + dp = [1] * n + parent = [-1] * n + for i in range(1, n): + for j in range(i): + if values[j] <= values[i] and dp[j] + 1 > dp[i]: + dp[i] = dp[j] + 1 + parent[i] = j + + idx = dp.index(max(dp)) + lis: list[int] = [] + while idx != -1: + lis.append(idx) + idx = parent[idx] + is_normal = [False] * n + for k in lis: + is_normal[k] = True + + result = [float(v) for v in values] + i = 0 + while i < n: + if is_normal[i]: + i += 1 + continue + j = i + while j < n and not is_normal[j]: + j += 1 + count = j - i + left = next((result[k] for k in range(i - 1, -1, -1) if is_normal[k]), None) + right = next((result[k] for k in range(j, n) if is_normal[k]), None) + if count <= 2: + for k in range(i, j): + if left is None: + result[k] = right + elif right is None: + result[k] = left + else: + result[k] = left if (k - (i - 1)) <= (j - k) else right + elif left is not None and right is not None: + step = (right - left) / (count + 1) + for k in range(i, j): + result[k] = left + step * (k - i + 1) + elif left is not None: + for k in range(i, j): + result[k] = left + elif right is not None: + for k in range(i, j): + result[k] = right + i = j + + return [int(r) for r in result] + + +def resolve_timestamp_token_id(tokenizer: Any, timestamp_token: str = TIMESTAMP_TOKEN) -> int: + """Look up the integer id of the timestamp special token.""" + convert = getattr(tokenizer, "convert_tokens_to_ids", None) + if not callable(convert): + raise RuntimeError("Aligner tokenizer has no convert_tokens_to_ids method.") + tid = convert(timestamp_token) + if isinstance(tid, list): + tid = tid[0] if tid else None + if tid is None or (isinstance(tid, int) and tid < 0): + raise RuntimeError( + f"Aligner tokenizer does not recognise {timestamp_token!r} (got id={tid}). " + "Check the model card; the marker token may use a different name." + ) + return int(tid) + + +def _reset_for_tests() -> None: + """Drop the cached official-processor probe so the next call re-checks.""" + global _official_processor, _official_processor_unavailable + _official_processor = None + _official_processor_unavailable = False diff --git a/vllm_omni/utils/tracking_parser.py b/vllm_omni/utils/tracking_parser.py index c6dc5fee798..64849a5aa09 100644 --- a/vllm_omni/utils/tracking_parser.py +++ b/vllm_omni/utils/tracking_parser.py @@ -9,6 +9,29 @@ UNSET = object() +def build_shadow_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: + """Build kwargs for the shadow argument with an ``UNSET`` default. + + Actions that mutate their default in place (append/extend/append_const/ + count) would crash on the bare ``UNSET`` sentinel, so they are remapped to + an equivalent store-style action; the shadow value only needs to flip away + from ``UNSET`` when the arg is passed explicitly. + """ + shadow_kwargs = {**kwargs, "default": UNSET} + action = kwargs.get("action") + + if action in ("append", "extend"): + shadow_kwargs["action"] = "store" + + elif action in ("append_const", "count"): + shadow_kwargs["action"] = "store_const" + + if action == "count": + shadow_kwargs["const"] = True + + return shadow_kwargs + + class TrackingNamespace(argparse.Namespace): """Proxy that wraps an argparse namespace with explicit keys, which can be filtered down to a dict containing only explicitly passed values. @@ -58,7 +81,7 @@ def __init__( def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action: """Add an argument to the real group and to the shadow group.""" action = self._real.add_argument(*args, **kwargs) - default_kwargs = {**kwargs, "default": UNSET} + default_kwargs = build_shadow_kwargs(kwargs) self._shadow.add_argument(*args, **default_kwargs) return action @@ -115,7 +138,7 @@ def __init__(self, *args, **kwargs): def add_argument(self, *args: Any, **kwargs: Any) -> argparse.Action: """Add an arg to the parser & the shadow, where the latter has UNSET for the default.""" action = super().add_argument(*args, **kwargs) - shadow_kwargs = {**kwargs, "default": UNSET} + shadow_kwargs = build_shadow_kwargs(kwargs) self._shadow.add_argument(*args, **shadow_kwargs) return action diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d29af7314fc..6cd2b35663e 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -7,7 +7,8 @@ from __future__ import annotations import gc -from collections.abc import Mapping +import threading +from collections.abc import Callable, Sequence from contextlib import nullcontext from copy import copy from dataclasses import replace @@ -21,6 +22,7 @@ from vllm.distributed.parallel_state import get_pp_group, get_tp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger +from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.dflash import DFlashProposer @@ -41,13 +43,199 @@ from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.outputs import OmniModelRunnerOutput -from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element +from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest +from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids logger = init_logger(__name__) +def _to_cpu_contiguous(tensor: torch.Tensor) -> torch.Tensor: + tensor = tensor.detach() + if tensor.device.type == "cpu": + return tensor.contiguous() + return tensor.to("cpu").contiguous() + + +def _clone_cuda_tensor_payload(value: Any, sources: list[torch.Tensor]) -> Any: + """Clone CUDA tensors on the current stream before async CPU copies. + + The clone protects async Omni output snapshots from CUDA graph output + buffers that may be reused by subsequent decode steps. CPU tensors are + cloned synchronously because they are already host-owned snapshots. + """ + if isinstance(value, torch.Tensor): + if value.device.type == "cuda": + cloned = value.detach().clone() + sources.append(cloned) + return cloned + return value.detach().clone() + if isinstance(value, dict): + return {k: _clone_cuda_tensor_payload(v, sources) for k, v in value.items()} + if isinstance(value, list): + return [_clone_cuda_tensor_payload(v, sources) for v in value] + if isinstance(value, tuple): + return tuple(_clone_cuda_tensor_payload(v, sources) for v in value) + return value + + +def _copy_tensor_payload_to_cpu(value: Any, pin_memory: bool) -> Any: + if isinstance(value, torch.Tensor): + if value.device.type != "cuda": + return value + cpu = torch.empty_like(value, device="cpu", pin_memory=pin_memory) + cpu.copy_(value, non_blocking=True) + return cpu + if isinstance(value, dict): + return {k: _copy_tensor_payload_to_cpu(v, pin_memory) for k, v in value.items()} + if isinstance(value, list): + return [_copy_tensor_payload_to_cpu(v, pin_memory) for v in value] + if isinstance(value, tuple): + return tuple(_copy_tensor_payload_to_cpu(v, pin_memory) for v in value) + return value + + +class _AsyncCPUPayloadSnapshot: + def __init__( + self, + payload: Any, + ready_event: torch.cuda.Event | None, + cuda_sources: list[torch.Tensor], + ) -> None: + self.payload = payload + self._ready_event = ready_event + self._cuda_sources = cuda_sources + self._waited = False + + def wait(self) -> None: + if self._waited: + return + if self._ready_event is not None: + self._ready_event.synchronize() + self._cuda_sources.clear() + self._waited = True + + +def _snapshot_tensor_payload_to_cpu_async( + value: Any, + *, + copy_stream: torch.cuda.Stream, + pin_memory: bool, +) -> _AsyncCPUPayloadSnapshot: + cuda_sources: list[torch.Tensor] = [] + cloned = _clone_cuda_tensor_payload(value, cuda_sources) + if not cuda_sources: + return _AsyncCPUPayloadSnapshot(cloned, None, cuda_sources) + + source_stream = torch.cuda.current_stream() + ready_event = torch.cuda.Event() + with torch.cuda.stream(copy_stream): + copy_stream.wait_stream(source_stream) + cpu_payload = _copy_tensor_payload_to_cpu(cloned, pin_memory) + ready_event.record(copy_stream) + return _AsyncCPUPayloadSnapshot(cpu_payload, ready_event, cuda_sources) + + +class _OmniOutputTensorSnapshot(NamedTuple): + hidden_states: torch.Tensor + staged_hidden_states_cpu: torch.Tensor | None + multimodal_outputs: Any + async_payload: _AsyncCPUPayloadSnapshot | None = None + + +class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): + def __init__( + self, + *, + model_runner_output_builder: Callable[[], OmniModelRunnerOutput], + cuda_device: torch.device | int | str | None = None, + **kwargs: Any, + ) -> None: + sampled_token_ids = kwargs.pop("sampled_token_ids") + logprobs_tensors = kwargs.pop("logprobs_tensors") + invalid_req_indices = kwargs.pop("invalid_req_indices") + async_output_copy_stream = kwargs.pop("async_output_copy_stream") + vocab_size = kwargs.pop("vocab_size") + routed_experts = kwargs.pop("routed_experts", None) + # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault + # for EP all2all fault tolerance (PR #43637). Omni doesn't use this + # feature but must consume the kwarg to prevent TypeError from stray + # kwargs and initialize the attribute so super().get_output() works. + kwargs.pop("check_ep_fault", False) + if kwargs: + raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}") + + self._model_runner_output = None + self._invalid_req_indices = invalid_req_indices + + self.async_copy_ready_event = torch.Event() + self._sampled_token_ids = sampled_token_ids + self.vocab_size = vocab_size + self._logprobs_tensors = logprobs_tensors + self._routed_experts = routed_experts + self._has_fault: torch.Tensor | None = None + + default_stream = torch.cuda.current_stream() + with torch.cuda.stream(async_output_copy_stream): + async_output_copy_stream.wait_stream(default_stream) + # Keep sampled-token feedback identical to upstream async + # scheduling. This tensor drives the next decode step, so avoid + # changing its host-copy allocation semantics while building Omni + # output asynchronously. + self.sampled_token_ids_cpu = self._sampled_token_ids.to("cpu", non_blocking=True) + self._logprobs_tensors_cpu = self._logprobs_tensors.to_cpu_nonblocking() if self._logprobs_tensors else None + self._routed_experts_cpu = ( + self._routed_experts.to_cpu_nonblocking() if self._routed_experts is not None else None + ) + self.async_copy_ready_event.record() + + self._model_runner_output_builder = model_runner_output_builder + self._background_exception: BaseException | None = None + self._background_thread: threading.Thread | None = None + self._cuda_device = cuda_device + self._background_thread = threading.Thread( + target=self._build_output_in_background, + daemon=True, + name="omni-async-output-builder", + ) + self._background_thread.start() + + def _build_model_runner_output_once(self) -> None: + if self._model_runner_output is not None: + return + with record_function_or_nullcontext("omni_async_output:get_output/build_model_runner_output"): + self._model_runner_output = self._model_runner_output_builder() + self._model_runner_output_builder = None + + def _build_output_in_background(self) -> None: + try: + if self._cuda_device is not None: + torch.cuda.set_device(self._cuda_device) + self._build_model_runner_output_once() + except BaseException as exc: # noqa: BLE001 - re-raised by get_output(). + self._background_exception = exc + + def get_output(self) -> OmniModelRunnerOutput: + background_thread = getattr(self, "_background_thread", None) + if background_thread is not None: + background_thread.join() + self._background_thread = None + background_exception = getattr(self, "_background_exception", None) + if background_exception is not None: + raise background_exception + self._build_model_runner_output_once() + # Upstream AsyncGPUModelRunnerOutput.get_output() accesses + # self._has_fault for EP all2all fault tolerance (PR #43637). + # Ensure the attribute exists even when __init__ was bypassed + # (e.g. unit tests using object.__new__). + if not hasattr(self, "_has_fault"): + self._has_fault = None + with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"): + return super().get_output() + + def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]: """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``. @@ -115,6 +303,7 @@ def __init__(self, *args, **kwargs): self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) # Initialize KV cache manager (preserve vllm_config fallback behavior) self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config) + self._async_chunk = getattr(self.model_config, "async_chunk", False) # Worker-connector init is gated by a per-`model_arch` allowlist # (covers both producer-side and consumer-side runners for the # arches below). Consumer-wait stages must be registered @@ -130,6 +319,7 @@ def __init__(self, *args, **kwargs): "Qwen3TTSCode2Wav", "CosyVoice3Model", "DyninOmniForConditionalGeneration", + "IndexTTS2TalkerForConditionalGeneration", } if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: self.init_omni_connectors( @@ -256,6 +446,23 @@ def _sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None: return None return [rid for rid in req_ids if isinstance(rid, str)] + @staticmethod + def _resolve_sparse_mm_routing( + *, + engine_output_type: str, + req_ids_output_copy: list[str], + downstream_req_ids: list[str], + multimodal_outputs: Any, + ) -> tuple[list[str], dict[str, int], bool]: + sparse_mm_req_ids = GPUARModelRunner._sparse_mm_req_ids(multimodal_outputs) + sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} + if engine_output_type != "audio" or sparse_mm_req_ids is None: + return downstream_req_ids, sparse_mm_index, False + + sparse_req_id_set = set(sparse_mm_req_ids) + sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] + return sparse_downstream_req_ids, sparse_mm_index, True + @staticmethod def _is_sparse_audio_marker(value: Any) -> bool: if isinstance(value, list): @@ -395,6 +602,93 @@ def _model_needs_full_prefix_hidden_states(self) -> bool: model = getattr(self, "model", None) return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)) + def _get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_reqs_padded: int, + num_scheduled_tokens_np: np.ndarray, + num_computed_tokens_cpu: torch.Tensor | np.ndarray, + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest | None: + # Models without this hook keep the normal runner path. VoxCPM2 uses + # it to ask the runner for padded FULL attention metadata while keeping + # graph policy and lifecycle in the model layer. + hook = getattr(self.model, "get_runner_assisted_full_attention_metadata_request", None) + if not callable(hook): + return None + num_computed_tokens = num_computed_tokens_cpu + if hasattr(num_computed_tokens, "numpy"): + num_computed_tokens = num_computed_tokens.numpy() + request = hook( + req_ids=req_ids, + num_reqs=num_reqs, + num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], + num_computed_tokens=num_computed_tokens, + max_num_scheduled_tokens=max_num_scheduled_tokens, + ) + if request is None: + return None + if not isinstance(request, RunnerAssistedFullAttentionMetadataRequest): + raise TypeError( + "runner-assisted full attention metadata hook must return " + "RunnerAssistedFullAttentionMetadataRequest or None, got " + f"{type(request).__name__}" + ) + padded_num_reqs = max( + num_reqs_padded, + min(int(request.num_reqs_padded), self.scheduler_config.max_num_seqs), + ) + return RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=padded_num_reqs, + for_cudagraph_capture=bool(request.for_cudagraph_capture), + ) + + def _refresh_runner_assisted_full_attention_metadata_buffers( + self, + *, + num_reqs: int, + num_reqs_padded: int, + num_scheduled_tokens_np: np.ndarray, + ) -> None: + num_computed_tokens = self.input_batch.num_computed_tokens_cpu[:num_reqs] + if hasattr(num_computed_tokens, "numpy"): + num_computed_tokens = num_computed_tokens.numpy() + np.add( + num_computed_tokens, + num_scheduled_tokens_np, + out=self.optimistic_seq_lens_cpu[:num_reqs].numpy(), + ) + self.optimistic_seq_lens_cpu[num_reqs:num_reqs_padded].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + + cum_num_tokens = self._get_cumsum_and_arange( + num_scheduled_tokens_np, + self.query_pos.np, + ) + self.query_start_loc.np[0] = 0 + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1]) + self.query_start_loc.copy_to_gpu() + + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + def _set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> bool: + hook = getattr(self.model, "set_runner_assisted_full_attention_metadata_context", None) + if not callable(hook): + return False + hook( + enabled=enabled, + num_reqs=num_reqs, + ) + return True + def _deferred_prefix_cache_mm_keys(self) -> set[str]: """Model-declared multimodal keys whose prefix-cache writes are deferred.""" model = getattr(self, "model", None) @@ -421,12 +715,6 @@ def _maybe_update_prefix_cache( if self.omni_prefix_cache is not None and is_last_pp_rank: # If this happens, it generally means the model is not following the correct # interface yet and is therefore currently not compatible with prefix cache. - if multimodal_outputs is not None and not isinstance(multimodal_outputs, Mapping): - logger.warning_once( - "prefix caching expects mm outputs to be a dict, but got %s", - type(multimodal_outputs), - ) - hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None # FIX: The .cpu attribute of slot_mapping is stale (not updated by the Triton # _compute_slot_mapping_kernel which only writes to .gpu). We must use .gpu and @@ -483,6 +771,171 @@ def _maybe_get_combined_prefix_cache_tensors( ) return combined_hidden_states, combined_multimodal_outputs + def _stage_deferred_prefix_cache_mm_outputs( + self, + *, + scheduler_output: SchedulerOutput, + multimodal_outputs: Any, + query_start_loc_cpu: Any, + ) -> None: + if self.omni_prefix_cache is None: + return + + deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() + if not deferred_mm_cache_keys: + return + + self.omni_prefix_cache.stage_deferred_mm_outputs( + query_start_loc=query_start_loc_cpu, + input_batch=self.input_batch, + multimodal_outputs=flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + deferred_mm_cache_keys=deferred_mm_cache_keys, + ) + + def _prepare_prefix_cache_pooler_payload_sources( + self, + *, + hidden_states: torch.Tensor, + staged_hidden_states_cpu: torch.Tensor | None, + multimodal_outputs: Any, + scheduler_output: SchedulerOutput, + needs_scheduled_hidden_payload: bool, + ) -> tuple[torch.Tensor | None, dict[str, torch.Tensor] | None, dict | None]: + hidden_states_cpu = None + if needs_scheduled_hidden_payload: + if staged_hidden_states_cpu is None: + raise RuntimeError("Prefix-cache hidden-state payload requires staged CPU hidden states.") + hidden_states_cpu = staged_hidden_states_cpu + + combined_hidden_states, combined_multimodal_outputs = self._maybe_get_combined_prefix_cache_tensors( + hidden_states, + staged_hidden_states_cpu, + multimodal_outputs, + scheduler_output.num_scheduled_tokens, + ) + return hidden_states_cpu, combined_hidden_states, combined_multimodal_outputs + + @staticmethod + def _build_combined_prefix_cache_mm_payload( + combined_multimodal_outputs: dict, + *, + rid: str, + idx: int, + ) -> dict[str, object]: + def _unwrap_lists(v): + if isinstance(v, list): + return v[idx] if idx < len(v) else v[0] + if isinstance(v, dict): + return {k: _unwrap_lists(sv) for k, sv in v.items()} + return v + + return { + mm_key: _unwrap_lists(combined_multimodal_outputs[mm_key][rid]) + for mm_key in combined_multimodal_outputs.keys() + } + + def _build_omni_mm_payload( + self, + *, + combined_multimodal_outputs: dict | None, + mm_cpu: dict[str, object] | None, + rid: str, + idx: int, + start: int, + end: int, + audio_sparse_output: bool, + sparse_mm_index: dict[str, int], + hidden_seq_len: int, + scheduled_seq_len: int, + ) -> dict[str, object]: + if combined_multimodal_outputs: + return self._build_combined_prefix_cache_mm_payload( + combined_multimodal_outputs, + rid=rid, + idx=idx, + ) + + mm_payload: dict[str, object] = {} + if not mm_cpu: + return mm_payload + + for mm_key, mm_val in mm_cpu.items(): + if mm_key in {"meta.req_id", "meta.sparse_audio"}: + continue + if audio_sparse_output and isinstance(mm_val, list): + sparse_idx = sparse_mm_index.get(rid) + if sparse_idx is None: + continue + if sparse_idx >= len(mm_val): + logger.warning( + "Sparse multimodal payload mismatch for request %s: index %d >= %d.", + rid, + sparse_idx, + len(mm_val), + ) + continue + sparse_val = mm_val[sparse_idx] + mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val + continue + mm_payload[mm_key] = to_payload_element( + element=mm_val, + idx=idx, + start=start, + end=end, + pass_lists_through=False, + seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, + ) + return mm_payload + + def _build_omni_pooler_payload( + self, + *, + rid: str, + idx: int, + start: int, + end: int, + hidden_states_cpu: torch.Tensor | None, + req_hidden_states_cpu: dict[str, torch.Tensor] | None, + combined_hidden_states: dict[str, torch.Tensor] | None, + combined_multimodal_outputs: dict | None, + mm_cpu: dict[str, object] | None, + audio_sparse_output: bool, + sparse_mm_index: dict[str, int], + hidden_seq_len: int, + scheduled_seq_len: int, + ) -> dict[str, object]: + payload: dict[str, object] = {} + if not audio_sparse_output: + if req_hidden_states_cpu is not None and combined_hidden_states is None: + req_hidden_states = req_hidden_states_cpu[rid] + else: + req_hidden_states = self._resolve_req_hidden_states( + hidden_states_cpu, + combined_hidden_states, + rid, + start, + end, + ) + if req_hidden_states is not None: + payload["hidden"] = req_hidden_states + + mm_payload = self._build_omni_mm_payload( + combined_multimodal_outputs=combined_multimodal_outputs, + mm_cpu=mm_cpu, + rid=rid, + idx=idx, + start=start, + end=end, + audio_sparse_output=audio_sparse_output, + sparse_mm_index=sparse_mm_index, + hidden_seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, + ) + payload.update(mm_payload) + return payload + @torch.inference_mode() def execute_model( self, @@ -500,6 +953,16 @@ def execute_model( if hasattr(self.model, "_clear_warmup_state"): self.model._clear_warmup_state() + # Async-write pipeline: apply any pending GPU->CPU prefix-cache writes + # whose copy event has already fired. Non-blocking — entries whose D2H + # is still in flight stay queued and will be picked up on the next + # step's drain. This guarantees any downstream read of + # ``omni_prefix_cache.hidden_states_cache`` / + # ``omni_prefix_cache.mm_outputs_cache`` in this step sees the + # state produced no later than the previous forward step. + if self.omni_prefix_cache is not None: + self.omni_prefix_cache.drain_ready_async_writes() + # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"): @@ -664,6 +1127,26 @@ def execute_model( ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu[:num_reqs] + + runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request( + req_ids=req_ids[:num_reqs], + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_scheduled_tokens_np=num_scheduled_tokens_np, + num_computed_tokens_cpu=num_computed_tokens_cpu, + max_num_scheduled_tokens=max_num_scheduled_tokens, + ) + runner_assisted_full_attn_capture = False + if runner_assisted_full_attn_request is not None: + num_reqs_padded = runner_assisted_full_attn_request.num_reqs_padded + runner_assisted_full_attn_capture = runner_assisted_full_attn_request.for_cudagraph_capture + num_tokens_padded = max( + num_tokens_padded, + num_reqs_padded * max_num_scheduled_tokens, + ) + runner_assisted_full_attn = runner_assisted_full_attn_request is not None + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( should_ubatch, num_scheduled_tokens_np, @@ -671,8 +1154,7 @@ def execute_model( num_reqs_padded, self.parallel_config.num_ubatches, ) - - pad_attn = cudagraph_mode == CUDAGraphMode.FULL + pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices @@ -694,11 +1176,18 @@ def execute_model( num_tokens_unpadded=num_tokens_unpadded, ubatch_slices=ubatch_slices_padded, ) + if runner_assisted_full_attn: + self._refresh_runner_assisted_full_attention_metadata_buffers( + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_scheduled_tokens_np=num_scheduled_tokens_np, + ) + attn_num_reqs = num_reqs_padded if runner_assisted_full_attn else num_reqs attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, - num_reqs=num_reqs, + num_reqs=attn_num_reqs, num_reqs_padded=num_reqs_padded if pad_attn else None, max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, @@ -707,6 +1196,7 @@ def execute_model( num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, slot_mappings=slot_mappings_by_group, + for_cudagraph_capture=runner_assisted_full_attn_capture, ) self._maybe_attach_attention_metadata_extensions( attn_metadata=attn_metadata, @@ -715,6 +1205,7 @@ def execute_model( max_query_len=max_num_scheduled_tokens, pad_attn=pad_attn, num_scheduled_tokens_np=num_scheduled_tokens_np, + for_cudagraph_capture=runner_assisted_full_attn_capture, ) ( @@ -728,14 +1219,15 @@ def execute_model( # Let the model adjust inputs before forward (e.g. restore input_ids # for multimodal position detection, fix decode position offsets). - if hasattr(self.model, "prepare_runner_inputs"): - input_ids, positions = self.model.prepare_runner_inputs( + prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None) + if callable(prepare_runner_inputs): + input_ids, positions = prepare_runner_inputs( input_ids=input_ids, positions=positions, inputs_embeds=inputs_embeds, req_ids=req_ids[:num_reqs], - num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)], - num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)], + num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs], + num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], input_ids_buffer=self.input_ids.gpu[:num_tokens_padded], ) @@ -744,46 +1236,60 @@ def execute_model( # with CUDA graph capture. if self.calculate_kv_scales: cudagraph_mode = CUDAGraphMode.NONE + runner_assisted_full_attn = False # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False + runner_assisted_context_enabled = False + if runner_assisted_full_attn: + runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context( + enabled=True, + num_reqs=num_reqs, + ) + runner_assisted_full_attn = runner_assisted_context_enabled + # Run the model. # Use persistent buffers for CUDA graphs. # When spec decode is enabled, defer connector finalization # (wait_for_save + clear metadata) until after draft model runs. defer_kv_connector_finalize = self.speculative_config is not None - with ( - nullcontext(), - set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens_padded, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_mode, - batch_descriptor=batch_desc, - ubatch_slices=ubatch_slices_padded, - slot_mapping=slot_mappings, # OMNI: required for KV cache operations - ), - record_function_or_nullcontext("gpu_model_runner: forward"), - self.maybe_get_kv_connector_output( - scheduler_output, - defer_finalize=defer_kv_connector_finalize, - ) as kv_connector_output, - ): - model_output = self._model_forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - sampling_metadata=self.input_batch.sampling_metadata, - logits_index=logits_indices, - sampler=self.sampler, - ) + try: + with ( + nullcontext(), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=(CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode), + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations + ), + record_function_or_nullcontext("gpu_model_runner: forward"), + self.maybe_get_kv_connector_output( + scheduler_output, + defer_finalize=defer_kv_connector_finalize, + ) as kv_connector_output, + ): + model_output = self._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, + ) - # [Omni] Map pending ropes metadata to req_ids. - if hasattr(self.model, "flush_pending_metadata"): - self.model.flush_pending_metadata(list(req_ids)) + # [Omni] Map pending ropes metadata to req_ids. + flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) + if callable(flush_pending_metadata): + flush_pending_metadata(req_ids[:num_reqs]) + finally: + if runner_assisted_context_enabled: + self._set_runner_assisted_full_attention_metadata_context(enabled=False) with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: @@ -796,18 +1302,30 @@ def execute_model( hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) hidden_states_cpu = None - if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: - hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() - # Cache hidden states & multimodal outputs if we've enabled hidden state - # prefix caching unless this isn't the last pipeline parallelism rank. - self._maybe_update_prefix_cache( - hidden_states=hidden_states, - hidden_states_cpu=hidden_states_cpu, - multimodal_outputs=multimodal_outputs, - num_tokens_unpadded=num_tokens_unpadded, - num_tokens_padded=num_tokens_padded, - ) + # Async-write pipeline (replaces the per-step blocking + # ``.to("cpu")`` + ``aten::index_put_`` on pageable host memory). + # Schedules non-blocking GPU->CPU copies on a dedicated stream; + # the actual CPU scatter into ``hidden_states_cache`` / + # ``mm_outputs_cache`` happens in ``drain_ready_async_writes`` + # at the top of subsequent execute_model() calls. + if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: + hs_for_cache = hidden_states if self._model_needs_full_prefix_hidden_states() else None + # Some models (e.g. qwen3-tts-talker) opt out of full-hidden-state + # prefix caching but the downstream pooler payload path still + # needs a CPU hidden-states view. Materialize it synchronously + # in that case; the legacy behavior is preserved. + if hs_for_cache is None and self._model_omni_pooler_payload_include_hidden(): + hidden_states_cpu = hidden_states[:num_tokens_unpadded].detach().to("cpu").contiguous() + slot_mapping_gpu = self.input_batch.block_table[0].slot_mapping.gpu + self.omni_prefix_cache.schedule_async_write( + hidden_states_gpu=hs_for_cache, + multimodal_outputs_gpu=(flatten_payload(multimodal_outputs) if multimodal_outputs else None), + slot_mapping_gpu=slot_mapping_gpu, + num_tokens_unpadded=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + skip_mm_cache_keys=self._deferred_prefix_cache_mm_keys(), + ) if not self.broadcast_pp_output: # Common case. @@ -887,7 +1405,7 @@ def execute_model( if deferred_state_corrections_fn: deferred_state_corrections_fn() - if self.routed_experts_initialized and hasattr(self, "_positions_cpu"): + if self._should_return_omni_routed_experts() and hasattr(self, "_positions_cpu"): self._omni_routed_experts_d2h(scheduler_output) return None @@ -948,8 +1466,8 @@ def _resolve_req_hidden_states( def _build_multimodal_outputs( self, - per_req_payloads: list[dict[str, object]] | None, - ) -> list[dict[str, torch.Tensor]] | None: + per_req_payloads: list[dict[str, object] | None] | None, + ) -> list[dict[str, torch.Tensor] | None] | None: """Build per-request multimodal output payloads (dedicated channel). Reuses the per-request payloads assembled by the pooler-payload loop @@ -964,7 +1482,373 @@ def _build_multimodal_outputs( return None if per_req_payloads is None: return None - return [_ensure_tensor_values(payload) if payload else {} for payload in per_req_payloads] + wire_payloads: list[dict[str, torch.Tensor] | None] = [] + for payload in per_req_payloads: + if not payload: + wire_payloads.append(None) + else: + wire_payloads.append(_ensure_tensor_values(payload)) + if all(item is None for item in wire_payloads): + return None + return wire_payloads + + def _snapshot_query_start_loc_cpu(self) -> Any: + query_start_loc_cpu = self.query_start_loc.cpu + if callable(query_start_loc_cpu): + query_start_loc_cpu = query_start_loc_cpu() + if isinstance(query_start_loc_cpu, torch.Tensor): + return query_start_loc_cpu.detach().cpu().clone() + if isinstance(query_start_loc_cpu, np.ndarray): + return query_start_loc_cpu.copy() + if isinstance(query_start_loc_cpu, list): + return list(query_start_loc_cpu) + return query_start_loc_cpu + + @staticmethod + def _snapshot_scheduler_output_for_async_omni_output( + scheduler_output: SchedulerOutput, + ) -> SchedulerOutput: + updates: dict[str, Any] = {} + for attr in ("num_scheduled_tokens", "scheduled_spec_decode_tokens"): + val = getattr(scheduler_output, attr, None) + if isinstance(val, dict): + updates[attr] = val.copy() + elif isinstance(val, list): + updates[attr] = list(val) + if not updates: + return scheduler_output + try: + return replace(scheduler_output, **updates) + except TypeError: + return scheduler_output + + def _should_return_omni_routed_experts(self) -> bool: + model_config = getattr(self, "model_config", None) + if model_config is None: + model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) + return bool(getattr(model_config, "enable_return_routed_experts", False)) and bool( + getattr(self, "routed_experts_initialized", False) + ) + + @staticmethod + def _model_omni_flag(model: Any, name: str, default: bool = False) -> bool: + return bool(getattr(model, name, default)) if model is not None else default + + def _runner_model_omni_flag(self, name: str, default: bool = False) -> bool: + return self._model_omni_flag(getattr(self, "model", None), name, default) + + def _model_omni_pooler_payload_include_hidden(self) -> bool: + return self._runner_model_omni_flag("omni_pooler_payload_include_hidden", default=True) + + def _should_use_async_omni_output(self) -> bool: + if not self.use_async_scheduling: + return False + if self.omni_prefix_cache is not None: + return False + if self.speculative_config is not None: + return False + + model_config = getattr(self, "model_config", None) + if model_config is None: + model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) + if not bool(getattr(model_config, "async_chunk", False)): + return False + if bool(getattr(model_config, "enable_return_routed_experts", False)): + return False + + model = getattr(self, "model", None) + if not self._model_omni_flag(model, "use_async_omni_output"): + return False + if self._model_omni_flag(model, "has_postprocess") and not self._model_omni_flag( + model, "eager_omni_postprocess_before_async_output" + ): + return False + + return True + + def _build_omni_async_snapshot_payload( + self, + *, + hidden_states: torch.Tensor, + staged_hidden_states_cpu: torch.Tensor | None, + multimodal_outputs: Any, + ) -> dict[str, Any]: + payload: dict[str, Any] = {"multimodal_outputs": multimodal_outputs} + if self._model_omni_pooler_payload_include_hidden(): + payload["hidden_states"] = hidden_states + payload["staged_hidden_states_cpu"] = staged_hidden_states_cpu + return payload + + def _snapshot_omni_output_tensors_for_async_output( + self, + *, + use_async_omni_output: bool, + hidden_states: torch.Tensor, + staged_hidden_states_cpu: torch.Tensor | None, + multimodal_outputs: Any, + ) -> _OmniOutputTensorSnapshot: + if not use_async_omni_output: + return _OmniOutputTensorSnapshot( + hidden_states=hidden_states, + staged_hidden_states_cpu=staged_hidden_states_cpu, + multimodal_outputs=multimodal_outputs, + ) + + with record_function_or_nullcontext("omni_async_output:snapshot_cpu_payload"): + async_payload_snapshot = _snapshot_tensor_payload_to_cpu_async( + self._build_omni_async_snapshot_payload( + hidden_states=hidden_states, + staged_hidden_states_cpu=staged_hidden_states_cpu, + multimodal_outputs=multimodal_outputs, + ), + copy_stream=self._get_or_create_omni_payload_copy_stream(), + # NOTE: vLLM v0.24.0's GPUModelRunner no longer exposes a + # ``self.pin_memory`` attribute (it uses a module-level + # ``PIN_MEMORY`` constant instead), so the old + # ``getattr(self, "pin_memory", False)`` silently fell back to + # False. That allocated the async D2H snapshot destination in + # *pageable* host memory, which turns ``copy_(non_blocking=True)`` + # into a fully synchronous, stream-stalling copy (~240 ms/step + # on the 17.5k-token Thinker prefill). Resolve pinning from the + # platform helper so the copy is a true async cudaMemcpyAsync. + pin_memory=is_pin_memory_available(), + ) + + payload = async_payload_snapshot.payload + hidden_states_snapshot = payload.get("hidden_states") + if hidden_states_snapshot is None: + # Models that omit hidden from the async snapshot only need + # multimodal payloads (for example, talker codes.audio). + hidden_states_snapshot = hidden_states[:0] + + return _OmniOutputTensorSnapshot( + hidden_states=hidden_states_snapshot, + staged_hidden_states_cpu=payload.get("staged_hidden_states_cpu"), + multimodal_outputs=payload["multimodal_outputs"], + async_payload=async_payload_snapshot, + ) + + def _maybe_run_eager_omni_postprocess_before_async_output( + self, + *, + hidden_states: torch.Tensor, + multimodal_outputs: Any, + num_scheduled_tokens_np: np.ndarray, + scheduler_output: SchedulerOutput, + req_ids_output_copy: list[str], + query_start_loc_cpu: Any, + ) -> bool: + """Apply model postprocess on live GPU tensors before async payload D2H.""" + model = getattr(self, "model", None) + if not self._model_omni_flag(model, "has_postprocess"): + return False + if not self._model_omni_flag(model, "eager_omni_postprocess_before_async_output"): + return False + + _, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) + if not downstream_req_ids: + return False + + with record_function_or_nullcontext("omni_output_builder:eager_postprocess"): + self._process_additional_information_updates( + hidden_states, + multimodal_outputs, + num_scheduled_tokens_np, + scheduler_output, + None, + None, + req_ids_filter=set(downstream_req_ids), + req_ids=req_ids_output_copy, + query_start_loc_cpu=query_start_loc_cpu, + ) + return True + + def _get_or_create_omni_payload_copy_stream(self) -> torch.cuda.Stream: + stream = getattr(self, "_omni_payload_copy_stream", None) + if stream is None: + stream = torch.cuda.Stream() + self._omni_payload_copy_stream = stream + return stream + + def _build_omni_model_runner_output_from_snapshot( + self, + *, + scheduler_output: SchedulerOutput, + hidden_states: torch.Tensor, + staged_hidden_states_cpu: torch.Tensor | None, + multimodal_outputs: Any, + req_ids_output_copy: list[str], + req_id_to_index_output_copy: dict[str, int], + valid_sampled_token_ids: list[list[int]], + logprobs_lists: Any, + prompt_logprobs_dict: dict[str, Any], + num_nans_in_logits: Any, + kv_connector_output: Any, + ec_connector_output: Any, + cudagraph_stats: Any, + kv_extracted_req_ids: list[str] | None, + num_scheduled_tokens_np: np.ndarray, + query_start_loc_cpu: Any, + postprocess_already_applied: bool = False, + ) -> OmniModelRunnerOutput: + combined_hidden_states = None + combined_multimodal_outputs = None + + engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) + downstream_req_ids, sparse_mm_index, audio_sparse_output = self._resolve_sparse_mm_routing( + engine_output_type=engine_output_type, + req_ids_output_copy=req_ids_output_copy, + downstream_req_ids=downstream_req_ids, + multimodal_outputs=multimodal_outputs, + ) + + needs_pooler_payload = len(downstream_req_ids) > 0 + downstream_req_id_set = set(downstream_req_ids) + hidden_states_cpu = None + req_hidden_states_cpu: dict[str, torch.Tensor] | None = None + include_hidden_payload = self._model_omni_pooler_payload_include_hidden() + needs_scheduled_hidden_payload = ( + include_hidden_payload + and needs_pooler_payload + and (self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states()) + ) + self._stage_deferred_prefix_cache_mm_outputs( + scheduler_output=scheduler_output, + multimodal_outputs=multimodal_outputs, + query_start_loc_cpu=query_start_loc_cpu, + ) + + if self.omni_prefix_cache is None and needs_scheduled_hidden_payload and not audio_sparse_output: + num_valid_tokens = min( + int(scheduler_output.total_num_scheduled_tokens), + int(hidden_states.shape[0]), + ) + if len(downstream_req_ids) == len(req_ids_output_copy): + with record_function_or_nullcontext("omni_output_builder:hidden_d2h/scheduled"): + hidden_states_cpu = _to_cpu_contiguous(hidden_states[:num_valid_tokens]) + else: + req_hidden_states_cpu = {} + with record_function_or_nullcontext("omni_output_builder:hidden_d2h/per_request"): + for rid in downstream_req_ids: + idx = req_id_to_index_output_copy[rid] + start = int(query_start_loc_cpu[idx]) + sched = int(num_scheduled_tokens_np[idx]) + end = start + sched + req_hidden_states_cpu[rid] = _to_cpu_contiguous(hidden_states[start:end]) + + # NOTE: pooler_output here is used only for the full-payload accumulation + # path (accumulate_full_payload_output) and is NOT passed on the wire via + # OmniModelRunnerOutput.pooler_output (which is set to None below). + # The actual multimodal wire transport uses multimodal_outputs instead. + pooler_output: list[dict[str, object]] | None = None + if needs_pooler_payload: + hidden_seq_len = int(hidden_states.shape[0]) + scheduled_seq_len = int(scheduler_output.total_num_scheduled_tokens) + mm_cpu = None + if self.omni_prefix_cache is not None: + ( + hidden_states_cpu, + combined_hidden_states, + combined_multimodal_outputs, + ) = self._prepare_prefix_cache_pooler_payload_sources( + hidden_states=hidden_states, + staged_hidden_states_cpu=staged_hidden_states_cpu, + multimodal_outputs=multimodal_outputs, + scheduler_output=scheduler_output, + needs_scheduled_hidden_payload=needs_scheduled_hidden_payload, + ) + if combined_multimodal_outputs is None: + with record_function_or_nullcontext("omni_output_builder:build_mm_cpu"): + mm_cpu = build_mm_cpu( + flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs + ) + + with record_function_or_nullcontext("omni_output_builder:process_additional_information"): + if not postprocess_already_applied: + self._process_additional_information_updates( + hidden_states, + multimodal_outputs, + num_scheduled_tokens_np, + scheduler_output, + combined_hidden_states, + combined_multimodal_outputs, + req_ids_filter=downstream_req_id_set, + req_ids=req_ids_output_copy, + query_start_loc_cpu=query_start_loc_cpu, + ) + + pooler_output = [] + with record_function_or_nullcontext("omni_output_builder:build_pooler_payloads"): + for rid in req_ids_output_copy: + if rid not in downstream_req_id_set: + pooler_output.append({}) + continue + idx = req_id_to_index_output_copy[rid] + start = int(query_start_loc_cpu[idx]) + sched = int(num_scheduled_tokens_np[idx]) + end = start + sched + payload = self._build_omni_pooler_payload( + rid=rid, + idx=idx, + start=start, + end=end, + hidden_states_cpu=hidden_states_cpu, + req_hidden_states_cpu=req_hidden_states_cpu, + combined_hidden_states=combined_hidden_states, + combined_multimodal_outputs=combined_multimodal_outputs, + mm_cpu=mm_cpu, + audio_sparse_output=audio_sparse_output, + sparse_mm_index=sparse_mm_index, + hidden_seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, + ) + pooler_output.append(flatten_payload(payload)) + + pooler_output = pooler_output or [] + if self._async_chunk: + pooler_inter, pooler_client = partition_payload_list(pooler_output) + else: + # Non-async-chunk still ships the full payload to the next stage (via + # accumulate_full_payload_output and the inter_stage_outputs field); only + # client mm keys are split out when async_chunk is enabled. #4527 set this + # to (None, pooler_output), which skipped accumulation and starved the + # downstream stage (300s connector-input timeout / empty audio). (PR #4792) + pooler_inter, pooler_client = pooler_output, pooler_output + + if pooler_inter and self._should_accumulate_full_payload_output(): + with record_function_or_nullcontext("omni_output_builder:accumulate_full_payload_output"): + for i, rid in enumerate(req_ids_output_copy): + req_state = self.requests.get(rid) + if req_state is not None and pooler_inter[i]: + self.accumulate_full_payload_output(rid, pooler_inter[i], req_state) + + with record_function_or_nullcontext("omni_output_builder:build_multimodal_outputs"): + inter_stage_outputs = self._build_multimodal_outputs(pooler_inter) + multimodal_outputs = self._build_multimodal_outputs(pooler_client) + + with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + routed_experts_lists = None + if self._should_return_omni_routed_experts(): + routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) + output = OmniModelRunnerOutput( + req_ids=req_ids_output_copy, + req_id_to_index=req_id_to_index_output_copy, + sampled_token_ids=valid_sampled_token_ids, + logprobs=logprobs_lists, + prompt_logprobs_dict=prompt_logprobs_dict, + pooler_output=None, + multimodal_outputs=multimodal_outputs, + inter_stage_outputs=inter_stage_outputs, + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, + num_nans_in_logits=num_nans_in_logits, + cudagraph_stats=cudagraph_stats, + ) + output.kv_extracted_req_ids = kv_extracted_req_ids + with record_function_or_nullcontext("omni_output_builder:get_omni_connector_output"): + output.omni_connector_output = self.get_omni_connector_output() + output.routed_experts = routed_experts_lists + return output @torch.inference_mode() def sample_tokens( @@ -974,13 +1858,6 @@ def sample_tokens( kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None) self.kv_extracted_req_ids = None - # Used for prefix cache - combined_hidden_states = None - combined_multimodal_outputs = None - # Used when we don't use prefix cache; prefix cache builds the payloads - # internally since it already needs to do this for the cached tensors - mm_cpu = {} - if self.execute_model_state is None: kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -1009,7 +1886,6 @@ def sample_tokens( slot_mappings, # OMNI: unpack slot_mappings for drafter ) = self.execute_model_state self.execute_model_state = None - seq_len = hidden_states.shape[0] # Apply structured output bitmasks if present. if grammar_output is not None: @@ -1023,6 +1899,15 @@ def sample_tokens( if self.input_batch.vocab_size > logits_vocab: smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + # Drop min-tokens stop ids the head cannot emit (e.g. the text + # tokenizer EOS folded into all_stop_token_ids on a narrow codec + # talker head); they would index_put_ out of bounds (#4962). + if logits is not None: + sanitize_min_tokens_stop_ids( + self.input_batch.sampling_metadata.logitsprocs, + logits.shape[-1], + ) + with record_function_or_nullcontext("gpu_model_runner: sample"): sampler_output = self._sample(logits, spec_decode_metadata) @@ -1052,10 +1937,7 @@ def propose_draft_token_ids(sampled_token_ids): spec_config = self.speculative_config propose_drafts_after_bookkeeping = False if spec_config is not None: - input_fits_in_drafter = spec_decode_common_attn_metadata is not None and ( - spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens - <= self.effective_drafter_max_model_len - ) + input_fits_in_drafter = self._input_fits_in_drafter(spec_decode_common_attn_metadata) use_gpu_toks = ( spec_config.use_eagle() or spec_config.uses_draft_model() or spec_config.uses_extract_hidden_states() ) and not spec_config.disable_padded_drafter_batch @@ -1121,204 +2003,94 @@ def propose_draft_token_ids(sampled_token_ids): kv_connector_output = self.kv_connector_output self.kv_connector_output = None - engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) - sparse_mm_req_ids = self._sparse_mm_req_ids(multimodal_outputs) - sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} - if engine_output_type == "audio" and sparse_mm_req_ids is not None: - sparse_req_id_set = set(sparse_mm_req_ids) - downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] - needs_pooler_payload = len(downstream_req_ids) > 0 - downstream_req_id_set = set(downstream_req_ids) - hidden_states_cpu = None - req_hidden_states_cpu: dict[str, torch.Tensor] | None = None - audio_sparse_output = engine_output_type == "audio" and sparse_mm_req_ids is not None - needs_scheduled_hidden_payload = needs_pooler_payload and ( - self.omni_prefix_cache is None or not self._model_needs_full_prefix_hidden_states() - ) - if needs_scheduled_hidden_payload and self.omni_prefix_cache is not None: - if staged_hidden_states_cpu is None: - raise RuntimeError("Prefix-cache hidden-state payload requires staged CPU hidden states.") - hidden_states_cpu = staged_hidden_states_cpu - elif needs_scheduled_hidden_payload: - num_valid_tokens = min( - int(scheduler_output.total_num_scheduled_tokens), - int(hidden_states.shape[0]), - ) - if audio_sparse_output: - pass - elif len(downstream_req_ids) == len(req_ids_output_copy): - hidden_states_cpu = hidden_states[:num_valid_tokens].detach().to("cpu").contiguous() - else: - req_hidden_states_cpu = {} num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) if num_scheduled_tokens_np is None: - req_ids = self.input_batch.req_ids num_scheduled_tokens_np = np.array( - [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], + [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids_output_copy], dtype=np.int32, ) - query_start_loc_cpu = self.query_start_loc.cpu - if callable(query_start_loc_cpu): - query_start_loc_cpu = query_start_loc_cpu() - if self.omni_prefix_cache is not None: - deferred_mm_cache_keys = self._deferred_prefix_cache_mm_keys() - if deferred_mm_cache_keys: - self.omni_prefix_cache.stage_deferred_mm_outputs( - query_start_loc=query_start_loc_cpu, - input_batch=self.input_batch, - multimodal_outputs=( - flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs - ), - num_scheduled_tokens=scheduler_output.num_scheduled_tokens, - deferred_mm_cache_keys=deferred_mm_cache_keys, - ) - - # NOTE: pooler_output here is used only for the full-payload accumulation - # path (accumulate_full_payload_output) and is NOT passed on the wire via - # OmniModelRunnerOutput.pooler_output (which is set to None below). - # The actual multimodal wire transport uses multimodal_outputs instead. - pooler_output: list[dict[str, object]] | None = None - if needs_pooler_payload: - mm_cpu = None - if self.omni_prefix_cache is not None: - ( - combined_hidden_states, - combined_multimodal_outputs, - ) = self._maybe_get_combined_prefix_cache_tensors( - hidden_states, - staged_hidden_states_cpu, - multimodal_outputs, - scheduler_output.num_scheduled_tokens, - ) - if self.omni_prefix_cache is None or combined_multimodal_outputs is None: - mm_cpu = build_mm_cpu(flatten_payload(multimodal_outputs) if multimodal_outputs else multimodal_outputs) + else: + num_scheduled_tokens_np = np.asarray(num_scheduled_tokens_np, dtype=np.int32).copy() + + query_start_loc_cpu = self._snapshot_query_start_loc_cpu() + scheduler_output_snapshot = self._snapshot_scheduler_output_for_async_omni_output(scheduler_output) + req_ids_output_snapshot = list(req_ids_output_copy) + req_id_to_index_output_snapshot = dict(req_id_to_index_output_copy) + valid_sampled_token_ids_snapshot = [list(token_ids) for token_ids in valid_sampled_token_ids] + logprobs_lists_snapshot = copy(logprobs_lists) if logprobs_lists is not None else None + prompt_logprobs_dict_snapshot = dict(prompt_logprobs_dict) if prompt_logprobs_dict is not None else {} + num_nans_in_logits_snapshot = ( + dict(num_nans_in_logits) if isinstance(num_nans_in_logits, dict) else num_nans_in_logits + ) - self._process_additional_information_updates( - hidden_states, - multimodal_outputs, - num_scheduled_tokens_np, - scheduler_output, - combined_hidden_states, - combined_multimodal_outputs, - req_ids_filter=downstream_req_id_set, + use_async_omni_output = self._should_use_async_omni_output() + omni_postprocess_already_applied = False + if use_async_omni_output: + omni_postprocess_already_applied = self._maybe_run_eager_omni_postprocess_before_async_output( + hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + num_scheduled_tokens_np=num_scheduled_tokens_np, + scheduler_output=scheduler_output, + req_ids_output_copy=req_ids_output_copy, + query_start_loc_cpu=query_start_loc_cpu, ) + output_tensor_snapshot = self._snapshot_omni_output_tensors_for_async_output( + use_async_omni_output=use_async_omni_output, + hidden_states=hidden_states, + staged_hidden_states_cpu=staged_hidden_states_cpu, + multimodal_outputs=multimodal_outputs, + ) - if req_hidden_states_cpu is not None and combined_hidden_states is None: - for rid in downstream_req_ids: - idx = req_id_to_index_output_copy[rid] - start = int(query_start_loc_cpu[idx]) - sched = int(num_scheduled_tokens_np[idx]) - end = start + sched - req_hidden_states_cpu[rid] = hidden_states[start:end].detach().to("cpu").contiguous() + def output_builder() -> OmniModelRunnerOutput: + if output_tensor_snapshot.async_payload is not None: + with record_function_or_nullcontext("omni_async_output:wait_cpu_payload"): + output_tensor_snapshot.async_payload.wait() + with record_function_or_nullcontext("omni_output_builder:total"): + return self._build_omni_model_runner_output_from_snapshot( + scheduler_output=scheduler_output_snapshot, + hidden_states=output_tensor_snapshot.hidden_states, + staged_hidden_states_cpu=output_tensor_snapshot.staged_hidden_states_cpu, + multimodal_outputs=output_tensor_snapshot.multimodal_outputs, + req_ids_output_copy=req_ids_output_snapshot, + req_id_to_index_output_copy=req_id_to_index_output_snapshot, + valid_sampled_token_ids=valid_sampled_token_ids_snapshot, + logprobs_lists=logprobs_lists_snapshot, + prompt_logprobs_dict=prompt_logprobs_dict_snapshot, + num_nans_in_logits=num_nans_in_logits_snapshot, + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output, + cudagraph_stats=cudagraph_stats, + kv_extracted_req_ids=kv_extracted_req_ids, + num_scheduled_tokens_np=num_scheduled_tokens_np, + query_start_loc_cpu=query_start_loc_cpu, + postprocess_already_applied=omni_postprocess_already_applied, + ) - pooler_output = [] - for rid in req_ids_output_copy: - if rid not in downstream_req_id_set: - pooler_output.append({}) - continue - idx = req_id_to_index_output_copy[rid] - start = int(query_start_loc_cpu[idx]) - sched = int(num_scheduled_tokens_np[idx]) - end = start + sched - payload: dict[str, object] = {} - if not audio_sparse_output: - if req_hidden_states_cpu is not None and combined_hidden_states is None: - req_hidden_states = req_hidden_states_cpu[rid] - else: - req_hidden_states = self._resolve_req_hidden_states( - hidden_states_cpu, - combined_hidden_states, - rid, - start, - end, - ) - if req_hidden_states is not None: - payload["hidden"] = req_hidden_states - - mm_payload: dict[str, object] = {} - if combined_multimodal_outputs or mm_cpu: - if combined_multimodal_outputs: - - def _unwrap_lists(v): - if isinstance(v, list): - return v[idx] if idx < len(v) else v[0] - if isinstance(v, dict): - return {k: _unwrap_lists(sv) for k, sv in v.items()} - return v - - for mm_key in combined_multimodal_outputs.keys(): - mm_payload[mm_key] = _unwrap_lists(combined_multimodal_outputs[mm_key][rid]) - else: - for mm_key, mm_val in mm_cpu.items(): - if mm_key in {"meta.req_id", "meta.sparse_audio"}: - continue - if audio_sparse_output and isinstance(mm_val, list): - sparse_idx = sparse_mm_index.get(rid) - if sparse_idx is None: - continue - if sparse_idx >= len(mm_val): - logger.warning( - "Sparse multimodal payload mismatch for request %s: index %d >= %d.", - rid, - sparse_idx, - len(mm_val), - ) - continue - sparse_val = mm_val[sparse_idx] - mm_payload[mm_key] = ( - sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val - ) - continue - mm_payload[mm_key] = to_payload_element( - element=mm_val, - idx=idx, - start=start, - end=end, - pass_lists_through=False, - seq_len=seq_len, - ) - payload.update(mm_payload) - pooler_output.append(flatten_payload(payload)) - - if pooler_output and self._should_accumulate_full_payload_output(): - for i, rid in enumerate(req_ids_output_copy): - req_state = self.requests.get(rid) - if req_state is not None and pooler_output[i]: - self.accumulate_full_payload_output(rid, pooler_output[i], req_state) - - multimodal_outputs = self._build_multimodal_outputs(pooler_output) - with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): - routed_experts_lists = None - if self.routed_experts_initialized: - routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) - output = OmniModelRunnerOutput( - req_ids=req_ids_output_copy, - req_id_to_index=req_id_to_index_output_copy, - sampled_token_ids=valid_sampled_token_ids, - logprobs=logprobs_lists, - prompt_logprobs_dict=prompt_logprobs_dict, - pooler_output=None, - multimodal_outputs=multimodal_outputs, - kv_connector_output=kv_connector_output, - ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, - num_nans_in_logits=num_nans_in_logits, - cudagraph_stats=cudagraph_stats, - ) - output.kv_extracted_req_ids = kv_extracted_req_ids - output.omni_connector_output = self.get_omni_connector_output() - output.routed_experts = routed_experts_lists + if not use_async_omni_output: + output = output_builder() - if not self.use_async_scheduling: - return output + if not self.use_async_scheduling: + return output with record_function_or_nullcontext("gpu_model_runner: AsyncGPUModelRunnerOutput"): - async_output = AsyncGPUModelRunnerOutput( - model_runner_output=output, + async_output_cls = OmniAsyncGPUModelRunnerOutput if use_async_omni_output else AsyncGPUModelRunnerOutput + async_output_kwargs = dict( sampled_token_ids=sampler_output.sampled_token_ids, logprobs_tensors=sampler_output.logprobs_tensors, invalid_req_indices=invalid_req_indices, async_output_copy_stream=self.async_output_copy_stream, vocab_size=self.input_batch.vocab_size, ) + if use_async_omni_output: + async_output = async_output_cls( + model_runner_output_builder=output_builder, + cuda_device=self.device, + **async_output_kwargs, + ) + else: + async_output = async_output_cls( + model_runner_output=output, + **async_output_kwargs, + ) with record_function_or_nullcontext("gpu_model_runner: set_async_sampled_token_ids"): # Save ref of sampled_token_ids CPU tensor if the batch contains # any requests with sampling params that require output ids. diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 08551ad6948..2afa595763c 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -12,6 +12,7 @@ from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -49,16 +50,32 @@ def init_device(self): # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK self.local_rank += dp_local_rank * tp_pp_world_size - assert self.local_rank < torch.accelerator.device_count(), ( - f"DP adjusted local rank {self.local_rank} is out of bounds. " + + # Publish the logical-to-physical mapping for topology queries + # such as NIC affinity and P2P checks (upstream PR #45026). + assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids + if assigned_physical_gpu_ids is not None: + from vllm.platforms.interface import set_assigned_physical_gpu_ids + + set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) + assert self.local_rank < len(assigned_physical_gpu_ids), ( + f"local_rank {self.local_rank} is out of bounds for " + f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" ) - visible_device_count = torch.accelerator.device_count() - assert self.parallel_config.local_world_size <= visible_device_count, ( - f"local_world_size ({self.parallel_config.local_world_size}) must " - f"be less than or equal to the number of visible devices " - f"({visible_device_count})." + if parallel_config.distributed_executor_backend not in ("ray", "external_launcher"): + assert self.parallel_config.local_world_size <= len(assigned_physical_gpu_ids), ( + f"local_world_size ({self.parallel_config.local_world_size}) " + "exceeds assigned_physical_gpu_ids count " + f"({len(assigned_physical_gpu_ids)})" + ) + else: + assert self.local_rank < torch.accelerator.device_count(), ( + f"DP adjusted local rank {self.local_rank} is out of " + f"bounds for {torch.accelerator.device_count()} devices." ) - self.device = torch.device(f"cuda:{self.local_rank}") + + visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank) + self.device = current_omni_platform.get_torch_device(visible_device_index) torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index bc3d6ea2d82..ec41900b145 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -37,6 +37,7 @@ from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.utils.mm_outputs import partition_payload_list from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState, _ensure_tensor_values from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin @@ -54,6 +55,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self._async_chunk = getattr(self.model_config, "async_chunk", False) # Mirrors the init allowlist in gpu_ar_model_runner.py. _OMNI_CONNECTOR_INIT_ARCHS = { "Qwen3OmniMoeForConditionalGeneration", @@ -64,6 +66,7 @@ def __init__(self, *args, **kwargs): "Qwen3TTSCode2Wav", "CosyVoice3Model", "DyninOmniForConditionalGeneration", + "IndexTTS2S2MelDecoder", } if getattr(self.model_config, "model_arch", None) in _OMNI_CONNECTOR_INIT_ARCHS: self.init_omni_connectors( @@ -150,6 +153,13 @@ def execute_model( if self.model_config.async_chunk and num_scheduled_tokens: self._update_request_states(scheduler_output) deferred_state_corrections_fn = self._update_states(scheduler_output) + + # Notify stateful models of finished requests before the + # zero-token early return. Request cleanup is a model lifecycle + # concern and does not depend on the inter-stage transfer mode. + if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): + self.model.on_requests_finished(scheduler_output.finished_req_ids) + if not scheduler_output.total_num_scheduled_tokens: return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) @@ -417,20 +427,20 @@ def sample_tokens( # Build per-request multimodal_outputs list (dedicated channel). # pooler_output is no longer used for multimodal data. - multimodal_outputs: list[dict[str, object]] = [] + per_req_payloads: list[dict[str, object]] = [] if isinstance(multimodal_outputs_raw, torch.Tensor): assert multimodal_outputs_raw.shape[0] == 1, ( "model should return a single tensor, to return multiple tensors, use a dict" ) assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs for i in range(self.input_batch.num_reqs): - multimodal_outputs.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) + per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs_raw, list): assert len(multimodal_outputs_raw) == 1, ( "model should return a single list, to return multiple lists, use a dict" ) for out in multimodal_outputs_raw: - multimodal_outputs.append( + per_req_payloads.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} ) elif isinstance(multimodal_outputs_raw, Mapping): @@ -449,20 +459,28 @@ def sample_tokens( mm_payload[key] = out.detach().to("cpu").contiguous() else: logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}") - multimodal_outputs.append(_ensure_tensor_values(mm_payload)) + per_req_payloads.append(_ensure_tensor_values(mm_payload)) else: raise RuntimeError("Unsupported diffusion output type") + + if self._async_chunk: + inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads) + else: + # See gpu_ar_model_runner: non-async-chunk ships the full payload to the next + # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792) + inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads + # [Omni] Copy req_id mappings to avoid async scheduling mutation. req_ids_output_copy = self.input_batch.req_ids.copy() req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy() routed_experts_lists = None if self.routed_experts_initialized: routed_experts_lists = self._omni_extract_routed_experts(scheduler_output) - if self._should_accumulate_full_payload_output(): + if inter_stage_outputs and self._should_accumulate_full_payload_output(): for i, rid in enumerate(req_ids_output_copy): req_state = self.requests.get(rid) - if req_state is not None and multimodal_outputs[i]: - self.accumulate_full_payload_output(rid, multimodal_outputs[i], req_state) + if req_state is not None and inter_stage_outputs[i]: + self.accumulate_full_payload_output(rid, inter_stage_outputs[i], req_state) output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, @@ -472,6 +490,7 @@ def sample_tokens( prompt_logprobs_dict={}, pooler_output=None, multimodal_outputs=multimodal_outputs, + inter_stage_outputs=inter_stage_outputs, kv_connector_output=kv_connector_output, num_nans_in_logits={}, cudagraph_stats=cudagraph_stats, @@ -777,6 +796,10 @@ def _dummy_run( if hasattr(self.model, "get_dummy_runtime_additional_information"): runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs) model_kwargs["runtime_additional_information"] = runtime_addi + # Generation-stage code2wav models split flat input_ids using this + # shared contract. Dummy/profile runs use padded synthetic inputs, + # so expose one synthetic segment covering the actual tensor length. + model_kwargs["seq_token_counts"] = [int(num_tokens_padded)] if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_tokens_padded] diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 2a85436a007..38a88c84c05 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -11,6 +11,7 @@ from vllm.v1.worker.gpu_worker import init_worker_distributed_environment from vllm.v1.worker.workspace import init_workspace_manager +from vllm_omni.platforms import current_omni_platform from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.memory_utils import request_memory_tolerant @@ -57,7 +58,7 @@ def init_device(self): f"be less than or equal to the number of visible devices " f"({visible_device_count})." ) - self.device = torch.device(f"cuda:{self.local_rank}") + self.device = current_omni_platform.get_torch_device(self.local_rank) torch.accelerator.set_device_index(self.device) current_platform.check_if_supports_dtype(self.model_config.dtype) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 349e1da0411..b816a36f734 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,4 +1,5 @@ import contextlib +import inspect from collections.abc import Callable from dataclasses import replace from typing import TYPE_CHECKING, Any, cast @@ -48,6 +49,30 @@ logger = init_logger(__name__) +def _filter_mrope_kwargs_for_model(model: object, kwargs: dict[str, Any]) -> dict[str, Any]: + """Return only M-RoPE kwargs accepted by the model implementation.""" + method = getattr(model, "get_mrope_input_positions") + try: + signature = inspect.signature(method) + except (TypeError, ValueError): + return kwargs + + if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in signature.parameters.values()): + return kwargs + + accepted = { + name + for name, param in signature.parameters.items() + if name not in {"self", "input_tokens"} + and param.kind + in { + inspect.Parameter.POSITIONAL_OR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + } + } + return {key: value for key, value in kwargs.items() if key in accepted} + + class OmniGPUModelRunner(GPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -293,6 +318,12 @@ def _init_mrope_positions(self, req_state: CachedRequestState): Raises: AssertionError: If the model does not support M-RoPE + + Note: + Upstream vLLM (commit 470229c37) added a fallback for the case + where ``req_state.prompt_token_ids`` is None but + ``req_state.prompt_embeds`` is available. Omni models always set + ``prompt_token_ids``, so this fallback is deliberately omitted. """ image_grid_thw = [] video_grid_thw = [] @@ -338,7 +369,7 @@ def _init_mrope_positions(self, req_state: CachedRequestState): kwargs["target_w"] = target_w req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions( req_state.prompt_token_ids, - **kwargs, + **_filter_mrope_kwargs_for_model(self.model, kwargs), ) else: req_state.mrope_positions, req_state.mrope_position_delta = MRotaryEmbedding.get_input_positions_tensor( @@ -570,9 +601,6 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None # Decode additional_information payloads (dictionary) try: if getattr(new_req_data, "additional_information", None) is not None: - logger.warning_once( - "additional_information on request data is deprecated, use model_intermediate_buffer" - ) info_dict = deserialize_additional_information(new_req_data.additional_information) if info_dict: self.model_intermediate_buffer[req_id] = info_dict @@ -1313,6 +1341,20 @@ def _build_model_kwargs_extra(self) -> dict: traceback.print_exc() + # Per-request (start, end) hidden-row spans so make_omni_output can map + # flat hidden rows to the right request in mixed prefill+decode steps, + # instead of assuming an equal rows-per-request split (which samples the + # wrong rows whenever per-request token counts differ). + nstp = self._omni_num_scheduled_tokens_np + if nstp is not None and len(nstp) == len(self.input_batch.req_ids): + try: + model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(nstp) + except Exception as e: + # Visible on purpose: the fallback is the equal rows-per-request + # split, which can re-introduce the cross-request corruption this + # plumbing fixes — a silent failure here must not pass unnoticed. + logger.warning("[OMNI] Failed to compute request_token_spans: %s", e) + if self._omni_query_start_loc_model_kwarg: try: num_reqs = len(self.input_batch.req_ids) @@ -1339,8 +1381,15 @@ def _process_additional_information_updates( combined_hidden_states: dict[str, torch.Tensor] | None = None, combined_multimodal_outputs: dict[str, object] | None = None, req_ids_filter: set[str] | None = None, + req_ids: list[str] | None = None, + query_start_loc_cpu: object | None = None, ) -> None: """Process model-provided per-request updates and merge into model_intermediate_buffer.""" + req_ids = req_ids if req_ids is not None else self.input_batch.req_ids + if query_start_loc_cpu is None: + query_start_loc_cpu = self.query_start_loc.cpu + if callable(query_start_loc_cpu): + query_start_loc_cpu = query_start_loc_cpu() try: # execute the custom postprocess function # TODO(Peiqi): do we have a more elegant way to do this? @@ -1348,7 +1397,7 @@ def _process_additional_information_updates( postprocess_uses_hidden_states = getattr(self.model, "postprocess_uses_hidden_states", True) postprocess_uses_multimodal_outputs = getattr(self.model, "postprocess_uses_multimodal_outputs", True) postprocess_uses_req_infos = getattr(self.model, "postprocess_uses_req_infos", True) - for req_index, req_id in enumerate(self.input_batch.req_ids): + for req_index, req_id in enumerate(req_ids): if req_ids_filter is not None and req_id not in req_ids_filter: continue req_infos = self.model_intermediate_buffer.get(req_id, {}) if postprocess_uses_req_infos else {} @@ -1357,7 +1406,7 @@ def _process_additional_information_updates( # Combined hidden states contains all hidden states for every request hidden_states_slice = combined_hidden_states[req_id] else: - start_offset = int(self.query_start_loc.cpu[req_index]) + start_offset = int(query_start_loc_cpu[req_index]) sched_tokens = int(num_scheduled_tokens_np[req_index]) s, e = start_offset, start_offset + sched_tokens # only consider to store data into update dict. @@ -1387,7 +1436,7 @@ def _process_additional_information_updates( ) self._update_intermediate_buffer(req_id, update_dict) except Exception as e: - logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}") + logger.error(f"Error merging for requests:{req_ids} additional information update: {e}") import traceback traceback.print_exc() @@ -1419,15 +1468,9 @@ def _update_additional_information(self, scheduler_output: "SchedulerOutput") -> for new_req in scheduler_output.scheduled_new_reqs: payload_info = getattr(new_req, "additional_information", None) if isinstance(payload_info, dict): - logger.warning_once( - "additional_information on request data is deprecated, use model_intermediate_buffer" - ) self._update_intermediate_buffer(new_req.req_id, payload_info) if hasattr(scheduler_output.scheduled_cached_reqs, "additional_information"): - logger.warning_once( - "additional_information on scheduled_cached_reqs is deprecated, use model_intermediate_buffer" - ) cached_infos = getattr(scheduler_output.scheduled_cached_reqs, "additional_information", {}) if isinstance(cached_infos, dict): for req_id, req_infos in cached_infos.items(): @@ -1478,7 +1521,16 @@ def _preprocess( num_input_tokens: int, intermediate_tensors: IntermediateTensors | None = None, ): - """Align with v0.14.0 preprocess and omni's additional information handling.""" + """Align with v0.14.0 preprocess and omni's additional information handling. + + Note: + Upstream vLLM (commit c621af169) added a conditional in the + ``supports_mm_inputs`` path to handle precomputed ``prompt_embeds`` + alongside multimodal inputs. Omni models that reach this override + always go through the omni-specific ``model.preprocess`` / + ``has_preprocess`` code path below, so the upstream change is + deliberately not ported. + """ num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens is_first_rank = get_pp_group().is_first_rank is_encoder_decoder = self.model_config.is_encoder_decoder @@ -1646,7 +1698,7 @@ def flush_decode_batch() -> None: self.text_step.gpu[dst].copy_(text_step) for req_id_b, update_dict_b in zip(req_ids_b, updates, strict=True): - self._merge_additional_information_update(req_id_b, update_dict_b) + self._update_intermediate_buffer(req_id_b, update_dict_b) decode_req_ids.extend(req_ids_b) decode_start_offsets.extend(start_offsets_b) @@ -1701,7 +1753,7 @@ def flush_decode_batch() -> None: decode_start_offsets.append(s) # TODO(Peiqi): the merge stage could move out from the critical path - self._merge_additional_information_update(req_id, update_dict) + self._update_intermediate_buffer(req_id, update_dict) # update the inputs_embeds and input_ids seg_len = min(span_len, req_embeds.shape[0]) @@ -1760,10 +1812,38 @@ def _talker_mtp_forward( def _explicit_talker_seed(req_id: str) -> int | None: sampling_params = getattr(self.requests[req_id], "sampling_params", None) extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None - seed = extra_args.get("qwen3_tts_request_seed") if isinstance(extra_args, dict) else None + seed = None + if isinstance(extra_args, dict): + seed = extra_args.get("tts_local_seed") return int(seed) if seed is not None else None - if decode_batch_size > 1 and any(_explicit_talker_seed(req_id) is not None for req_id in decode_req_ids): + def _row_generator(req_id: str) -> torch.Generator | None: + seed = _explicit_talker_seed(req_id) + if seed is None: + return None + cache = getattr(self, "_talker_mtp_generators", None) + if cache is None: + cache = {} + self._talker_mtp_generators = cache + generator = cache.get(req_id) + if generator is None or generator.device != req_input_ids.device: + generator = torch.Generator(device=req_input_ids.device) + generator.manual_seed(seed) + cache[req_id] = generator + return generator + + row_generators = [_row_generator(req_id) for req_id in decode_req_ids] + cache = getattr(self, "_talker_mtp_generators", None) + if cache: + # Generators live as long as their request; drop finished ones. + for stale_id in [rid for rid in cache if rid not in self.requests]: + del cache[stale_id] + + if ( + decode_batch_size > 1 + and any(generator is not None for generator in row_generators) + and not getattr(self.model, "talker_mtp_accepts_per_row_generators", False) + ): # A torch.Generator is a single stream. Using one generator for a # multi-row batch would make explicitly-seeded requests depend on # other rows in the same scheduler step, so keep that path scalar. @@ -1786,28 +1866,22 @@ def _explicit_talker_seed(req_id: str) -> int | None: self.text_step.gpu[:decode_batch_size].copy_(saved_text) return - generator = None - if decode_req_ids: - first_req_id = decode_req_ids[0] - seed = _explicit_talker_seed(first_req_id) - if seed is not None: - generators = getattr(self, "_talker_mtp_generators", None) - if generators is None: - generators = {} - self._talker_mtp_generators = generators - generator = generators.get(first_req_id) - if generator is None or generator.device != req_input_ids.device: - generator = torch.Generator(device=req_input_ids.device) - generator.manual_seed(int(seed)) - generators[first_req_id] = generator talker_kwargs = { "do_sample": subtalker_params.get("do_sample"), "temperature": subtalker_params.get("temperature"), "top_k": subtalker_params.get("top_k"), "top_p": subtalker_params.get("top_p"), } - if generator is not None: - talker_kwargs["generator"] = generator + if decode_batch_size == 1: + if row_generators[0] is not None: + talker_kwargs["generator"] = row_generators[0] + elif any(generator is not None for generator in row_generators): + talker_kwargs["generators"] = row_generators + if getattr(self.model, "talker_mtp_accepts_req_infos", False): + talker_kwargs["req_ids"] = decode_req_ids + talker_kwargs["req_infos"] = [ + self.model_intermediate_buffer.setdefault(req_id, {}) for req_id in decode_req_ids + ] with current_omni_platform.set_forward_context( None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc ): @@ -1827,8 +1901,9 @@ def _explicit_talker_seed(req_id: str) -> int | None: start_offsets = [int(self.query_start_loc.cpu[id_to_index[req_id]]) for req_id in decode_req_ids] for idx, (req_id, start_offset) in enumerate(zip(decode_req_ids, start_offsets, strict=True)): inputs_embeds[start_offset : start_offset + 1] = req_embeds[idx : idx + 1] - update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} - self._merge_additional_information_update(req_id, update_dict) + if code_predictor_codes is not None: + update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}} + self._update_intermediate_buffer(req_id, update_dict) def _model_forward( self, @@ -1902,10 +1977,6 @@ def _update_intermediate_buffer(self, req_id: str, upd: dict) -> None: # Backward compatible: mirror to old setattr location setattr(req_state, "additional_information_cpu", existing) - def _merge_additional_information_update(self, req_id, upd): - logger.warning_once("_merge_additional_information_update is deprecated, use _update_intermediate_buffer") - return self._update_intermediate_buffer(req_id, upd) - def _update_streaming_input_additional_info(self, req_id): # For streaming input prefill case only. Set num processed tokens = 0 for new segment input cached_additional_info = self.model_intermediate_buffer.get(req_id, {}) diff --git a/vllm_omni/worker/omni_connector_model_runner_mixin.py b/vllm_omni/worker/omni_connector_model_runner_mixin.py index 205a28c6d02..b173f81acc1 100644 --- a/vllm_omni/worker/omni_connector_model_runner_mixin.py +++ b/vllm_omni/worker/omni_connector_model_runner_mixin.py @@ -110,7 +110,7 @@ def init_omni_connectors( self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config) self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg() - logger.info( + logger.debug( "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s", self._stage_id, self._async_chunk, @@ -1154,7 +1154,7 @@ def send_chunk( connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}" if chunk_id == 0: - logger.info( + logger.debug( "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s", self._stage_id, request_id, @@ -1252,7 +1252,7 @@ def receive_multi_kv_cache( active_requests = getattr(self, "requests", None) if active_requests is not None and request_id not in active_requests: - logger.info("Skip receiving KV cache for inactive request %s", request_id) + logger.debug("Skip receiving KV cache for inactive request %s", request_id) return False primary_ok = False @@ -1272,7 +1272,7 @@ def receive_multi_kv_cache( if cfg_kvs and hasattr(req, "sampling_params") and req.sampling_params is not None: for key, value in cfg_kvs.items(): setattr(req.sampling_params, key, value) - logger.info("Applied CFG KV caches: %s", list(cfg_kvs.keys())) + logger.debug("Applied CFG KV caches: %s", list(cfg_kvs.keys())) except Exception: logger.exception("Failed to collect CFG KV caches for %s", request_id) @@ -1657,7 +1657,7 @@ def _recv_loop(self) -> None: _recv_poll_count += 1 if _recv_poll_count % 5000 == 1: - logger.info( + logger.debug( "[Stage-%s] _recv_loop: polling %s pending reqs: %s (poll#%s)", self._stage_id, len(pending_ids), diff --git a/vllm_omni/worker/runner_assisted_metadata.py b/vllm_omni/worker/runner_assisted_metadata.py new file mode 100644 index 00000000000..a71180a253a --- /dev/null +++ b/vllm_omni/worker/runner_assisted_metadata.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Sequence +from typing import NamedTuple, Protocol + + +class RunnerAssistedFullAttentionMetadataRequest(NamedTuple): + num_reqs_padded: int + for_cudagraph_capture: bool + + +class RunnerAssistedAttentionMetadataProvider(Protocol): + """Optional model hook for runner-built full attention metadata.""" + + def get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_scheduled_tokens: Sequence[int], + num_computed_tokens: Sequence[int], + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest | None: ... + + def set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> None: ... diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py new file mode 100644 index 00000000000..c38a9bbdce4 --- /dev/null +++ b/vllm_omni/worker/sampling_utils.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sampling-state guards shared by the GPU and NPU AR model runners.""" + +import torch +from vllm.logger import init_logger +from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsProcessor + +logger = init_logger(__name__) + +__all__ = ["sanitize_min_tokens_stop_ids"] + + +def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: int) -> None: + """Drop stop ids the model head cannot emit from min-tokens masking state. + + vLLM's input processor unconditionally folds the stage tokenizer's EOS id + into ``SamplingParams.all_stop_token_ids``. AR stages whose lm_head is + narrower than the tokenizer vocabulary (codec talkers such as Qwen3-TTS: + 3072 logits vs text EOS 151645) then crash on any ``min_tokens >= 1``: + ``MinTokensLogitsProcessor.apply`` writes ``-inf`` at an out-of-range + index and ``index_put_`` triggers a CUDA device-side assert (#4962). + + Out-of-range ids are unreachable for the head, so dropping them never + changes sampling or stopping behavior. The per-request stop-id set is + mutated in place (it is shared with the request's ``SamplingParams``), + so each request is sanitized at most once; the processor's device-side + mask slice is rebuilt only when an out-of-range id was found. + """ + for proc in logitsprocs.non_argmax_invariant: + if not isinstance(proc, MinTokensLogitsProcessor): + continue + min_toks = getattr(proc, "min_toks", None) + if not min_toks: + continue + needs_rebuild = False + for _, _, stop_tok_ids in min_toks.values(): + oob = [tok for tok in stop_tok_ids if tok >= logits_vocab] + if not oob: + continue + stop_tok_ids.difference_update(oob) + needs_rebuild = True + logger.warning_once( + "min_tokens: dropped stop token ids %s that exceed the logits vocabulary (%d); " + "the model head cannot emit them.", + str(sorted(oob)), + logits_vocab, + ) + if needs_rebuild: + reqs: list[int] = [] + tok_ids: list[int] = [] + for index, (_, _, stop_tok_ids) in min_toks.items(): + reqs.extend([index] * len(stop_tok_ids)) + tok_ids.extend(stop_tok_ids) + proc.logits_slice = ( + proc._device_tensor(reqs, torch.int32), + proc._device_tensor(tok_ids, torch.int32), + )